samtools-iteratee-0.1: src/BamCount.hs

module Main
       where

import Control.Applicative
import Control.Monad.Reader
import qualified Data.ByteString.Char8 as BS
import Data.Int (Int64)
import Data.List
import Data.Maybe

import System.Console.GetOpt
import System.Environment
import System.IO

import qualified Data.Iteratee as Iter
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import qualified Bio.SamTools.Bam as Bam
import qualified Bio.SamTools.Iteratee as Bam

main :: IO ()
main = getArgs >>= handleOpt . getOpt RequireOrder optDescrs
    where handleOpt (_,    _,         errs@(_:_)) = usage (unlines errs)
          handleOpt (args, [bam], []) = either usage (doBamCount bam) $ argsToConf args
          handleOpt (_,    _,     []) = usage "Specify exactly one BAM input"
          usage errs = do prog <- getProgName
                          hPutStr stderr $ usageInfo prog optDescrs
                          hPutStrLn stderr errs

doBamCount :: FilePath -> Conf -> IO ()
doBamCount bam conf = Bam.withBamInFile bam $ \hin -> do
  tctio <- UM.replicate (Bam.nTargets . Bam.inHeader $ hin) 0
  let countBam b = flip (maybe (return ())) (Bam.targetID b) $ \tid -> do 
        n0 <- UM.read tctio tid
        UM.write tctio tid $! succ n0
      bamiter = Iter.joinI $ Iter.filter (wanted conf) $ Iter.mapM_ countBam
  Bam.enumInHandle hin bamiter >>= Iter.run
  U.freeze tctio >>= writeFile (confOutput conf) . targetTable (Bam.inHeader hin)
  
targetTable :: Bam.Header -> U.Vector Int64 -> String
targetTable h = unlines . U.ifoldr addTargetLine []
  where addTargetLine idx n = (targetLine :)
          where targetLine = intercalate "\t" [ BS.unpack $ Bam.targetSeqName h idx
                                              , show $ Bam.targetSeqLen h idx
                                              , show n
                                              ]

wanted :: Conf -> Bam.Bam1 -> Bool
wanted conf | confPerfect conf = maybe False (== 0) . Bam.nMismatch
            | otherwise = const True
                                    
data Conf = Conf { confOutput :: !FilePath
                 , confPerfect :: !Bool
                 } deriving (Show)

data Arg = ArgOutput { unArgOutput :: !String }
         | ArgPerfect
         deriving (Show, Read, Eq, Ord)

argOutput :: Arg -> Maybe String
argOutput (ArgOutput del) = Just del
argOutput _ = Nothing

optDescrs :: [OptDescr Arg]
optDescrs = [ Option ['o'] ["output"]  (ReqArg ArgOutput "OUTFILE") "Output filename"
            , Option []    ["perfect"] (NoArg ArgPerfect)           "Filter perfect (NM:i:0) reads"
            ]

argsToConf :: [Arg] -> Either String Conf
argsToConf = runReaderT conf
    where conf = Conf <$> 
                 findOutput <*>
                 (ReaderT $ return . elem ArgPerfect)
          findOutput = ReaderT $ maybe (Left "No output file") return . listToMaybe . mapMaybe argOutput