flower-0.1: src/Flower.hs

-- FlowEr - FLOWgram ExtractoR

module Main (main) where

import Bio.Sequence.SFF
import Bio.Sequence.Fasta
import Bio.Sequence.FastQ

import System.IO (stdout)
import System.Environment (getArgs)
import Numeric (showFFloat)
import Data.List (intersperse, partition)
import Data.ByteString.Char8 (unpack,ByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString as B1
import qualified Data.ByteString.Lazy.Char8 as LB

import Data.Array.Unboxed
import Data.Array.ST
import Control.Monad.ST

main :: IO ()
main = do
  args <- getArgs
  let (opts,files) = partition (\p -> case p of ('-':_) -> True; _ -> False) args
  case opts of 
    ["-r"] -> mapM_ (\f -> hWriteFasta stdout . sffToSequence =<< readSFF f) files
    ["-q"] -> mapM_ (\f -> hWriteFastQ stdout . sffToSequence =<< readSFF f) files
    ["-f"] -> LB.putStrLn . LB.fromChunks . intersperse (B.pack "\n") . concat =<< mapM showflow files
    ["-h"] -> mapM_ (\f -> (putStr . sffToHistogram) =<< readSFF f) files

    _ -> error ("Usage: flower -[f|q|r] <file.sff> [<file2.sff> ..]\n"
                ++"  -r  output reads in Fasta format\n"
                ++"  -q  output in FastQ format\n"
                ++"  -f  output the flowgram in tabular format\n"
                ++"  -h  output a histogram table of flow values")
tab :: ByteString
tab = B.pack "\t"

-- these are clumsy, since we just might need the file name
showflow :: FilePath -> IO [ByteString]
showflow f = return . {- map (\s -> B.concat [B.pack f,t,s]) . -} showrun =<< readSFF f

fi :: Flow -> ByteString
fi = (!) farray 

farray :: Array Flow ByteString
farray = listArray (0,10000) [B.pack (showFFloat (Just 2) i "") | i <- [0,0.01..99.99::Double]]

showrun :: SFF -> [ByteString]
showrun (SFF h rs) = concatMap (showread h) rs

showread :: CommonHeader -> ReadBlock -> [ByteString]
showread h rd = let rn = read_name $ read_header rd
                    qgroups = qgroup (B1.unpack $ flow_index rd) (B1.unpack $ quality rd)
                    format p c v q = B.concat [rn,tab,B.pack (show p),tab,B.pack [c],tab,fi v,tab,B.pack (show q)]
                in zipWith4 format [(1::Int)..] (unpack $ flow h) (flowgram rd) qgroups

zipWith4 :: (a -> b -> c -> d -> e) -> [a] -> [b] -> [c] -> [d] -> [e]
zipWith4 f (a:as) (b:bs) (c:cs) (d:ds) =  f a b c d : zipWith4 f as bs cs ds
zipWith4 _ _ _ _ _ = []

-- | Take the unpacked index_offsets and quality values, and return 
--   a list of groups of quality values, each group corresponding to a flow value. 
--   Flow values < 0.5 result in empty groups.
qgroup :: [Index] -> [Qual] -> [[Qual]]
qgroup [] []       = let rest = []:rest in rest
qgroup is@(1:_) qs = let (iz,irest) = span (==0) (tail is)
                         (q1,qrest) = splitAt (length iz+1) qs
                     in q1 : qgroup irest qrest
qgroup (i:is) qs = [] : qgroup (i-1:is) qs

sffToHistogram :: SFF -> String
sffToHistogram (SFF h rs) = showHist . histogram (B.unpack $ flow h) . map flowgram $ rs

type Hist = UArray Flow Int

histogram :: String -> [[Flow]] -> (Hist,Hist,Hist,Hist)
histogram fl scores = runST $ do 
  let zero = newArray (0,9999) 0 :: ST s (STUArray s Flow Int)
  a <- zero
  c <- zero
  g <- zero
  t <- zero
  let ins1 ('A',i) = bump a i
      ins1 ('C',i) = bump c i
      ins1 ('G',i) = bump g i
      ins1 ('T',i) = bump t i
      ins1 (x,_)   = error ("Illegal character "++show x++" in flow!")
      bump ar i = readArray ar i >>= \x -> writeArray ar i (x+1)
  mapM_ ins1 (zip (cycle fl) (concat scores))
  a' <- unsafeFreeze a
  c' <- unsafeFreeze c
  g' <- unsafeFreeze g
  t' <- unsafeFreeze t
  return (a',c',g',t')

showHist :: (Hist,Hist,Hist,Hist) -> String
showHist (as,cs,gs,ts) = "Score\tA\tC\tG\tT\tsum\n" ++ 
    unlines [concat $ intersperse "\t" $ (showFFloat (Just 2) (fromIntegral sc/100::Double) "") : map show [as!sc,cs!sc,gs!sc,ts!sc, as!sc+cs!sc+gs!sc+ts!sc]
                 | sc <- [0..9999]]