{- TooManyCells.Program.Spatial
Gregory W. Schwartz

Spatial entry point for command line program.
-}

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ViewPatterns #-}

module TooManyCells.Program.Spatial where

-- Remote
import BirchBeer.Load (loadLabelData)
import BirchBeer.Types (LabelFile (..), Delimiter (..), LabelMap (..), Sample (..), Feature (..), Label (..))
import Control.Concurrent.Async.Pool (withTaskGroup, mapTasks)
import Control.Monad (join, forM_, guard, unless, when, void)
import Data.Bool (bool)
import Data.Maybe (isJust, fromMaybe)
import GHC.Conc (getNumCapabilities)
import Language.R.Instance as R
import Language.R.QQ
import System.Directory (getTemporaryDirectory)
import System.IO (hPutStrLn, stderr)
import Text.Read (readMaybe)
import qualified Control.Foldl as Fold
import qualified Control.Lens as L
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified H.Prelude as H
import qualified System.FilePath as FP
import qualified Turtle as TU

-- Local
import TooManyCells.Matrix.Load (loadProjectionMap)
import TooManyCells.Program.LoadMatrix (loadAllSSM)
import TooManyCells.Program.Options
import TooManyCells.Spatial.AnnoSpat (scToAnnoSpatFile, runAnnoSpat)
import TooManyCells.Spatial.ProjectionPlot (plotSpatialProjection)
import TooManyCells.Spatial.SummaryPlot (plotSummary)
import TooManyCells.Spatial.Relationships (spatialRelationshipsR)
import TooManyCells.Spatial.Utility (subsampleProjectionMap, markToText)
import qualified TooManyCells.File.Types as Too
import qualified TooManyCells.Matrix.Types as Too
import qualified TooManyCells.Program.Utility as Too
import qualified TooManyCells.Spatial.Types as Too

-- | General call for spatialRelationshipsR.
spatialRelationshipsCall :: Too.OutputDirectory
                         -> Too.SkipFinishedFlag
                         -> Too.PCFCrossFlag
                         -> Too.ProjectionMap
                         -> Maybe LabelMap
                         -> Too.SingleCells
                         -> Maybe Sample
                         -> Maybe Too.TopDistances
                         -> [Too.Mark]
                         -> IO ()
