content: {- TooManyCells.Spatial.SummaryPlot
  Gregory W. Schwartz
  
  Collects functions pertaining to plotting the summary of statistics for all
  samples.
  -}
  
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE PackageImports #-}
  
  module TooManyCells.Spatial.SummaryPlot
      ( plotSummary
      ) where
  
  import Control.Exception (SomeException (..), try, evaluate)
  import Control.Monad (when, mfilter, guard)
  import Data.Bool (bool)
  import Data.Function (on)
  import Data.List (sortBy, groupBy, foldl', sort, maximumBy, zipWith4)
  import Data.Maybe (fromMaybe, mapMaybe, isNothing)
  import Ploterific.Plot.Plot (labelColorScale)
  import Ploterific.Plot.Types (ColorLabel (..))
  import Safe (headMay)
  import System.Environment (getArgs)
  import System.IO (hPutStrLn, stderr)
  import TextShow (showt)
  import qualified BirchBeer.Types as Birch
  import qualified Control.Foldl as Fold
  import qualified Control.Lens as L
  import qualified Data.ByteString.Lazy.Char8 as BL
  import qualified Data.ByteString.Char8 as B
  import qualified Data.Csv.Streaming as CSVStream
  import qualified Data.Csv as CSV
  import qualified Data.Foldable as F
  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.Read as T
  import qualified Data.Vector.Unboxed as VU
  import qualified Data.Vector as V
  import qualified Graphics.Vega.VegaLite as VL
  import qualified Graphics.Vega.VegaLite.Theme as VL
  import qualified Statistics.Test.KruskalWallis as S
  import qualified "statistics" Statistics.Quantile as S
  import qualified "statistics" Statistics.Types as S
  import qualified Filesystem.Path as FP
  import qualified Turtle as TU
  
  import qualified TooManyCells.File.Types as Too
  import qualified TooManyCells.Spatial.Types as Too
  
  -- | Get all pairwise comparisons.
  pairwise :: (Eq a, Ord a) => [a] -> [[a]]
  pairwise xs = Set.toList
              . Set.fromList
              . filter (\[x,y] -> x /= y)
              . fmap sort
              $ (\x y -> [x,y]) <$> xs <*> xs
  
  -- | Safer Kruskal-Wallis test.
  safeKWTest :: (VU.Unbox a, Ord a) => [VU.Vector a] -> IO Double
  safeKWTest xs = do
    res <- try ( evaluate
               . maybe 1 (S.pValue . S.testSignificance)
               $ S.kruskalWallisTest xs
               ) :: IO (Either SomeException Double)
    res' <- case res of
              Left ex -> do
                hPutStrLn stderr $ "Caught exception, continuing with p = 1: " <> show ex
                return 1
              Right val -> return val
    return res'
  
  -- | Add on the relevant p-values
  pVal :: Birch.Feature
       -> [Map.Map T.Text T.Text]
       -> IO ([Map.Map T.Text T.Text], [Map.Map T.Text T.Text])
  pVal (Birch.Feature feature) ms = do
    statRows' <- statRows
    statSummary' <- statSummary
    return (statRows', statSummary')
    where
      statRows = do
        tests' <- tests
        pairwiseTests' <- pairwiseTests
        return
          . mconcat
          . zipWith3
              (\ x xs ys
              -> fmap (\ m
                      -> foldl' (\acc (!a, !b) -> Map.insert a b acc) m (("kwPValue", x):xs)
                      )
              . concatMap Too.unLabelList
              . Too.unVarList
              $ ys
              )
              tests'
              (fmap Too.unVarList pairwiseTests')
          $ grouped
      statSummary = do
        tests' <- tests
        pairwiseTests' <- pairwiseTests
        return
          . mconcat
          . zipWith4
                (\ s x xs ys
                -> fmap (\ m
                        -> foldl'
                            (\acc (!a, !b) -> Map.insert a b acc)
                            m
                            (("highestLabels", s):("kwPValue", x):xs)
                        )
                . concatMap Too.unLabelList
                . Too.unVarList
                $ ys
                )
                highestLabels
                tests'
                (fmap Too.unVarList pairwiseTests')
            $ summarizedVarList
      highestLabels :: [T.Text]
      highestLabels = fmap ( T.intercalate "/"
                           . fmap showt
                           . sortBy ((flip compare) `on` fst)  -- Descending
                           . concatMap
                               ( fmap
                                   (\ !x
                                   -> ( maybe
                                         (error "No median found")
                                         (either error fst . T.double)
                                      $ Map.lookup "median" x
                                      , Map.findWithDefault
                                         (error "No label found")
                                         "statLabel"
                                         x
                                      )
                                   )
                                   . Too.unLabelList
                               )
                           . Too.unVarList
                           )
                    $ summarizedVarList
      summarizedVarList :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
      summarizedVarList = Too.VarList
                        . fmap ( Too.LabelList
                               . (:[])
                               . getStatSummaries
                               . Too.unLabelList
                               )
                        . Too.unVarList
                      <$> grouped
      getStatSummaries !ms = Map.fromList
                           $ ("median", T.pack . show $ getMedianFromMaps ms)
                           : ("feature", feature)
                           : foldl'
                              (\ acc x
                              -> ( x
                                 , fromMaybe "" $ Map.lookup x =<< headMay ms
                                 )
                                 : acc
                              )
                              []
                              ["statLabel", "Var1", "Var2"]
      getMedianFromMaps = S.median S.s . VU.fromList . fmap getFeature
      kruskalOrBust = fmap (T.pack . show) . safeKWTest
      tests :: IO [T.Text]
      tests = mapM ( kruskalOrBust
                   . fmap (VU.fromList . fmap getFeature . Too.unLabelList)
                   . Too.unVarList
                   )
            $ grouped
      pairwiseTests :: IO [Too.VarList (T.Text, T.Text)]
      pairwiseTests = mapM ( fmap Too.VarList
                           . mapM
                              ( sequence
                              . L.over
                                  L._2
                                  ( kruskalOrBust
                                  . fmap
                                      ( VU.fromList
                                      . fmap getFeature
                                      . Too.unLabelList
                                      )
                                  )
                                  . Too.unCompare
                              )
                           . setupComparisonVar
                           )
                     $ grouped
      setupComparisonVar :: Too.VarList (Too.LabelList (Map.Map T.Text T.Text))
                         -> [Too.Compare (Too.LabelList (Map.Map T.Text T.Text))]
      setupComparisonVar = fmap ( Too.Compare
                                . L.over L._1 (T.intercalate "/")
                                . unzip
                                )
                         . (\x -> pairwise x :: [[(T.Text, Too.LabelList (Map.Map T.Text T.Text))]])
                         . fmap ( L.over L._2 Too.LabelList
                                . L.over L._1 (fromMaybe "" . headMay)
                                . unzip
                                . fmap (\ !x -> (fromMaybe "" $ compLabel x, x))
                                . Too.unLabelList
                                )
                         . Too.unVarList
      getFeature = maybe 0 (either error fst . T.double) . Map.lookup feature
      grouped :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
      grouped = fmap ( Too.VarList
                     . fmap Too.LabelList
                     . groupBy ((==) `on` compLabel)
                     . sortBy (compare `on` compLabel)
                     )
              . groupBy ((==) `on` compVar)
              . sortBy (compare `on` compVar)
              $ ms
      compLabel = Map.lookup "statLabel"
      compVar x = (Map.lookup "Var1" x, Map.lookup "Var2" x)
  
  -- -- | Add marksOther-markOther label if present to every group to compare with
  -- -- all others.
  -- addOthersOthersToEachGroup
  --   :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  --   -> [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  -- addOthersOthersToEachGroup varLists = fmap addToVarList noOthersVarList
  --   where
  --     noOthersVarList =
  --       filter (not . (\x -> isOthersList (&&) x && isOthersList (||) x)) varLists
  --     addToVarList xs = xs <> mconcat others
  --     others :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  --     others = filter (isOthersList (&&)) varLists
  --     isOthersList :: (Bool -> Bool -> Bool)
  --                  -> Too.VarList (Too.LabelList (Map.Map T.Text T.Text))
  --                  -> Bool
  --     isOthersList f xs = fromMaybe False $ do
  --       var1 <- getVarFromVarList "Var1" xs
  --       var2 <- getVarFromVarList "Var2" xs
  --       return $ f (var1 == "marksOther") (var2 == "marksOther")
  
  -- -- | Add markOther label if present to every group to compare with
  -- -- all others.
  -- addOthersToGroup :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  --                  -> [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  -- addOthersToGroup varLists = fmap addToVarList varLists
  --   where
  --     addToVarList xs = fromMaybe xs $ do
  --       var1 <- getVarFromVarList "Var1" xs
  --       others <- Map.lookup var1 otherMap
  --       return $ xs <> others
  --     otherMap = otherComparisonsMap varLists
  
  -- -- | Get marksOther-marksOther label if exists.
  -- otherComparisonsMap
  --   :: [Too.VarList (Too.LabelList (Map.Map T.Text T.Text))]
  --   -> Map.Map T.Text (Too.VarList (Too.LabelList (Map.Map T.Text T.Text)))
  -- otherComparisonsMap = Map.unions . mapMaybe getVarListMap
  --   where
  --     getVarListMap xs = do
  --       var1 <- getVarFromVarList "Var1" xs
  --       var2 <- getVarFromVarList "Var2" xs
  --       guard (var2 == "marksOther")
  --       return $ Map.singleton var1 xs
  
  -- | Return true if this is an other comparison row
  isOther :: Map.Map T.Text T.Text -> Bool
  isOther m = fromMaybe False $ do
        var1 <- Map.lookup "Var1" m
        var2 <- Map.lookup "Var2" m
        return (var1 == "marksOther" || var2 == "marksOther")
  
  -- | Get a Var column value from a VarList
  getVarFromVarList :: T.Text
                    -> Too.VarList (Too.LabelList (Map.Map T.Text T.Text))
                    -> Maybe T.Text
  getVarFromVarList var xs = (headMay . Too.unVarList $ xs)
                         >>= headMay . Too.unLabelList
                         >>= Map.lookup var
  -- | Parse a file into CSV.
  parseFile :: Maybe Too.StateLabelMap -> TU.FilePath -> IO [Map.Map T.Text T.Text]
  parseFile slm file = do
    contents <- BL.readFile . T.unpack . TU.format TU.fp $ file
    -- Force handles to close
    let !rows = either error ( fmap (insertSampleLabel slm file)
                             . insertOldVar2
                             . F.toList
                             . snd
                             )
              . CSVStream.decodeByName
              $ contents
    return rows
  
  -- | Insert samples and labels in a row. Labels are for plotting niceness.
  insertSampleLabel :: Maybe Too.StateLabelMap
                    -> TU.FilePath
                    -> Map.Map T.Text T.Text
                    -> Map.Map T.Text T.Text
  insertSampleLabel slm file m =
    Map.insert "label" (fromMaybe "No label" label)
      . Map.insert "statLabel" statLabel
      . Map.insert "sample" sample
      $ m
    where
      statLabel = maybe (fromMaybe "No label" alternateLabel) id label
      label
        | isOldVar2OtherOnly m = (\x y -> x <> "-" <> y)
                             <$> state
                             <*> (Just "marksOther")
        | otherwise = state
        where
          state = fmap Birch.unLabel
                . (=<<) (Map.lookup (Birch.Id sample) . Too.unStateLabelMap)
                $ slm
      alternateLabel = (\x y -> x <> "-" <> y)
                   <$> Map.lookup "Var1" m
                   <*> Map.lookup "Var2" m
      sample = Birch.unSample $ getSample file
  
  -- | Insert the original Var2 for reference in the other comparison.
  insertOldVar2 :: [Map.Map T.Text T.Text] -> [Map.Map T.Text T.Text]
  insertOldVar2 ms = mapMaybe insertOldVar2 ms
    where
      insertOldVar2 m = do
        var2 <- Map.lookup "Var2" m
        relevantVar <- getRelevantVar m
        return
          . bool id (Map.insert "Var2" relevantVar) (isVar2OtherOnly m)
          . Map.insert "oldVar2" var2
          $ m
      getRelevantVar m = do
        var1 <- Map.lookup "Var1" m
        var2 <- Map.lookup "Var2" m
        if (var1 == var2)
          then return var1
          else do
            -- If Nothing, then it is Var1 vs. Var1 Other, so use Var1.
            oldVar2 <- Just . fromMaybe var1 . headMay . filter (/= var1) $ vars
            return oldVar2
      vars = filter (/= "marksOther")
           . mapMaybe (Map.lookup "Var2")
           $ ms
  
  -- | Check if Var2 is marksOther.
  isVar2OtherOnly :: Map.Map T.Text T.Text -> Bool
  isVar2OtherOnly m = fromMaybe False $ do
    var1 <- Map.lookup "Var1" m
    var2 <- Map.lookup "Var2" m
    return (var1 /= "marksOther" && var2 == "marksOther")
  
  -- | Check if oldVar2 is marksOther.
  isOldVar2OtherOnly :: Map.Map T.Text T.Text -> Bool
  isOldVar2OtherOnly m = fromMaybe False $ do
    var1 <- Map.lookup "Var1" m
    var2 <- Map.lookup "oldVar2" m
    return (var1 /= "marksOther" && var2 == "marksOther")
  
  -- | Check if Var1 is marksOther.
  isVar1Other :: Map.Map T.Text T.Text -> Bool
  isVar1Other m = fromMaybe False $ do
    var1 <- Map.lookup "Var1" m
    return (var1 == "marksOther")
  
  -- | Get sample from file path.
  getSample :: TU.FilePath -> Birch.Sample
  getSample = Birch.Sample . TU.format TU.fp . TU.dirname . p . p . p
    where
      p = TU.parent
  
  -- | Remove columns with "marksOther".
  removeOthers :: [Map.Map T.Text T.Text] -> [Map.Map T.Text T.Text]
  removeOthers =
    fmap (Map.filterWithKey (\k _ -> not . T.isInfixOf "marksOther" $ k))
  
  -- | Remove rows with "marksOther".
  removeOthersRows :: [Map.Map T.Text T.Text] -> [Map.Map T.Text T.Text]
  removeOthersRows = filter (\ x -> (not . isOther $ x)
                                 && (Map.lookup "oldVar2" x /= Just "marksOther")
                            )
  
  -- | Make sure all columns are represented in a row
  fillColumns :: [Map.Map T.Text T.Text] -> [Map.Map T.Text T.Text]
  fillColumns xs = fmap addDummy xs
    where
      addDummy = flip Map.union dummyColMap  -- Prefer original value, not dummy.
      dummyColMap = Map.fromList
                  . fmap (\ !x -> (x, "NA"))
                  . Set.toList
                  . Set.fromList
                  . concatMap Map.keys
                  $ xs
  
  plotSummary :: Too.OutputDirectory
              -> Too.IncludeOthersFlag
              -> Maybe Too.StateLabelMap
              -> Birch.Feature
              -> IO ()
  plotSummary (Too.OutputDirectory inDir) iof slm feature = do
    let inDir' = TU.fromText . T.pack $ inDir
    statsFiles <- TU.reduce Fold.list . TU.find (TU.suffix "stats.csv") $ inDir'
    statsParsed <- mapM (parseFile slm) statsFiles
  
    when (null statsParsed) $ print "No stats.csv files found, first generate spatial results"
  
    let outDir = inDir' FP.</> "summary"
  
    TU.mktree outDir
  
    let outputBasename = "summary_" <> Birch.unFeature feature
        outputPlot = outDir FP.</> TU.fromText (outputBasename <> ".html")
        outputStats = outDir FP.</> TU.fromText (outputBasename <> ".csv")
        rows' = bool
                  removeOthersRows
                  (filter (not . isVar1Other))
                  (Too.unIncludeOthersFlag iof)
              . mconcat
              $ statsParsed
        noLabelFlag = isNothing slm
  
    (rows, stats) <- pVal feature rows'
  
    let pValComparisons = fmap (T.intercalate "/" . fmap (fromMaybe ""))
                        . pairwise
                        . Set.toList
                        . Set.fromList
                        . fmap (Map.lookup "label")
                        $ rows'
        colorLabels = fmap ColorLabel
                    . Set.toAscList
                    . Set.fromList
                    $ fmap (fromMaybe "" . Map.lookup "label") rows
        feature' = Birch.unFeature feature
        dataSet = VL.dataFromColumns []
                . VL.dataColumn
                    "var1"
                    (VL.Strings . fmap (Map.findWithDefault "" "Var1") $ rows)
                . VL.dataColumn
                    "var2"
                    (VL.Strings . fmap (Map.findWithDefault "" "Var2") $ rows)
                . VL.dataColumn
                    "label"
                    (VL.Strings . fmap (Map.findWithDefault "" "label") $ rows)
                . VL.dataColumn
                    "sample"
                    (VL.Strings . fmap (Map.findWithDefault "" "sample") $ rows)
                . VL.dataColumn
                    "kwPValue"
                    ( VL.Numbers
                    . fmap
                        ( maybe 0 (either error fst . T.double)
                        . Map.lookup "kwPValue"
                        )
                    $ rows
                    )
                . (\ y
                  -> foldl'
                      (\ acc x
                      -> VL.dataColumn
                          x
                          ( VL.Numbers
                          . fmap ( maybe 0 (either error fst . T.double)
                                 . Map.lookup x
                                 )
                          $ rows
                          ) acc
                      )
                      y
                      pValComparisons
                  )
                $ VL.dataColumn
                    feature'
                    ( VL.Numbers
                    . fmap
                        ( maybe 0 (either error fst . T.double)
                        . Map.lookup feature'
                        )
                    $ rows
                    )
                    []
        transDens = transMax
                  . VL.density feature' [ VL.DnGroupBy ["var1", "var2", "label"]]
        transHref = VL.calculateAs ("'../' + datum.sample + '/projections/' + datum.sample + '_projection.html'") "url"
        transMax = VL.joinAggregate
                    [ VL.opAs VL.Max feature' "max" ]
                    [ VL.WGroupBy ["var1", "var2"] ]
                 . VL.calculateAs ("datum." <> feature' <> "/datum.max") "maxTrans"
                 . transHref
        picked = "picked"
        sel = VL.selection
            . VL.select
                picked
                VL.Interval
                [ VL.Encodings [ VL.ChX, VL.ChY ], VL.BindScales ]
        axisSize = VL.Axis [ VL.TickWidth 0.756, VL.DomainWidth 0.756 ]
        densitySpec = [ VL.mark VL.Area [ VL.MOpacity 0.5]
                      , sel []
                      , VL.encoding
                      . VL.position
                          VL.Y
                          [ VL.PName "density"
                          , VL.PmType VL.Quantitative
                          ]
                      . VL.position
                          VL.X
                          [ VL.PName "value"
                          , VL.PmType VL.Quantitative
                          ]
                      . VL.color
                          [ VL.MName "label"
                          , VL.MmType VL.Nominal
                          , labelColorScale colorLabels
                          ]
                      $ []
                      , VL.transform $ transDens []
                      ]
        tickSpec = [ VL.mark VL.Tick []
                   , VL.transform $ transMax []
                   , VL.encoding
                   . VL.position
                      VL.X
                      [VL.PName feature', VL.PmType VL.Quantitative]
                   . VL.tooltips ( fmap
                                    (\x ->  [VL.TName x, VL.TmType VL.Nominal])
                                    ( [ feature'
                                      , "sample"
                                      , "label"
                                      , "kwPValue"
                                      ]
                                   <> pValComparisons
                                    )
                                 )
                   . VL.color [ VL.MName "label"
                              , VL.MmType VL.Nominal
                              , labelColorScale colorLabels
                              ]
                   . VL.hyperlink [ VL.HName "url", VL.HmType VL.Nominal ]
                   $ []
                   ]
        boxSpec = [ VL.mark
                      VL.Boxplot
                      [ VL.MNoOutliers
                      , VL.MOpacity 0.5
                      , VL.MMedian [ VL.MColor "black" ]
                      ]
                  , VL.transform $ transMax []
                  , VL.encoding
                  . VL.position
                      VL.X
                      [VL.PName feature', VL.PmType VL.Quantitative]
                  . VL.position VL.Y [ VL.PNumber 30 ]
                  . VL.color
                      [ VL.MName "label"
                      , VL.MmType VL.Nominal
                      , labelColorScale colorLabels
                      ]
                  $ []
                  ]
        res = VL.resolve
            . VL.resolution ( VL.RAxis
                                [ (VL.ChX, VL.Independent)
                                , (VL.ChY, VL.Independent)
                                ]
                            )
        p = VL.toVegaLite [ dataSet
                          , VL.specification
                          $ VL.asSpec
                              [ VL.layer
                                  (fmap VL.asSpec [densitySpec, tickSpec, boxSpec])
                              ]
                          , VL.facet [ VL.RowBy
                                        [VL.FName "var1", VL.FmType VL.Nominal]
                                     , VL.ColumnBy
                                        [VL.FName "var2", VL.FmType VL.Nominal]
                                     ]
                          , res []
                          , VL.theme
                              ( VL.defaultConfig
                                  { VL.configHeight = Just 170
                                  , VL.configWidth = Just 200
                                  , VL.configFontSize = Just 8
                                  , VL.configTitleFontSize = Just 9.33333015
                                  }
                              )
                              (VL.configuration axisSize [])
                          ]
  
    -- Output plot
    VL.toHtmlFile (T.unpack . TU.format TU.fp $ outputPlot) p
  
    -- Output stats
    let header = V.fromList
               . fmap (B.pack . T.unpack)
               . Set.toList
               . Set.fromList
               . concatMap Map.keys
               $ stats
    BL.writeFile (T.unpack . TU.format TU.fp $ outputStats)
      . CSV.encodeByName header
      . fillColumns
      $ stats
  