{-# LANGUAGE CPP
           , GADTs
           , KindSignatures
           , TypeOperators
           , TypeFamilies
           , EmptyCase
           , DataKinds
           , PolyKinds
           , ExistentialQuantification
           , FlexibleContexts
           , OverloadedStrings
           #-}

{-# OPTIONS_GHC -Wall -fwarn-tabs #-}

module Language.Hakaru.Sample where

import           Numeric.SpecFunctions            (logFactorial)
import qualified Data.Number.LogFloat             as LF
import qualified Math.Combinatorics.Exact.Binomial as EB
-- import qualified Numeric.Integration.TanhSinh     as TS
import qualified System.Random.MWC                as MWC
import qualified System.Random.MWC.CondensedTable as MWC
import qualified System.Random.MWC.Distributions  as MWCD

import qualified Data.Vector                      as V
import           Data.STRef
import           Data.Sequence (Seq)
import qualified Data.Foldable                    as F
import qualified Data.List.NonEmpty               as L
import           Data.List.NonEmpty               (NonEmpty(..))
import           Data.Maybe                       (fromMaybe)

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative   (Applicative(..), (<$>))
#endif
import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.Identity
import           Control.Monad.Trans.Maybe
import           Control.Monad.State.Strict
import qualified Data.IntMap                      as IM

import Data.Number.Nat     (fromNat)
import Data.Number.Natural (fromNatural, fromNonNegativeRational, Natural, unsafeNatural)
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.Sing
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Syntax.TypeOf
import Language.Hakaru.Syntax.Value
import Language.Hakaru.Syntax.Reducer
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.DatumCase
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.ABT

data EAssoc =
    forall a. EAssoc {-# UNPACK #-} !(Variable a) !(Value a)

newtype Env = Env (IM.IntMap EAssoc)

emptyEnv :: Env
emptyEnv = Env IM.empty

updateEnv :: EAssoc -> Env -> Env
updateEnv v@(EAssoc x _) (Env xs) =
    Env $ IM.insert (fromNat $ varID x) v xs

updateEnvs
    :: List1 Variable xs
    -> List1 Value xs
    -> Env
    -> Env
updateEnvs Nil1         Nil1         env = env
updateEnvs (Cons1 x xs) (Cons1 y ys) env =
    updateEnvs xs ys (updateEnv (EAssoc x y) env)

lookupVar :: Variable a -> Env -> Maybe (Value a)
lookupVar x (Env env) = do
    EAssoc x' e' <- IM.lookup (fromNat $ varID x) env
    Refl         <- varEq x x'
    return e'

---------------------------------------------------------------

-- Makes use of Atkinson's algorithm as described in:
-- Monte Carlo Statistical Methods pg. 55
--
-- Further discussion at:
-- http://www.johndcook.com/blog/2010/06/14/generating-poisson-random-values/
poisson_rng :: Double -> MWC.GenIO -> IO Int
poisson_rng lambda g' = make_poisson g'
    where
    smu   = sqrt lambda
    b     = 0.931 + 2.53*smu
    a     = -0.059 + 0.02483*b
    vr    = 0.9277 - 3.6224/(b - 2)
    arep  = 1.1239 + 1.1368/(b - 3.4)
    lnlam = log lambda

    make_poisson :: MWC.GenIO -> IO Int
    make_poisson g = do
        u <- MWC.uniformR (-0.5,0.5) g
        v <- MWC.uniformR (0,1) g
        let us = 0.5 - abs u
            k = floor $ (2*a / us + b)*u + lambda + 0.43
        case () of
            () | us >= 0.07 && v <= vr -> return k
            () | k < 0                 -> make_poisson g
            () | us <= 0.013 && v > us -> make_poisson g
            () | accept_region us v k  -> return k
            _                          -> make_poisson g

    accept_region :: Double -> Double -> Int -> Bool
    accept_region us v k =
        log (v * arep / (a/(us*us)+b))
        <=
        -lambda + fromIntegral k * lnlam - logFactorial k

normalize :: [Value 'HProb] -> (LF.LogFloat, Double, [Double])
normalize []          = (0, 0, [])
normalize [(VProb x)] = (x, 1, [1])
normalize xs          = (m, y, ys)
    where
    xs' = map (\(VProb x) -> x) xs
    m   = maximum xs'
    ys  = [ LF.fromLogFloat (x/m) | x <- xs' ]
    y   = sum ys

normalizeVector
    :: Value ('HArray 'HProb) -> (LF.LogFloat, Double, V.Vector Double)
normalizeVector (VArray xs) =
    let xs' = V.map (\(VProb x) -> x) xs in
    case V.length xs of
    0 -> (0, 0, V.empty)
    1 -> (V.unsafeHead xs', 1, V.singleton 1)
    _ ->
        let m   = V.maximum xs'
            ys  = V.map (\x -> LF.fromLogFloat (x/m)) xs'
            y   = V.sum ys
        in (m, y, ys)

---------------------------------------------------------------

runEvaluate
    :: (ABT Term abt)
    => abt '[] a
    -> Value a
runEvaluate prog = evaluate prog emptyEnv

evaluate
    :: (ABT Term abt)
    => abt '[] a
    -> Env
    -> Value a
evaluate e env = caseVarSyn e (evaluateVar env) (flip evaluateTerm env)

evaluateVar :: Env -> Variable a -> Value a
evaluateVar env v =
    case lookupVar v env of
    Nothing -> error "variable not found!"
    Just a  -> a

evaluateTerm
    :: (ABT Term abt)
    => Term abt a
    -> Env
    -> Value a
evaluateTerm t env =
    case t of
    o :$          es -> evaluateSCon    o es    env
    NaryOp_  o    es -> evaluateNaryOp  o es    env
    Literal_ v       -> evaluateLiteral v
    Empty_   _       -> evaluateEmpty
    Array_   n    es -> evaluateArray   n es    env
    ArrayLiteral_ es -> VArray . V.fromList $ map (flip evaluate env) es
    Bucket b e    rs -> evaluateBucket  b e  rs env
    Datum_   d       -> evaluateDatum   d       env
    Case_    o    es -> evaluateCase    o es    env
    Superpose_    es -> evaluateSuperpose es    env
    Reject_  _       -> VMeasure $ \_ _ -> return Nothing

evaluateSCon
    :: (ABT Term abt)
    => SCon args a
    -> SArgs abt args
    -> Env
    -> Value a
evaluateSCon Lam_ (e1 :* End) env =
    caseBind e1 $ \x e1' ->
        VLam $ \v -> evaluate e1' (updateEnv (EAssoc x v) env)
evaluateSCon App_ (e1 :* e2 :* End) env =
    case evaluate e1 env of
    VLam f -> f (evaluate e2 env)
evaluateSCon Let_ (e1 :* e2 :* End) env =
    let v = evaluate e1 env
    in caseBind e2 $ \x e2' ->
        evaluate e2' (updateEnv (EAssoc x v) env)
evaluateSCon (CoerceTo_   c) (e1 :* End) env =
    coerceTo c $ evaluate e1 env
evaluateSCon (UnsafeFrom_ c) (e1 :* End) env =
    coerceFrom c $ evaluate e1 env
evaluateSCon (PrimOp_ o)     es env = evaluatePrimOp    o es env
evaluateSCon (ArrayOp_ o)    es env = evaluateArrayOp   o es env
evaluateSCon (MeasureOp_  m) es env = evaluateMeasureOp m es env
evaluateSCon Dirac           (e1 :* End) env =
    VMeasure $ \p _ -> return $ Just (evaluate e1 env, p)
evaluateSCon MBind (e1 :* e2 :* End) env =
    case evaluate e1 env of
    VMeasure m1 -> VMeasure $ \ p g -> do
        x <- m1 p g
        case x of
            Nothing -> return Nothing
            Just (a, p') ->
                caseBind e2 $ \x' e2' ->
                    case evaluate e2' (updateEnv (EAssoc x' a) env) of
                    VMeasure y -> y p' g

evaluateSCon Plate (n :* e2 :* End) env =
    case evaluate n env of
    VNat n' -> caseBind e2 $ \x e' ->
        VMeasure $ \(VProb p) g -> runMaybeT $ do
            (v', ps) <- fmap V.unzip . V.mapM (performMaybe g) $
                V.generate (fromInteger $ fromNatural n') $ \v ->
                    evaluate e' $
                    updateEnv (EAssoc x . VNat $ intToNatural v) env
            return
                ( VArray v'
                , VProb $ p * V.product (V.map (\(VProb y) -> y) ps)
                )
    where
    performMaybe
        :: MWC.GenIO
        -> Value ('HMeasure a)
        -> MaybeT IO (Value a, Value 'HProb)
    performMaybe g (VMeasure m) = MaybeT $ m (VProb 1) g

evaluateSCon Chain (n :* s :* e :* End) env =
    case (evaluate n env, evaluate s env) of
    (VNat n', start) ->
        caseBind e $ \x e' ->
            let s' = VLam $ \v -> evaluate e' (updateEnv (EAssoc x v) env) in
            VMeasure (\(VProb p) g -> runMaybeT $ do
                (evaluates, sout) <- runStateT (replicateM (unsafeInt n') $ convert g s') start
                let (v', ps) = unzip evaluates
                    bodyType :: Sing ('HMeasure (HPair a b)) -> Sing ('HArray a)
                    bodyType = SArray . fst . sUnPair . sUnMeasure
                return
                    ( VDatum $ dPair_ (bodyType $ caseBind e (const typeOf)) (typeOf s)
                        (VArray . V.fromList $ v') sout
                    , VProb $ p * product (map (\(VProb y) -> y) ps)
                    ))
    where
    convert
        :: MWC.GenIO
        -> Value (s ':-> 'HMeasure (HPair a s))
        -> StateT (Value s) (MaybeT IO) (Value a, Value 'HProb)
    convert g (VLam f) = StateT $ \s' ->
        case f s' of
        VMeasure f' -> do
            (as'', p') <- MaybeT (f' (VProb 1) g)
            let (a, s'') = unPair as''
            return ((a, p'), s'')

    unPair :: Value (HPair a b) -> (Value a, Value b)
    unPair (VDatum (Datum "pair" _typ
        (Inl (Et (Konst a)
            (Et (Konst b) Done))))) = (a, b)
    unPair x = case x of {}

evaluateSCon (Summate hd hs) (e1 :* e2 :* e3 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (lo, hi) ->
        caseBind e3 $ \x e3' ->
            foldl (\t i ->
                   evalOp (Sum  hs) t $
                     evaluate e3' (updateEnv (EAssoc x i) env))
                  (identityElement $ Sum hs)
                  (enumFromUntilValue hd lo hi)

evaluateSCon (Product hd hs) (e1 :* e2 :* e3 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (lo, hi) ->
        caseBind e3 $ \x e3' ->
            foldl (\t i ->
                   evalOp (Prod hs) t $
                     evaluate e3' (updateEnv (EAssoc x i) env))
                  (identityElement $ Prod hs)
                  (enumFromUntilValue hd lo hi)

evaluateSCon s _ _ = error $ "TODO: evaluateSCon{" ++ show s ++ "}"

evaluatePrimOp
    ::  ( ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
    => PrimOp typs a
    -> SArgs abt args
    -> Env
    -> Value a
evaluatePrimOp Not (e1 :* End) env =
    case evaluate e1 env of
      VDatum a -> if a == dTrue
                  then VDatum dFalse
                  else VDatum dTrue

evaluatePrimOp Pi  End         _   = VProb . LF.logFloat $ pi
evaluatePrimOp Cos (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . cos $ v1

evaluatePrimOp Sin (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . sin $ v1

evaluatePrimOp Tan (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VReal . tan $ v1

evaluatePrimOp RealPow (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
      (VProb v1, VReal v2) -> VProb $ LF.pow v1 v2

evaluatePrimOp Choose (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
      (VNat v1, VNat v2) -> VNat $ EB.choose v1 v2

evaluatePrimOp Exp (e1 :* End) env =
    case evaluate e1 env of
      VReal v1 -> VProb . LF.logToLogFloat $ v1

evaluatePrimOp Log (e1 :* End) env =
    case evaluate e1 env of
      VProb v1 -> VReal . LF.logFromLogFloat $ v1

evaluatePrimOp (Infinity h) End _ =
    case h of
      HIntegrable_Nat  -> error "Can not evaluate infinity for natural numbers"
      HIntegrable_Prob -> VProb $ LF.logFloat LF.infinity

evaluatePrimOp (Equal _) (e1 :* e2 :* End) env = (VDatum . dBool) $ evaluate e1 env == evaluate e2 env

evaluatePrimOp (Less _) (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (VNat  v1, VNat  v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VInt  v1, VInt  v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VProb v1, VProb v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    (VReal v1, VReal v2) -> VDatum $ if v1 < v2 then dTrue else dFalse
    _                    -> error "TODO: evaluatePrimOp{Less}"
evaluatePrimOp (NatPow _) (e1 :* e2 :* End) env =
    case evaluate e2 env of
    VNat  v2 ->
        let v2' = fromNatural v2 in
        case evaluate e1 env of
          VNat  v1 -> VNat  (v1 ^ v2')
          VInt  v1 -> VInt  (v1 ^ v2')
          VProb v1 -> VProb (v1 ^ v2')
          VReal v1 -> VReal (v1 ^ v2')
          _        -> error "NatPow should always return some kind of number"
evaluatePrimOp (Negate _) (e1 :* End) env =
    case evaluate e1 env of
    VInt  v -> VInt  (negate v)
    VReal v -> VReal (negate v)
    v       -> case v of {}
evaluatePrimOp (Abs   _) (e1 :* End) env =
    case evaluate e1 env of
    VInt  v -> VNat  . unsafeNatural   $ abs v
    VReal v -> VProb . LF.logFloat $ abs v
    v       -> case v of {}
evaluatePrimOp (Recip _) (e1 :* End) env =
    case evaluate e1 env of
    VProb v -> VProb (recip v)
    VReal v -> VReal (recip v)
    v       -> case v of {}
evaluatePrimOp (NatRoot _) (e1 :* e2 :* End) env =
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VNat v2) -> VProb $ LF.pow v1 (recip . fromIntegral $ v2)
    v                   -> case v of {}

evaluatePrimOp (Floor) (e1 :* End) env =
    case (evaluate e1 env) of
    VProb v1 -> VNat (floor (LF.fromLogFloat v1))

evaluatePrimOp prim _ _ =
    error ("TODO: evaluatePrimOp{" ++ show prim ++ "}")

evaluateArrayOp
    :: ( ABT Term abt
       , typs ~ UnLCs args
       , args ~ LCs typs)
    => ArrayOp typs a
    -> SArgs abt args
    -> Env
    -> Value a
evaluateArrayOp (Index _) = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VArray v, VNat n) -> v V.! unsafeInt n

evaluateArrayOp (Size _) = \(e1 :* End) env ->
    case evaluate e1 env of
    VArray v -> VNat . intToNatural $ V.length v

evaluateArrayOp (Reduce _) = \(e1 :* e2 :* e3 :* End) env ->
    case ( evaluate e1 env
         , evaluate e2 env
         , evaluate e3 env) of
    (f, a, VArray v) -> V.foldl' (lam2 f) a v

evaluateMeasureOp
    :: ( ABT Term abt
       , typs ~ UnLCs args
       , args ~ LCs typs)
    => MeasureOp typs a
    -> SArgs abt args
    -> Env
    -> Value ('HMeasure a)

evaluateMeasureOp Lebesgue = \(e1 :* e2 :* End) env ->
  case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VReal v2) | v1 < v2 ->
      VMeasure $ \(VProb p) g ->
        case (isInfinite v1, isInfinite v2) of
          (False, False) -> do
            x <- MWC.uniformR (v1, v2) g
            return $ Just (VReal $ x,
                           VProb $ p * LF.logFloat (v2 - v1))
          (False, True) -> do
            u <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ v1 + n,
                           VProb $ p * LF.logToLogFloat n)
          (True, False) -> do
            u <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ v2 - n,
                           VProb $ p * LF.logToLogFloat n)
          (True, True) -> do
            (u,b) <- MWC.uniform g
            let l = log u
            let n = -l
            return $ Just (VReal $ if b then n else l,
                           VProb $ p * 2 * LF.logToLogFloat n)
    (VReal _, VReal _) -> error "Lebesgue with length 0 or flipped endpoints"

evaluateMeasureOp Counting = \End _ ->
    VMeasure $ \(VProb p) g -> do
        let success = LF.logToLogFloat (-3 :: Double)
        let pow x y = LF.logToLogFloat (LF.logFromLogFloat x *
                                       (fromIntegral y :: Double))
        u' <- MWCD.geometric0 (LF.fromLogFloat success) g
        let u = toInteger u'
        b <- MWC.uniform g
        return $ Just
            ( VInt  $ if b then -1-u else u
            , VProb $ p * 2 / pow (1-success) u / success)

evaluateMeasureOp Categorical = \(e1 :* End) env ->
    VMeasure $ \p g -> do
        let (_,y,ys) = normalizeVector (evaluate e1 env)
        if not (y > (0::Double)) -- TODO: why not use @y <= 0@ ??
        then error "Categorical needs positive weights"
        else do
            u <- MWC.uniformR (0, y) g
            return $ Just
                ( VNat
                . intToNatural
                . fromMaybe 0
                . V.findIndex (u <=)
                . V.scanl1' (+)
                $ ys
                , p)

evaluateMeasureOp Uniform = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VReal v2) -> VMeasure $ \p g -> do
        x <- MWC.uniformR (v1, v2) g
        return $ Just (VReal x, p)

evaluateMeasureOp Normal = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VReal v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.normal v1 (LF.fromLogFloat v2) g
        return $ Just (VReal x, p)

evaluateMeasureOp Poisson = \(e1 :* End) env ->
    case evaluate e1 env of
    VProb v1 -> VMeasure $ \ p g -> do
        x <- MWC.genFromTable (MWC.tablePoisson (LF.fromLogFloat v1)) g
        return $ Just (VNat $ intToNatural x, p)

evaluateMeasureOp Gamma = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.gamma (LF.fromLogFloat v1) (LF.fromLogFloat v2) g
        return $ Just (VProb $ LF.logFloat x, p)

evaluateMeasureOp Beta = \(e1 :* e2 :* End) env ->
    case (evaluate e1 env, evaluate e2 env) of
    (VProb v1, VProb v2) -> VMeasure $ \ p g -> do
        x <- MWCD.beta (LF.fromLogFloat v1) (LF.fromLogFloat v2) g
        return $ Just (VProb $ LF.logFloat x, p)

evaluateNaryOp
    :: (ABT Term abt)
    => NaryOp a -> Seq (abt '[] a) -> Env -> Value a
evaluateNaryOp s es =
    F.foldr (evalOp s) (identityElement s) . mapEvaluate es

identityElement :: NaryOp a -> Value a
identityElement And                   = VDatum dTrue
identityElement (Sum HSemiring_Nat)   = VNat  0
identityElement (Sum HSemiring_Int)   = VInt  0
identityElement (Sum HSemiring_Prob)  = VProb 0
identityElement (Sum HSemiring_Real)  = VReal 0
identityElement (Prod HSemiring_Nat)  = VNat  1
identityElement (Prod HSemiring_Int)  = VInt  1
identityElement (Prod HSemiring_Prob) = VProb 1
identityElement (Prod HSemiring_Real) = VReal 1
identityElement (Max  HOrd_Prob)      = VProb 0
identityElement (Max  HOrd_Real)      = VReal LF.negativeInfinity
identityElement (Min  HOrd_Prob)      = VProb (LF.logFloat LF.infinity)
identityElement (Min  HOrd_Real)      = VReal LF.infinity
identityElement _                     = error "Missing identity elements?"

evalOp
    :: NaryOp a -> Value a -> Value a -> Value a
evalOp And (VDatum a) (VDatum b)
    | a == dTrue && b == dTrue = VDatum dTrue
    | otherwise = VDatum dFalse
evalOp (Sum  HSemiring_Nat)  (VNat  a) (VNat  b) = VNat  (a + b)
evalOp (Sum  HSemiring_Int)  (VInt  a) (VInt  b) = VInt  (a + b)
evalOp (Sum  HSemiring_Prob) (VProb a) (VProb b) = VProb (a + b)
evalOp (Sum  HSemiring_Real) (VReal a) (VReal b) = VReal (a + b)
evalOp (Prod HSemiring_Nat)  (VNat  a) (VNat  b) = VNat  (a * b)
evalOp (Prod HSemiring_Int)  (VInt  a) (VInt  b) = VInt  (a * b)
evalOp (Prod HSemiring_Prob) (VProb a) (VProb b) = VProb (a * b)
evalOp (Prod HSemiring_Real) (VReal a) (VReal b) = VReal (a * b)
evalOp (Max  HOrd_Prob)      (VProb a) (VProb b) = VProb (max a b)
evalOp (Max  HOrd_Real)      (VReal a) (VReal b) = VReal (max a b)
evalOp (Min  HOrd_Prob)      (VProb a) (VProb b) = VProb (min a b)
evalOp (Min  HOrd_Real)      (VReal a) (VReal b) = VReal (min a b)

evalOp op                    _          _        =
    error ("TODO: evalOp{" ++ show op ++ "}")

mapEvaluate
    :: (ABT Term abt)
    => Seq (abt '[] a) -> Env -> Seq (Value a)
mapEvaluate es env = fmap (flip evaluate env) es

evaluateLiteral :: Literal a -> Value a
evaluateLiteral (LNat  n) = VNat  . fromInteger $ fromNatural n -- TODO: catch overflow errors
evaluateLiteral (LInt  n) = VInt  $ fromInteger n -- TODO: catch overflow errors
evaluateLiteral (LProb n) = VProb . fromRational $ fromNonNegativeRational n
evaluateLiteral (LReal n) = VReal $ fromRational n

evaluateEmpty :: Value ('HArray a)
evaluateEmpty = VArray V.empty

evaluateArray
    :: (ABT Term abt)
    => (abt '[] 'HNat)
    -> (abt '[ 'HNat ] a)
    -> Env
    -> Value ('HArray a)
evaluateArray n e env =
    case evaluate n env of
    VNat n' -> caseBind e $ \x e' ->
        VArray $ V.generate (unsafeInt n') $ \v ->
            let v' = VNat $ intToNatural v in
            evaluate e' (updateEnv (EAssoc x v') env)

evaluateBucket
    :: (ABT Term abt)
    => abt '[] 'HNat
    -> abt '[] 'HNat
    -> Reducer abt '[] a
    -> Env
    -> Value a
evaluateBucket b e rs env =
    case (evaluate b env, evaluate e env) of
      (VNat b', VNat e') -> runST $ do
          s' <- init Nil1 rs env
          mapM_ (\i -> accum (VNat i) Nil1 rs s' env) [b' .. e' - 1]
          done s'
    where init :: (ABT Term abt)
               => List1 Value xs
               -> Reducer abt xs a
               -> Env
               -> ST s (VReducer s a)
          init ix (Red_Fanout r1 r2)    env  =
              VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env
          init ix (Red_Index  n  _  mr) env' =
              let (vars, n') = caseBinds n in
              case evaluate n' (updateEnvs vars ix env') of
                VNat n'' -> VRed_Array <$> V.generateM (fromIntegral n'')
                            (\bb -> init (Cons1 (vnat bb) ix) mr env')
          init ix (Red_Split _ r1 r2)   env' =
              VRed_Pair (type_ r1) (type_ r2) <$> init ix r1 env <*> init ix r2 env'
          init _  Red_Nop               _    = return VRed_Unit
          init _  (Red_Add h _) _ = VRed_Num <$> newSTRef (identityElement (Sum h))

          type_ = typeOfReducer

          vnat :: Int -> Value 'HNat
          vnat  = VNat . fromIntegral

          accum :: (ABT Term abt)
                => Value 'HNat
                -> List1 Value xs
                -> Reducer abt xs a
                -> VReducer s a
                -> Env
                -> ST s ()
          accum n ix (Red_Fanout r1 r2)   (VRed_Pair _ _ v1 v2) env' =
              accum n ix r1 v1 env >> accum n ix r2 v2 env'
          accum n ix (Red_Index n' a1 r2) (VRed_Array v)          env' =
              caseBind a1 $ \i a1' ->
              let (vars, a1'') = caseBinds a1'
                  VNat ov = evaluate a1''
                            (updateEnv (EAssoc i n) (updateEnvs vars ix env'))
                  ov' = fromIntegral ov in
              accum n (Cons1 (VNat ov) ix) r2 (v V.! ov') env
          accum n ix (Red_Split bb r1 r2) (VRed_Pair _ _ v1 v2) env' =
              caseBind bb $ \i b' ->
                  let (vars, b'') = caseBinds b' in
                  case evaluate b''
                       (updateEnv (EAssoc i n) (updateEnvs vars ix env')) of
                  VDatum bb -> if bb == dTrue then
                                   accum n ix r1 v1 env'
                               else
                                   accum n ix r2 v2 env'
          accum n ix (Red_Add h ee) (VRed_Num s) env' =
              caseBind ee $ \i e' ->
                  let (vars, e'') = caseBinds e'
                      v = evaluate e''
                          (updateEnv (EAssoc i n) (updateEnvs vars ix env')) in
                  modifySTRef' s (evalOp (Sum h) v)
          accum _ _ Red_Nop _ _ = return ()
          accum _ _ _ _ _ = error "Some impossible combinations happened?"

          done :: VReducer s a -> ST s (Value a)
          done (VRed_Num s)            = readSTRef s
          done VRed_Unit               = return (VDatum dUnit)
          done (VRed_Pair s1 s2 v1 v2) = do
            v1' <- done v1
            v2' <- done v2
            return (VDatum $ dPair_ s1 s2 v1' v2')
          done (VRed_Array v)          = VArray <$> V.sequence (V.map done v)

evaluateDatum
    :: (ABT Term abt)
    => Datum (abt '[]) (HData' a)
    -> Env
    -> Value (HData' a)
evaluateDatum d env = VDatum (fmap11 (flip evaluate env) d)

evaluateCase
    :: forall abt a b
    .  (ABT Term abt)
    => abt '[] a
    -> [Branch a abt b]
    -> Env
    -> Value b
evaluateCase o es env =
    case runIdentity $ matchBranches evaluateDatum' (evaluate o env) es of
    Just (Matched rho b) ->
        evaluate b (extendFromMatch (fromAssocs rho) env)
    _ -> error "Missing cases in match expression"
    where
    extendFromMatch :: [Assoc Value] -> Env -> Env
    extendFromMatch []                env' = env'
    extendFromMatch (Assoc x v : xvs) env' =
        extendFromMatch xvs (updateEnv (EAssoc x v) env')

    evaluateDatum' :: DatumEvaluator Value Identity
    evaluateDatum' = return . Just . getVDatum

    getVDatum :: Value (HData' a) -> Datum Value (HData' a)
    getVDatum (VDatum a) = a

evaluateSuperpose
    :: (ABT Term abt)
    => NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
    -> Env
    -> Value ('HMeasure a)
evaluateSuperpose ((q, m) :| []) env =
    case evaluate m env of
    VMeasure m' ->
        let VProb q' = evaluate q env
        in  VMeasure (\(VProb p) g -> m' (VProb $ p * q') g)

evaluateSuperpose pms@((_, m) :| _) env =
    case evaluate m env of
    VMeasure m' ->
        let pms'     = L.toList pms
            weights  = map ((flip evaluate env) . fst) pms'
            (x,y,ys) = normalize weights
        in VMeasure $ \(VProb p) g ->
            if not (y > (0::Double)) then return Nothing else do
            u <- MWC.uniformR (0, y) g
            case [ m1 | (v,(_,m1)) <- zip (scanl1 (+) ys) pms', u <= v ] of
                m2 : _ ->
                    case evaluate m2 env of
                    VMeasure m2' -> m2' (VProb $ p * x * LF.logFloat y) g
                []     -> m' (VProb $ p * x * LF.logFloat y) g

----------------------------------------------------------------

-- Useful 'short-hand'
intToNatural :: Int -> Natural
intToNatural = unsafeNatural . toInteger

unsafeInt :: Natural -> Int
unsafeInt = fromInteger . fromNatural
----------------------------------------------------------- fin.


-- |
-- Module      : Copernicus
-- Description : Contains the core of the Copernicus project
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental
-- Portability : POSIX (not sure)
--
-- Jonatan H Sundqvist
-- November 25 2014
--

-- TODO | - Model planetary motion and physics in 2D
--        - Cartoon earth with gravitational field
--        - Flexible event handling
--        - App typeclass (run, manage events, window properties, etc.)
--        - Lenses
--        - Options (eg. toggle grid) (cf. 'when')
--        - UI
--        - Use types to encode units (eg. SI, radians)
--        - Move to Cairo (branch?)
--        - Polymorphic types (not just Floats) (?)

-- SPEC | -
--        -

module Copernicus where

---------------------------------------------------------------------------------------------------
-- We'll need these
---------------------------------------------------------------------------------------------------
import Data.Complex
-- import Control.Monad (when)

import Graphics.Gloss.Geometry.Angle (degToRad, radToDeg, normaliseAngle)

---------------------------------------------------------------------------------------------------
-- Types
---------------------------------------------------------------------------------------------------
-- Make typeclass (mesh, body, bounding box, collision, etc) (?)
-- type Number a = Floating a -- Real number
type Vector = Complex -- TODO: Polymorphic (cf. related TODO item)
data Body f = Body (Vector f) (Vector f) (Vector f) deriving Show -- Add argument

---------------------------------------------------------------------------------------------------
-- Functions
---------------------------------------------------------------------------------------------------
clampAngle :: Float -> Float
clampAngle = radToDeg . normaliseAngle . degToRad

--
-- TODO: Make them change colour when bouncing (?)
animate :: (RealFloat f, Floating f) => f -> Body f -> Body f
animate t (Body p v a) = collide ground $ Body (parabola t p v a) (v + (t:+0)*a) a
	where ground = 0.0 -- 30+30/2-540/2 (previous hard-coded value)

--
parabola :: (RealFloat f, Floating f) => f -> Vector f -> Vector f -> Vector f -> Vector f
parabola t p v a = let (px:+py) = p
                       (vx:+vy) = v
                       (ax:+ay) = a
                   in (px + vx*t + 0.5*ax*t**2) :+ (py + vy*t + 0.5*ay*t**2)

-- collide
-- Very primitive for now
-- TODO: Use 'contains' function (Range -> Value -> Bool)
-- TODO: Don't hard-code bounds (left, right)
-- TODO: Take bounds of Body into account (don't hard-code that either)
collide :: (RealFloat f, Floating f) => f -> Body f -> Body f
collide gnd (Body (px:+py) (vx:+vy) a) = Body (px:+py) ((invertIf (\ _ -> (px <= left) || ( px >= right)) vx) :+ (invertIf (\ v -> (v < 0) && (py <= gnd)) vy)) a
	where invertIf p v | p v 	   = -v
	                   | otherwise =  v
	      (left, right) = (-5, 5) --(15-720/2, 720/2-30/2)

-- | ETA (estimated time of arrival)
-- TODO: Rename (eg. timeUntil, solveForT, etc)
-- TODO: Parabola type (eg. Parabola a v x)
-- eta :: Acceleration

-- Utilities (should eventually be moved to separate module or library ----------------------------
-- Python-style String formatting (eg. keyword interpolation, {0}, customisation, format specs.)
-- Parsec

---------------------------------------------------------------------------------------------------
-- Entry point
---------------------------------------------------------------------------------------------------


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric, TupleSections, RecordWildCards, TemplateHaskell, RankNTypes, FlexibleContexts #-}

module BayesStack.Models.Topic.CitationInfluenceNoTopics
  ( -- * Primitives
    NetData
  , dHypers, dArcs, dItems, dNodeItems, dCitingNodes, dCitedNodes
  , netData
  , HyperParams(..)
  , MState
  , stGammas, stOmegas, stPsis, stCiting, stLambdas
  , CitingUpdateUnit
  , ItemSource(..)
  , CitedNode(..), CitedNodeItem(..)
  , CitingNode(..), CitingNodeItem(..)
  , Citing(..), Cited(..)
  , Item(..), Topic(..), NodeItem(..), Node(..), Arc(..)
  , setupNodeItems
    -- * Initialization
  , verifyNetData, cleanNetData
  , ModelInit
  , randomInitialize
  , model
  , updateUnits
    -- * Diagnostics
  , modelLikelihood
  ) where

import qualified Data.Vector as V
import Statistics.Sample (mean)

import           Prelude hiding (mapM_, sum)
import           Data.Maybe (fromMaybe)

import           Control.Lens hiding (Setting)
import           Data.Set (Set)
import qualified Data.Set as S

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import           Data.Foldable hiding (product)
import           Control.Applicative ((<$>), (<*>))
import           Control.Monad (when)
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict

import           Data.Random
import           Data.Random.Lift (lift)
import           Data.Random.Distribution.Categorical (categorical)
import           Numeric.Log hiding (sum)

import           BayesStack.Types
import           BayesStack.Gibbs
import           BayesStack.DirMulti
import           BayesStack.TupleEnum ()
import           BayesStack.Models.Topic.Types

import           GHC.Generics (Generic)
import           Data.Binary (Binary)
import           Control.DeepSeq

at' :: At m => Index m -> IndexedLens' (Index m) m (IxValue m)
at' i = at i . _fromMaybe
  where _fromMaybe = iso (fromMaybe $ error "at': Unexpected Nothing") Just

data ItemSource = Shared | Own deriving (Show, Eq, Enum, Ord, Generic)
instance Binary ItemSource
instance NFData ItemSource

newtype Citing a = Citing a deriving (Show, Eq, Enum, Ord, Generic, NFData)
newtype Cited a = Cited a deriving (Show, Eq, Enum, Ord, Generic, NFData)
instance Binary a => Binary (Citing a)
instance Binary a => Binary (Cited a)

type CitingNode = Citing Node
type CitedNode = Cited Node
type CitingNodeItem = Citing NodeItem
type CitedNodeItem = Cited NodeItem

-- ^ A directed edge
data Arc = Arc { citingNode :: !CitingNode, citedNode :: !CitedNode }
            deriving (Show, Eq, Ord, Generic)
instance Binary Arc

data HyperParams = HyperParams
                   { _alphaPsi         :: Double
                   , _alphaLambda      :: Double
                   , _alphaOmega       :: Double
                   , _alphaGammaShared :: Double
                   , _alphaGammaOwn    :: Double
                   , _alphaBetaFG      :: Double
                   , _alphaBetaBG      :: Double
                   }
                 deriving (Show, Eq, Generic)
instance Binary HyperParams
makeLenses ''HyperParams

data NetData = NetData { _dHypers             :: !(HyperParams)
                       , _dArcs               :: !(Set Arc)
                       , _dItems              :: !(Map Item Double)
                       , _dNodeItems          :: !(Map NodeItem (Node, Item))
                       , _dCitingNodes        :: !(Map CitingNode (Set CitedNode))
                         -- ^ Maps each citing node to the set of nodes cited by it
                       , _dCitedNodes         :: !(Map CitedNode (Set CitingNode))
                         -- ^ Maps each cited node to the set of nodes citing it
                       }
              deriving (Show, Eq, Generic)
instance Binary NetData
makeLenses ''NetData

netData :: HyperParams -> Set Arc -> Map Item Double -> Map NodeItem (Node,Item) -> NetData
netData hypers arcs items nodeItems =
    NetData { _dHypers       = hypers
            , _dArcs         = arcs
            , _dItems        = items
            , _dNodeItems    = nodeItems
            , _dCitingNodes  = M.unionsWith S.union
                               $ map (\(Arc a b)->M.singleton a $ S.singleton b)
                               $ S.toList arcs
            , _dCitedNodes   = M.unionsWith S.union
                               $ map (\(Arc a b)->M.singleton b $ S.singleton a)
                               $ S.toList arcs
            }

dCitingNodeItems :: NetData -> Map CitingNodeItem (CitingNode, Item)
dCitingNodeItems nd =
    M.mapKeys Citing
    $ M.map (\(n,i)->(Citing n, i))
    $ M.filter (\(n,i)->Citing n `M.member` (nd^.dCitingNodes))
    $ nd^.dNodeItems

itemsOfCitingNode :: NetData -> CitingNode -> [Item]
itemsOfCitingNode d (Citing u) =
    map snd $ M.elems $ M.filter (\(n,_)->n==u) $ d^.dNodeItems

connectedNodes :: Set Arc -> Set Node
connectedNodes arcs =
    S.map ((\(Cited n)->n) . citedNode) arcs `S.union` S.map ((\(Citing n)->n) . citingNode) arcs

cleanNetData :: NetData -> NetData
cleanNetData d =
    let nodesWithItems = S.fromList $ map fst $ M.elems $ d^.dNodeItems
        nodesWithArcs = connectedNodes $ d^.dArcs
        keptNodes = nodesWithItems `S.intersection` nodesWithArcs
        keepArc (Arc (Citing citing) (Cited cited)) =
            citing `S.member` keptNodes && cited `S.member` keptNodes
        go = do dArcs %= S.filter keepArc
                dNodeItems %= M.filter (\(n,i)->n `S.member` keptNodes)
    in execState go d

verifyNetData :: (Node -> String) -> NetData -> [String]
verifyNetData showNode d = execWriter $ do
    let nodesWithItems = S.fromList $ map fst $ M.elems $ d^.dNodeItems
    forM_ (d^.dArcs) $ \(Arc (Citing citing) (Cited cited))->do
        when (cited `S.notMember` nodesWithItems)
            $ tell [showNode cited++" has arc yet has no items"]
        when (citing `S.notMember` nodesWithItems)
            $ tell [showNode citing++" has arc yet has no items"]

-- Citing Update unit (Shared Taste-like)
data CitingUpdateUnit = CitingUpdateUnit { _uuNI    :: CitingNodeItem
                                         , _uuN     :: CitingNode
                                         , _uuX     :: Item
                                         , _uuCites :: Set CitedNode
                                         , _uuItemWeight :: Double
                                         }
                      deriving (Show, Generic)
instance Binary CitingUpdateUnit
makeLenses ''CitingUpdateUnit

citingUpdateUnits :: NetData -> [CitingUpdateUnit]
citingUpdateUnits d =
    map (\(ni,(n,x))->CitingUpdateUnit { _uuNI      = ni
                                       , _uuN       = n
                                       , _uuX       = x
                                       , _uuCites   = d^.dCitingNodes . at' n
                                       , _uuItemWeight = (d ^. dItems . at' x)
                                       }
        ) $ M.assocs $ dCitingNodeItems d

updateUnits :: NetData -> [WrappedUpdateUnit MState]
updateUnits d = map WrappedUU (citingUpdateUnits d)

-- | Model State
data CitingSetting = OwnSetting
                   | SharedSetting !CitedNode
                   deriving (Show, Eq, Generic)
instance Binary CitingSetting
instance NFData CitingSetting where
    rnf (OwnSetting)      = ()
    rnf (SharedSetting c) = rnf c `seq` ()

data MState = MState { -- Citing model state
                       _stGammas   :: !(Map CitingNode (Multinom Int ItemSource))
                     , _stOmegas   :: !(Map CitingNode (Multinom Int Item))
                     , _stPsis     :: !(Map CitingNode (Multinom Int CitedNode))

                     , _stCiting   :: !(Map CitingNodeItem CitingSetting)

                     -- Cited model state
                     , _stLambdas  :: !(Map CitedNode (Multinom Int Item))
                     }
            deriving (Show, Generic)
instance Binary MState
makeLenses ''MState

-- | Model initialization
type ModelInit = Map CitingNodeItem (Setting CitingUpdateUnit)

modify' :: Monad m => (a -> a) -> StateT a m ()
modify' f = do x <- get
               put $! f x

randomInitializeCiting :: NetData -> ModelInit -> RVar ModelInit
randomInitializeCiting d init = execStateT doInit init
    where doInit :: StateT ModelInit RVar ()
          doInit = let unset = M.keysSet (dCitingNodeItems d) `S.difference` M.keysSet init
                   in mapM_ (randomInitCitingUU d) (S.toList unset)

randomInitCitingUU :: NetData -> CitingNodeItem -> StateT ModelInit RVar ()
randomInitCitingUU d cni@(Citing ni) =
    let (n,_) = d ^. dNodeItems . at' ni
    in case d ^. dCitingNodes . at' (Citing n) of
           a | S.null a -> do
               modify' $ M.insert cni OwnSetting

           citedNodes -> do
               s <- lift $ randomElement [Shared, Own]
               c <- lift $ randomElement $ toList citedNodes
               modify' $ M.insert cni $
                   case s of Shared -> SharedSetting c
                             Own    -> OwnSetting

randomInitialize :: NetData -> RVar ModelInit
randomInitialize d = randomInitializeCiting d M.empty

model :: NetData -> ModelInit -> MState
model d citingInit =
    let citingNodes = M.keys $ d^.dCitingNodes
        hp = d^.dHypers
        s = MState { -- Citing model
                     _stPsis = let dist n = case d ^. dCitingNodes . at' n . to toList of
                                                []    -> M.empty
                                                nodes -> M.singleton n
                                                         $ symDirMulti (hp^.alphaPsi) nodes
                               in foldMap dist citingNodes
                   , _stGammas = let dist = multinom [ (Shared, hp^.alphaGammaShared)
                                                     , (Own, hp^.alphaGammaOwn) ]
                                 in foldMap (\t->M.singleton t dist) citingNodes
                   , _stOmegas = let dist = symDirMulti (hp^.alphaOmega) (M.keys $ d^.dItems)
                                 in foldMap (\t->M.singleton t dist) citingNodes
                   , _stCiting = M.empty

                   -- Cited model
                   , _stLambdas = let dist = symDirMulti (hp^.alphaLambda) (M.keys $ d^.dItems)
                                      lambdas0 = foldMap (\n->M.singleton n dist) $ M.keys $ d^.dCitedNodes
                                  in foldl' (\dms (n,x)->M.adjust (incMultinom x) (Cited n) dms) lambdas0 (M.elems $ d^.dNodeItems)
                   }

        initCitingUU :: CitingUpdateUnit -> State MState ()
        initCitingUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given\n"
                s = maybe err id $ M.lookup (uu^.uuNI) citingInit
            modify' $ setCitingUU uu (Just s)

    in execState (mapM_ initCitingUU $ citingUpdateUnits d) s

modelLikelihood :: MState -> Probability
modelLikelihood model =
    product $ (model ^.. stGammas  . folded . to likelihood)
           ++ (model ^.. stLambdas . folded . to likelihood)
           ++ (model ^.. stOmegas  . folded . to likelihood)
           ++ (model ^.. stPsis    . folded . to likelihood)

instance UpdateUnit CitingUpdateUnit where
    type ModelState CitingUpdateUnit = MState
    type Setting CitingUpdateUnit = CitingSetting
    fetchSetting uu ms = ms ^. stCiting . at' (uu^.uuNI)
    evolveSetting ms uu = categorical $ citingFullCond (setCitingUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitingUU uu (Just s') . setCitingUU uu Nothing

citingProb :: MState -> CitingUpdateUnit -> Setting CitingUpdateUnit -> Double
citingProb st (CitingUpdateUnit {_uuN=n, _uuX=x}) setting =
    let gamma = st ^. stGammas . at' n
        omega = st ^. stOmegas . at' n
        psi = st ^. stPsis . at' n
    in case setting of
        SharedSetting c   -> let lambda = st ^. stLambdas . at' c
                             in sampleProb gamma Shared
                              * sampleProb psi c
                              * sampleProb lambda x
        OwnSetting        ->  sampleProb gamma Own
                            * sampleProb omega x

citingFullCond :: MState -> CitingUpdateUnit -> [(Double, Setting CitingUpdateUnit)]
citingFullCond ms uu = map (\s->(citingProb ms uu s, s)) $ citingDomain ms uu

citingDomain :: MState -> CitingUpdateUnit -> [Setting CitingUpdateUnit]
citingDomain ms uu = do
    s <- [Own, Shared]
    case s of
        Shared -> do c <- uu ^. uuCites . to S.toList
                     return $ SharedSetting c
        Own    -> do return $ OwnSetting

setCitingUU :: CitingUpdateUnit -> Maybe (Setting CitingUpdateUnit) -> MState -> MState
setCitingUU uu@(CitingUpdateUnit {_uuNI=ni, _uuN=n, _uuX=x}) setting ms = execState go ms
  where
    set = maybe Unset (const Set) setting
    go = case maybe (fetchSetting uu ms) id setting of
           SharedSetting c    -> do stPsis .    at' n %= setMultinom set c
                                    stLambdas . at' c %= setMultinom set x
                                    stGammas .  at' n %= setMultinom set Shared
                                    stCiting .  at ni .= setting

           OwnSetting         -> do stOmegas .  at' n %= setMultinom set x
                                    stGammas .  at' n %= setMultinom set Own
                                    stCiting .  at ni .= setting


-- |
-- Copyright: (c) 2020 Tom Westerhout
-- SPDX-License-Identifier: BSD-3-Clause
-- Maintainer: Tom Westerhout <14264576+twesterhout@users.noreply.github.com>
--
-- High-level wrapper around [PRIMME C
-- library](https://github.com/primme/primme). /Quote from README/:
--
-- @
--   PRIMME, pronounced as prime, is a high-performance library for computing a few
--   eigenvalues/eigenvectors, and singular values/vectors. PRIMME is especially
--   optimized for large, difficult problems. Real symmetric and complex Hermitian
--   problems, standard @A x = λ x@ and generalized @A x = λ B x@, are supported.
--   Besides, standard eigenvalue problems with a normal matrix are supported. It
--   can find largest, smallest, or interior singular/eigenvalues, and can use
--   preconditioning to accelerate convergence.
-- @
module Numeric.PRIMME
  ( -- * Defining the matrix

    -- | One of the great things about PRIMME library is that it works with
    -- block matrix-vector products (i.e. matrix-matrix products). Following the
    -- example of [vector](https://hackage.haskell.org/package/vector) library,
    -- we define two types for mutable and immutable dense blocks.
    MBlock (..),
    Block (..),
    -- | Now we are can define the "operator" we want to diagonalize. Since this
    -- library is meant to be used with rather large matrices, it might very
    -- well be the case that the matrix does not fit into the memory of your
    -- computer (even in sparse format such as CSR). Sometimes, however, we can
    -- define the "operator" implicitly, namely, by defining its action on a
    -- vector (or in case of PRIMME on a block of vectors). This is done with
    -- 'PrimmeOperator' type.
    PrimmeOperator,

    -- * Choosing what to compute

    -- | Having defined a 'PrimmeOperator' we now need to tell PRIMME what to
    -- compute. This is done by constructing a 'PrimmeOptions' object.
    PrimmeOptions (..),
    primmeDefaults,
    finalizeOptions,
    PrimmeTarget (..),

    -- * Diagonalizing
    eigh,
    eigh',

    -- * Logging
    PrimmeMonitor (..),
    PrimmeInfo (..),
    PrimmeEventInfo (..),
    PrimmeStats (..),
    primmePrettyInfo,

    -- * Dense Matrices

    -- | The most trivial example of an "operator" is of course a square dense
    -- matrix. We thus provide a function which constructs an operator from a
    -- matrix.
    primmeFromDense,

    -- * Misc
    getPrimmeVersion,
    PrimmeException (..),
    BlasDatatype (blasTag),
    BlasRealPart,
    BlasDatatypeTag (..),
  )
where

import Control.Exception.Safe (throw)
import Control.Monad (forM_, unless)
import Control.Monad.ST (RealWorld)
import Data.Complex
import Data.Proxy
import Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import Data.Vector.Storable.Mutable (MVector)
import qualified Data.Vector.Storable.Mutable as MV
import Foreign.C.Types (CInt (..))
import Foreign.Ptr (Ptr)
import Foreign.Storable
import qualified Language.C.Inline as C
import qualified Language.C.Inline.Unsafe as CU
import Numeric.PRIMME.Context
import Numeric.PRIMME.Dense
import Numeric.PRIMME.Monitor
import Numeric.PRIMME.Options
import Numeric.PRIMME.Types
import Prelude hiding (init)

C.context (C.baseCtx <> primmeCtx)
C.include "<primme.h>"

-- | Diagonalize the operator to find the first few eigenpairs.
eigh ::
  forall a.
  BlasDatatype a =>
  PrimmeOptions ->
  PrimmeOperator a ->
  IO (Vector (BlasRealPart a), Block a, Vector (BlasRealPart a))
eigh options matrix = do
  let dim = pDim options
      numEvals = pNumEvals options
  (evals :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
  (evecs :: MVector RealWorld a) <- MV.new (dim * numEvals)
  (rnorms :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
  status <-
    MV.unsafeWith evals $ \evalsPtr ->
      MV.unsafeWith evecs $ \evecsPtr ->
        MV.unsafeWith rnorms $ \rnormsPtr ->
          withOptions options matrix $ \optionsPtr ->
            cPrimme evalsPtr evecsPtr rnormsPtr optionsPtr
  unless (status == 0) $ throwPrimmeError status
  evals' <- V.unsafeFreeze evals
  evecs' <- Block (dim, numEvals) dim <$> V.unsafeFreeze evecs
  rnorms' <- V.unsafeFreeze rnorms
  return (evals', evecs', rnorms')

mkEmptyBlock :: Storable a => Int -> Int -> IO (MBlock RealWorld a)
mkEmptyBlock rows columns
  | rows <= 0 || columns <= 0 = error $ "invalid shape specified: " <> show (rows, columns)
  | otherwise = MBlock (rows, columns) rows <$> MV.new (rows * columns)

copyBlock :: Storable a => MBlock RealWorld a -> Block a -> IO ()
copyBlock dest@(MBlock (r, c) _ _) src@(Block (r', c') _ _)
  | r /= r' || c /= c' = error $ "block shape mismatch: " <> show (r, c) <> " != " <> show (r', c')
  | otherwise = do
    src' <- unsafeThaw src
    forM_ [0 .. (c - 1)] $ \i ->
      MV.copy (getColumn i dest) (getColumn i src')
  where
    getColumn :: Storable a => Int -> MBlock RealWorld a -> MVector RealWorld a
    getColumn i (MBlock (rows, _) stride v) = MV.slice (i * stride) rows v

unsafeThaw :: Storable a => Block a -> IO (MBlock RealWorld a)
unsafeThaw (Block shape stride v) = MBlock shape stride <$> V.unsafeThaw v

unsafeFreeze :: Storable a => MBlock RealWorld a -> IO (Block a)
unsafeFreeze (MBlock shape stride v) = Block shape stride <$> V.unsafeFreeze v

sliceBlock1 :: Storable a => Int -> Int -> MBlock s a -> MBlock s a
sliceBlock1 i n (MBlock (r, c) stride v)
  | i < 0 || n < 0 || i + n > c = error $ "invalid slice: [" <> show i <> ", " <> show (i + n) <> ")"
  | otherwise = MBlock (r, n) stride $ MV.slice (i * stride) (n * stride) v

eigh' ::
  forall a.
  BlasDatatype a =>
  PrimmeOptions ->
  MBlock RealWorld a ->
  PrimmeOperator a ->
  IO (Vector (BlasRealPart a), Block a, Vector (BlasRealPart a))
eigh' options init@(MBlock (dim', initSize) _ _) matrix
  | pDim options /= dim' =
    error $
      "'init' has wrong shape: "
        <> show (dim', initSize)
        <> "; expected a block with "
        <> show (pDim options)
        <> " rows"
  | otherwise = do
    let dim = pDim options
        numEvals = pNumEvals options
    evecs@(MBlock (_, _) evecsStride evecsData) <- case initSize < numEvals of
      True -> do
        temp <- mkEmptyBlock dim numEvals
        copyBlock (sliceBlock1 0 initSize temp) =<< unsafeFreeze init
        return temp
      False -> return init
    (evals :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
    (rnorms :: MVector RealWorld (BlasRealPart a)) <- MV.new numEvals
    status <-
      MV.unsafeWith evals $ \evalsPtr ->
        MV.unsafeWith evecsData $ \evecsPtr ->
          MV.unsafeWith rnorms $ \rnormsPtr ->
            withOptions options matrix $ \optionsPtr -> do
              let c_initSize = fromIntegral initSize
                  c_evecsStride = fromIntegral evecsStride
              [C.block| void {
                $(primme_params* optionsPtr)->initSize = $(int c_initSize);
                $(primme_params* optionsPtr)->ldevecs = $(PRIMME_INT c_evecsStride);
              } |]
              cPrimme evalsPtr evecsPtr rnormsPtr optionsPtr
    unless (status == 0) $ throwPrimmeError status
    evals' <- V.unsafeFreeze evals
    evecs' <- unsafeFreeze $ sliceBlock1 0 numEvals evecs
    rnorms' <- V.unsafeFreeze rnorms
    return (evals', evecs', rnorms')

throwPrimmeError :: CInt -> IO a
throwPrimmeError c = case (- c) of
  0 -> error "no error"
  3 -> throw PrimmeMaximumIterationsReached
  40 -> throw PrimmeLapackFailure
  41 -> throw PrimmeUserFailure
  _ -> throw $ PrimmeOtherFailure c

foreign import ccall "sprimme"
  sprimme :: Ptr Float -> Ptr Float -> Ptr Float -> Ptr Cprimme_params -> IO CInt

foreign import ccall "dprimme"
  dprimme :: Ptr Double -> Ptr Double -> Ptr Double -> Ptr Cprimme_params -> IO CInt

foreign import ccall "cprimme"
  cprimme :: Ptr Float -> Ptr (Complex Float) -> Ptr Float -> Ptr Cprimme_params -> IO CInt

foreign import ccall "zprimme"
  zprimme :: Ptr Double -> Ptr (Complex Double) -> Ptr Double -> Ptr Cprimme_params -> IO CInt

cPrimme ::
  forall a.
  (BlasDatatype a) =>
  Ptr (BlasRealPart a) ->
  Ptr a ->
  Ptr (BlasRealPart a) ->
  Ptr Cprimme_params ->
  IO CInt
cPrimme = case blasTag (Proxy :: Proxy a) of
  FloatTag -> sprimme
  DoubleTag -> dprimme
  ComplexFloatTag -> cprimme
  ComplexDoubleTag -> zprimme

getPrimmeVersion :: (Int, Int)
getPrimmeVersion =
  ( fromIntegral [CU.pure| int { PRIMME_VERSION_MAJOR } |],
    fromIntegral [CU.pure| int { PRIMME_VERSION_MINOR } |]
  )


{-# LANGUAGE AllowAmbiguousTypes, CPP, ConstraintKinds, DataKinds,
             FlexibleInstances, GADTs, MultiParamTypeClasses, PolyKinds,
             QuantifiedConstraints, TypeFamilyDependencies, TypeOperators,
             UndecidableInstances #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | The class defining dual components of dual numbers and related classes,
-- type families, constraints and instances. This is a low-level API
-- used to define types and operations in "HordeAd.Core.DualNumber"
-- that is the high-level API.
module HordeAd.Core.DualClass
  ( IsPrimalWithScalar, IsPrimalAndHasFeatures, IsScalar, HasDelta
  , DMode(..), Dual, IsPrimal(..), HasRanks(..)
  , HasVariables(..)  -- use sparringly
  ) where

import Prelude

import qualified Data.Array.Convert
import qualified Data.Array.Dynamic as OTB
import qualified Data.Array.DynamicS as OT
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.MonoTraversable (Element, MonoFunctor)
import           Data.Proxy (Proxy)
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, natVal, type (+))
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Internal.Delta

-- * Abbreviations to export (not used anywhere below)

-- | The intended semantics (not fully enforced by the constraint in isolation)
-- is that the second type is the primal component of a dual number type
-- at an unknown rank, with the given differentiation mode
-- and underlying scalar.
type IsPrimalWithScalar (d :: DMode) a r =
  (ScalarOf a ~ r, IsPrimal d a, HasVariables a)

-- | A shorthand for a useful set of constraints.
type IsPrimalAndHasFeatures (d :: DMode) a r =
  (IsPrimalWithScalar d a r, RealFloat a, MonoFunctor a, Element a ~ r)

-- | A mega-shorthand for a bundle of connected type constraints.
-- The @Scalar@ in the name means that the second argument is the underlying
-- scalar type of a well behaved (wrt the differentiation mode in the first
-- argument) collection of primal and dual components of dual numbers.
type IsScalar (d :: DMode) r =
  ( HasRanks d r, Ord r, Numeric r, Show r
  , IsPrimalAndHasFeatures d r r
  , IsPrimalAndHasFeatures d (Vector r) r
  , IsPrimalAndHasFeatures d (Matrix r) r
  , IsPrimalAndHasFeatures d (OT.Array r) r
  -- This fragment is for @OS.Array@ and it's irregular, because we can't
  -- mention @sh@ and so fully apply the type constructor.
  , IsPrimalS d r  -- TODO: Floating (OS.Array sh r), MonoFunctor
  )

-- | Is a scalar and will be used to compute gradients via delta-expressions.
type HasDelta r = ( IsScalar 'DModeGradient r
                  , Dual 'DModeGradient r ~ Delta0 r )

-- * Class definitions

-- | The enumeration of all possible differentiation (and more generally,
-- computation with dual numbers) schemes.
data DMode =
    DModeGradient
  | DModeDerivative

-- | The type family that enumerates all possible "ranks"
-- for each differentiation mode.
-- The second type argument is meant to be the primal component
-- of dual numbers. The result is the dual component.
--
-- Rank 0 is special because, in derivatives mode, the dual component
-- is not the primal component wrapped in a datatype or newtype constructor.
-- This makes impossible a representation of primal and dual components as
-- the primal plus the type constructor for creating the dual.
--
-- Rank S is special, because of the extra type parameter @sh@ representing
-- a shape. This is another obstacle to a dual number representation via
-- a single-argument type constructor.
type family Dual (d :: DMode) a = result | result -> d a where
  Dual 'DModeGradient Double = Delta0 Double
  Dual 'DModeGradient Float = Delta0 Float
  Dual 'DModeGradient (Vector r) = Delta1 r
  Dual 'DModeGradient (Matrix r) = Delta2 r
  Dual 'DModeGradient (OT.Array r) = DeltaX r
  Dual 'DModeGradient (OS.Array sh r) = DeltaS sh r
-- not injective:  Dual 'DModeDerivative r = r
  Dual 'DModeDerivative Double = Double
  Dual 'DModeDerivative Float = Float
  Dual 'DModeDerivative (Vector r) = Vector r
  Dual 'DModeDerivative (Matrix r) = Matrix r
  Dual 'DModeDerivative (OT.Array r) = OT.Array r
  Dual 'DModeDerivative (OS.Array sh r) = OS.Array sh r

-- | The underlying scalar of a given primal component of a dual number.
-- A long name to remember not to use, unless necessary, and not to export.
type family ScalarOf a where
  ScalarOf Double = Double
  ScalarOf Float = Float
  ScalarOf (Vector r) = r
  ScalarOf (Matrix r) = r
  ScalarOf (OT.Array r) = r
  ScalarOf (OS.Array sh r) = r

-- | Second argument is a primal component of dual numbers at some rank
-- wrt the differentiation mode given in the first argument.
class IsPrimal d a where
  dZero :: Dual d a
  dScale :: a -> Dual d a -> Dual d a
  dAdd :: Dual d a -> Dual d a -> Dual d a
  dDelay :: Dual d a -> Dual d a

-- | Part 1/2 of a hack to squeeze the shaped tensors rank,
-- with its extra @sh@ parameter, into the 'IsPrimal' class.
class IsPrimalS d r where
  dZeroS :: forall sh. OS.Shape sh => Dual d (OS.Array sh r)
  dScaleS :: forall sh. OS.Shape sh
          => OS.Array sh r -> Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)
  dAddS :: forall sh. OS.Shape sh
        => Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)
        -> Dual d (OS.Array sh r)
  dDelayS :: forall sh. OS.Shape sh
             => Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)

-- | Part 2/2 of a hack to squeeze the shaped tensors rank,
-- with its extra @sh@ parameter, into the 'IsPrimal' class.
instance (IsPrimalS d r, OS.Shape sh) => IsPrimal d (OS.Array sh r) where
  dZero = dZeroS
  dScale = dScaleS
  dAdd = dAddS
  dDelay = dDelayS

-- | Assuming that the first argument is the primal component of dual numbers
-- with the underyling scalar in the second argument and with differentiation
-- mode `DModeGradient`, it additionally admits delta-variable
-- introduction and binding as defined by the methods of the class.
class HasVariables a where
  dVar :: DeltaId a -> Dual 'DModeGradient a
  bindInState :: Dual 'DModeGradient a
              -> DeltaState (ScalarOf a)
              -> (DeltaState (ScalarOf a), DeltaId a )
  dOutline :: CodeOut -> [a] -> [Dual 'DModeGradient a]
           -> Dual 'DModeGradient a

-- | The class provides methods required for the second type parameter
-- to be the underlying scalar of a well behaved collection of dual numbers
-- of various ranks wrt the differentation mode given in the first argument.
class HasRanks (d :: DMode) r where
  dSumElements0 :: Dual d (Vector r) -> Int -> Dual d r
  dIndex0 :: Dual d (Vector r) -> Int -> Int -> Dual d r
  dDot0 :: Vector r -> Dual d (Vector r) -> Dual d r
  dFromX0 :: Dual d (OT.Array r) -> Dual d r
  dFromS0 :: Dual d (OS.Array '[] r) -> Dual d r

  dSeq1 :: Data.Vector.Vector (Dual d r) -> Dual d (Vector r)
  dKonst1 :: Dual d r -> Int -> Dual d (Vector r)
  dAppend1 :: Dual d (Vector r) -> Int -> Dual d (Vector r) -> Dual d (Vector r)
  dSlice1 :: Int -> Int -> Dual d (Vector r) -> Int -> Dual d (Vector r)
  dSumRows1 :: Dual d (Matrix r) -> Int -> Dual d (Vector r)
  dSumColumns1 :: Dual d (Matrix r) -> Int -> Dual d (Vector r)
  dM_VD1 :: Matrix r -> Dual d (Vector r) -> Dual d (Vector r)
  dMD_V1 :: Dual d (Matrix r) -> Vector r -> Dual d (Vector r)
  dFromX1 :: Dual d (OT.Array r) -> Dual d (Vector r)
  dFromS1 :: KnownNat len
          => Dual d (OS.Array '[len] r) -> Dual d (Vector r)
  dReverse1 :: Dual d (Vector r) -> Dual d (Vector r)
  dFlatten1 :: Int -> Int -> Dual d (Matrix r) -> Dual d (Vector r)
  dFlattenX1 :: OT.ShapeL -> Dual d (OT.Array r) -> Dual d (Vector r)
  dFlattenS1 :: OS.Shape sh
             => Dual d (OS.Array sh r) -> Dual d (Vector r)

  dFromRows2 :: Data.Vector.Vector (Dual d (Vector r)) -> Dual d (Matrix r)
  dFromColumns2 :: Data.Vector.Vector (Dual d (Vector r)) -> Dual d (Matrix r)
  dKonst2 :: Dual d r -> (Int, Int) -> Dual d (Matrix r)
  dTranspose2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dM_MD2 :: Matrix r -> Dual d (Matrix r) -> Dual d (Matrix r)
  dMD_M2 :: Dual d (Matrix r) -> Matrix r -> Dual d (Matrix r)
  dRowAppend2 :: Dual d (Matrix r) -> Int -> Dual d (Matrix r) -> Dual d (Matrix r)
  dColumnAppend2 :: Dual d (Matrix r) -> Int -> Dual d (Matrix r) -> Dual d (Matrix r)
  dRowSlice2 :: Int -> Int -> Dual d (Matrix r) -> Int -> Dual d (Matrix r)
  dColumnSlice2 :: Int -> Int -> Dual d (Matrix r) -> Int -> Dual d (Matrix r)
  dAsRow2 :: Dual d (Vector r) -> Dual d (Matrix r)
  dAsColumn2 :: Dual d (Vector r) -> Dual d (Matrix r)
  dFromX2 :: Dual d (OT.Array r) -> Dual d (Matrix r)
  dFromS2 :: (KnownNat rows, KnownNat cols)
          => Dual d (OS.Array '[rows, cols] r) -> Dual d (Matrix r)

  dFlipud2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dFliprl2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dReshape2 :: Int -> Dual d (Vector r) -> Dual d (Matrix r)
  dConv2 :: Matrix r -> Dual d (Matrix r) -> Dual d (Matrix r)

  dKonstX :: Dual d r -> OT.ShapeL -> Dual d (OT.Array r)
  dAppendX :: Dual d (OT.Array r) -> Int -> Dual d (OT.Array r) -> Dual d (OT.Array r)
  dSliceX :: Int -> Int -> Dual d (OT.Array r) -> Int -> Dual d (OT.Array r)
  dIndexX :: Dual d (OT.Array r) -> Int -> Int -> Dual d (OT.Array r)
  dRavelFromListX :: [Dual d (OT.Array r)] -> Dual d (OT.Array r)
  dReshapeX :: OT.ShapeL -> OT.ShapeL -> Dual d (OT.Array r) -> Dual d (OT.Array r)
  dFrom0X :: Dual d r -> Dual d (OT.Array r)
  dFrom1X :: Dual d (Vector r) -> Dual d (OT.Array r)
  dFrom2X :: Dual d (Matrix r) -> Int -> Dual d (OT.Array r)
  dFromSX :: OS.Shape sh
          => Dual d (OS.Array sh r) -> Dual d (OT.Array r)

  dKonstS :: OS.Shape sh
          => Dual d r -> Dual d (OS.Array sh r)
  dAppendS :: (OS.Shape sh, KnownNat m, KnownNat n)
           => Dual d (OS.Array (m ': sh) r) -> Dual d (OS.Array (n ': sh) r)
           -> Dual d (OS.Array ((m + n) ': sh) r)
  dSliceS :: (KnownNat i, KnownNat n, KnownNat k, OS.Shape rest)
          => Proxy i -> Proxy n -> Dual d (OS.Array (i + n + k ': rest) r)
          -> Dual d (OS.Array (n ': rest) r)
  dIndexS :: (KnownNat ix, KnownNat k, OS.Shape rest)
          => Dual d (OS.Array (ix + 1 + k ': rest) r)
          -> Proxy ix -> Dual d (OS.Array rest r)
  dRavelFromListS :: (KnownNat k, OS.Shape rest)
                  => [Dual d (OS.Array rest r)]
                  -> Dual d (OS.Array (k : rest) r)
  dReshapeS :: (OS.Shape sh, OS.Shape sh', OS.Size sh ~ OS.Size sh')
            => Dual d (OS.Array sh r) -> Dual d (OS.Array sh' r)
  dFrom0S :: Dual d r -> Dual d (OS.Array '[] r)
  dFrom1S :: KnownNat n => Dual d (Vector r) -> Dual d (OS.Array '[n] r)
  dFrom2S :: (KnownNat rows, KnownNat cols)
          => Proxy cols
          -> Dual d (Matrix r) -> Dual d (OS.Array '[rows, cols] r)
  dFromXS :: OS.Shape sh => Dual d (OT.Array r) -> Dual d (OS.Array sh r)

-- * Backprop gradient method instances

instance IsPrimal 'DModeGradient Double where
  dZero = Zero0
  dScale = Scale0
  dAdd = Add0
  dDelay = Delay0

instance IsPrimal 'DModeGradient Float where
  -- Identical as above:
  dZero = Zero0
  dScale = Scale0
  dAdd = Add0
  dDelay = Delay0

instance IsPrimal 'DModeGradient (Vector r) where
  dZero = Zero1
  dScale = Scale1
  dAdd = Add1
  dDelay = Delay1

instance IsPrimal 'DModeGradient (Matrix r) where
  dZero = Zero2
  dScale = Scale2
  dAdd = Add2
  dDelay = Delay2

instance IsPrimal 'DModeGradient (OT.Array r) where
  dZero = ZeroX
  dScale = ScaleX
  dAdd = AddX
  dDelay = DelayX

instance IsPrimalS 'DModeGradient r where
  dZeroS = ZeroS
  dScaleS = ScaleS
  dAddS = AddS
  dDelayS = DelayS

instance HasVariables Double where
  dVar = Var0
  {-# INLINE bindInState #-}
  bindInState = bindInState0
  dOutline = Outline0

instance HasVariables Float where
  dVar = Var0
  {-# INLINE bindInState #-}
  bindInState = bindInState0
  dOutline = Outline0

instance HasVariables (Vector r) where
  dVar = Var1
  {-# INLINE bindInState #-}
  bindInState = bindInState1
  dOutline = Outline1

instance HasVariables (Matrix r) where
  dVar = Var2
  {-# INLINE bindInState #-}
  bindInState = bindInState2
  dOutline = Outline2

instance HasVariables (OT.Array r) where
  dVar = VarX
  {-# INLINE bindInState #-}
  bindInState = bindInStateX
  dOutline = OutlineX

instance OS.Shape sh => HasVariables (OS.Array sh r) where
  dVar = VarS
  {-# INLINE bindInState #-}
  bindInState u' st = let (st2, did) = bindInStateX (FromSX u') st
                      in (st2, convertDeltaId did)
  dOutline = OutlineS

instance Dual 'DModeGradient r ~ Delta0 r
         => HasRanks 'DModeGradient r where
  dSumElements0 = SumElements0
  dIndex0 = Index0
  dDot0 = Dot0
  dFromX0 = FromX0
  dFromS0 = FromS0
  dSeq1 = Seq1
  dKonst1 = Konst1
  dAppend1 = Append1
  dSlice1 = Slice1
  dSumRows1 = SumRows1
  dSumColumns1 = SumColumns1
  dM_VD1 = M_VD1
  dMD_V1 = MD_V1
  dFromX1 = FromX1
  dFromS1 = FromS1
  dReverse1 = Reverse1
  dFlatten1 = Flatten1
  dFlattenX1 = FlattenX1
  dFlattenS1 = FlattenS1
  dFromRows2 = FromRows2
  dFromColumns2 = FromColumns2
  dKonst2 = Konst2
  dTranspose2 = Transpose2
  dM_MD2 = M_MD2
  dMD_M2 = MD_M2
  dRowAppend2 = RowAppend2
  dColumnAppend2 = ColumnAppend2
  dRowSlice2 = RowSlice2
  dColumnSlice2 = ColumnSlice2
  dAsRow2 = AsRow2
  dAsColumn2 = AsColumn2
  dFromX2 = FromX2
  dFromS2 = FromS2
  dFlipud2 = Flipud2
  dFliprl2 = Fliprl2
  dReshape2 = Reshape2
  dConv2 = Conv2
  dKonstX = KonstX
  dAppendX = AppendX
  dSliceX = SliceX
  dIndexX = IndexX
  dRavelFromListX = RavelFromListX
  dReshapeX = ReshapeX
  dFrom0X = From0X
  dFrom1X = From1X
  dFrom2X = From2X
  dFromSX = FromSX
  dKonstS = KonstS
  dAppendS = AppendS
  dSliceS = SliceS
  dIndexS = IndexS
  dRavelFromListS = RavelFromListS
  dReshapeS = ReshapeS
  dFrom0S = From0S
  dFrom1S = From1S
  dFrom2S = From2S
  dFromXS = FromXS

-- * Alternative instances: forward derivatives computed on the spot

instance IsPrimal 'DModeDerivative Double where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id  -- no delaying

instance IsPrimal 'DModeDerivative Float where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

-- These constraints force @UndecidableInstances@.
instance Num (Vector r)
         => IsPrimal 'DModeDerivative (Vector r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance Num (Matrix r)
         => IsPrimal 'DModeDerivative (Matrix r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance Num (OT.Array r)
         => IsPrimal 'DModeDerivative (OT.Array r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance (Numeric r, Num (Vector r))
         => IsPrimalS 'DModeDerivative r where
  dZeroS = 0
  dScaleS k d = k * d
  dAddS d e = d + e
  dDelayS = id

instance ( Numeric r, Num (Vector r)
         , Dual 'DModeDerivative r ~ r )
         => HasRanks 'DModeDerivative r where
  dSumElements0 vd _ = HM.sumElements vd
  dIndex0 d ix _ = d V.! ix
  dDot0 = (HM.<.>)
  dFromX0 = OT.unScalar
  dFromS0 = OS.unScalar
  dSeq1 = V.convert
  dKonst1 = HM.konst
  dAppend1 d _k e = d V.++ e
  dSlice1 i n d _len = V.slice i n d
  dM_VD1 = (HM.#>)
  dMD_V1 = (HM.#>)
  dSumRows1 dm _cols = V.fromList $ map HM.sumElements $ HM.toRows dm
  dSumColumns1 dm _rows = V.fromList $ map HM.sumElements $ HM.toColumns dm
  dFromX1 = OT.toVector
  dFromS1 = OS.toVector
  dReverse1 = V.reverse
  dFlatten1 _rows _cols = HM.flatten
  dFlattenX1 _sh = OT.toVector
  dFlattenS1 = OS.toVector
  dFromRows2 = HM.fromRows . V.toList
  dFromColumns2 = HM.fromColumns . V.toList
  dKonst2 = HM.konst
  dTranspose2 = HM.tr'
  dM_MD2 = (HM.<>)
  dMD_M2 = (HM.<>)
  dAsRow2 = HM.asRow
  dAsColumn2 = HM.asColumn
  dRowAppend2 d _k e = d HM.=== e
  dColumnAppend2 d _k e = d HM.||| e
  dRowSlice2 i n d _rows = HM.takeRows n $ HM.dropRows i d
  dColumnSlice2 i n d _cols = HM.takeColumns n $ HM.dropColumns i d
  dFromX2 d = case OT.shapeL d of
    [_rows, cols] -> HM.reshape cols $ OT.toVector d
    _ -> error "dFromX2: wrong tensor dimensions"
  dFromS2 d = case OS.shapeL d of
    [_rows, cols] -> HM.reshape cols $ OS.toVector d
    _ -> error "dFromS2: wrong tensor dimensions"
  dFlipud2 = HM.flipud
  dFliprl2 = HM.fliprl
  dReshape2 = HM.reshape
  dConv2 = HM.conv2
  dKonstX d sz = OT.constant sz d
  dAppendX d _k e = d `OT.append` e
  dSliceX i n d _len = OT.slice [(i, n)] d
  dIndexX d ix _len = OT.index d ix
  dRavelFromListX ld =
    let sh = case ld of
          d : _ -> length ld : OT.shapeL d
          [] -> []
    in OT.ravel $ OTB.fromList sh ld
  dReshapeX _sh = OT.reshape
  dFrom0X = OT.scalar
  dFrom1X d = OT.fromVector [V.length d] d
  dFrom2X d cols = OT.fromVector [HM.rows d, cols] $ HM.flatten d
  dFromSX = Data.Array.Convert.convert
  dKonstS = OS.constant
  dAppendS = OS.append
  dSliceS (_ :: Proxy i) (_ :: Proxy n) = OS.slice @'[ '(i, n) ]
  dIndexS d proxyIx = OS.index d (fromInteger $ natVal proxyIx)
  dRavelFromListS = OS.ravel . OSB.fromList
  dReshapeS = OS.reshape
  dFrom0S = OS.scalar
  dFrom1S = OS.fromVector
  dFrom2S _ = OS.fromVector . HM.flatten
  dFromXS = Data.Array.Convert.convert