content: {- Copyright 2014-2019 NGLess Authors
   - License: MIT
   -}
  {-# LANGUAGE ScopedTypeVariables, FlexibleContexts #-}
  module Data.Sam
      ( SamLine(..)
      , SamGroup
      , samLength
      , readSamGroupsC'
      , readSamLine
      , encodeSamLine
      , isAligned
      , isPositive
      , isNegative
      , isFirstInPair
      , isSecondInPair
      , isSamHeaderString
      , hasSequence
      , matchSize
      , matchSize'
      , matchIdentity
  
      , samStatsC
      , samStatsC'
      , samIntTag
      ) where
  
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Builder as BB
  import qualified Data.ByteString.Char8 as B8
  import qualified Data.Conduit.List as CL
  import qualified Data.Conduit.Lift as C
  import qualified Data.Conduit as C
  import           Data.Conduit.Algorithms.Async as CAlg
  import           Data.Conduit.Algorithms.Utils (awaitJust)
  import qualified Data.Conduit.Combinators as CC
  import           Data.Conduit ((.|))
  import qualified Data.Vector as V
  import qualified Data.Vector.Mutable as VM
  import           Data.Strict.Tuple (Pair(..))
  import Data.Bits (testBit)
  import Data.List (intersperse)
  import Control.Error (note)
  import Control.DeepSeq
  
  import Data.Maybe
  import           Control.Monad.IO.Class
  import Control.Monad.Except
  import NGLess.NGError
  import Utils.Utils
  import Utils.Conduit (ByteLine(..))
  
  
  type SamGroup = [SamLine]
  
  data SamLine = SamLine
              { samQName :: !B.ByteString
              , samFlag :: {-# UNPACK #-} !Int
              , samRName :: {-# UNPACK #-} !B.ByteString
              , samPos :: {-# UNPACK #-} !Int
              , samMapq :: {-# UNPACK #-} !Int
              , samCigar :: {-# UNPACK #-} !B.ByteString
              , samRNext :: {-# UNPACK #-} !B.ByteString
              , samPNext :: {-# UNPACK #-} !Int
              , samTLen :: {-# UNPACK #-} !Int
              , samSeq :: {-# UNPACK #-} !B.ByteString
              , samQual :: {-# UNPACK #-} !B.ByteString
              , samExtra :: {-# UNPACK #-} !B.ByteString
              } | SamHeader !B.ByteString
               deriving (Eq, Show, Ord)
  
  
  instance NFData SamLine where
      rnf SamLine{} = ()
      rnf (SamHeader !_) = ()
  
  isHeader SamHeader{} = True
  isHeader SamLine{} = False
  {-# INLINE isHeader #-}
  
  samLength :: SamLine -> Int
  samLength = B8.length . samSeq
  {-# INLINE samLength #-}
  
  -- log 2 of N
  -- 4 -> 2
  isAligned :: SamLine -> Bool
  isAligned = not . (`testBit` 2) . samFlag
  {-# INLINE isAligned #-}
  
  -- 16 -> 4
  isNegative :: SamLine -> Bool
  isNegative = (`testBit` 4) . samFlag
  {-# INLINE isNegative #-}
  
  -- all others
  isPositive :: SamLine -> Bool
  isPositive = not . isNegative
  {-# INLINE isPositive #-}
  
  isFirstInPair :: SamLine -> Bool
  isFirstInPair = (`testBit` 6) . samFlag
  {-# INLINE isFirstInPair #-}
  
  isSecondInPair :: SamLine -> Bool
  isSecondInPair = (`testBit` 7) . samFlag
  {-# INLINE isSecondInPair #-}
  
  isSamHeaderString :: B.ByteString -> Bool
  isSamHeaderString s = not (B.null s) && (B.head s == 64) -- 64 is '@'
  {-# INLINE isSamHeaderString #-}
  
  hasSequence :: SamLine -> Bool
  hasSequence SamHeader{} = False
  hasSequence SamLine{samSeq=s} = s /= "*"
  {-# INLINE hasSequence #-}
  
  newtype SimpleParser a = SimpleParser { runSimpleParser :: B.ByteString -> Maybe (Pair a B.ByteString) }
  instance Functor SimpleParser where
      fmap f p = SimpleParser $ \b -> do
                                  (v :!: rest) <- runSimpleParser p b
                                  return $! (f v :!: rest)
  
  instance Applicative SimpleParser where
      pure v = SimpleParser (\b -> Just (v :!: b))
      f <*> g = SimpleParser (\b -> do
                                  (f' :!: rest) <- runSimpleParser f b
                                  (g' :!: rest') <- runSimpleParser g rest
                                  return $! (f' g' :!: rest'))
  
  encodeSamLine :: SamLine -> BB.Builder
  encodeSamLine (SamHeader b) = BB.byteString b
  encodeSamLine samline = mconcat . intersperse (BB.char7 '\t') $ map ($ samline)
      [ BB.byteString . samQName
      , BB.intDec     . samFlag
      , BB.byteString . samRName
      , BB.intDec     . samPos
      , BB.intDec     . samMapq
      , BB.byteString . samCigar
      , BB.byteString . samRNext
      , BB.intDec .     samPNext
      , BB.intDec     . samTLen
      , BB.byteString . samSeq
      , BB.byteString . samQual
      , BB.byteString . samExtra
      ]
  
  readSamLine :: B.ByteString -> Either NGError SamLine
  readSamLine line
      | B.null line = throwDataError "Unexpected empty line"
      | B8.head line == '@' = return (SamHeader line)
      | otherwise = case runSimpleParser samP line of
          Just (v :!: _) -> return v
          Nothing -> throwDataError ("Could not parse sam line "++show line)
  
  tabDelim :: SimpleParser B.ByteString
  tabDelim = SimpleParser $ \input -> do
      ix <- B8.elemIndex '\t' input
      return $! (B.take ix input :!: B.drop (ix+1) input)
  
  tabDelimOpts :: SimpleParser B.ByteString
  tabDelimOpts = SimpleParser $ \input ->
      case B8.elemIndex '\t' input of
           Just ix -> return $! (B.take ix input :!: B.drop (ix+1) input)
           Nothing -> return $! (B.empty :!: input)
  
  readIntTab = SimpleParser $ \b -> do
          (v,rest) <- B8.readInt b
          return $! (v :!: B.tail rest)
  restParser = SimpleParser $ \b -> Just (b :!: B.empty)
  samP = SamLine
      <$> tabDelim
      <*> readIntTab
      <*> tabDelim
      <*> readIntTab
      <*> readIntTab
      <*> tabDelim
      <*> tabDelim
      <*> readIntTab
      <*> readIntTab
      <*> tabDelim
      <*> tabDelimOpts
      <*> restParser
  
  {--
  Op     Description
  M alignment match (can be a sequence match or mismatch).
  I insertion to the reference
  D deletion from the reference.
  N skipped region from the reference.
  S soft clipping (clipped sequences present inSEQ)
  H hard clipping (clipped sequences NOT present inSEQ)
  P padding (silent deletion from padded reference).
  = sequence match.
  X sequence mismatch.
  --}
  
  matchSize :: Bool -> SamLine -> Either NGError Int
  matchSize includeI = matchSize' includeI False . samCigar
  
  matchSize' includeI includeSoft cigar
      | B8.null cigar = return 0
      | otherwise = case B8.readInt cigar of
          Nothing -> throwDataError ("could not parse cigar '"++B8.unpack cigar ++"'")
          Just (n,code_rest) -> do
              let code = B8.head code_rest
                  rest = B8.tail code_rest
                  n' = case code of
                              'M' -> n
                              '=' -> n
                              'X' -> n
                              'S'
                                  | includeSoft -> n
                              'I'
                                  | includeI -> n
                              _ -> 0
              r <- matchSize' includeI includeSoft rest
              return (n' + r)
  
  matchIdentity :: Bool -> SamLine -> Either NGError Double
  matchIdentity useNewer samline = do
      let errmsg = "Could not get NM tag for samline " ++ B8.unpack (samQName samline) ++ ", extra tags were: "++ B8.unpack (samExtra samline)
      errors <- note (NGError DataError errmsg) $ samIntTag samline "NM"
      len <- matchSize useNewer samline
      let toDouble = fromInteger . toInteger
          mid = toDouble (len - errors) / toDouble len
      return mid
  
  samIntTag :: SamLine -> B.ByteString -> Maybe Int
  samIntTag samline tname
      | isHeader samline = Nothing
      | otherwise = listToMaybe . mapMaybe gettag . B8.split '\t' . samExtra $ samline
      where
          gettag match
              | B.take 2 match == tname
                      && (fst <$> B8.uncons (B.drop 3 match)) == Just 'i' = fst <$> B8.readInt (B.drop 5 match)
              | otherwise = Nothing
  
  -- | take in *ByteLines* and transform them into groups of SamLines all
  -- refering to the same read
  --
  -- This works in chunks (vectors) for efficiency
  --
  -- Header lines are silently ignored
  --
  -- When respectPairs is False, then the two mates of the same fragment will be
  -- considered grouped in different blocks
  readSamGroupsC' :: forall m . (MonadError NGError m, MonadIO m) =>
              Int -- ^ number of threads
              -> Bool -- respectPairs
              -> C.ConduitT (V.Vector ByteLine) (V.Vector [SamLine]) m ()
  readSamGroupsC' mapthreads respectPairs = do
          CC.dropWhileE (isSamHeaderString . unwrapByteLine)
          CC.filter (not . V.null)
              .| asyncMapEitherC mapthreads (fmap groupByName . V.mapM (readSamLineCheck . unwrapByteLine))
              -- the groups may not be aligned on the group boundary, thus we need to fix them
              .| fixSamGroups
      where
          samQNameMate :: SamLine -> (B.ByteString, Bool)
          samQNameMate s
              | respectPairs = (samQName s, True)
              | otherwise = (samQName s, isFirstInPair s)
          readSamLineCheck sm
              | isSamHeaderString sm = throwDataError "SAM header appears after SAM data (malformed file?)"
              | otherwise = readSamLine sm
          groupByName :: V.Vector SamLine -> V.Vector [SamLine]
          groupByName vs = V.unfoldr groupByName' (0, error "null value", [])
              where
                  groupByName' :: (Int, (B.ByteString, Bool), [SamLine]) -> Maybe ([SamLine], (Int, (B.ByteString, Bool), [SamLine]))
                  groupByName' (ix, name, acc)
                          | ix == V.length vs = if null acc
                                                  then Nothing
                                                  else Just (reverse acc, (ix, error "null value", []))
                          | null acc = groupByName' (ix + 1, cur_tag, [cur])
                          | cur_tag == name = groupByName' (ix + 1, name, cur:acc)
                          | otherwise = Just (reverse acc, (ix, error "null value", []))
                      where
                          cur = vs V.! ix
                          cur_tag = samQNameMate cur
          fixSamGroups :: C.ConduitT (V.Vector [SamLine]) (V.Vector [SamLine]) m ()
          fixSamGroups = awaitJust fixSamGroups'
          fixSamGroups' :: V.Vector [SamLine] -> C.ConduitT (V.Vector [SamLine]) (V.Vector [SamLine]) m ()
          fixSamGroups' prev = C.await >>= \case
              Nothing -> C.yield prev
              Just cur -> case (V.last prev, V.head cur) of
                      (lastprev:_, curfirst:_)
                          | samQNameMate lastprev /= samQNameMate curfirst -> do
                              -- lucky case, the groups align with the vector boundary
                              C.yield prev
                              fixSamGroups' cur
                          | otherwise -> do
                              C.yield (V.slice 0 (V.length prev - 1) prev)
                              cur' <- liftIO $ injectLast (V.last prev) cur
                              fixSamGroups' cur'
                      _ -> throwShouldNotOccur "This should have been an impossible situation in `fixSamGroups`"
          injectLast :: [SamLine] -> V.Vector [SamLine] -> IO (V.Vector [SamLine])
          injectLast prev gs = do
              gs' <- V.unsafeThaw gs
              VM.modify gs' (prev ++ ) 0
              V.unsafeFreeze gs'
  
  samStatsC :: (MonadIO m) => C.ConduitT (V.Vector ByteLine) C.Void m (NGLess (Int, Int, Int))
  samStatsC = C.runExceptC $ readSamGroupsC' 1 True .| samStatsC'
  
  samStatsC' :: forall m. (MonadError NGError m) => C.ConduitT (V.Vector SamGroup) C.Void m (Int, Int, Int)
  samStatsC' = CL.foldM summarizeV (0, 0, 0)
      where
          add1if !v True = v+1
          add1if !v False = v
  
          summarizeV :: (Int, Int, Int) -> V.Vector SamGroup -> m (Int, Int, Int)
          summarizeV c vs = V.foldM summarize c vs
  
          summarize :: (Int, Int, Int) -> SamGroup -> m (Int, Int, Int)
          summarize c [] = return c
          summarize c (samline:_)
              | isHeader samline = return c
          summarize (!t, !al, !u) g = let
                      aligned = any isAligned g
                      sameRName = allSame (samRName <$> g)
                      unique = aligned && sameRName
              in return
                  (t + 1
                  ,add1if al aligned
                  ,add1if  u unique
                  )
  