content: {- Copyright 2015-2020 NGLess Authors
   - License: MIT
   -}
  
  {-# LANGUAGE RankNTypes, FlexibleContexts, TypeFamilies #-}
  
  module Interpretation.Select
      ( executeSelect
      , executeMappedReadMethod
      , fixCigar
      , reinjectSequences
      ) where
  
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as BL
  import qualified Data.ByteString.Char8 as B8
  import qualified Data.ByteString.Builder as BB
  import qualified Data.Vector as V
  import qualified Data.Conduit as C
  import qualified Data.Conduit.Combinators as CC
  import           Data.Conduit ((.|))
  import qualified Data.Conduit.List as CL
  import qualified Data.Text as T
  import qualified Data.Text.Encoding as TE
  import           Data.Bits (Bits(..))
  import           Control.Monad.Except (throwError)
  import           Data.Either.Combinators (fromRight)
  import           Data.List (foldl')
  import           Data.Either.Extra (eitherToMaybe)
  import           Data.Tuple.Extra (fst3)
  import           Data.Ratio (Ratio, (%))
  
  import Data.Maybe
  
  import Data.Sam
  import FileManagement
  import FileOrStream (asSamStream, FileOrStream(..))
  import Language
  import Output
  
  import NGLess.NGLEnvironment
  import NGLess
  
  import Utils.Conduit
  
  data SelectCondition = SelectMapped | SelectUnmapped | SelectUnique
      deriving (Eq, Show)
  
  data MatchCondition = KeepIf [SelectCondition] | DropIf [SelectCondition]
      deriving (Eq, Show)
  
  _parseConditions :: KwArgsValues -> NGLessIO MatchCondition
  _parseConditions args = do
          keep_if <- lookupSymbolListOrScriptErrorDef (return []) "arguments to select" "keep_if" args
          drop_if <- lookupSymbolListOrScriptErrorDef (return []) "arguments to select" "drop_if" args
          keep_if' <- mapM asSC keep_if
          drop_if' <- mapM asSC drop_if
          case (keep_if', drop_if') of
              (cs, []) -> return $! KeepIf cs
              ([], cs) -> return $! DropIf cs
              (_, _) -> throwScriptError "To select, you cannot use both keep_if and drop_if"
      where
          asSC :: T.Text -> NGLessIO SelectCondition
          asSC "mapped" = return SelectMapped
          asSC "unmapped" = return SelectUnmapped
          asSC "unique" = return SelectUnique
          asSC c = throwShouldNotOccur ("Check failed.  Should not have seen this condition: '" ++ show c ++ "'")
  
  -- Notes on "Sequence reinjection"
  --
  -- Here is the case where it is necessary:
  --  1) the aligner includes the sequence in SamLine 1 (typically, the best hit, by its standards)
  --  2) we filter out SamLine 1, but keep SamLine 2 (this is a rare case, a few in a million)
  --  3) SamLine 2 does not have sequence information, so we need to reinject it
  --
  -- Here is why we need to rewrite the CIGAR string:
  --  1) there are two forms of trimming in CIGAR: hard & soft. Hard means that
  --  the bases were removed from the reads, whilst soft means that they are
  --  kept. If one of the lines uses hard trimming and the other one uses soft
  --  trimming, the read length may not match correctly
  --  2) we want to keep the full sequence, so we want to use soft trimming (if
  --  at all possible)
  matchConditions :: Bool -> MatchCondition -> [(SamLine,B.ByteString)] -> NGLess [(SamLine, B.ByteString)]
  matchConditions doReinject conds sg =
      let sg' = matchConditions' conds sg
          in if doReinject
              then reinjectSequencesIfNeeded sg sg'
              else pure sg'
  
  toStrictBS :: BB.Builder -> B.ByteString
  toStrictBS = BL.toStrict . BB.toLazyByteString
  
  
  reinjectSequencesIfNeeded :: [(SamLine,B.ByteString)] -> [(SamLine,B.ByteString)] -> NGLess [(SamLine,B.ByteString)]
  reinjectSequencesIfNeeded _ filtered@((SamHeader{},_):_) = return filtered
  reinjectSequencesIfNeeded orig filtered =
          if needsReinject (fmap fst filtered)
              then do
                  filtered' <- reinjectSequences (fmap fst orig) (fmap fst filtered)
                  return $ fmap (\s -> (s, toStrictBS $ encodeSamLine s)) filtered'
              else return filtered
      where
          needsReinject :: [SamLine] -> Bool
          needsReinject fs = let (fs1, fs2, fs0) = splitSamlines3 fs
                              in needsReinject' fs1 || needsReinject' fs2 || needsReinject' fs0
          needsReinject' :: [SamLine] -> Bool
          needsReinject' [] = False
          needsReinject' xs = not (any hasSequence xs)
  
  -- See note above on "Sequence reinjection" about why this function is necessary
  reinjectSequences :: [SamLine] -> [SamLine] -> NGLess [SamLine]
  reinjectSequences original filtered = case (splitSamlines3 original, splitSamlines3 filtered) of
      ((o1, o2, os), (f1, f2, fs)) -> do
          r1 <- reinjectSequences' o1 f1
          r2 <- reinjectSequences' o2 f2
          ss <- reinjectSequences' os fs
          return (r1 ++ r2 ++ ss)
  reinjectSequences' :: [SamLine] -> [SamLine] -> NGLess [SamLine]
  reinjectSequences' original f@(s@SamLine{}:rs)
      | not (any hasSequence f) = let
                  fixed = flip map (filter hasSequence original) $ \s' -> do
                      cigar <- fixCigar (samCigar s) (samLength s')
                      return (s { samSeq = samSeq s', samQual = samQual s', samCigar = cigar }:rs)
                  asum' [] = return f
                  asum' (x:xs) = case x of
                      Right{} -> x
                      Left{} -> asum' xs
              in asum' fixed
  reinjectSequences' _ f = return f
  -- See note above on "Sequence reinjection" about why this function is necessary
  fixCigar :: B.ByteString -> Int -> NGLess B.ByteString
  fixCigar prev n = do
      prevM <- matchSize' True True prev
      if prevM == n
          then return prev
          else do
              let prev' = B8.map (\c -> if c == 'H' then 'S' else c) prev
              prevM' <- matchSize' True True prev'
              if prevM' == n
                  then return prev'
                  else throwDataError ("Cannot fix CIGAR string \"" ++ B8.unpack prev ++ "\" to represent a sequence of length " ++ show n)
  
  matchConditions' :: MatchCondition -> [(SamLine,B.ByteString)] -> [(SamLine, B.ByteString)]
  matchConditions' _ r@[(SamHeader _,_)] = r
  matchConditions' (DropIf []) slines = slines
  matchConditions' (DropIf (c:cs)) slines = matchConditions' (DropIf cs) (drop1 c slines)
  
  matchConditions' (KeepIf []) slines = slines
  matchConditions' (KeepIf (c:cs)) slines = matchConditions' (KeepIf cs) (keep1 c slines)
  
  drop1 SelectUnmapped g = filter (isAligned . fst) g
  drop1 SelectMapped g = if any (isAligned . fst) g
                              then []
                              else g
  drop1 SelectUnique g = if isGroupUnique (map fst g) then [] else g
  
  keep1 SelectMapped g = filter (isAligned . fst) g
  keep1 SelectUnmapped g = if any (isAligned . fst) g
                              then []
                              else g
  keep1 SelectUnique g = if isGroupUnique (map fst g) then g else []
  
  -- readSamGroupsAsConduit :: (MonadIO m, MonadResource m) => FileOrStream -> C.Producer m [(SamLine, B.ByteString)]
  -- The reason we cannot just use readSamGroupsC is that we want to get the unparsed ByteString on the side
  readSamGroupsAsConduit istream paired =
          istream
              .| CC.concat
              .| readSamLineOrDie
              .| CL.groupBy groupLine
      where
          readSamLineOrDie = C.awaitForever $ \(ByteLine line) ->
              case readSamLine line of
                  Left err -> throwError err
                  Right parsed -> C.yield (parsed,line)
          groupLine (SamHeader _,_) _ = False
          groupLine _ (SamHeader _,_) = False
          groupLine (s0,_) (s1,_)
              | paired = samQName s0 == samQName s1
              | otherwise = samQName s0 == samQName s1 && isFirstInPair s0 == isFirstInPair s1
  
  
  executeSelect :: NGLessObject -> KwArgsValues -> NGLessIO NGLessObject
  executeSelect (NGOMappedReadSet name istream ref) args = do
      paired <- lookupBoolOrScriptErrorDef (return True) "select" "paired" args
      conditions <- _parseConditions args
      Just lno <- ngleLno <$> nglEnvironment
      doReinject <- do
          v <- ngleVersion <$> nglEnvironment
          if v < NGLVersion 0 8
              then do
                  outputListLno' WarningOutput ["Select changed behaviour (for the better) in ngless 0.8. If possible, upgrade your version statement."]
                  return False
              else return True
      let (fpsam, istream') = asSamStream istream
          stream =
              readSamGroupsAsConduit istream' paired
                  .| CL.mapM (runNGLess . matchConditions doReinject conditions)
                  .| streamedSamStats lno ("select_"++T.unpack name) ("select.lno"++show lno)
                  .| CL.map (V.fromList . map (ByteLine . snd))
  
          out = Stream [istream] ("selected_" ++ takeBaseNameNoExtensions fpsam ++ ".sam") stream
      return $! NGOMappedReadSet name out ref
  executeSelect o _ = throwShouldNotOccur ("NGLESS type checking error (Select received " ++ show o ++ ")")
  
  streamedSamStats lno ifile ref = C.passthroughSink (CL.map (V.singleton . map fst) .| samStatsC') $ \(total, aligned, unique) ->
      outputMappedSetStatistics (MappingInfo lno ifile ref total aligned unique)
  
  
  splitSamlines3 = reverse3 . foldl' add1 ([],[],[])
      where
          reverse3 (g1, g2, gs) = (reverse g1, reverse g2, reverse gs)
          add1 (g1,g2,gs) s
              | isFirstInPair s = (s:g1,g2,gs)
              | isSecondInPair s = (g1, s:g2, gs)
              | otherwise = (g1, g2, s:gs)
  
  data FilterAction = FADrop | FAUnmatch
      deriving (Eq)
  
  data SelectGroupOptions = SelectGroupOptions
      !Double -- ^ min id
      !Int -- ^ minMatchSize
      !Int -- ^ maxTrim
      !Bool -- ^ reverse
      !FilterAction
  
  applySelect :: Bool -> SelectGroupOptions -> SamGroup -> SamGroup
  applySelect useNewer (SelectGroupOptions minID minMatchSize (-1) rev act) =
                      case act of
                          FADrop -> filter passTest
                          FAUnmatch -> map (unmatchUnless passTest)
      where
          okID
              | minID < 0.0 = const True
              | otherwise = \s -> fromRight 0.0 (matchIdentity useNewer s) >= minID
          okSize :: SamLine -> Bool
          okSize
              | minMatchSize == -1 = const True
              | otherwise = \s -> fromRight 0 (matchSize useNewer s) >= minMatchSize
          rawTest s = okID s && okSize s
          passTest
              | rev = not . rawTest
              | otherwise = rawTest
  applySelect useNewer (SelectGroupOptions minID minMatch maxTrim rev act) = \samlines ->
                      let samlines' = applySelect useNewer (SelectGroupOptions minID minMatch (-1) rev act) samlines
                          (g1,g2,gs) = splitSamlines3 samlines'
                          (s1,s2,ss) = (seqSize g1, seqSize g2, seqSize gs)
                          okTrim :: Int -> SamLine -> Bool
                          okTrim seqlen = \s -> (seqlen - fromRight 0 (matchSize useNewer s) <= maxTrim)
                          in case act of
                              FADrop -> filter (okTrim s1) g1 ++ filter (okTrim s2) g2 ++ filter (okTrim ss) gs
                              FAUnmatch -> map (unmatchUnless $ okTrim s1) g1 ++ map (unmatchUnless $ okTrim s2) g2 ++ map (unmatchUnless $ okTrim ss) gs
  
  unmatchUnless c s
      | not (c s) = unmatch s
      | otherwise = s
  unmatch samline = samline { samFlag = (samFlag samline .|. 4) `clearBit` 1, samRName = "*", samCigar = "*" }
  seqSize :: SamGroup -> Int
  seqSize = maximum . map (B.length . samSeq)
  
  executeMappedReadMethod :: MethodName -> [SamLine] -> Maybe NGLessObject -> KwArgsValues -> NGLess NGLessObject
  executeMappedReadMethod (MethodName "flag") samlines (Just (NGOSymbol flag)) [] = do
          f <- getFlag flag
          return (NGOBool $ f samlines)
      where
          getFlag :: T.Text -> NGLess ([SamLine] -> Bool)
          getFlag "mapped" = return (any isAligned)
          getFlag "unmapped" = return (not . any isAligned)
          getFlag ferror = throwScriptError ("Flag " ++ show ferror ++ " is unknown for method flag")
  executeMappedReadMethod (MethodName "some_match") samlines (Just (NGOString target)) [] = return . NGOBool $ any ismatch samlines
      where
          ismatch :: SamLine -> Bool
          ismatch = (==target') . samRName
          target' = TE.encodeUtf8 target
  executeMappedReadMethod (MethodName "pe_filter") samlines Nothing [] = return . NGOMappedRead . filterPE $ samlines
  executeMappedReadMethod (MethodName "filter") samlines Nothing kwargs = do
      minID <- lookupIntegerOrScriptErrorDef (return (-1)) "filter method" "min_identity_pc" kwargs
      minMatchSize <- fromInteger <$> lookupIntegerOrScriptErrorDef (return (-1)) "filter method" "min_match_size" kwargs
      maxTrim <- fromInteger <$> lookupIntegerOrScriptErrorDef (return (-1)) "filter method" "max_trim" kwargs
      reverseTest <- lookupBoolOrScriptErrorDef (return False) "filter method" "reverse" kwargs
      useNewer <- lookupBoolOrScriptErrorDef (return False) "filter method" "__version11_or_higher" kwargs
      action <- lookupSymbolOrScriptErrorDef (return "drop") "filter method" "action" kwargs >>= \case
          "drop" -> return FADrop
          "unmatch" -> return FAUnmatch
          other -> throwScriptError ("unknown action in filter(): `" ++ T.unpack other ++"`.\nAllowed values are:\n\tdrop\n\tunmatch\n\tkeep\n")
      let samlines' = applySelect useNewer (SelectGroupOptions (fromInteger minID / 100.0) minMatchSize maxTrim reverseTest action) samlines
      return (NGOMappedRead samlines')
  executeMappedReadMethod (MethodName "unique") samlines Nothing [] = return . NGOMappedRead . mUnique $ samlines
  executeMappedReadMethod (MethodName "allbest") samlines Nothing kwargs = do
      useNewer <- lookupBoolOrScriptErrorDef (return False) "filter method" "__version11_or_higher" kwargs
      return . NGOMappedRead . mBesthit useNewer $ samlines
  executeMappedReadMethod m self arg kwargs = throwShouldNotOccur ("Method " ++ show m ++ " with self="++show self ++ " arg="++ show arg ++ " kwargs="++show kwargs ++ " is not implemented")
  
  filterPE :: [SamLine] -> [SamLine]
  filterPE slines = (filterPE' . filter isAligned) slines
      where
          filterPE' [] = []
          filterPE' (sl:sls)
              | isPositive sl = case findMatch sl slines of
                      Just sl2 -> sl:sl2:filterPE' sls
                      Nothing -> filterPE' sls
              | otherwise = filterPE' sls
          findMatch target = listToMaybe . filter (isMatch target)
          isMatch target other = isNegative other && samRName target == samRName other
  
  mUnique :: [SamLine] -> [SamLine]
  mUnique slines
      | isGroupUnique slines = slines
      | otherwise = []
  
  mBesthit :: Bool -> [SamLine] -> [SamLine]
  mBesthit _ [] = []
  mBesthit _ sl@[_] = sl
  mBesthit useNewer slines = let (g1,g2,gs) = splitSamlines3 slines
                      in mBesthit' useNewer g1 ++ mBesthit' useNewer g2 ++ mBesthit' useNewer gs
  
  mBesthit' :: Bool -> [SamLine] -> [SamLine]
  mBesthit' _ [] = []
  mBesthit' _ sl@[_] = sl
  mBesthit' useNewer samlines = case mapMaybe extract samlines of
          [] -> samlines
          extracted ->
              let
                  ms = maximum (fst3 <$> extracted)
                  -- fractional distance
                  ds :: Ratio Int
                  ds = minimum ((\(_, e, _) -> e % ms) <$> extracted)
              in mapMaybe (\(_, e, sl) ->
                          if e % ms <= ds
                              then Just sl
                              else Nothing) extracted
  
      where
          extract :: SamLine -> Maybe (Int, Int, SamLine)
          extract sl = do
              dist <- samIntTag sl "NM"
              size <- eitherToMaybe (matchSize useNewer sl)
              return (size, dist, sl)
  
  isGroupUnique :: [SamLine] -> Bool
  isGroupUnique [] = True
  isGroupUnique [_] = True
  isGroupUnique [f,s] = (isFirstInPair f /= isFirstInPair s) &&
                              (not (isAligned f && isAligned s) || (samRName f == samRName s))
  isGroupUnique _ = False
  
  