spatialRelationshipsCall (Too.OutputDirectory outDir) skipFlag' pcfCrossFlag' pm lm sc sample td marks = do
  let markString = T.unpack . T.intercalate "_" . fmap markToText $ marks
      folderName (Just (Sample x)) = T.unpack x <> "_" <> markString
      folderName Nothing = markString
      outDir' = Too.OutputDirectory
              $ outDir
         FP.</> folderName sample
      finalFile = TU.fromText
                . T.pack
                $ Too.unOutputDirectory outDir' FP.</> "stats.csv"

  -- If stats.csv exists and is not empty, then run everything
  !skip <- (&&)
       <$> (pure $ Too.unSkipFinishedFlag skipFlag')
       <*> (TU.testfile finalFile)
  when (not skip)
    $ spatialRelationshipsR outDir' pcfCrossFlag' pm lm sc td marks

-- | Helper for spatial relationships function.
spatialRelationshipsHelper :: Too.OutputDirectory
                           -> Too.SkipFinishedFlag
                           -> Too.PCFCrossFlag
                           -> Too.ProjectionMap
                           -> Maybe LabelMap
                           -> Too.SingleCells
                           -> Maybe Sample
                           -> Maybe Too.TopDistances
                           -> [Too.Mark]
                           -> IO ()
spatialRelationshipsHelper _ _ _ _ _ _ _ _ [] = pure ()
spatialRelationshipsHelper outDir skipFlag' pcfCrossFlag' pm lm sc sample td (fmap markToText -> ["ALL"]) = do
  cores <- getNumCapabilities
  withTaskGroup cores $ \workers ->
    void
      . mapTasks workers  -- mapTasks_ not working
      . fmap (\(!x, !y) -> spatialRelationshipsCall outDir skipFlag' pcfCrossFlag' pm lm sc sample td [x, y])
      $ (,) <$> getMarks lm <*> getMarks lm
  where
    getMarks Nothing = fmap Too.MarkFeature . V.toList . L.view Too.colNames $ sc
    getMarks (Just lm) =
      fmap Too.MarkLabel . Set.toList . Set.fromList . Map.elems . unLabelMap $ lm
spatialRelationshipsHelper outDir skipFlag' pcfCrossFlag' pm lm sc sample td marks =
  spatialRelationshipsCall outDir skipFlag' pcfCrossFlag' pm lm sc sample td marks

-- | Helper for plotting spatial summary.
spatialSummary :: Too.OutputDirectory
               -> Delimiter
               -> Too.IncludeOthersFlag
               -> Maybe Too.StateLabelsFile
               -> IO ()
spatialSummary outputDir' delimiter' iof stateLabelsFile' = do
  stateLabelMap <-
    mapM
      ( fmap (Too.StateLabelMap . unLabelMap)
      . loadLabelData delimiter'
      . LabelFile
      . Too.unStateLabelsFile
      )
      stateLabelsFile'
  let featureList = fmap
                      Feature
                      [ "meanCorr"
                      , "maxCorr"
                      , "minCorr"
                      , "topMaxCorr"
                      , "topMinCorr"
                      , "topMeanCorr"
                      , "negSwap"
                      , "posSwap"
                      , "longestPosLength"
                      , "longestNegLength"
                      , "maxPosWithVal"
                      , "logMaxPosWithVal"
                      , "maxPos"
                      , "minPos"
                      ]
  mapM_ (plotSummary outputDir' iof stateLabelMap) featureList

-- | Spatial path.
spatialMain :: Subcommand -> IO ()
spatialMain sub@(SpatialCommand opts) = H.withEmbeddedR R.defaultConfig $ do
  let onlySummaryFlag' = onlySummaryFlag opts
      outputDir' = Too.OutputDirectory . (output :: Spatial -> String) $ opts
      stateLabelsFile' = fmap Too.StateLabelsFile . stateLabelsFile $ opts
      delimiter' = Delimiter
                 . (delimiter :: LoadMatrixOptions -> Char)
                 . (loadMatrixOptions :: Spatial -> LoadMatrixOptions)
                 $ opts
      skipFlag' = Too.SkipFinishedFlag $ skipFinishedFlag opts
      iof' = Too.IncludeOthersFlag $ includeOthersFlag opts
      td' = topDistances opts

  unless onlySummaryFlag' $ do

    let readOrErr err = fromMaybe (error err) . readMaybe
        delimiter'        = Delimiter
                          . (delimiter :: LoadMatrixOptions -> Char)
                          . (loadMatrixOptions :: Spatial -> LoadMatrixOptions)
                          $ opts
        projectionFile' =
            maybe (error "--projection-file required") Too.ProjectionFile
              . (projectionFile :: Spatial -> Maybe String)
              $ opts
        annoSpatMarkerFile' = fmap Too.AnnoSpatMarkerFile . annoSpatMarkerFile $ opts
        annoSpatCommand' = Too.AnnoSpatCommand . annoSpatCommand $ opts

    scRes <- fmap (fromMaybe (error "Requires --matrix-path"))
          . loadAllSSM sub
          $ (loadMatrixOptions :: Spatial -> LoadMatrixOptions) opts
    let processedSc = fst scRes
        customLabelMap = snd scRes

    projectionMap <- loadProjectionMap projectionFile'

    let labelsFileError = hPutStrLn stderr "Warning: Problem in AnnoSpat, skipping label generation ..."
    labelsFile' <-
      case annoSpatMarkerFile' of
        Nothing ->
          pure . fmap LabelFile . (labelsFile :: Spatial -> Maybe String) $ opts
        (Just mf) -> (=<<) (maybe (labelsFileError >> pure Nothing) pure) . TU.reduce Fold.head $ do
          let annoOutDir = Too.OutputDirectory
                        $ Too.unOutputDirectory outputDir' FP.</> "AnnoSpat_out"
          tmpDir <- TU.liftIO $ fmap (TU.fromText . T.pack) getTemporaryDirectory
          tmpOut <- fmap Too.TempPath $ TU.mktempfile tmpDir "AnnoSpat_input.csv"
          startEndCols <- TU.liftIO $ scToAnnoSpatFile projectionMap processedSc tmpOut

          case startEndCols of
            Nothing -> pure Nothing
            (Just (startCol, endCol)) ->
              TU.liftIO
                $ runAnnoSpat
                    annoSpatCommand'
                    tmpOut
                    mf
                    annoOutDir
                    startCol
                    endCol

    labelMap <- if isJust labelsFile'
                  then mapM (loadLabelData delimiter') $ labelsFile'
                  else return customLabelMap

    let marks' = fmap
                  ( bool (Too.MarkFeature . Feature) (Too.MarkLabel . Label)
                  . isJust
                  $ labelMap
                  )
              . mark
              $ opts

    pcfCrossFlag' <- case (isJust labelMap, pcfCrossFlag opts) of
                      (False, True) -> do
                        hPutStrLn stderr "Warning: Continuous feature marks detected but pcfcross requested, ignoring pcfcross request ..."
                        return $ Too.PCFCrossFlag False
                      _ ->
                        return . Too.PCFCrossFlag . pcfCrossFlag $ opts

    let samples = Set.toList
                . Set.fromList
                . fmap fst
                . Map.elems
                . Too.unProjectionMap
                $ projectionMap

    forM_ samples $ \s -> do
      let s' = fromMaybe (Sample "total") s
          sOutLabel = T.unpack $ unSample s'
          projectionOutput = Too.OutputDirectory
                            . (FP.</> sOutLabel FP.</> "projections")
                            . Too.unOutputDirectory
                            $ outputDir'
          relationshipsOutput = Too.OutputDirectory
                              . (FP.</> sOutLabel FP.</> "relationships")
                              . Too.unOutputDirectory
                              $ outputDir'
          subPm = subsampleProjectionMap s projectionMap
      plotSpatialProjection projectionOutput labelMap subPm processedSc s'

      case marks' of
        [] -> pure ()
        m -> spatialRelationshipsHelper relationshipsOutput skipFlag' pcfCrossFlag' subPm labelMap processedSc s td' m

  spatialSummary outputDir' delimiter' iof' stateLabelsFile'
spatialMain _ = error "Wrong path in spatial, contact Gregory Schwartz for this error."