content: {- TooManyCells.MakeTree.ClusterPeaks
  Gregory W. Schwartz
  
  Collects functions pertaining to printing cluster fragments and obtaining peaks
  per cluster.
  -}
  
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TupleSections #-}
  
  module TooManyCells.Peaks.ClusterPeaks
      ( saveClusterFragments
      , peakCallFiles
      ) where
  
  -- Remote
  import BirchBeer.Types (Cluster (..), Label, LabelMap (..), Id (..))
  import Codec.Compression.GZip (compress, decompress)
  import Control.Concurrent.Async.Pool (withTaskGroup, mapTasks)
  import Control.Monad (unless, when, void)
  import Control.Monad.Trans.Maybe (runMaybeT, MaybeT (..))
  import Control.Monad.Trans.Resource (runResourceT, MonadResource)
  import Data.Bool (bool)
  import Data.List (intercalate)
  import Data.List.Split (splitOn)
  import Data.Maybe (mapMaybe, fromMaybe, isNothing, isJust, fromJust)
  import Data.Streaming.Zlib (WindowBits (..))
  import Data.Tuple (swap)
  import GHC.Conc (getNumCapabilities)
  import Safe (headMay, atMay)
  import System.IO (hPutStrLn, stderr)
  import System.IO.Temp (withSystemTempFile)
  import System.Process (callCommand)
  import Text.Read (readMaybe)
  import TextShow (showt)
  import Turtle
  import Turtle.Line
  import qualified Control.Foldl as Fold
  import qualified Control.Lens as L
  import qualified Data.ByteString.Lazy.Char8 as B
  import qualified Streaming.ByteString.Char8 as BS
  import qualified Data.HashMap.Strict as HMap
  import qualified Data.IntMap.Strict as IMap
  import qualified Data.IntSet as ISet
  import qualified Data.Map.Strict as Map
  import qualified Data.Set as Set
  import qualified Text.Printf as TP
  import qualified Data.Text as T
  import qualified Data.Text.Encoding as T
  import qualified Streaming as S
  import qualified Streaming.Prelude as S
  import qualified System.Directory as FP
  import qualified System.FilePath as FP
  import qualified Filesystem.Path as FSP
  import qualified System.IO as IO
  import qualified Turtle.Bytes as TB
  
  -- Local
  import TooManyCells.File.Types
  import TooManyCells.MakeTree.Types
  import TooManyCells.Matrix.Types
  import TooManyCells.Matrix.Utility
  import TooManyCells.Peaks.Types
  
  -- | External sort fragment file.
  externalFragmentsSort :: IO.FilePath -> IO ()
  externalFragmentsSort file =
    void . callCommand $ TP.printf  "sort -k1,1 -k2,2n -o %s %s" file file
  
  -- | External sort compressed fragment file.
  externalCompressedFragmentsSort :: IO.FilePath -> IO ()
  externalCompressedFragmentsSort file = do
    void . callCommand $ TP.printf  "gzip -d %s" file
    void
      . callCommand
      $ TP.printf  "cat %s | sort -k1,1 -k2,2n | gzip -c > %s" (FP.dropExtension file) file
    void . callCommand $ TP.printf  "rm %s" (FP.dropExtension file)
  
  -- | Get cluster fragment sizes from a split stream.
  getClusterTotalMap :: (MonadResource m)
                     => HMap.HashMap T.Text [Int]
                     -> S.Stream (S.Of [T.Text]) m r
                     -> m (IMap.IntMap Double)
  getClusterTotalMap cellClusterMap =
    S.fold_ (\acc x -> IMap.insertWith (+) x 1 acc) mempty id
      . S.concat
      . S.mapMaybe getCluster
    where
      getCluster :: [T.Text] -> Maybe [Int]
      getCluster row = atMay row 3 >>= flip HMap.lookup cellClusterMap
  
  -- | Get the genome coverage bedgraph.
  toGenomecov :: GenomecovCommand
              -> GenomeFile
              -> IMap.IntMap Double
              -> IO.FilePath
              -> IO.FilePath
              -> IO (Maybe ())
  toGenomecov (GenomecovCommand gcc) (GenomeFile gf) clusterTotalMap outputPath file = runMaybeT $ do
    output <- MaybeT
            . return
            . fmap ((FP.</>) outputPath . (<> "_cov.bdg"))
            . headMay
            . splitOn "."
            . FP.takeFileName
            $ file
    wigOut <- MaybeT
            . return
            . fmap ((FP.</>) outputPath . (<> "_cov.bw"))
            . headMay
            . splitOn "."
            . FP.takeFileName
            $ file
  
    cluster <- MaybeT
             . return
             $ readMaybe
           =<< (flip atMay 1 . splitOn "_" . FP.takeFileName $ file)
  
    scale <-
      MaybeT . return . fmap (1000000 /) $ IMap.lookup cluster clusterTotalMap
  
    _ <- MaybeT . fmap Just $ externalCompressedFragmentsSort file
    _ <- MaybeT . fmap Just . callCommand $ TP.printf gcc file gf scale output
    -- _ <- MaybeT . fmap Just $ externalFragmentsSort output  -- Removing as sorting ruins track line
    _ <- MaybeT . fmap Just . callCommand $ TP.printf "bedGraphToBigWig %s %s %s" output gf wigOut
    return ()
  
  -- | Write the line of the Fragment or BedGraph file, depending on
  -- normalization.
  writeLine :: (MonadResource m) => HMap.HashMap T.Text [Int]
                                 -> ClusterHandleMap
                                 -> [T.Text]
                                 -> m (Maybe ())
  writeLine cellClusterMap (ClusterHandleMap chm) all@(chrom:start:end:barcode:duplicateCount:_) = S.liftIO . runMaybeT $ do
    clusters <- MaybeT . return $ HMap.lookup barcode cellClusterMap
  
    let outputRow = compress
                  . flip B.append "\n"
                  . B.intercalate "\t"
                  $ fmap (B.fromStrict . T.encodeUtf8) all
        write :: Int -> MaybeT IO ()
        write c = do
          h <- MaybeT . return $ IMap.lookup c chm
          MaybeT . fmap Just $ B.hPutStr h outputRow
  
    -- Fragment file
    mapM_ write clusters
  -- writeLine (Just ctm) cellClusterMap outputPath all@(chrom:start:end:barcode:_:_) = S.liftIO . runMaybeT $ do
  --   cluster <- MaybeT . return $ HMap.lookup barcode cellClusterMap
  --   clusterSize <- MaybeT . return $ IMap.lookup cluster ctm
  
  --   let val = (1 / clusterSize) * 1000000 -- per million
  --       file = outputPath FP.</> T.unpack ("cluster_" <> showt cluster <> "_fragments.bdg")
  
  --   -- Write header if file does not exist.
  --   outputPathExists <- MaybeT . fmap Just $ FP.doesFileExist file
  --   MaybeT . fmap Just . unless outputPathExists $
  --     B.appendFile file
  --       . B.fromStrict
  --       . T.encodeUtf8
  --       $ "track type=bedGraph name=" <> "cluster_" <> showt cluster <> "\n"
  
  --   -- BedGraph file
  --   MaybeT
  --     . fmap Just
  --     . B.appendFile file
  --     . B.intercalate "\t"
  --     $ (fmap (B.fromStrict . T.encodeUtf8) [chrom, start, end, showt val])
  --    <> ["\n"]
  writeLine _ _ xs = error $ "writeLine: Unexpected number of columns, did you use the fragments.tsv.gz 10x format? Input: " <> show xs
  
  -- | Get cluster handle map to have several handles open.
  getClusterHandleMap :: IO.FilePath -> [Int] -> IO ClusterHandleMap
  getClusterHandleMap outputPath =
    fmap (ClusterHandleMap . IMap.fromList)
      . mapM (\x -> IO.openBinaryFile (getFileName x) IO.AppendMode >>= return . (x,))
    where
      getFileName cluster =
        outputPath FP.</> T.unpack ("cluster_" <> showt cluster <> "_fragments.tsv.gz")
  
  -- | Get the list of cluster output files.
  getClusterFiles :: IO.FilePath -> [Int] -> [IO.FilePath]
  getClusterFiles outputPath = fmap getFileName
    where
      getFileName cluster =
        outputPath FP.</> T.unpack ("cluster_" <> showt cluster <> "_fragments.tsv.gz")
  
  -- | Save the fragments per cluster. Loads in the original fragments file to
  -- stream out cluster specific fragments.
  saveClusterFragments :: OutputDirectory
                       -> BedGraphFlag
                       -> GenomeFile
                       -> GenomecovCommand
                       -> Maybe LabelMap
                       -> AllNodesFlag
                       -> PeakNodes
                       -> PeakNodesLabels
                       -> ClusterResults
                       -> Either MatrixFileType MatrixFileType
                       -> IO ()
  saveClusterFragments
    output
    (BedGraphFlag bgf)
    genome
    gcc
    lm
    (AllNodesFlag anf)
    (PeakNodes pkns)
    (PeakNodesLabels pknls)
    cr
    (Left (CompressedFragments (FragmentsFile file))) = do
    -- Where to place output files
    let outputPath = unOutputDirectory output FP.</> "cluster_fragments"
  
    restartDir outputPath
  
    -- Check for occurrence of label map.
    when (isNothing lm && (not $ IMap.null pknls)) $
      hPutStrLn IO.stderr "--labels-file argument required for --peak-node-labels. Ignoring filter ..."
  
    let checkLabel :: CellInfo -> Int -> Maybe Bool
        checkLabel cellInfo cluster = do
          validLabs <- IMap.lookup cluster pknls
          label <- join
                 $ fmap
                    ( Map.lookup (Id . unCell $ L.view barcode cellInfo)
                    . unLabelMap
                    )
                    lm
          return $ Set.member label validLabs
        clusterAssocList =
          fmap
              (\ (!cellInfo, !c)
              -> ( unCell $ _barcode cellInfo
                 , bool (filter (fromMaybe True . checkLabel cellInfo)) id (IMap.null pknls)  -- Specific node label selection.
                 . bool (filter (flip ISet.member pkns)) id (ISet.null pkns)  -- Specific node selection.
                 . fmap unCluster
                 $ bool (take 1 c) c anf
                 )
              )
            . _clusterList
            $ cr
        cellClusterMap :: HMap.HashMap T.Text [Int]
        cellClusterMap = HMap.fromList clusterAssocList
        processedStream = S.map (T.splitOn "\t")
                        . S.filter ((/=) (Just '#') . fmap fst . T.uncons)  -- No commented
                        . S.map T.decodeUtf8
                        . S.mapped BS.toStrict
                        . BS.lines
                        . decompressStreamAll (WindowBits 31) -- For gunzip
                        . BS.readFile
                        $ file
        clusters = Set.toList
                 . Set.fromList
                 . mconcat
                 . HMap.elems
                 $ cellClusterMap
  
    -- Open the cluster file handles.
    clusterHandleMap <- getClusterHandleMap outputPath clusters
  
    hPutStrLn IO.stderr "Streaming fragments to cluster fragment files ..."
    -- Stream in file for fragments.
    runResourceT
      . S.effects
      . S.mapMaybeM (writeLine cellClusterMap clusterHandleMap)
      . S.filter (maybe False (not . T.null) . headMay)
      $ processedStream
  
    hPutStrLn IO.stderr "Closing cluster fragment files ..."
    -- Close the cluster file handles.
    mapM_ IO.hClose . unClusterHandleMap $ clusterHandleMap
  
    -- Sort cluster files.
    hPutStrLn IO.stderr "Sorting cluster fragment files ..."
    mapM_ externalCompressedFragmentsSort . getClusterFiles outputPath $ clusters
  
    -- Stream in file for bedgraph.
    when bgf $ do
      hPutStrLn IO.stderr "Creating bedgraphs ..."
      let bedPath = unOutputDirectory output FP.</> "cluster_bedgraphs"
  
      restartDir bedPath
  
      clusterTotalMap <-
        runResourceT . getClusterTotalMap cellClusterMap $ processedStream
  
      files <- fmap (fmap (outputPath FP.</>)) . FP.listDirectory $ outputPath
      cores <- getNumCapabilities
      withTaskGroup cores $ \workers ->
        void . mapTasks workers . fmap (toGenomecov gcc genome clusterTotalMap bedPath) $ files -- mapTasks_ not working
      -- runResourceT
      --   . S.effects
      --   . S.mapMaybeM (writeLine (Just clusterTotalMap) cellClusterMap bedPath)
      --   $ processedStream
  saveClusterFragments _ _ _ _ _ _ _ _ _ _ = error "saveClusterFragments: Does not supported this matrix type. See too-many-cells -h for each entry point for more information"
  
  -- | Run the specified command on a file. Command should have \"%s\" where the
  -- output file and then input file (order is important) should be.
  -- For instance: `command --output %s --other --args --input %s`.
  peakCallFile :: PeakCommand
               -> GenomeFile
               -> IO.FilePath
               -> IO.FilePath
               -> IO IO.FilePath
  peakCallFile (PeakCommand c) (GenomeFile gf) outputDir file =
    withSystemTempFile "temp_fragments.tsv" $ \temp h -> do
      let output = (FP.</>) outputDir
                 . intercalate ("_")
                 . (<> ["peaks"])
                 . take 1
                 . splitOn "_fragments.tsv.gz"
                 . FP.takeFileName
                 $ file
          wigIn = output FP.</> (name <> "_treat_pileup.bdg")
          wigOut = output FP.</> (name <> "_treat_pileup.bw")
          name = fromMaybe (error "peakCallFile: Bad file name for fragments.")
               . headMay
               . splitOn "_fragments.tsv.gz"
               . FP.takeFileName
               $ file
      B.readFile file >>= B.hPut h . decompress >> IO.hClose h  -- Decompress file in temporary directory
      _ <- callCommand $ TP.printf c temp name output
      _ <- externalFragmentsSort wigIn
      -- _ <- callCommand $ TP.printf "bedGraphToBigWig %s %s %s" wigIn gf wigOut
      return $ output FP.</> (name <> "_peaks.narrowPeak")
  
  -- | Convert between filepaths.
  toTurtleFile :: IO.FilePath -> Turtle.FilePath
  toTurtleFile = Turtle.fromText . T.pack
  
  -- | Convert narrowPeak to bdg format.
  narrowPeakToBdg :: Turtle.FilePath -> IO Turtle.FilePath
  narrowPeakToBdg file = do
    let outFile = FSP.replaceExtension file "bdg"
  
    sh $
      output outFile
        . fmap fromJust
        . mfilter isJust
        . fmap ( join
               . fmap (textToLine . T.intercalate "\t")
               . (\xs -> mapM (atMay xs) [0, 1, 2, 4])
               . T.splitOn "\t"
               . lineToText
               )
        . input
        $ file
  
    return outFile
  
  -- | Return the union of peaks from all requested nodes.
  unionPeaks :: OutputDirectory -> [IO.FilePath] -> IO IO.FilePath
  unionPeaks (OutputDirectory path) files = do
    let outputPath = toTurtleFile $ path
        outputFile = outputPath Turtle.</> "union.bdg"
        union [x] _ = input x
        union xs stdin = inproc "bedtools" (["unionbedg", "-i"] <> fmap (format fp) xs) stdin
  
    bdgFiles <- mapM (narrowPeakToBdg . Turtle.fromText . T.pack) files
  
    sh $ do
      mktree outputPath
  
      output outputFile
        . inproc "bedtools" ["merge", "-i", "stdin"]
        . union bdgFiles
        $ mempty
  
    return . T.unpack . format fp $ outputFile
  
  -- | Return the original fragments file with unioned peaks.
  unionPeaksFragments :: OutputDirectory
                      -> IO.FilePath
                      -> [Either MatrixFileType MatrixFileType]
                      -> IO ()
  unionPeaksFragments (OutputDirectory path) unionFile files = sh $ do
    let outputFile = toTurtleFile $ path FP.</> "union_fragments.tsv.gz"
        getInputFile (Left (CompressedFragments (FragmentsFile file))) = file
        getInputFile _ = error "unionPeaksFragments: Does not supported this matrix type. See too-many-cells -h for each entry point for more information"
  
    TB.output outputFile
      . TB.compress 6 (TB.WindowBits 31)
      . TB.fromUTF8
      . fmap ( flip T.append "\n"
             . T.intercalate "\t"
             . fromMaybe (error "unionPeaksFragments: Not enough columns.")
             . (\xs -> mapM (atMay xs) [0, 1, 2, 6, 7])
             . T.splitOn "\t"
             . lineToText
             )
      . inproc "bedtools" ( ["intersect", "-wa", "-wb", "-a", T.pack unionFile, "-b"]
                           <> fmap (T.pack . getInputFile) files
                          )
      $ mempty
  
  -- | Run a specified command on all cluster files asynchronously.
  peakCallFiles :: PeakCommand
                -> GenomeFile
                -> OutputDirectory
                -> [Either MatrixFileType MatrixFileType]
                -> IO ()
  peakCallFiles pc gf (OutputDirectory path) fragmentsFiles = do
    let output = path FP.</> "cluster_peaks"
        fragmentsPath = path FP.</> "cluster_fragments"
  
    restartDir output
  
    files <- fmap (fmap (fragmentsPath FP.</>)) . FP.listDirectory $ fragmentsPath
  
    cores <- getNumCapabilities
    peakFiles <- withTaskGroup cores $ \workers ->
      mapTasks workers . fmap (peakCallFile pc gf output) $ files -- mapTasks_ not working
  
    unionFile <- unionPeaks (OutputDirectory output) peakFiles
    unionPeaksFragments (OutputDirectory output) unionFile fragmentsFiles
  
  -- | Delete a directory if it exists.
  restartDir :: IO.FilePath -> IO ()
  restartDir dir = do
    dirExists <- FP.doesDirectoryExist dir
    when dirExists $ FP.removeDirectoryRecursive dir
    FP.createDirectoryIfMissing True dir
    return ()
  
  -- -- | Filter a BED file by labels.
  -- labelFilterBed :: LabelMap
  --                -> Set.Set Label
  --                -> Either MatrixFileType MatrixFileType
  --                -> FP.FilePath
  --                -> IO FP.FilePath
  -- labelFilterBed lm labels (Left (CompressedFragments (FragmentsFile file))) inFile = sh $ do
  --   let validCells = Set.fromList
  --                  . fmap fst
  --                  . filter (flip Set.member labels . snd)
  --                  . Map.toAscList
  --                  . unLabelMap
  --                  $ lm
  
  --   out <- mktempfile ".temp" "temp.bed"
  
  --   output out
  --     . inproc "bedtools" ["intersect", "-a", "stdin", "-b", inFile]
  --     . mfilter (maybe False (flip Set.member validCells) . atMay 3 . T.splitOn "\t")
  --     . toLines
  --     . TB.decompress (WindowBits 31)
  --     . TB.input
  --     $ file
  