content: {- Copyright 2013-2022 NGLess Authors
   - License: MIT
   -}
  {-# LANGUAGE TemplateHaskell, RecordWildCards, CPP, FlexibleContexts #-}
  
  module Output
      ( OutputType(..)
      , MappingInfo(..)
      , AutoComment(..)
      , buildComment
      , commentC
      , outputListLno
      , outputListLno'
      , lookupNrSeqs
      , traceStatus
      , outputFQStatistics
      , outputMappedSetStatistics
      , writeOutputJS
      , writeOutputTSV
      , writeOutputTo
      , outputConfiguration
      , writeTransientMsg
      ) where
  
  import           Text.Printf (printf)
  import           System.IO (hFlush, hPutStr, hPutStrLn, hIsTerminalDevice, stdout, stderr, Handle)
  import           System.IO.Unsafe (unsafePerformIO)
  import           System.IO.SafeWrite (withOutputFile)
  import           Data.Maybe (maybeToList, fromMaybe, isJust)
  import           Data.IORef (IORef, newIORef, modifyIORef, readIORef, writeIORef)
  import           Data.List (sort)
  import           Data.Aeson ((.=))
  import qualified Data.Aeson as Aeson
  import           System.FilePath ((</>))
  import           Data.Aeson.TH (deriveToJSON, defaultOptions, Options(..))
  import           Data.Time (getZonedTime, ZonedTime(..))
  import           Data.Time.Format (formatTime, defaultTimeLocale)
  import qualified System.Console.ANSI as ANSI
  import           Control.Monad (forM_, when, unless)
  import           Control.Monad.IO.Class (liftIO)
  import           Control.Monad.Extra (whenJust, whenM)
  import           Numeric (showFFloat)
  import           Control.Arrow (first)
  import qualified Data.Text as T
  import qualified Data.Text.Encoding as TE
  import qualified Data.Conduit as C
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy.Char8 as BL8
  import qualified Data.ByteString.Lazy as BL
  
  import           System.Environment (lookupEnv)
  
  
  import           Data.FastQ (FastQEncoding(..), encodingName)
  import qualified Data.FastQ as FQ
  import Configuration
  import CmdArgs (Verbosity(..))
  import NGLess.NGLEnvironment
  import NGLess.NGError
  
  data AutoComment = AutoScript | AutoDate | AutoResultHash
                          deriving (Eq, Show)
  
  buildComment :: Maybe T.Text -> [AutoComment] -> T.Text -> NGLessIO [T.Text]
  buildComment c ac rh = do
          ac' <- mapM buildAutoComment ac
          return $ maybeToList c ++ concat ac'
      where
          buildAutoComment :: AutoComment -> NGLessIO [T.Text]
          buildAutoComment AutoDate = liftIO $ do
              t <- formatTime defaultTimeLocale "%a %d-%m-%Y %R" <$> getZonedTime
              return . (:[]) $ T.concat ["Script run on ", T.pack t]
          buildAutoComment AutoScript = (("Output generated by:":) . map addInitialIndent . T.lines . ngleScriptText) <$> nglEnvironment
          buildAutoComment AutoResultHash = return [T.concat ["Output hash: ", rh]]
          addInitialIndent t = T.concat ["    ", t]
  
  -- Output a comment as a conduit
  commentC :: Monad m => B.ByteString -> [T.Text] -> C.ConduitT () B.ByteString m ()
  commentC cmarker cs = forM_ cs $ \c -> do
                                  C.yield cmarker
                                  C.yield (TE.encodeUtf8 c)
                                  C.yield "\n"
  
  
  data OutputType = TraceOutput | DebugOutput | InfoOutput | ResultOutput | WarningOutput | ErrorOutput
      deriving (Eq, Ord)
  
  instance Show OutputType where
      show TraceOutput = "trace"
      show DebugOutput = "debug"
      show InfoOutput = "info"
      show ResultOutput = "result"
      show WarningOutput = "warning"
      show ErrorOutput = "error"
  
  data OutputLine = OutputLine !Int !OutputType !ZonedTime !String
  
  instance Aeson.ToJSON OutputLine where
      toJSON (OutputLine lno ot t m) = Aeson.object
                                          ["lno" .= lno
                                          , "time" .=  formatTime defaultTimeLocale "%a %d-%m-%Y %T" t
                                          , "otype" .= show ot
                                          , "message" .= m
                                          ]
  
  
  data BPosInfo = BPosInfo
                      { _mean :: !Int
                      , _median :: !Int
                      , _lowerQuartile :: !Int
                      , _upperQuartile :: !Int
                      } deriving (Show)
  $(deriveToJSON defaultOptions{fieldLabelModifier = drop 1} ''BPosInfo)
  
  data FQInfo = FQInfo
                  { fileName :: String
                  , scriptLno :: !Int
                  , gcContent :: !Double
                  , nonATCGFrac :: !Double
                  , encoding :: !String
                  , numSeqs :: !Int
                  , numBasepairs :: !Integer
                  , seqLength :: !(Int,Int)
                  , perBaseQ :: [BPosInfo]
                  } deriving (Show)
  
  $(deriveToJSON defaultOptions ''FQInfo)
  
  type TaggedFQInfo = (FilePath, FQInfo)
  
  data MappingInfo = MappingInfo
                  { mi_lno :: Int
                  , mi_inputFile :: FilePath
                  , mi_reference :: String
                  , mi_totalReads :: !Int
                  , mi_totalAligned :: !Int
                  , mi_totalUnique :: !Int
                  } deriving (Show)
  
  $(deriveToJSON defaultOptions{fieldLabelModifier = drop 3} ''MappingInfo)
  
  data SavedOutput = SavedOutput
          { outOutput :: [OutputLine]
          , fqOutput :: [TaggedFQInfo]
          , mapOutput :: [MappingInfo]
          }
  
  savedOutput :: IORef SavedOutput
  {-# NOINLINE savedOutput #-}
  savedOutput = unsafePerformIO (newIORef (SavedOutput [] [] []))
  
  lineClearNeeded :: IORef Bool
  {-# NOINLINE lineClearNeeded #-}
  lineClearNeeded = unsafePerformIO (newIORef False)
  
  lookupNrSeqs :: FilePath -> NGLessIO (Maybe Int)
  lookupNrSeqs fname = do
      SavedOutput{..} <- liftIO $ readIORef savedOutput
      return $ (numSeqs <$> lookup fname fqOutput)
  
  addOutputLine :: OutputLine -> SavedOutput -> SavedOutput
  addOutputLine !oline so@(SavedOutput ells _ _) = so { outOutput = oline:ells }
  
  addFQOutput !fq so@(SavedOutput _ fqs _) = so { fqOutput = fq:fqs }
  addMapOutput !mp so@(SavedOutput _ _ maps) = so { mapOutput = mp:maps }
  
  outputReverse (SavedOutput a b c) = SavedOutput (reverse a) (reverse b) (reverse c)
  
  -- | See `outputListLno'`, which is often the right function to use
  outputListLno :: OutputType      -- ^ Level at which to output
                      -> Maybe Int -- ^ Line number (in script). Use 'Nothing' for global messages
                      -> [String]
                      -> NGLessIO ()
  outputListLno ot lno ms = output ot (fromMaybe 0 lno) (concat ms)
  
  -- | Output a message using the global line number.
  -- This function takes a list as it is often a more convenient interface
  outputListLno' :: OutputType      -- ^ Level at which to output
                      -> [String]   -- ^ output. Will be 'concat' together (without any spaces or similar in between)
                      -> NGLessIO ()
  outputListLno' !ot ms = do
      lno <- ngleLno <$> nglEnvironment
      outputListLno ot lno ms
  
  shouldPrint :: Bool -- ^ is terminal
                  -> OutputType
                  -> Verbosity
                  -> Bool
  shouldPrint _ TraceOutput _  = False
  shouldPrint _      _ Loud    = True
  shouldPrint False ot Quiet   = ot == ErrorOutput
  shouldPrint False ot Normal  = ot > InfoOutput
  shouldPrint True  ot Quiet   = ot >= WarningOutput
  shouldPrint True  ot Normal  = ot >= InfoOutput
  
  traceStatus :: String -> NGLessIO ()
  traceStatus s = do
      outputListLno' TraceOutput [s]
      traceSet <- nConfTrace <$> nglConfiguration
      unless traceSet $ do
          writeTransientMsg s
          doColor <- isDoColor
          verb <- nConfVerbosity <$> nglConfiguration
          when (shouldPrint True InfoOutput verb) $ do
              let c = colorFor InfoOutput
                  st = if doColor
                          then ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull c]
                          else ""
                  rst = if doColor
                          then ANSI.setSGRCode [ANSI.Reset]
                          else ""
              writeTransientMsg (st ++ s ++ rst)
  
  isDoColor :: NGLessIO Bool
  isDoColor = do
      colorOpt <- nConfColor <$> nglConfiguration
      case colorOpt of
          ForceColor -> return True
          NoColor -> return False
          AutoColor -> do
              outputTo <- nConfOutputTo <$> nglConfiguration
              let outputHandle = case outputTo of
                      NGLOutStdout -> stdout
                      NGLOutStderr -> stderr
              isTerm <- liftIO $ hIsTerminalDevice outputHandle
              hasNOCOLOR <- isJust <$> liftIO (lookupEnv "NO_COLOR")
              return (isTerm && not hasNOCOLOR)
  
  output :: OutputType -> Int -> String -> NGLessIO ()
  output !ot !lno !msg = do
      outputTo <- nConfOutputTo <$> nglConfiguration
      let outputHandle = case outputTo of
              NGLOutStdout -> stdout
              NGLOutStderr -> stderr
      isTerm <- liftIO $ hIsTerminalDevice outputHandle
      verb <- nConfVerbosity <$> nglConfiguration
      traceSet <- nConfTrace <$> nglConfiguration
      doColor <- isDoColor
      liftIO $ do
          t <- getZonedTime
          modifyIORef savedOutput (addOutputLine $ OutputLine lno ot t msg)
          when (traceSet || shouldPrint isTerm ot verb) $ do
              let c = colorFor ot
                  st = if doColor
                          then ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull c]
                          else ""
                  rst = if doColor
                          then ANSI.setSGRCode [ANSI.Reset]
                          else ""
                  tformat = if traceSet -- when trace is set, output seconds
                                  then "%a %d-%m-%Y %T"
                                  else "%a %d-%m-%Y %R"
                  tstr = formatTime defaultTimeLocale tformat t
                  lineStr = if lno > 0
                                  then printf " Line %s" (show lno)
                                  else "" :: String
              clearTransientIfNeeded outputHandle
              hPutStrLn outputHandle $ printf "%s[%s]%s: %s%s" st tstr lineStr msg rst
              writeIORef lineClearNeeded False
  
  writeTransientMsg :: String -> NGLessIO ()
  writeTransientMsg m = do
      outputTo <- nConfOutputTo <$> nglConfiguration
      let outputHandle = case outputTo of
              NGLOutStdout -> stdout
              NGLOutStderr -> stderr
      liftIO $ whenM (hIsTerminalDevice outputHandle) $ do
          clearTransientIfNeeded outputHandle
          hPutStr outputHandle m
          ANSI.hSetCursorColumn outputHandle 0
          hFlush outputHandle
          writeIORef lineClearNeeded True
  
  clearTransientIfNeeded :: Handle -> IO ()
  clearTransientIfNeeded h =
      whenM (readIORef lineClearNeeded) $
          ANSI.hClearLine h
  
  
  colorFor :: OutputType -> ANSI.Color
  colorFor TraceOutput   = ANSI.White
  colorFor DebugOutput   = ANSI.White
  colorFor InfoOutput    = ANSI.Blue
  colorFor ResultOutput  = ANSI.Black
  colorFor WarningOutput = ANSI.Yellow
  colorFor ErrorOutput   = ANSI.Red
  
  
  encodeBPStats :: FQ.FQStatistics -> [BPosInfo]
  encodeBPStats res = map encode1 (FQ.qualityPercentiles res)
      where encode1 (mean, median, lq, uq) = BPosInfo mean median lq uq
  
  outputFQStatistics :: FilePath -> FilePath -> FQ.FQStatistics -> FastQEncoding -> NGLessIO ()
  outputFQStatistics rawpath fname stats enc = do
      lno' <- ngleLno <$> nglEnvironment
      let enc'    = encodingName enc
          sSize'  = FQ.seqSize stats
          nSeq'   = FQ.nSeq stats
          gc'     = FQ.gcFraction stats
          nATGC   = FQ.nonATCGFrac stats
          st      = encodeBPStats stats
          lno     = fromMaybe 0 lno'
          nbps    = FQ.nBasepairs stats
          binfo   = FQInfo fname lno gc' nATGC enc' nSeq' nbps sSize' st
      let p s0 s1  = outputListLno' DebugOutput [s0, s1]
      p "Simple Statistics completed for: " fname
      p "Number of base pairs: "      (show $ length (FQ.qualCounts stats))
      p "Encoding is: "               (show enc)
      p "Number of sequences: "   (show $ FQ.nSeq stats)
      liftIO $ modifyIORef savedOutput (addFQOutput (rawpath, binfo))
  
  outputMappedSetStatistics :: MappingInfo -> NGLessIO ()
  outputMappedSetStatistics mi@(MappingInfo _ _ ref total aligned unique) = do
          lno <- ngleLno <$> nglEnvironment
          let out = outputListLno' ResultOutput
          out ["Mapped readset stats (", ref, "):"]
          out ["Total reads: ", show total]
          out ["Total reads aligned: ", showNumAndPercentage aligned]
          out ["Total reads Unique map: ", showNumAndPercentage unique]
          out ["Total reads Non-Unique map: ", showNumAndPercentage (aligned - unique)]
          liftIO $ modifyIORef savedOutput (addMapOutput $ mi { mi_lno = fromMaybe 0 lno })
      where
          showNumAndPercentage :: Int -> String
          showNumAndPercentage v = concat [show v, " [", showFFloat (Just 2) ((fromIntegral (100*v) / fromIntegral total') :: Double) "", "%]"]
          total' = if total /= 0 then total else 1
  
  
  data InfoLink = HasQCInfo !Int
                  | HasStatsInfo !Int
      deriving (Eq, Show)
  instance Aeson.ToJSON InfoLink where
      toJSON (HasQCInfo lno) = Aeson.object
                                  [ "info_type" .= ("has_QCInfo" :: String)
                                  , "lno" .= show lno
                                  ]
      toJSON (HasStatsInfo lno) = Aeson.object
                                  [ "info_type" .= ("has_StatsInfo" :: String)
                                  , "lno" .= show lno
                                  ]
  
  data ScriptInfo = ScriptInfo String String [(Maybe InfoLink,T.Text)] deriving (Show, Eq)
  instance Aeson.ToJSON ScriptInfo where
     toJSON (ScriptInfo a b c) = Aeson.object [ "name" .= a,
                                              "time" .= b,
                                              "script" .= Aeson.toJSON c ]
  
  wrapScript :: [(Int, T.Text)] -> [FQInfo] -> [Int] -> [(Maybe InfoLink, T.Text)]
  wrapScript script tags stats = first annotate <$> script
      where
          annotate i
              | i `elem` (scriptLno <$> tags) = Just (HasQCInfo i)
              | i `elem` stats = Just (HasStatsInfo i)
              | otherwise =  Nothing
  
  writeOutputTo :: Handle -> IO ()
  writeOutputTo h = do
      SavedOutput fullOutput _ _ <- outputReverse <$> readIORef savedOutput
      forM_ fullOutput $ \(OutputLine lno ot t' msg) -> do
          let tformat = "%a %d-%m-%Y %T"
              tstr = formatTime defaultTimeLocale tformat t'
              lineStr = if lno > 0
                              then printf " Line %s" (show lno)
                              else "" :: String
          hPutStrLn h $ printf "[%s]%s (%s): %s" tstr lineStr (show ot) msg
  
  writeOutputJS :: FilePath -> FilePath -> T.Text -> NGLessIO ()
  writeOutputJS odir scriptName script = liftIO $ do
      SavedOutput fullOutput taggedFQStats mapStats <- outputReverse <$> readIORef savedOutput
      let fqStats = map snd taggedFQStats
      t <- getZonedTime
      let script' = zip [1..] (T.lines script)
          sInfo = ScriptInfo (odir </> "output.js") (show t) (wrapScript script' fqStats (mi_lno <$> mapStats))
      withOutputFile (odir </> "output.js") $ \hout ->
          BL.hPutStr hout (Aeson.encode $ Aeson.object
                          [ "output" .= fullOutput
                          , "processed" .= sInfo
                          , "fqStats" .= fqStats
                          , "mapStats" .= mapStats
                          , "scriptName" .= scriptName
                          ])
  
  
  -- | Writes QC stats to the given filepaths.
  writeOutputTSV :: Bool -- ^ whether to transpose matrix
                  -> Maybe FilePath -- ^ FastQ statistics
                  -> Maybe FilePath -- ^ Mapping statistics
                  -> NGLessIO ()
  writeOutputTSV transpose fqStatsFp mapStatsFp = liftIO $ do
          SavedOutput _ taggedFQStats mapStats <- outputReverse <$> readIORef savedOutput
          let fqStats = map snd taggedFQStats
          whenJust fqStatsFp $ \fp ->
              withOutputFile fp $ \hout ->
                  BL.hPut hout  . formatTSV $ encodeFQStats <$> fqStats
          whenJust mapStatsFp $ \fp ->
              withOutputFile fp $ \hout ->
                  BL.hPutStr hout . formatTSV $ encodeMapStats <$> mapStats
      where
          formatTSV :: [[(BL.ByteString, String)]] -> BL.ByteString
          formatTSV [] = BL.empty
          formatTSV contents@(h:_)
              | transpose = BL.concat ("\tstats\n":(formatTSV1 <$> zip [0..] contents))
              | otherwise = BL.concat [BL8.intercalate "\t" (fst <$> h), "\n",
                                      BL8.intercalate "\n" (asTSVline . fmap snd <$> contents), "\n"]
          formatTSV1 :: (Int, [(BL.ByteString, String)]) -> BL.ByteString
          formatTSV1 (i, contents) = BL.concat [BL8.concat [BL8.concat [BL8.pack . show $ i, ":", h], "\t", BL8.pack c, "\n"]
                                                                          | (h, c) <- contents]
          asTSVline = BL8.intercalate "\t" . map BL8.pack
  
          encodeFQStats FQInfo{..} = sort
                              [ ("file", fileName)
                              , ("encoding", encoding)
                              , ("numSeqs", show numSeqs)
                              , ("numBasepairs", show numBasepairs)
                              , ("minSeqLen", show (fst seqLength))
                              , ("maxSeqLen", show (snd seqLength))
                              , ("gcContent", show gcContent)
                              , ("nonATCGFraction", show nonATCGFrac)
                              ]
          encodeMapStats MappingInfo{..} = sort
                  [ ("inputFile", mi_inputFile)
                  , ("lineNumber", show mi_lno)
                  , ("reference", mi_reference)
                  , ("total", show mi_totalReads)
                  , ("aligned", show mi_totalAligned)
                  , ("unique", show mi_totalUnique)
                  ]
  
  outputConfiguration :: NGLessIO ()
  outputConfiguration = do
      cfg <- ngleConfiguration <$> nglEnvironment
      outputListLno' DebugOutput ["# Configuration"]
      outputListLno' DebugOutput ["\tdownload base URL: ", nConfDownloadBaseURL cfg]
      outputListLno' DebugOutput ["\tglobal data directory: ", nConfGlobalDataDirectory cfg]
      outputListLno' DebugOutput ["\tuser directory: ", nConfUserDirectory cfg]
      outputListLno' DebugOutput ["\tuser data directory: ", nConfUserDataDirectory cfg]
      outputListLno' DebugOutput ["\ttemporary directory: ", nConfTemporaryDirectory cfg]
      outputListLno' DebugOutput ["\tkeep temporary files: ", show $ nConfKeepTemporaryFiles cfg]
      outputListLno' DebugOutput ["\tcreate report: ", show $ nConfCreateReportDirectory cfg]
      outputListLno' DebugOutput ["\treport directory: ", nConfReportDirectory cfg]
      outputListLno' DebugOutput ["\tcolor setting: ", show $ nConfColor cfg]
      outputListLno' DebugOutput ["\tprint header: ",  show $ nConfPrintHeader cfg]
      outputListLno' DebugOutput ["\tsubsample: ", show $ nConfSubsample cfg]
      outputListLno' DebugOutput ["\tverbosity: ", show $ nConfVerbosity cfg]
      forM_ (nConfIndexStorePath cfg) $ \p ->
          outputListLno' DebugOutput ["\tindex storage path: ", p]
      outputListLno' DebugOutput ["\tsearch path:"]
      forM_ (nConfSearchPath cfg) $ \p ->
          outputListLno' DebugOutput ["\t\t", p]
  
  
  