content: {- TooManyCells.Spatial.ProjectionPlot
  Gregory W. Schwartz
  
  Collects functions pertaining to plotting interactive figures of point proximity
  on the left with cumulative distribution functions of features on the right.
  -}
  
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE DeriveGeneric #-}
  
  module TooManyCells.Spatial.ProjectionPlot
      ( plotSpatialProjection
      ) where
  
  -- Remote
  import BirchBeer.Types (Feature (..), LabelMap (..), Label (..), getColNames, Id (..), Sample (..))
  import Data.Bool (bool)
  import Data.Colour.Palette.BrewerSet (Kolor, brewerSet, ColorCat (..) )
  import Data.Colour.Palette.Harmony (colorRamp)
  import Data.Colour.SRGB (sRGB24show)
  import Data.List (genericLength)
  import Data.Maybe (fromMaybe, isJust, catMaybes)
  import qualified Control.Foldl as Fold
  import qualified Control.Lens as L
  import qualified Data.ByteString.Char8 as BS
  import qualified Data.ByteString.Lazy.Char8 as B
  import qualified Data.Foldable as F
  import qualified Data.Map.Strict as Map
  import qualified Data.Set as Set
  import qualified Data.Text as T
  import qualified Data.Text.Read as T
  import qualified Data.Vector as V
  import qualified Graphics.Vega.VegaLite as VL
  import qualified Graphics.Vega.VegaLite.Theme as VL
  import qualified System.FilePath as FP
  import qualified Turtle as TU
  
  -- Local
  import TooManyCells.File.Types (OutputDirectory (..))
  import TooManyCells.Matrix.Types (SingleCells (..), ProjectionMap (..), X (..), Y (..))
  import TooManyCells.Spatial.Types (ColorMap (..), Range (..))
  import TooManyCells.Spatial.Utility
  
  -- | Get the minimum and maximum values for a projection map.
  getMinMax :: ProjectionMap -> ((Double, Double), (Double, Double))
  getMinMax =
    fromMaybe ((0, 0), (0, 0))
      . (L.sequenceOf L.both :: (Maybe (Double, Double), Maybe (Double, Double)) -> Maybe ((Double, Double), (Double, Double)))  -- m ((a, b), (c, d))
      . (L.over L.both (L.sequenceOf L.both) :: ((Maybe Double, Maybe Double), (Maybe Double, Maybe Double)) -> (Maybe (Double, Double), Maybe (Double, Double)))  -- (m (a, b), m (b, c))
      . (L.over L.both (Fold.fold ((,) <$> Fold.minimum <*> Fold.maximum)) :: ([Double], [Double]) -> ((Maybe Double, Maybe Double), (Maybe Double, Maybe Double))) -- ((m a, m b), (m c, m d))
      . (L.over (L._2 . L.mapped) unY :: ([Double], [Y]) -> ([Double], [Double]) )
      . L.over (L._1 . L.mapped) unX
      . unzip
      . fmap snd
      . Map.elems
      . unProjectionMap
  
  -- | Get the color mapping from feature to color.
  getColorMap :: [Feature] -> ColorMap
  getColorMap xs =
    ColorMap . Map.fromList . zip xs . colorRamp (length xs) . brewerSet Set1 $ 9
  
  -- | Get the feature selection of a distribution.
  featureSel :: Feature -> [VL.SelectSpec] -> VL.PropertySpec
  featureSel (Feature f) = VL.selection
                         . VL.select
                             ("pick_" <> f)
                             VL.Interval [VL.Encodings [VL.ChX], VL.Empty]
  
  -- | Color by features.
  colorByFeatures
    :: Maybe LabelMap -> ColorMap -> [Feature] -> VL.BuildEncodingSpecs
  colorByFeatures lm colorMap features =
          VL.color ( maybe  -- Choose whether to color by labels or features
                      [ VL.MDataCondition
                          featureSels
                          [ VL.MString "white" ]
                      ]
                      (\ x
                      -> [ VL.MName "label", VL.MmType VL.Nominal, labelColorScale x]
                      )
                      lm
                    )
    where
      featureSels = fmap
                      (\ f
                      -> ( VL.Selection ("pick_" <> unFeature f)
                         ,  [ VL.MString
                                ( maybe "white" (T.pack . sRGB24show)
                                . Map.lookup f
                                $ unColorMap colorMap
                                )
                            ]
                          )
                      )
                      features
  
  
  -- | Pseudo-legend for selection.
  legendSpec :: LabelMap -> VL.VLSpec
  legendSpec lm =
    VL.asSpec [ VL.mark VL.Circle []
              , selection []
              , encoding []
              ]
    where
      selection = VL.selection
                . VL.select
                    "pick_legend"
                    VL.Multi
                    [ VL.Fields [ "label" ], VL.Empty ]
      labels = Set.toAscList
             . Set.fromList
             . fmap unLabel
             . Map.elems
             . unLabelMap
             $ lm
      encoding =
        VL.encoding
          . VL.position VL.Y [ VL.PName "label", VL.PmType VL.Nominal ]
          . VL.color [ VL.MName "label"
                     , VL.MmType VL.Nominal
                     , labelColorScale lm
                     ,  VL.MLegend []
                     ]
  
  -- | Get the circle spec for plotting.
  getCircleSpec :: Maybe LabelMap -> Range -> ColorMap -> [Feature] -> VL.VLSpec
  getCircleSpec lm range colorMap features =
    VL.asSpec [ VL.mark VL.Circle []
              , VL.width plotWidth
              , VL.height plotHeight
              , circleEnc []
              , circleFilterTrans []
              ]
    where
      aspectRatio = (maxY range - minY range) / (maxX range - minX range)
      plotWidth = 800
      plotHeight = aspectRatio * plotWidth
      circleEnc =
        VL.encoding
          . VL.position VL.X [ VL.PName "x"
                              , VL.PmType VL.Quantitative
                              , VL.PAxis [ VL.AxTitle "X Axis" ]
                              , VL.PScale
                                  [ VL.SDomain (VL.DNumbers [minX range, maxX range])
                                  ]
                              ]
          . VL.position VL.Y [ VL.PName "y"
                              , VL.PmType VL.Quantitative
                              , VL.PAxis [ VL.AxTitle "Y Axis" ]
                              , VL.PScale
                                  [ VL.SDomain (VL.DNumbers [minY range, maxY range])
                                  ]
                              ]
          . colorByFeatures lm colorMap features
      circleFilterTrans =
        VL.transform
          . VL.filter
              ( VL.FCompose
                  ( bool (VL.Expr "false") (VL.Selection "pick_legend") (isJust lm)
            `VL.Or` F.foldl'
                    (\acc x -> VL.Or acc (VL.Selection $ "pick_" <> unFeature x))
                    (VL.Expr "false")
                    features
                  )
              )
  
  -- | Get a window spec of a feature for plotting.
  getWindowSpec :: ColorMap -> Feature -> VL.VLSpec
  getWindowSpec colorMap (Feature feature) =
    VL.asSpec [ VL.title feature []
              , VL.height 40
              , windowEnc []
              , windowTrans []
              , featureSel (Feature feature) []
              , VL.mark VL.Area []
              ]
    where
      windowTrans =
        VL.transform
          . VL.filter (VL.FExpr $ "datum." <> feature <> " > 0")
          . VL.window
              [([VL.WOp VL.CumeDist, VL.WField feature], "window_" <> feature)]
              [VL.WFrame Nothing (Just 0), VL.WSort [VL.WAscending feature]]
      windowEnc = VL.encoding
                 . VL.position VL.X [ VL.PName $  feature
                                    , VL.PmType VL.Quantitative
                                    , VL.PAxis [ VL.AxTitle "Expression"]
                                    ]
                 . VL.position VL.Y [ VL.PName $ "window_" <> feature
                                    , VL.PmType VL.Quantitative
                                    , VL.PAxis [ VL.AxTitle "Probability"]
                                    ]
                 . VL.color
                    [ VL.MString ( maybe "white" (T.pack . sRGB24show)
                                 . Map.lookup (Feature feature)
                                 $ unColorMap colorMap
                                 )
                    ]
  
  -- | The color scheme for the label field.
  labelColorScale :: LabelMap -> VL.MarkChannel
  labelColorScale lm = VL.MScale [ VL.SDomain (VL.DStrings labels)
                                 , VL.SRange (VL.RStrings colors)
                                 ]
    where
      labels =
        Set.toAscList . Set.fromList . fmap unLabel . Map.elems . unLabelMap $ lm
      colors =
        fmap (T.pack . sRGB24show) . colorRamp (length labels) . brewerSet Set1 $ 9
  
  plotSpatialProjection :: OutputDirectory
                        -> Maybe LabelMap
                        -> ProjectionMap
                        -> SingleCells
                        -> Sample
                        -> IO ()
  plotSpatialProjection outputDir' labelMap' pm sc sample = do
    let labelMap = LabelMap . Map.insert (Id "") (Label "NA") . unLabelMap
               <$> labelMap'  -- Insert dummy cell to account for NA labels.
        dataSet = scToVLData labelMap pm sc
        features = fmap Feature . V.toList . getColNames $ sc
        ((miX, maX), (miY, maY)) = getMinMax pm
        range = Range miX maX miY maY
        numWindowCols = ceiling . sqrt . genericLength $ features
  
        colorMap = getColorMap features
        allSelections =
          ( VL.FilterOp
          $ VL.FCompose
              ( F.foldl'
                (\acc x -> VL.And acc (VL.Selection $ "pick_" <> unFeature x))
                (VL.Expr "true")
                features
              )
          , [VL.MString "white"]
          )
  
        circleSpec = getCircleSpec labelMap range colorMap features
        windowSpecs = fmap (getWindowSpec colorMap) features
        allSpec = VL.hConcat
                $ [ maybe
                      circleSpec
                      (\ lm
                      -> VL.asSpec
                          [VL.columns 2, VL.vlConcat [circleSpec, legendSpec lm]]
                      )
                      labelMap
                  , VL.asSpec [VL.columns numWindowCols, VL.vlConcat windowSpecs]
                  ]
  
        p = VL.toVegaLite
              $ [ dataSet
                , VL.theme VL.defaultConfig []
                , allSpec
                ]
  
    TU.mktree . TU.fromText . T.pack $ unOutputDirectory outputDir'
    let fileName (Sample x) = T.unpack x <> "_projection.html"
        outputPath = unOutputDirectory outputDir' FP.</> fileName sample
    VL.toHtmlFile outputPath p
  