content: {- TooManyCells.Classify.Classify
  Gregory W. Schwartz
  
  Collects functions pertaining to classifying cells from reference populations.
  -}
  
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE BangPatterns #-}
  
  module TooManyCells.Classify.Classify
      ( classifyCells
      , classifyMat
      ) where
  
  -- Remote
  import BirchBeer.Types (getMatrix, Label (..))
  import Data.Function (on)
  import Data.List (maximumBy)
  import qualified Control.Lens as L
  import qualified Data.Sparse.Common as S
  import qualified Data.Vector as V
  import qualified Math.Clustering.Spectral.Sparse as Spectral
  
  -- Local
  import TooManyCells.Matrix.Types
  import TooManyCells.Matrix.Utility
  import TooManyCells.MakeTree.Adjacency
  
  -- | Classify cells in a SingleCells type from a list of references.
  classifyCells :: SingleCells -> [AggReferenceMat] -> [(Cell, (Label, Double))]
  classifyCells sc refs = zip (V.toList . L.view rowNames $ sc')
                        . fmap (getMaxLabel . zip refNames . S.toDenseListSV)
                        . S.toRowsL
                        $ classifyMat sc' refs'
    where
      (sc', refs') = unifyFeatures sc refs
      refNames = fmap (Label . unCell)
               . V.toList
               . L.view rowNames
               . unAggSingleCells
               $ unAggReferenceMat refs'
      getMaxLabel = maximumBy (compare `on` snd)
  
  -- | Classify matrices using cosine similarity. Results in a matrix with scores
  -- (cosine similarity) for each column and observations per row.
  classifyMat :: SingleCells -> AggReferenceMat -> S.SpMatrix Double
  classifyMat sc refs = scMat S.## S.transpose refMat
    where
      refMat = normScMat . unAggSingleCells . unAggReferenceMat $ refs
      scMat = normScMat sc
      normScMat = Spectral.unB . Spectral.b2ToB . Spectral.B2 . getMatrix
  
  -- | Force unification of features with reference features.
  unifyFeatures :: SingleCells -> [AggReferenceMat] -> (SingleCells, AggReferenceMat)
  unifyFeatures sc refs = (newSc, newRefs)
    where
      newSc = L.over rowNames (V.take nRowsSc)
            . L.over matrix (extractObsRows 0 (nRowsSc - 1))
            $ unified
      newRefs = AggReferenceMat
              . AggSingleCells
              . L.over rowNames (V.drop nRowsSc)
              . L.over matrix (extractObsRows nRowsSc (nRows - 1))
              $ unified
      nRefs = length refs
      (nRows, nCols) = S.dim . getMatrix $ unified
      (nRowsSc, nColsSc) = S.dim . getMatrix $ sc
      unified = mconcat
              $ sc : fmap (unAggSingleCells . unAggReferenceMat) refs
      extractObsRows lb ub (MatObsRow x) =
        MatObsRow $ S.extractSubmatrixRebalanceKeys x (lb, ub) (0, nCols - 1)
  
  -- | Get the name of a reference from an AggReferenceMat.
  getRefName :: AggReferenceMat -> Label
  getRefName = maybe (error "unifyFeatures: no reference data") (Label . unCell)
             . flip (V.!?) 0
             . L.view rowNames
             . unAggSingleCells
             . unAggReferenceMat
  
  -- | Get the reference vector from an AggReferenceMat.
  getRefVec :: AggReferenceMat -> S.SpVector Double
  getRefVec = flip S.extractRow 0
            . getMatrix
            . unAggSingleCells
            . unAggReferenceMat
  