{- TooManyCells.Matrix.Utility
Gregory W. Schwartz

Collects helper functions in the program.
-}

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE QuasiQuotes #-}

module TooManyCells.Matrix.Utility
    ( matToRMat
    , scToRMat
    , sparseToHMat
    , sparseMatToSparseRMat
    , hToSparseMat
    , matToHMat
    , matToSpMat
    , spMatToMat
    , loadMatrixMarket
    , extractSc
    , writeMatrixLike
    , isCsvFile
    , getMatrixOutputType
    , matrixValidity
    , hashNub
    , decompressStreamAll
    , transposeSc
    , extractCellSc
    , extractCellsSc
    , removeCellsSc
    , subsetCellsSc
    , extractCellV
    , aggSc
    ) where

-- Remote
import BirchBeer.Types
import Codec.Compression.GZip (compress)
import Control.Monad (liftM2)
import Control.Monad.Managed (runManaged)
import Control.Monad.State (MonadState (..), State (..), evalState, execState, modify)
import Control.Monad.IO.Class (MonadIO)
import Data.Bool (bool)
import Data.Char (toUpper)
import Data.Function (on)
import Data.Hashable (Hashable)
import Data.List (maximumBy, isInfixOf, foldl')
import Data.Maybe (fromMaybe, catMaybes)
import Data.Matrix.MatrixMarket (Matrix(RMatrix, IntMatrix), Structure (..), writeMatrix')
import Data.Scientific (toRealFloat, Scientific)
import Data.Streaming.Zlib (WindowBits)
import Language.R as R
import Language.R.QQ (r)
import System.FilePath ((</>))
import TextShow (showt)
import qualified Control.Foldl as Fold
import qualified Control.Lens as L
import qualified Data.Attoparsec.Text as A
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as BL
import qualified Streaming.ByteString.Char8 as BS
import qualified Data.Clustering.Hierarchical as HC
import qualified Data.Graph.Inductive as G
import qualified Data.HashMap.Lazy as HMap
import qualified Data.HashSet as HSet
import qualified Data.IntMap.Strict as IMap
import qualified Data.IntervalMap.Strict as IntervalMap
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Sparse.Common as S
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL
import qualified Data.Text.Lazy.IO as TL
import qualified Data.Text.Lazy.Read as TL
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra as H
import qualified Streaming as Stream
import qualified Streaming.Cassava as Stream
import qualified Streaming.Prelude as Stream
import qualified Streaming.With.Lifted as SW
import qualified Streaming.Zip as S
import qualified System.Directory as FP

-- Local
import TooManyCells.File.Types
import TooManyCells.MakeTree.Types
import TooManyCells.Matrix.Types

-- | Convert a mat to an RMatrix.
matToRMat :: MatObsRow -> R s (RMatObsRow s)
matToRMat (MatObsRow m) = do
    [r| library(jsonlite) |]

    let mString = show . H.toLists . sparseToHMat $ m

    -- We want rows as observations and columns as features.
    mat <- [r| as.matrix(fromJSON(mString_hs)) |]
    return . RMatObsRow $ mat

-- | Convert a mat to an RMatrix.
sparseMatToSparseRMat :: SingleCells -> R s (RMatObsRow s)
sparseMatToSparseRMat sc = do
    [r| library(Matrix) |]

    let rowNamesR = fmap (T.unpack . unCell) . V.toList . L.view rowNames $ sc
        colNamesR = fmap (T.unpack . unFeature) . V.toList . L.view colNames $ sc
        idxs = S.toListSM . unMatObsRow . L.view matrix $ sc
        d :: [Double]
        d = (\(!m, !n) -> fmap fromIntegral [m, n])
          . S.dimSM
          . unMatObsRow
          . L.view matrix
          $ sc
        is :: [Double]
        is = fmap (fromIntegral . (\x -> x + 1) . L.view L._1) idxs
        js :: [Double]
        js = fmap (fromIntegral . (\x -> x + 1) . L.view L._2) idxs
        xs :: [Double]
        xs = fmap (L.view L._3) idxs

    namedMat <- [r| mat = sparseMatrix(i = is_hs, j = js_hs, x = xs_hs, dims = d_hs)
                    rownames(mat) = rowNamesR_hs
                    colnames(mat) = colNamesR_hs
                    mat
                |]

    return . RMatObsRow $ namedMat

-- | Convert a sc structure to an RMatrix.
scToRMat :: SingleCells -> R s (RMatObsRow s)
scToRMat sc = do
    [r| library(Matrix) |]

    let rowNamesR = fmap (T.unpack . unCell) . V.toList . _rowNames $ sc
        colNamesR = fmap (T.unpack . unFeature) . V.toList . _colNames $ sc

    mat <- fmap unRMatObsRow . matToRMat . _matrix $ sc

    -- namedMat <- [r| rownames(mat_hs) = rowNamesR_hs
    --                 colnames(mat_hs) = colNamesR_hs
    --             |]

    return . RMatObsRow $ mat

-- | Convert a sparse matrix to an hmatrix.
sparseToHMat :: S.SpMatrix Double -> H.Matrix H.R
sparseToHMat mat = H.assoc (S.dimSM mat) 0
                 . fmap (\(!x, !y, !z) -> ((x, y), z))
                 . S.toListSM
                 $ mat

-- | Convert a sparse matrix to an hmatrix.
hToSparseMat :: H.Matrix H.R -> S.SpMatrix Double
hToSparseMat =
    S.transposeSM . S.sparsifySM . S.fromColsL . fmap S.vr . H.toLists

-- | Convert a Matrix to an hmatrix Matrix. Assumes matrix market is 1 indexed.
matToHMat :: Matrix Scientific -> H.Matrix H.R
matToHMat (RMatrix size _ _ xs) =
    H.assoc size 0
        . fmap (\(!x, !y, !z) -> ((fromIntegral x - 1, fromIntegral y - 1), toRealFloat z))
        $ xs
matToHMat (IntMatrix size _ _ xs) =
    H.assoc size 0
        . fmap (\(!x, !y, !z) -> ((fromIntegral x - 1, fromIntegral y - 1), fromIntegral z))
        $ xs
matToHMat _ = error "\nInput matrix is not a Real matrix."

-- | Convert a Matrix to a sparse matrix.
matToSpMat :: Matrix Scientific -> S.SpMatrix Double
matToSpMat (RMatrix size _ _ xs) =
    S.fromListSM size
        . fmap (\(!x, !y, !z) -> (fromIntegral x - 1, fromIntegral y - 1, toRealFloat z))
        $ xs
matToSpMat (IntMatrix size _ _ xs) =
    S.fromListSM size
        . fmap (\(!x, !y, !z) -> (fromIntegral x - 1, fromIntegral y - 1, fromIntegral z))
        $ xs
matToSpMat _ = error "\nInput matrix is not a Real matrix."

-- | Convert a sparse matrix to a Matrix.
spMatToMat :: S.SpMatrix Double -> Matrix Double
spMatToMat mat = RMatrix (S.dimSM mat) (S.nzSM mat) General
               . fmap (\(!x, !y, !z) -> (x + 1, y + 1, z)) . S.toListSM
               $ mat

-- | Load a matrix market format.
loadMatrixMarket :: MatrixFileFolder -> IO (S.SpMatrix Double)
loadMatrixMarket (MatrixFile file) = do
    let toDouble [r, c, v] =
            ( either error fst . TL.decimal $ r
            , either error fst . TL.decimal $ c
            , either error fst . TL.double $ v
            )
        toDouble _ = error "Matrix market contains non-standard coordinate."

    (rows, cols, _):cs <- fmap
                            ( fmap (toDouble . TL.words)
                            . dropWhile ((== '%') . TL.head)
                            . TL.lines
                            )
                        . TL.readFile
                        $ file

    return
        . S.fromListSM (rows, cols)
        . fmap (\(!r, !c, !v) -> (r - 1, c - 1, v))
        $ cs

-- | Determine presence of matrix.
extractSc :: Maybe SingleCells -> SingleCells
extractSc = fromMaybe (error "Need to provide matrix in --matrix-path for this functionality.")

-- | Write a matrix to a file.
writeSparseMatrixLike :: MatrixLike a => MatrixTranspose -> MatrixFileFolder -> a -> IO ()
writeSparseMatrixLike (MatrixTranspose mt) (MatrixFolder folder) mat = do
  -- Where to place output files.
  FP.createDirectoryIfMissing True folder

  let writeCompressedFile x = BL.writeFile x . compress

  (=<<) (writeCompressedFile (folder </> "matrix.mtx.gz"))
    . writeMatrix'
    . spMatToMat
    . (bool S.transposeSM id mt) -- whether to transpose
    . getMatrix
    $ mat
  writeCompressedFile (folder </> "features.tsv.gz")
    . TL.encodeUtf8
    . TL.fromStrict
    . T.unlines
    . fmap (\x -> T.intercalate "\t" [x, x, "NA"])
    . V.toList
    . getColNames
    $ mat
  writeCompressedFile (folder </> "barcodes.tsv.gz")
    . TL.encodeUtf8
    . TL.fromStrict
    . T.unlines
    . V.toList
    . getRowNames
    $ mat

  return ()

-- | Print a dense matrix to a streaming string.
printDenseMatrixLike :: (MatrixLike a, Monad m)
                     => MatrixTranspose
                     -> a
                     -> BS.ByteStream m ()
printDenseMatrixLike (MatrixTranspose mt) mat =
  Stream.encode (Just $ Stream.header header)
    . Stream.zipWith (:) rowN
    . Stream.each
    . fmap (fmap showt . S.toDenseListSV)
    . S.toRowsL
    . (bool S.transposeSM id mt) -- To have cells as columns
    . getMatrix
    $ mat
  where
    header = (B.empty :)
           . fmap T.encodeUtf8
           . V.toList
           . getRowNames -- To have rows as columns
           $ mat
    rowN = Stream.each . V.toList . getColNames $ mat -- To have columns as rows

-- | Write a matrix to a dense file.
writeDenseMatrixLike :: (MatrixLike a)
                     => MatrixTranspose
                     -> MatrixFileFolder
                     -> a
                     -> IO ()
writeDenseMatrixLike mt (MatrixFile file) =
  runManaged . SW.writeBinaryFile file . printDenseMatrixLike mt

-- | Write a MatrixLike to a file (dense) or folder (sparse).
writeMatrixLike :: (MatrixLike a)
                => MatrixTranspose
                -> MatrixFileFolder
                -> a
                -> IO ()
writeMatrixLike mt o@(MatrixFolder _) = writeSparseMatrixLike mt o
writeMatrixLike mt o@(MatrixFile _) = writeDenseMatrixLike mt o

-- | Check if a name ends with .csv
isCsvFile :: FilePath -> Bool
isCsvFile = (== ".CSV") . fmap toUpper . reverse . take 4 . reverse

-- | Get matrix output format from input name.
getMatrixOutputType :: FilePath -> MatrixFileFolder
getMatrixOutputType x = bool (MatrixFolder x) (MatrixFile x) . isCsvFile $ x

-- | Check validity of matrix.
matrixValidity :: (MatrixLike a) => a -> Maybe String
matrixValidity mat
  | rows /= numCells || cols /= numFeatures =
      Just $ "Warning: mismatch in number of (features, cells) ("
                           <> show numFeatures
                           <> ","
                           <> show numCells
                           <> ") with matrix (rows, columns) ("
                           <> show cols
                           <> ","
                           <> show rows
                           <> "), will probably result in error."
  | otherwise = Nothing
  where
    (rows, cols) = S.dimSM . getMatrix $ mat
    numCells = V.length . getRowNames $ mat
    numFeatures = V.length . getColNames $ mat

-- | Same as Control.Foldl.nub but using HashSet.
hashNub :: (Hashable a, Eq a) => Fold.Fold a [a]
hashNub = Fold.Fold step (HSet.empty, id) fin
  where
    step (!s, !r) a =
      if HSet.member a s
        then (s, r)
        else (HSet.insert a s, r . (a :))
    fin (_, !r) = r []
{-# INLINABLE hashNub #-}

-- | Keep decompressing a compressed bytestream until exhaused.
decompressStreamAll :: (MonadIO m) => WindowBits -> BS.ByteStream m r -> BS.ByteStream m r
decompressStreamAll w bs = S.decompress' w bs >>= go
  where
    go (Left bs) = S.decompress' w bs >>= go
    go (Right r) = return r
{-# INLINABLE decompressStreamAll #-}

-- | Transpose a SingleCells type.
transposeSc :: SingleCells -> SingleCells
transposeSc (SingleCells (MatObsRow mat) rows cols) =
  SingleCells
    (MatObsRow . S.transposeSM $ mat)
    (fmap (Cell . unFeature) cols)
    (fmap (Feature . unCell) rows)

-- | Get a single cell from a SingleCells type.
extractCellSc :: Cell -> SingleCells -> SingleCells
extractCellSc cell sc = L.set rowNames (V.singleton cell)
                    . L.set matrix newMatrix
                    $ sc
  where
    newMatrix = MatObsRow
              . S.fromRowsL
              . (:[])
              . flip S.extractRow (getCellIdx sc cell)
              . unMatObsRow
              . L.view matrix
              $ sc

-- | Get a list of single cells for each single cell from a SingleCells type.
extractCellsSc :: SingleCells -> [SingleCells]
extractCellsSc sc = zipWith
                      (\ v c -> sc { _rowNames = V.singleton c
                                   , _matrix = MatObsRow $ S.fromRowsL [v]
                                   }
                      )
                      (S.toRowsL . unMatObsRow . L.view matrix $ sc)
                  . V.toList
                  . L.view rowNames
                  $ sc

-- | Get a single cell vector from a SingleCells type.
extractCellV :: Cell -> SingleCells -> S.SpVector Double
extractCellV cell sc =
  flip S.extractRow (getCellIdx sc cell) . unMatObsRow . L.view matrix $ sc

-- | Get a single cell vector from a SingleCells type using hashes.
extractCellHashV ::
  HMap.HashMap T.Text Int -> SingleCells -> Cell -> Maybe (S.SpVector Double)
extractCellHashV idxMap sc (Cell cell) =
  S.extractRow (unMatObsRow . L.view matrix $ sc) <$> idx
  where
    idx = HMap.lookup cell idxMap

-- | Remove single cells from a SingleCells type.
removeCellsSc :: [Cell] -> SingleCells -> SingleCells
removeCellsSc cells sc = L.set rowNames newRowNames
                     . L.set matrix newMatrix
                     $ sc
  where
    blacklistCells = Set.fromList cells
    isValidV = fmap (not . flip Set.member blacklistCells) . L.view rowNames $ sc
    newRowNames = fmap snd . V.filter fst . V.zip isValidV . L.view rowNames $ sc
    newMatrix = MatObsRow
              . S.fromRowsV
              . fmap snd
              . V.filter fst
              . V.zip isValidV
              . V.fromList
              . S.toRowsL
              . unMatObsRow
              . L.view matrix
              $ sc

-- | Get ordered subset of single cells from a SingleCells type (if they exist).
subsetCellsSc :: [Cell] -> SingleCells -> SingleCells
subsetCellsSc cells sc = L.set rowNames newRowNames
                       . L.set matrix newMatrix
                       $ sc
  where
    idxMap = HMap.fromList
           . flip zip [0..]
           . fmap unCell
           . V.toList
           . L.view rowNames
           $ sc
    newRowNames = V.fromList
                . catMaybes
                . zipWith (liftM2 const) (fmap Just cells)
                $ validRows
    newMatrix = MatObsRow . S.fromRowsL . catMaybes $ validRows
    validRows = fmap (extractCellHashV idxMap sc) cells

-- | Get the index of a cell in a SingleCells type.
getCellIdx :: SingleCells -> Cell -> Int
getCellIdx sc x =
  fromMaybe (error $ "extractCellIdx: cell not found in matrix: " <> show x)
    . V.findIndex (== x)
    . L.view rowNames
    $ sc

-- | Aggregate SingleCells using average.
aggSc :: SingleCells -> AggSingleCells
aggSc sc = AggSingleCells
         . L.over matrix (MatObsRow . avgMat . unMatObsRow)
         . L.over rowNames ( fmap (Cell . (<> " Aggregate") . unCell)
                           . V.take 1
                           )
         $ sc
  where
    avgMat = S.fromListDenseSM numRows
           . fmap ((/ fromIntegral numCols) . foldl' (+) 0)
           . S.toRowsL
    (numRows, numCols) = S.dim . unMatObsRow . L.view matrix $ sc