-- | An implementation of ResourceHandle-based variables.
--
-- The main difference between this and 'Ref'-based variables is
-- that reads are explicit, via the 'readValue' op.
--
-- TODO: given that distinction, figure out a good story around
-- gradients and save/restore.  Then, merge this module into
-- TensorFlow.Ops.
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
module TensorFlow.Variable
    ( Variable
    , variable
    , variable'
    , readValue
    , initializedValue
    , initializedVariable
    , initializedVariable'
    , zeroInitializedVariable
    , zeroInitializedVariable'
    , assign
    , assign'
    , assignAdd
    , assignAdd'
    , resourceApplyAdam
    , resourceApplyAdam'
    ) where

import qualified Data.Complex
import qualified Data.Int
import qualified Data.Word
import Data.Text.Encoding (encodeUtf8)
import Lens.Family2 ((.~), (&))
import TensorFlow.Core
import TensorFlow.Build (opDef)
import TensorFlow.BuildOp (buildInputs, pureOp, OpParams)
import TensorFlow.Output (opInputs, unNodeName)
import TensorFlow.Tensor (Rendered(..), ToTensor(..), renderValue, tensorNodeName)
import TensorFlow.Types (tensorType)
import qualified TensorFlow.GenOps.Core as CoreOps
import TensorFlow.Ops (zeros)

data Variable a = Variable
    { variableHandle   :: Tensor Value ResourceHandle
    , initializedValue :: Maybe (Tensor Value a)
      -- ^ The initial value of a 'Variable' created with 'initializedVariable'.
    }

instance Rendered Variable where
    renderedOutput = renderedOutput . variableHandle

instance ToTensor Variable where
    toTensor = readValue

-- | Creates a new, uninitialized variable.
variable :: (MonadBuild m, TensorType a) => Shape -> m (Variable a)
variable = variable' id

variable' :: forall m a . (MonadBuild m, TensorType a)
                    => OpParams -> Shape -> m (Variable a)
variable' params s = build $ do
    -- Each variable needs a unique "shared_name".  Use MonadFix to
    -- set the attribute to the same name as the variable itself, without
    -- exposing more internals of the Build module.
    rec t <- CoreOps.varHandleOp' (params . (opAttr "shared_name" .~ n))
                                    (tensorType (undefined :: a)) s
        let n = encodeUtf8 $ unNodeName $ tensorNodeName t
    return $ Variable t Nothing

-- | Creates a variable initialized to the given value.
-- Initialization happens next time session runs.
initializedVariable :: (MonadBuild m, TensorType a)
                    => Tensor v a -> m (Variable a)
initializedVariable = initializedVariable' id

initializedVariable' :: forall a m v . (MonadBuild m, TensorType a)
                    => OpParams -> Tensor v a -> m (Variable a)
initializedVariable' params initializer = do
    -- The shape is not known initially.
    (Variable h Nothing :: Variable a) <- variable' params (Shape [])
    initializer' <- renderValue initializer
    i <- CoreOps.assignVariableOp h initializer'
    addInitializer =<< group i
    return (Variable h (Just initializer'))

-- | Creates a zero-initialized variable with the given shape.
zeroInitializedVariable
  :: (MonadBuild m, TensorType a, Num a) => Shape -> m (Variable a)
zeroInitializedVariable = zeroInitializedVariable' id

zeroInitializedVariable'
  :: (MonadBuild m, TensorType a, Num a) => OpParams -> Shape -> m (Variable a)
zeroInitializedVariable' params = initializedVariable' params . zeros

-- | Gets the value stored in a variable.
--
-- Note that this op is stateful since it depends on the value of the variable;
-- however, it may be CSE'd with other reads in the same context.  The context can
-- be fixed by using 'render' along with (for example) 'withControlDependencies'.
-- For example:
--
-- >   runSession $ do
-- >     v <- variable []
-- >     a <- assign v 24
-- >     r <- withControlDependencies a $ render $ readValue v + 18
-- >     result <- run r
-- >     liftIO $ (42 :: Float) @=? unScalar result
--
--
readValue :: TensorType a => Variable a -> Tensor Build a
readValue = readValue' id

readValue' :: forall a . TensorType a
    => OpParams -> Variable a -> Tensor Build a
readValue' params (Variable h _)
    = pureOp [] $ do
        os <- buildInputs h
        pure $ opDef "ReadVariableOp"
                & (params
                    . (opAttr "dtype" .~ tensorType (undefined :: a))
                    . (opInputs .~ os))

-- | Sets the value of a variable.
assign :: (MonadBuild m, TensorType a)
    => Variable a -> Tensor v a -> m ControlNode
assign = assign' id

assign' :: (MonadBuild m, TensorType a)
    => OpParams -> Variable a -> Tensor v a -> m ControlNode
assign' params (Variable h _) v = CoreOps.assignVariableOp' params h v

-- | Increments the value of a variable.
assignAdd :: (MonadBuild m, TensorType a)
    => Variable a -> Tensor v a -> m ControlNode
assignAdd = assignAdd' id

assignAdd' :: (MonadBuild m, TensorType a)
    => OpParams -> Variable a -> Tensor v a -> m ControlNode
assignAdd' params (Variable h _) v = CoreOps.assignAddVariableOp' params h v

-- | Update '*var' according to the Adam algorithm.
--
-- lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-- m_t <- beta1 * m_{t-1} + (1 - beta1) * g_t
-- v_t <- beta2 * v_{t-1} + (1 - beta2) * g_t * g_t
-- variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon)
resourceApplyAdam ::
    (MonadBuild m,
     OneOf '[(Data.Complex.Complex Double),
             (Data.Complex.Complex Float),
             Data.Int.Int16,
             Data.Int.Int32,
             Data.Int.Int64, Data.Int.Int8,
             Data.Word.Word16,
             Data.Word.Word8, Double,
             Float] t)
    => Variable t -- ^ __var__: Should be from a Variable().
    -> Variable t -- ^ __m__: Should be from a Variable().
    -> Variable t -- ^ __v__: Should be from a Variable().
    -> Tensor v1 t -- ^ __beta1_power__: Must be a scalar.
    -> Tensor v2 t -- ^ __beta2_power__: Must be a scalar.
    -> Tensor v3 t -- ^ __lr__: Scaling factor. Must be a scalar.
    -> Tensor v4 t -- ^ __beta1__: Momentum factor. Must be a scalar.
    -> Tensor v5 t -- ^ __beta2__: Momentum factor. Must be a scalar.
    -> Tensor v6 t -- ^ __epsilon__: Ridge term. Must be a scalar.
    -> Tensor v7 t -- ^ __grad__: The gradient.
    -> m (ControlNode)
