content: {- TooManyCells.MakeTree.Print
  Gregory W. Schwartz
  
  Collects the functions pertaining to the printing of information for the graph.
  -}
  
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module TooManyCells.MakeTree.Print
      ( printClusterDiversity
      , printClusterInfo
      , printNodeInfo
      , printLabelMap
      , saveFragments
      ) where
  
  -- Remote
  import BirchBeer.Types
  import BirchBeer.Utility (getGraphLeaves, getGraphLeavesWithParents)
  import Control.Monad (join)
  import Data.List (genericLength, intercalate, foldl')
  import Data.Maybe (fromMaybe, mapMaybe, catMaybes, isJust)
  import Data.Monoid ((<>))
  import Data.Streaming.Zlib (WindowBits (..))
  import Safe (headMay)
  import TextShow (showt)
  import Turtle hiding (Size)
  import qualified Control.Lens as L
  import qualified Data.ByteString.Lazy.Char8 as B
  import qualified Data.Csv as CSV
  import qualified Data.Foldable as F
  import qualified Data.Graph.Inductive as G
  import qualified Data.HashMap.Strict as HMap
  import qualified Data.Map.Strict as Map
  import qualified Data.Sequence as Seq
  import qualified Data.Set as Set
  import qualified Data.Text as T
  import qualified Data.Text.Encoding as T
  import qualified Data.Vector as V
  import qualified Filesystem.Path as FP
  import qualified Turtle.Bytes as TB
  
  -- Local
  import TooManyCells.Diversity.Types
  import TooManyCells.File.Types
  import TooManyCells.MakeTree.Types
  import TooManyCells.MakeTree.Cluster
  import TooManyCells.Matrix.Types
  
  -- | Print the diversity of each leaf cluster.
  printClusterDiversity :: [(Cluster, Diversity, Size)] -> B.ByteString
  printClusterDiversity =
      (<>) "cluster,diversity,size\n" . CSV.encode
          . fmap ( L.over L._3 unSize
                 . L.over L._2 unDiversity
                 . L.over L._1 unCluster
                 )
  
  -- | Get the size and modularity path of each leaf cluster path. Modularity
  -- starts from the parent of the cluster to the root for modularity.
  clusterInfo :: (TreeItem a) => ClusterGraph a -> [(Cluster, [Double], [Int])]
  clusterInfo (ClusterGraph gr) =
      F.toList
          . fmap (\ !x -> (Cluster . fst . snd $ x, getQs x, getSizes x))
          . getGraphLeavesWithParents gr
          $ 0
    where
      getSizes :: ([G.Node], a) -> [Int]
      getSizes = fmap getSize . fst
      getSize :: G.Node -> Int
      getSize = foldl' (+) 0 . fmap (maybe 0 Seq.length . snd) . getGraphLeaves gr
      getQs :: ([G.Node], a) -> [Double]
      getQs = mapMaybe getQ . fst
      getQ :: G.Node -> Maybe Double
      getQ  = join . fmap (L.view (L._2. edgeDistance)) . headMay . G.lsuc gr
  
  -- | Get the information of each leaf cluster path. Modularity
  -- starts from the parent of the cluster to the root for modularity.
  printClusterInfo :: (TreeItem a) => ClusterGraph a -> B.ByteString
  printClusterInfo =
      (<>) "cluster,modularity,size\n" . CSV.encode
          . fmap ( L.over L._3 (T.intercalate "/" . fmap showt)
                 . L.over L._2 (T.intercalate "/" . fmap showt)
                 . L.over L._1 unCluster
                 )
          . clusterInfo
  
  -- | Get the counts and frequencies of all cell labels in a node.
  getComposition ::
       (TreeItem a) => LabelMap -> ClusterGraph a -> G.Node -> LabelCompositions
  getComposition (LabelMap lm) (ClusterGraph gr) n =
      LabelCompositions
          . Map.elems
          . Map.mapWithKey
              (\ !k !v -> LabelComposition k v (fromIntegral v / total))
          $ counts
    where
      counts = Map.fromListWith (+) . flip zip [1,1..] $ labels
      total  = genericLength labels
      labels = catMaybes
             . fmap (flip Map.lookup lm . getId)
             . concatMap (maybe [] F.toList . snd)
             . F.toList
             $ getGraphLeaves gr n
  
  -- | Get various properties about the nodes in the ClusterGraph.
  nodeInfo :: (TreeItem a) => Maybe LabelMap -> ClusterGraph a -> [NodeInfo]
  nodeInfo lm (ClusterGraph gr) = fmap getNodeInfo . G.nodes $ gr
    where
      getNodeInfo x = NodeInfo
                          x
                          (getSize x)
                          (getProportion x)
                          (getQ x)
                          (getSignificance x)
                          (fmap (\a -> getComposition a (ClusterGraph gr) x) lm)
                          (getNodeChildren gr x)
      getSize :: G.Node -> Int
      getSize = foldl' (+) 0 . fmap (maybe 0 Seq.length . snd) . getGraphLeaves gr
      getQ :: G.Node -> Maybe Double
      getQ  = join . fmap (L.view edgeDistance . snd) . headMay . G.lsuc gr
      getSignificance :: G.Node -> Maybe Double
      getSignificance =
        join . fmap (L.view edgeSignificance . snd) . headMay . G.lsuc gr
      getProportion :: G.Node -> Maybe Double
      getProportion n =
          case fmap fst $ G.lsuc gr n of
              [x, y] -> Just $ fromIntegral (getSize x) / fromIntegral (getSize y)
              []     -> Nothing
      getNodeChildren gr = Set.toAscList
                         . Set.fromList
                         . concatMap fst
                         . F.toList
                         . getGraphLeavesWithParents gr
  
  -- | Print the node information to a string.
  printNodeInfo ::
       (TreeItem a) => Maybe LabelMap -> ClusterGraph a -> B.ByteString
  printNodeInfo lm =
      (<>) "node,size,proportion,modularity,significance,composition,subtree\n"
          . CSV.encode
          . fmap (\ (NodeInfo n s p m sig c ch)
                 -> ( n
                    , s
                    , maybe "" show p
                    , maybe "" show m
                    , maybe "" show sig
                    , maybe "" show c
                    , intercalate "/" . fmap show $ ch
                    )
                 )
          . nodeInfo lm
  
  -- | Print the label map to a string.
  printLabelMap :: LabelMap -> B.ByteString
  printLabelMap = (<>) "item,label\n" . CSV.encode 
                . fmap ( L.over L._2 unLabel
                       . L.over L._1 unId
                       )
                . Map.toAscList
                . unLabelMap
  
  -- | Print the fragments to a file. Will reassign original barcodes from the
  -- original file to those in the single cell matrix by splitting the barcode on
  -- "BARCODE-LABEL". Mainly for use with --custom-label.
  saveFragments :: OutputDirectory
                -> SingleCells
                -> [Either MatrixFileType MatrixFileType]
                -> IO ()
  saveFragments (OutputDirectory output') sc files = sh $ do
    let out = (fromText . T.pack $ output') FP.</> "fragments.tsv.gz"
        newBarcodes = HMap.fromList
                    . fmap (\ (Cell !x) -> (fst . T.breakOn "-" $ x, x))
                    . V.toList
                    . L.view rowNames
                    $ sc
        assignRow :: T.Text -> Maybe T.Text
        assignRow = fmap ((<> "\n") . T.intercalate "\t")
                  . ( (L.ix 3)
                L.%%~ (flip HMap.lookup newBarcodes . fst . T.breakOn "-")
                    )
                  . T.splitOn "\t"
        getFile (Left (CompressedFragments (FragmentsFile file))) = file
        getFile x = error $ "saveFragments: Not a fragments.tsv.gz file: "
                 <> show x
        loadFile =
          TB.decompress (WindowBits 31) . TB.input . fromText . T.pack . getFile
  
    TB.output out
      . TB.compress 6 (WindowBits 31)
      . fmap (maybe "" T.encodeUtf8)
      . mfilter isJust
      . fmap (assignRow . lineToText)
      . toLines
      . TB.toUTF8
      . msum
      . fmap loadFile
      $ files
  