resourceApplyAdam = resourceApplyAdam' id

resourceApplyAdam' ::
    (MonadBuild m,
     OneOf '[(Data.Complex.Complex Double),
             (Data.Complex.Complex Float),
             Data.Int.Int16, Data.Int.Int32,
             Data.Int.Int64, Data.Int.Int8,
             Data.Word.Word16, Data.Word.Word8, Double,
             Float] t)
    => OpParams
    -> Variable t -- ^ __var__: Should be from a Variable().
    -> Variable t -- ^ __m__: Should be from a Variable().
    -> Variable t -- ^ __v__: Should be from a Variable().
    -> Tensor v1 t -- ^ __beta1_power__: Must be a scalar.
    -> Tensor v2 t -- ^ __beta2_power__: Must be a scalar.
    -> Tensor v3 t -- ^ __lr__: Scaling factor. Must be a scalar.
    -> Tensor v4 t -- ^ __beta1__: Momentum factor. Must be a scalar.
    -> Tensor v5 t -- ^ __beta2__: Momentum factor. Must be a scalar.
    -> Tensor v6 t -- ^ __epsilon__: Ridge term. Must be a scalar.
    -> Tensor v7 t -- ^ __grad__: The gradient.
    -> m (ControlNode)
resourceApplyAdam' params (Variable var _) (Variable m _) (Variable v _) =
    CoreOps.resourceApplyAdam' params var m v


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
module Taiji.Utils
    ( BBIndex
    , openBBs
    , queryBB
    , readExpression
    , lp

    , readGenesValidated

    , computeRAS
    , computeSS
    , computeCDF

    , module Taiji.Utils.Matrix
    , module Taiji.Utils.Clustering
    ) where

import Data.BBI.BigBed
import Bio.Data.Experiment
import Bio.Data.Bed
import Bio.Utils.Functions (scale)
import           Data.CaseInsensitive    (mk)
import qualified Data.HashMap.Strict as HM
import qualified Data.Matrix.Unboxed as MU
import qualified Data.Matrix as Mat
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import System.Random.MWC.Distributions
import System.Random.MWC
import Statistics.Sample (mean, varianceUnbiased)
import Bio.RealWorld.GENCODE (readGenes, Gene(..))

import Taiji.Utils.Clustering
import Taiji.Utils.Matrix
import qualified Taiji.Utils.DataFrame as DF
import Taiji.Prelude

readGenesValidated :: FilePath -> IO [Gene]
readGenesValidated = fmap validate . readGenes
  where
    validate xs
        | null xs = error "No gene was found in the annotation file"
        | all (null . geneTranscripts) xs = error "No transcript was found in the annotation file"
        | otherwise = xs
{-# INLINE readGenesValidated #-}

-- | Open bigbed files.
openBBs :: [(B.ByteString, Maybe (File '[] 'BigBed))]
        -> IO BBIndex
openBBs xs = fmap (HM.fromList . catMaybes) $ forM xs $ \(chr, x) -> case x of
    Nothing -> return Nothing
    Just fl -> do
        bb <- openBBedFile $ fl^.location
        return $ Just (chr, bb)

queryBB :: BEDLike b => b -> BBIndex -> ConduitT () TFBS IO ()
queryBB bed idx = case HM.lookup (bed^.chrom) idx of
    Nothing -> return ()
    Just i -> query (bed^.chrom, bed^.chromStart, bed^.chromEnd) i .| mapC f
  where
    f (chr, s, e, rest) = BEDExt (asBed chr s e) site
      where
        site = SiteInfo
            { _tf_name = mk $ head $ B.split '+' f1
            , _site_affinity = toSiteAffinity $ readInt f2
            , _peak_affinity = toPeakAffinity 100 }
        (f1:f2:_) = B.split '\t' rest

lp :: Int -> [Double] -> Double
lp p = (**(1/fromIntegral p)) . foldl' (+) 0 . map (**fromIntegral p)
{-# INLINE lp #-}

-- | Read RNA expression data
readExpression :: Double    -- ^ Threshold to call a gene as non-expressed
               -> B.ByteString  -- ^ cell type
               -> FilePath
               -> IO (HM.HashMap GeneName (Double, Double)) -- ^ absolute value and z-score
readExpression thres ct fl = do
    c <- B.readFile fl
    let ((_:header):dat) = map (B.split '\t') $ B.lines c
        rowNames = map (mk . head) dat
        dataTable = map (U.fromList . map ((+pseudoCount) . readDouble) . tail) dat
        dataTable' = MU.fromRows $ zipWith U.zip dataTable $
            map computeZscore dataTable :: MU.Matrix (Double, Double)
        idx = fromMaybe (error $ "Cell type:" ++ B.unpack ct ++ " not found!") $
            lookup ct $ zip header [0..]
    return $ HM.fromList $ zip rowNames $ U.toList $ dataTable' `MU.takeColumn` idx
  where
    pseudoCount = 0.1
    computeZscore xs
        | U.length xs == 1 = U.map log xs
        | U.all (<thres) xs = U.replicate (U.length xs) 0
        | U.length xs == 2 = let fc = log $ U.head xs / U.last xs
                             in U.fromList [fc, negate fc]
        | U.all (== U.head xs) xs = U.replicate (U.length xs) 0
        | otherwise = scale xs
{-# INLINE readExpression #-}

computeRAS :: FilePath -> IO (U.Vector Double)
computeRAS fl = do
    mat <- mkSpMatrix readInt fl
    fmap accScore $ runResourceT $ runConduit $
        streamRows mat .| sink (_num_col mat)
  where
    sink n = do
        v <- lift $ UM.replicate n 0.1
        mapM_C $ \(_, xs) -> forM_ xs $ \(i, x) ->
            UM.unsafeModify v (+fromIntegral x) i
        lift $ U.unsafeFreeze v
    accScore xs = U.map (\x -> x * 1000000 / s) xs
      where
        s = U.sum xs

-- | Compute Specificity Score (SS).
computeSS :: DF.DataFrame Double -> DF.DataFrame Double
computeSS df = df{ DF._dataframe_data = mat}
  where
    mat = Mat.fromRows $ map (q_t . normalize) $ Mat.toRows $
        DF._dataframe_data df
    q_t xs = V.map (\x -> e - logBase 2 x) xs
      where
        e = negate $ V.sum $ V.map (\p -> p * logBase 2 p) xs
    normalize xs = V.map (/s) xs
      where
        s = V.sum xs
{-# INLINE computeSS #-}

computeCDF :: DF.DataFrame Double -> IO (V.Vector Double, Double, Double)
computeCDF df = do
    gen <- create
    let std = let rows = filter ((>1) . V.maximum) $ Mat.toRows $ DF._dataframe_data df
                  f xs = (xs, 1 / entropy (normalize xs))
                  n = truncate $ (0.3 :: Double) * fromIntegral (length rows)
                  getFold xs = let m = mean xs in V.map (logBase 2 . (/m)) xs
                  entropy xs = negate $ V.sum $ V.map (\p -> p * logBase 2 p) xs
              in sqrt $ varianceUnbiased $ V.concat $ map (getFold . fst) $
                  take n $ sortBy (comparing snd) $ map f $ rows
    runConduit $ replicateMC nSample (V.toList <$> mkSample gen std) .|
        concatC .| mkCDF 0.001
  where
    ncol = Mat.cols $ DF._dataframe_data df
    nSample = 500000
    mkCDF res = do
        vec <- VM.replicate n 0
        mapM_C $ \x -> do
          let i = min (n-1) $ truncate $ x / res
          VM.unsafeModify vec (+1) i
        v <- V.scanl1' (+) <$> V.unsafeFreeze vec
        return (V.map (/(V.last v)) v, res, fromIntegral $ nSample*ncol)
      where
        n = truncate $ 2 * logBase 2 (fromIntegral ncol) / res
    -- | https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2396180/
    normalize xs = V.map (/s) xs
      where
        s = V.sum xs
    mkSample gen !std = do
        folds <- replicateM ncol $ normal 0 std gen
        return $ q_t $ normalize $ V.fromList $ map (\x -> 2**x) folds
    q_t xs = V.map (\x -> e - logBase 2 x) xs
      where
        e = negate $ V.sum $ V.map (\p -> p * logBase 2 p) xs
{-# INLINE computeCDF #-}