{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

{-| This module defines `Distributive` for a fixed `Monad`
    instead of an arbitrary `Functor`.
-}

module Data.Distributive.Monadic (
  -- * Type class
    DistributiveM(..)

  -- ** distributem and collectm in terms of each other
  , distributemByCollectm
  , collectmByDistributem

  -- *  Utilities defined using `DistributiveM`
  , wrapFoldM

  -- ** ConstT
  , ConstT(..)

  -- *  MonadTrans instances
  , distributemTrans
  , collectmTrans

  -- *  Generics
  , gdistributem
  , gcollectm
  , GDistributiveM(..)
  ) where

import Control.Comonad (Comonad(..))
import Control.Comonad.Cofree (Cofree(..))
import Control.Comonad.Trans.Cofree (CofreeT(..))
import qualified Control.Comonad.Trans.Cofree as Trans (CofreeF(..))
import Control.Comonad.Trans.Coiter (CoiterT(..))
import Control.Comonad.Trans.Env (EnvT(..), ask, lowerEnvT)
import Control.Comonad.Trans.Store (StoreT(..), pos)
import Control.Comonad.Trans.Traced (TracedT(..))
import Control.Foldl (FoldM(..))
import Control.Foldl.Utils (extractM, innerJoinFoldM, stepFoldM, unwrapFoldM)
import Control.Monad (join, liftM2)
import Control.Monad.Free (Free(..))
import Control.Monad.Free.Church (F(..), fromF, toF)
import Control.Monad.Trans.Accum (AccumT(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.Const
import Control.Monad.Trans.Cont (ContT(..))
import Control.Monad.Trans.Except (ExceptT(..))
import Control.Monad.Trans.Free (FreeT(..))
import qualified Control.Monad.Trans.Free as Trans (FreeF(..))
import Control.Monad.Trans.Free.Church (FT(..), fromFT, toFT)
import Control.Monad.Trans.Identity (IdentityT(..))
import Control.Monad.Trans.Iter (IterT(..))
import Control.Monad.Trans.Maybe (MaybeT(..))
import Control.Monad.Trans.RWS (RWST(..))
import qualified Control.Monad.Trans.RWS.Strict as Strict (RWST(..))
import Control.Monad.Trans.Reader (ReaderT(..))
import Control.Monad.Trans.State (StateT(..))
import qualified Control.Monad.Trans.State.Strict as Strict (StateT(..))
import Control.Monad.Trans.Writer (WriterT(..))
import qualified Control.Monad.Trans.Writer.Strict as Strict (WriterT(..))
import Data.Bifunctor (Bifunctor(..))
import Data.Complex (Complex)
import Data.Distributive (Distributive(..))
import Data.Functor.Compose (Compose(..))
import Data.Functor.Identity (Identity(..))
import Data.Proxy (Proxy(..))
import Data.Semigroup (Dual, Product, Sum)
import GHC.Generics
  ( (:*:)(..)
  , (:.:)(..)
  , Generic1(..)
  , M1(..)
  , Par1(..)
  , Rec1(..)
  , U1(..)
  )

-- | This is the categorical dual of `Traversable`, for a `Monad` fixed by
-- our choice of @f@.
--
-- If you define one of `distributem`, `collectm` you may define the other
-- using `collectmByDistributem`, `distributemByCollectm`, respectively.
--
-- All `MonadTrans`formers that result in a `Monad` are instances of
-- `DistributiveM`, see `collectmTrans`.
class Monad m => DistributiveM m f | f -> m where
  -- | The dual of `sequence`, for a fixed `Monad`
  --
  -- @
  --  `distributem` = `collectm` `id`
  -- @
  --
  distributem :: m (f a) -> f (m a)
  default distributem :: (Generic1 f, GDistributiveM m (Rep1 f)) => m (f a) -> f (m a)
  distributem = gdistributem

  -- | Map then `distributem`:
  --
  -- @
  --  `collectm` f = `distributem` . `fmap` f
  -- @
  --
  collectm :: (a -> f b) -> m a -> f (m b)
  default collectm :: (Generic1 f, GDistributiveM m (Rep1 f)) => (a -> f b) -> m a -> f (m b)
  collectm = gcollectm

instance Monad m => DistributiveM m (FoldM m a) where
  distributem f = FoldM (fmap unwrapFoldM . flip stepFoldM) f (return . extractM)

  collectm = collectmByDistributem

-- | If there were no functional dependencies, this would be the "other"
-- instance for `Identity`:
--
-- @
--  instance `Monad` m => `DistributiveM` m `Identity` where
--    `distributem` :: m (`Identity` a) -> `Identity` (m a)
--    `distributem` = `Identity` . `fmap` `runIdentity`
-- @
--
-- But it'd violate the functional dependency since `Identity` does not
-- determine @m@.
--
-- By ignoring the chosen `Monad`, we can recreate the instance.
--
-- See `wrappedCofreeFToProd` for an example.
instance Monad m => DistributiveM m (ConstT m) where
  distributem = ConstT . fmap runConstT
  collectm f = ConstT . fmap (runConstT . f)

-- distributem and collectm in terms of each other

-- | `distributem` in terms of `collectm`
distributemByCollectm :: DistributiveM m f => m (f a) -> f (m a)
distributemByCollectm = collectm id

-- | `collectm` in terms of `distributem`
collectmByDistributem :: DistributiveM m f => (a -> f b) -> m a -> f (m b)
collectmByDistributem f = distributem . fmap f

-- Utilities defined using `DistributiveM`

-- | Implement `wrap` using `distributem` and `innerJoinFoldM`
wrapFoldM :: Monad m => m (FoldM m a b) -> FoldM m a b
wrapFoldM = innerJoinFoldM . distributem

-- GHC.Generics instances

instance (DistributiveM m f, DistributiveM m g) =>
         DistributiveM m (f :*: g)

instance (DistributiveM m f, DistributiveM m g, Functor f) =>
         DistributiveM m (f :.: g)

instance (DistributiveM m f, DistributiveM m g, Functor f) =>
         DistributiveM m (Compose f g)

instance DistributiveM Identity Par1

instance DistributiveM m f => DistributiveM m (M1 i c f)

instance DistributiveM m f => DistributiveM m (Rec1 f)

instance DistributiveM Identity U1

-- Lifted instances

instance DistributiveM Identity Identity

instance DistributiveM Identity Complex

instance DistributiveM Identity Dual

instance DistributiveM Identity Sum

instance DistributiveM Identity Product

instance DistributiveM Identity Proxy

-- MonadTrans instances

instance Monad m => DistributiveM m (ReaderT a m) where
  distributem = distributemTrans -- ReaderT . distribute . fmap runReaderT
  collectm = collectmTrans -- f xs = ReaderT $ \x -> flip (runReaderT . f) x <$> xs

instance (Monoid w, Monad m) => DistributiveM m (AccumT w m) where
  distributem = distributemTrans -- AccumT . fmap (>>= fmap (first return)) . collect runAccumT
  collectm = collectmTrans -- f xs = AccumT $ \x -> xs >>= fmap (first return) . flip (runAccumT . f) x

instance Monad m => DistributiveM m (ContT r m) where
  distributem = distributemTrans -- xs = ContT $ \f -> xs >>= ($ f . return) . runContT
  collectm = collectmTrans -- f xs = ContT $ \g -> xs >>= \y -> (runContT . f) y (g . return)

instance Monad m => DistributiveM m (ExceptT e m) where
  distributem = distributemTrans -- ExceptT . (>>= fmap (fmap return) . runExceptT)
  collectm = collectmTrans -- f xs = ExceptT $ xs >>= fmap (fmap return) . runExceptT . f

instance Monad m => DistributiveM m (IdentityT m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance Monad m => DistributiveM m (MaybeT m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance (Monoid w, Monad m) => DistributiveM m (RWST r w s m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance (Monoid w, Monad m) => DistributiveM m (Strict.RWST r w s m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance Monad m => DistributiveM m (StateT s m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance Monad m => DistributiveM m (Strict.StateT s m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance (Monoid w, Monad m) => DistributiveM m (WriterT w m) where
  distributem = distributemTrans
  collectm = collectmTrans

instance (Monoid w, Monad m) => DistributiveM m (Strict.WriterT w m) where
  distributem = distributemTrans
  collectm = collectmTrans

-- | `distributem` for any `MonadTrans` that's a `Monad` when its base functor is.
distributemTrans :: (Monad m, MonadTrans t, Monad (t m)) => m (t m a) -> t m (m a)
distributemTrans = fmap return . join . lift

-- | `collectm` for any `MonadTrans` that's a `Monad` when its base functor is.
--
-- We `lift` the `m` then @(`>>=`)@ the function to the result.
-- Finally, we `return` inside of the transformer.
collectmTrans :: (Monad m, MonadTrans t, Monad (t m)) => (a -> t m b) -> m a -> t m (m b)
collectmTrans f = fmap return . (>>= f) . lift

-- Comonad transformers

instance (Comonad m, Monad m) => DistributiveM m (EnvT e m) where
  distributem = liftM2 EnvT (ask . extract) (fmap lowerEnvT)
  collectm = collectmByDistributem

instance (Comonad m, Monad m) => DistributiveM m (StoreT s m) where
  distributem = liftM2 StoreT (fmap (\(~(StoreT acc _)) -> distribute acc)) (pos . extract)
  collectm = collectmByDistributem

instance (Monoid w, Monad m) => DistributiveM m (TracedT w m) where
  distributem = TracedT . fmap (distribute . runTracedT)
  collectm f = TracedT . (>>= runTracedT . fmap return . f)

-- Free

instance (Comonad m, Monad m) => DistributiveM m (Cofree m) where
  distributem xs =
    fmap extract xs :< (xs >>= (\(~(_ :< ys)) -> fmap return <$> ys))
  collectm = collectmByDistributem

-- | A wrapped representation of `CofreeF` without the outermost
-- `Comonad`ic layer
newtype WrappedCofreeF f w a = WrappedCofreeF
  { runWrappedCofreeF :: Trans.CofreeF f a (CofreeT f w a)
  }

instance (Functor f, Functor w) => Functor (WrappedCofreeF f w) where
  fmap f = WrappedCofreeF . bimap f (fmap f) . runWrappedCofreeF

-- | Convert `CofreeT` to our wrapped representation
toWrappedCofreeF :: Functor w => CofreeT f w a -> w (WrappedCofreeF f w a)
toWrappedCofreeF = fmap WrappedCofreeF . runCofreeT

-- | Convert `WrappedCofreeF` to a product, which has a suitable automatically
-- derived `DistributiveM` instance.
{-# INLINE wrappedCofreeFToProd #-}
wrappedCofreeFToProd ::
     WrappedCofreeF f w a -> (ConstT f :*: Compose (IdentityT f) (CofreeT f w)) a
wrappedCofreeFToProd ~(WrappedCofreeF (x Trans.:< xs)) =
  ConstT x :*: Compose (IdentityT xs)

-- | Convert `WrappedCofreeF` from a product
{-# INLINE wrappedCofreeFFromProd #-}
wrappedCofreeFFromProd ::
     (ConstT f :*: Compose (IdentityT f) (CofreeT f w)) a -> WrappedCofreeF f w a
wrappedCofreeFFromProd ~(ConstT x :*: Compose (IdentityT xs)) =
  WrappedCofreeF $ x Trans.:< xs

instance (DistributiveM m f, Functor f) =>
         DistributiveM m (WrappedCofreeF m f) where
  distributem = wrappedCofreeFFromProd . collectm wrappedCofreeFToProd
  collectm f = wrappedCofreeFFromProd . collectm (wrappedCofreeFToProd . f)

instance (DistributiveM m f, Functor f) => DistributiveM m (CofreeT m f) where
  distributem =
    CofreeT .
    fmap runWrappedCofreeF . getCompose . collectm (Compose . toWrappedCofreeF)
  collectm f =
    CofreeT .
    fmap runWrappedCofreeF .
    getCompose . collectm (Compose . toWrappedCofreeF . f)

instance Monad m => DistributiveM m (CoiterT m) where
  distributem = CoiterT . (>>= fmap (bimap return (fmap return)) . runCoiterT)
  collectm f =
    CoiterT . (>>= fmap (bimap return (fmap return)) . runCoiterT . f)

-- | Deconstruct a `Free`
{-# INLINE free #-}
free :: (a -> b) -> (f (Free f a) -> b) -> Free f a -> b
free f _ (Pure x) = f x
free _ g ~(Free x) = g x

instance Monad m => DistributiveM m (Free m) where
  distributem = Free . fmap (fmap return)
  collectm f =
    Free . (>>= free (return . return . return) (fmap return <$>) . f)

instance Monad m => DistributiveM m (F m) where
  distributem = toF . collectm fromF
  collectm = collectmByDistributem

-- | A wrapped representation of `FreeF` without the outermose `Monad`ic layer
newtype WrappedFreeF f m a = WrappedFreeF
  { runWrappedFreeF :: Trans.FreeF f a (FreeT f m a)
  }

instance (Functor f, Monad m) => Functor (WrappedFreeF f m) where
  fmap f = WrappedFreeF . bimap f (fmap f) . runWrappedFreeF

-- | Convert `FreeT` to our wrapped representation
toWrappedFreeF :: Functor m => FreeT f m a -> m (WrappedFreeF f m a)
toWrappedFreeF = fmap WrappedFreeF . runFreeT

-- | Deconstruct a `FreeF`
{-# INLINE freeF #-}
freeF :: (a -> c) -> (f b -> c) -> Trans.FreeF f a b -> c
freeF f _ (Trans.Pure x) = f x
freeF _ g ~(Trans.Free x) = g x

instance (Monad m, Monad n) => DistributiveM m (WrappedFreeF m n) where
  distributem =
    WrappedFreeF .
    Trans.Free .
    (>>= freeF (return . return . return) (fmap return <$>) . runWrappedFreeF)
  collectm f =
    WrappedFreeF .
    Trans.Free .
    (>>= freeF (return . return . return) (fmap return <$>) .
         runWrappedFreeF . f)

instance (DistributiveM m n, Monad n) => DistributiveM m (FreeT m n) where
  distributem =
    FreeT .
    fmap runWrappedFreeF . getCompose . collectm (Compose . toWrappedFreeF)
  collectm f =
    FreeT .
    fmap runWrappedFreeF . getCompose . collectm (Compose . toWrappedFreeF . f)

instance (DistributiveM m n, Monad n) => DistributiveM m (FT m n) where
  distributem = toFT . collectm fromFT
  collectm = collectmByDistributem

instance Monad m => DistributiveM m (IterT m) where
  distributem = IterT . (>>= fmap (bimap return (fmap return)) . runIterT)
  collectm f = IterT . (>>= fmap (bimap return (fmap return)) . runIterT . f)

-- Generics

-- | A default implementation of `distributem` for `Generic1` types
gdistributem :: (GDistributiveM m (Rep1 f), Generic1 f) => m (f a) -> f (m a)
gdistributem = to1 . gcollectm' from1

-- | A default implementation of `collectm` for `Generic1` types
gcollectm :: (GDistributiveM m (Rep1 f), Generic1 f) => (a -> f b) -> m a -> f (m b)
gcollectm f = to1 . gcollectm' (from1 . f)

-- | The typeclass implementing `gdistributem` (as `gdistributem'`)
-- for `Generic1` types
class Monad m => GDistributiveM (m :: * -> *) (f :: * -> *) | f -> m where
  gdistributem' :: m (f a) -> f (m a)
  gdistributem' = gcollectm' id

  gcollectm' :: (a -> f b) -> m a -> f (m b)

instance GDistributiveM m f => GDistributiveM m (M1 i c f) where
  gcollectm' f = M1 . gcollectm' (unM1 . f)

instance GDistributiveM Identity U1 where
  gcollectm' _ _ = U1

instance GDistributiveM Identity Par1 where
  gcollectm' f = Par1 . fmap (unPar1 . f)

instance (GDistributiveM m f, GDistributiveM m g) =>
         GDistributiveM m (f :*: g) where
  gcollectm' f xs = gcollectm' leftF xs :*: gcollectm' rightF xs
    where
      leftF = (\(~(ys :*: _)) -> ys) . f
      rightF = (\(~(_ :*: zs)) -> zs) . f

instance DistributiveM m f => GDistributiveM m (Rec1 f) where
  gcollectm' f xs = Rec1 $ collectm (unRec1 . f) xs

instance (DistributiveM m f, GDistributiveM m g, Functor f) =>
         GDistributiveM m (f :.: g) where
  gcollectm' f xs = Comp1 $ gdistributem' <$> collectm (unComp1 . f) xs


{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes, GADTs, TypeFamilies, TypeOperators, ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
{-# LANGUAGE DataKinds, PolyKinds, ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Language.Grappa.Distribution where

import Control.Monad
import Control.Monad.Trans
import GHC.Exts
import Data.Typeable
import qualified Data.Vector.Storable as V
--import Data.Number.LogFloat hiding (sum, product, log)

import Numeric.LinearAlgebra hiding (R, Uniform, (<>), (<#))
import qualified Numeric.LinearAlgebra as M

import qualified Numeric.Log as Log
import qualified Numeric.SpecFunctions as Gamma

import qualified Numeric.AD.Mode.Forward as ADF
import qualified Numeric.AD.Internal.Forward as ADF
import qualified Numeric.AD.Mode.Reverse as ADR
import qualified Numeric.AD.Internal.Reverse as ADR
import qualified Numeric.AD.Internal.Identity as ADR
import qualified Numeric.AD.Jacobian as ADR
import qualified Data.Reflection as ADR (Reifies)

-- | The R type alias
type R = Double

----------------------------------------------------------------------
-- * Supporting the Gamma and Beta Functions
----------------------------------------------------------------------

-- | Class for computing @ln (gamma x)@
class HasGamma a where
  logGamma :: a -> a
  digamma :: a -> a

instance HasGamma Double where
  logGamma = Gamma.logGamma
  digamma = Gamma.digamma

-- | The 'trigamma' function = the derivative of 'digamma' (which itself is the
-- derivative of 'logGamma'). This implementation was cribbed from
-- <https://stackoverflow.com/questions/2978979/how-to-improve-performance-of-this-numerical-computation-in-haskell/2983578 this>
-- Stack Overflow post
trigamma :: Double -> Double
trigamma x = go 0 (x + 5) $ computeP $ x + 6
  where
    go :: Int -> Double -> Double -> Double
    go !i !z !p
        | i >= 6    = p
        | otherwise = go (i+1) (z-1) (1 / (z*z) + p)
    invSq z = 1 / (z * z)
    computeP z = (((((5/66*p-1/30)*p+1/42)*p-1/30)*p+1/6)*p+1)/z+0.5*p
      where p = invSq z

----------------------------------------------------------------------
-- * Helper Definitions for Automatic Differentiation
----------------------------------------------------------------------

-- | Convert between two 'AD.Forward' types, using a conversion function which
-- should numerically be the identity, i.e., it should the representation of @r@
-- in type @a@ to the closest representable number to @r@ in @b@.
{-
unsafeConvertAD :: (a -> b) -> AD.Forward a -> AD.Forward b
unsafeConvertAD f (AD.Forward x y) =
  AD.Forward (f x) (f y)
unsafeConvertAD f (AD.Lift x) =
  AD.Lift (f x)
unsafeConvertAD _ AD.Zero = AD.Zero
-}

-- | Apply a function to an 'ADF.Forward' object that has a more efficient or
-- accurate version for the underlying type. For instance, log1p = log . (1+)
-- has a more accurate version in C code, but we still need to apply the above
-- mathematical definition to compute the derivative.
optimizedADOp :: Num a => (ADF.Forward a -> ADF.Forward a) -> (a -> a) ->
                 ADF.Forward a -> ADF.Forward a
optimizedADOp f f_opt x =
  ADF.bundle (f_opt $ ADF.primal x) (ADF.tangent $ f x)

-- Needed to use @'ADF.Forward' 'Double'@ Log
instance Log.Precise (ADF.Forward Double) where
  log1p = optimizedADOp Log.log1p (log . (1 +))
  expm1 = optimizedADOp Log.expm1 ((\x -> x - 1) . exp)

-- Needed to use @'ADR.Reverse' 'Double'@ Log
instance ADR.Reifies rs ADR.Tape => Log.Precise (ADR.Reverse rs Double) where
  log1p x = ADR.unary Log.log1p (ADR.Id (1 / (ADR.primal x + 1))) x
  expm1 x = ADR.unary Log.expm1 (ADR.Id ((ADR.primal x) * exp (ADR.primal x))) x

instance ADR.Reifies s ADR.Tape => HasGamma (ADR.Reverse s Double) where
  logGamma x =
    ADR.unary Gamma.logGamma (ADR.Id $ Gamma.digamma $ ADR.primal x) x
  digamma x =
    ADR.unary Gamma.digamma (ADR.Id $ trigamma $ ADR.primal x) x

----------------------------------------------------------------------
-- * Probabilities
----------------------------------------------------------------------

-- | We use the shorthand @Prob@ for the type of probabilities
newtype Prob = Prob { fromProb :: Log.Log Double }
             deriving (Num,Eq,Ord,Real,Fractional,Typeable)

instance Show Prob where
  show (Prob x) = show x

-- | Specialized method for testing if a 'Prob' is zero
probIsZero :: Prob -> Bool
probIsZero (Prob x) = x == 0

-- | Specialized method for testing if a 'Prob' is a NaN
probIsNaN :: Prob -> Bool
probIsNaN (Prob x) = isNaN $ Log.ln x

-- | Convert a (non-negative) real to a 'Prob'
rToProb :: R -> Prob
rToProb = Prob . Log.Exp . log

-- | Convert a real that is already in log-space to a 'Prob'
logRToProb :: R -> Prob
logRToProb = Prob . Log.Exp

-- | Convert a 'Prob' to a real (not in log space)
probToR :: Prob -> R
probToR = exp . Log.ln . fromProb

-- | Convert a 'Prob' to a real in log space
probToLogR :: Prob -> R
probToLogR = Log.ln . fromProb

-- | Lifting the 'Log.sum' method to 'Prob'
sumProb :: [Prob] -> Prob
sumProb ps = Prob $ Log.sum $ map fromProb ps

----------------------------------------------------------------------
-- * Matrices of reals and probabilities
----------------------------------------------------------------------

-- | A bundled up type for unboxed real-valued vectors
newtype RVector = RVector { unRVector :: Vector Double } deriving (Num,Eq,Show)

-- | Return the length of an 'RVector'
lengthV :: RVector -> Int
lengthV (RVector v) = V.length v

-- | Return the @i@th element of a 'RVector'
atV :: RVector -> Int -> Double
atV (RVector v) i = v V.! i

-- | Generate a probability vector of the given size
generateV :: Int -> (Int -> Double) -> RVector
generateV n f = RVector $ V.generate n f

-- | Map over an 'RVector'
mapV :: (R -> R) -> RVector -> RVector
mapV f (RVector v) = RVector (V.map f v)

-- | Fold over an 'RVector'
foldrV :: (R -> b -> b) -> b -> RVector -> b
foldrV f b (RVector v) = V.foldr' f b v

-- | Take the sum of an 'RVector'
sumV :: RVector -> Double
sumV (RVector v) = V.foldl' (+) 0 v

-- | Generate a vector from a list of its elements
fromListV :: [Double] -> RVector
fromListV = RVector . V.fromList

-- | Extract the elements of a vector as a list
toListV :: RVector -> [Double]
toListV = V.toList . unRVector

-- | The type of real-valued matrices
newtype RMatrix = RMatrix { unRMatrix :: Matrix Double } deriving (Num,Eq,Show)

-- | Return the number of rows of an 'RMatrix'
rowsM :: RMatrix -> Int
rowsM (RMatrix m) = rows m

-- | Return the number of columns of a probability matrix
colsM :: RMatrix -> Int
colsM (RMatrix m) = M.cols m

-- | Return the @(i,j)@th element of a 'RMatrix'
atM :: RMatrix -> Int -> Int -> Double
atM (RMatrix m) i j = atIndex m (i,j)

-- | Build a matrix from a list of its rows
fromRowsM :: [RVector] -> RMatrix
fromRowsM vs = RMatrix $ fromRows $ map unRVector vs

-- | Build a matrix of probabilities from a list of its columns
fromColsM :: [RVector] -> RMatrix
fromColsM vs = RMatrix $ fromColumns $ map unRVector vs

-- | Get the rows of a matrix as a list of vectors
toRowsM :: RMatrix -> [RVector]
toRowsM (RMatrix m) = map RVector $ toRows m

-- | Get the columns of a matrix as a list of vectors
toColsM :: RMatrix -> [RVector]
toColsM (RMatrix m) = map RVector $ toColumns m

-- | Matrix multiplication
mulM :: RMatrix -> RMatrix -> RMatrix
mulM (RMatrix m1) (RMatrix m2) = RMatrix (m1 M.<> m2)

-- | Matrix-vector multiplication
mulMV :: RMatrix -> RVector -> RVector
mulMV (RMatrix m) (RVector v) = RVector (m #> v)

-- | Vector-matrix multiplication
mulVM :: RVector -> RMatrix -> RVector
mulVM (RVector v) (RMatrix m) = RVector (v M.<# m)

-- | The type of (efficient storable) vectors of values in log space
newtype ProbVector = ProbVector { unProbVector :: Vector Double }
                   deriving (Eq, Show)

-- | Return the length of a 'ProbVector'
lengthPV :: ProbVector -> Int
lengthPV (ProbVector v) = V.length v

-- | Return the @i@th element of a 'ProbVector'
atPV :: ProbVector -> Int -> Prob
atPV (ProbVector v) i = Prob $ Log.Exp $ v V.! i

-- | Generate a probability vector of the given size
generatePV :: Int -> (Int -> Prob) -> ProbVector
generatePV n f = ProbVector $ V.generate n (Log.ln . fromProb . f)

-- | Generate a probability vector from a list of its elements
fromListPV :: [Prob] -> ProbVector
fromListPV = ProbVector . V.fromList . map probToLogR

-- | Extract the elements of a probability vector as a list
toListPV :: ProbVector -> [Prob]
toListPV = map logRToProb . V.toList . unProbVector

-- | Map over a 'ProbVector'
mapPV :: (Prob -> Prob) -> ProbVector -> ProbVector
mapPV f (ProbVector v) = ProbVector (V.map (probToLogR . f . logRToProb) v)

-- | Fold over a 'ProbVector'
foldrPV :: (Prob -> b -> b) -> b -> ProbVector -> b
foldrPV f b (ProbVector v) = V.foldr' (f . logRToProb) b v

-- | Apply a binary operator pointwise to two probability vectors
binOpPV :: (Prob -> Prob -> Prob) -> ProbVector -> ProbVector -> ProbVector
binOpPV f v1 v2 =
  generatePV (min (lengthPV v1) (lengthPV v2)) (\i -> f (atPV v1 i) (atPV v2 i))

instance Num ProbVector where
  (+) = binOpPV (+)
  (-) = binOpPV (-)
  (*) = binOpPV (*)
  abs = mapPV abs
  signum = mapPV signum
  fromInteger i = generatePV 1 (const $ fromInteger i)

-- | Take the sum of a 'ProbVector'. The algorithm for this is adapted from
-- 'Log.sum', though that function requires a 'Foldable' instance.
sumPV :: ProbVector -> Prob
sumPV (ProbVector v) | V.length v == 0 = Prob $ Log.Exp $ log 0
sumPV (ProbVector v) =
  let max_v = V.foldl1' max v in
  Prob $ Log.Exp $
  if isInfinite max_v then max_v else
    max_v + Log.log1p (V.foldl' (\r x -> r + Log.expm1 (x - max_v)) 0 v
                       + fromIntegral (V.length v - 1))

-- | The type of matrices of values in log space
newtype ProbMatrix = ProbMatrix (Matrix Double) deriving (Eq, Show)

-- | Return the number of rows of a probability matrix
rowsPM :: ProbMatrix -> Int
rowsPM (ProbMatrix m) = rows m

-- | Return the number of columns of a probability matrix
colsPM :: ProbMatrix -> Int
colsPM (ProbMatrix m) = M.cols m

-- | Return the @(i,j)@th element of a 'ProbMatrix'
atPM :: ProbMatrix -> Int -> Int -> Prob
atPM pm@(ProbMatrix m) i j =
  if i >= rowsPM pm then
    error ("atPM index out of bounds: row " ++ show i ++ " >= " ++ show (rowsPM pm))
    else
    if j >= colsPM pm then
      error ("atPM index out of bounds: column " ++ show j
             ++ " >= " ++ show (colsPM pm))
    else
      Prob $ Log.Exp $ atIndex m (i,j)

-- | Build a matrix of probabilities from a list of its rows
fromRowsPM :: [ProbVector] -> ProbMatrix
fromRowsPM vs = ProbMatrix $ fromRows $ map unProbVector vs

-- | Build a matrix of probabilities from a list of its columns
fromColsPM :: [ProbVector] -> ProbMatrix
fromColsPM vs = ProbMatrix $ fromColumns $ map unProbVector vs

-- | Get the rows of a matrix as a list of vectors
toRowsPM :: ProbMatrix -> [ProbVector]
toRowsPM (ProbMatrix m) = map ProbVector $ toRows m

-- | Get the columns of a matrix as a list of vectors
toColsPM :: ProbMatrix -> [ProbVector]
toColsPM (ProbMatrix m) = map ProbVector $ toColumns m

-- | Matrix-vector multiplication for probability matrices
mulPMV :: ProbMatrix -> ProbVector -> ProbVector
mulPMV m v =
  generatePV (rowsPM m) $ \i ->
  sumPV $ generatePV (min (colsPM m) (lengthPV v)) $ \j ->
  (atPM m i j) * (atPV v j)

-- | Vector-matrix multiplication
mulPVM :: ProbVector -> ProbMatrix -> ProbVector
mulPVM v m =
  generatePV (colsPM m) $ \j ->
  sumPV $ generatePV (min (rowsPM m) (lengthPV v)) $ \i ->
  (atPM m i j) * (atPV v i)

-- | Matrix-matrix multiplication for probability matrices
mulPM :: ProbMatrix -> ProbMatrix -> ProbMatrix
mulPM m1 (ProbMatrix m2) =
  ProbMatrix $ fromColumns $
  map (unProbVector . mulPMV m1 . ProbVector) $ toColumns m2

----------------------------------------------------------------------
-- * Distributions
----------------------------------------------------------------------

-- | This defines the support type of a distribution
type family Support (d :: *) :: *

-- | The base class for distributions, stating that we can evaluate the density
-- of any given distribution at a particular value.  For debugging purposes, we
-- also require distributions to be 'Show'able.
class PDFDist d where
  distDensity :: d -> Support d -> Prob

-- | This type class states that monad @m@ supports distribution type @d@ by
-- allowing distributions of type @d@ to be randomly sampled in @m@
class Monad m => SampleableIn m d where
  distSample :: d -> m (Support d)

instance (SampleableIn m d,
          Monad (t m), MonadTrans t) => SampleableIn (t m) d where
  distSample = lift . distSample

-- | This type family states that all distribution types in @ds@ satisfy @c@
type family DistsIn (ds :: [*]) (c :: * -> Constraint) :: Constraint where
  DistsIn '[] c = ()
  DistsIn (d ': ds) c = (c d, DistsIn ds c)

-- | This constraint says that the distribution type @d@ is continuous, i.e.,
-- isomorphic to the reals.
--
-- The class functions take @d@, not @Proxy d@, since e.g. we need
-- different implementations for @uniform 0 1 :: Uniform@ and @uniform
-- 1 2 :: Uniform@.
class Continuous d where
  -- | Note that 'toReal' need not be defined on all of @Support d@,
  -- and implementation should raise an error on out of bounds
  -- values. E.g., for @uniform 0 1@, we have @Support Uniform ~
  -- R@, but in fact 'toReal' is only defined on @[0,1]@. We
  -- don't use @Maybe R@ as the return type, since passing an out
  -- of bounds value to 'toReal' is probably a bug.
  toReal :: d -> Support d -> R
  fromReal :: d -> Double -> Support d

{- FIXME: this does not seem to be the derivative in the below, and is in fact
just fromReal in log-space. Figure this out or remove it!

  -- | The derivative of 'fromReal'.
  --
  -- In the multivariable case the derivative becomes the determinant
  -- of the Jacobian (matrix of partial derivatives).
  --
  -- We can factor this into a separate class if we end up with some
  -- continuous distributions for which we don't have Jacobian
  -- determinants.
  fromReal' :: d -> Double -> LogFloat
-}

-- | The PDF of a continuous dist transformed to be a dist on the
-- whole real line.
--
-- All of our single variable transforms are increasing, so I don't
-- think we have to worry about taking absolute values; indeed, the
-- @LogFloat@ is not even defined for negative values. (compare
-- formulas for single dimensional change of variables which uses
-- derivative and multiple dimension change of variables where the
-- *absolute value* of the Jacobian is used:
-- https://en.wikipedia.org/wiki/Integration_by_substitution).
{-
fromRealDistDensity :: (Continuous d, PDFDist d) => d -> Double -> Prob
fromRealDistDensity d x = distDensity d (fromReal d x) * Prob (fromReal' d x)
-}

----------------------------------------------------------------
-- * Helper functions for defining 'Continuous' instances
----------------------------------------------------------------

-- TODO(conathan): revisit these bijections once we have a better
-- understanding of any additional properties we want from them. We
-- plan to use these bijections to work in the reals and not have to
-- worry about bounds.
--
-- Potential trouble: consider the open interval @(lb,ub)@ bijection
-- below in @toRealLbUb@ and @fromRealLbUb@. The logit function @\p ->
-- log (p / (1 - p))@ is small for most of the reals, so if we do
-- anything where we choose random reals over a large range, nearly
-- all of them will cluster in the ends of underlying interval
-- (i.e. near @lb@ and @ub@). If we want more uniform choices in the
-- reals to map to more uniform choices in our underlying interval,
-- then we'll need a bijection that spreads the interval much wider.
--
-- Some bijections to consider (these need to be scaled to work with
-- @(lb,ub)@):
--
-- - @\x -> log (x / (1 - x))@:
--   https://www.wolframalpha.com/input/?i=plot+log+(x+/+(1+-+x))
--
-- - @\x -> (1 / (1 - x)) - (1 / x)@:
--   https://www.wolframalpha.com/input/?i=plot+(1/(1-x)+-+1/x)
--
-- - @\x -> tan x@:
--   https://www.wolframalpha.com/input/?i=plot+tan(x)
--
-- All of these can be made more uniform by scaling them by a
-- constant, e.g. @\x -> 100000 * tan x@.

{- FIXME: this is only used for fromReal'; see comments for it above
-- | The derivative of the identify function, returning @LogFloat@.
fromRealId' :: Double -> LogFloat
fromRealId' = logFloat
-}

-- | Biject the open interval @(lb,ub)@ onto the reals.
toRealLbUb :: Floating a => a -> a -> a -> a
toRealLbUb lb ub x = r
  where
    -- Biject @(lb,ub)@ onto @(0,1)@
    p = (x - lb) / (ub - lb)
    -- Biject @(0,1)@ onto the reals
    --r = log (p / (1 - p))
    r = (2*p - 1) / (p - p*p)

-- | Biject the reals on the open interval @(lb,ub)@
fromRealLbUb :: Floating a => a -> a -> a -> a
fromRealLbUb lb ub r = x
  where
    -- Inverse of @\p -> log (p / 1 - p)@, except that we have a special case
    -- for infinite values of exp(-r), to prevent us getting NaN values for AD
    -- exp_neg_r = exp (-r)
    -- p = if isInfinite exp_neg_r then 0 else 1 / (1 + exp_neg_r)
    p = (r - 2 + sqrt (r*r + 4)) / (2*r)
    -- Inverse of @\x -> (x - lb) / (ub - lb)@.
    x = p * (ub - lb) + lb

{-
-- | Derivative of 'fromRealLbUb'.
fromRealLbUb' :: R -> R -> R -> LogFloat
fromRealLbUb' lb ub r =
  -- The below is an expansion of
  -- @log ((ub - lb) * (exp (-r) / (1 + exp (-r))**2))@.
  --
  -- Here we use @log1p@ -- an implementation of @\p -> log (1 + p)@
  -- that's optimized for small @p@ -- from
  -- @Data.Numeric.LogFloat@. This implementation of @log1p@ will use
  -- an optimized C implementation if compiled with the @useffi@ flag,
  -- which is supposed to be enabled by default.
  --
  -- The @log-domain@ package that provides an native Haskell
  -- implementation of @log1p@, I think.
  logToLogFloat $ log (ub - lb) - r - 2 * log1p (exp (- r))
-}

-- | Biject the open interval @(0,\infty)@ onto the reals, using the piecewise
-- function that maps @x@ in @(0,1)@ to @log x@ and otherwise maps @x@ to @x-1@.
toRealZeroInftyPW :: Double -> Double
toRealZeroInftyPW x = if x < 1 then log x else x - 1

-- | Biject the reals on the open interval @(0,\infty)@, using the piecewise
-- function that maps negative @x@ to @exp x@ and otherwise maps @x@ to @x+1@.
fromRealZeroInftyPW :: Double -> Double
fromRealZeroInftyPW x = if x < 0 then exp x else x + 1

-- | Biject the open interval @(lb,\infty)@ onto the reals.
toRealLbInfty :: Double -> Double -> R
toRealLbInfty lb x = log (x - lb)

-- | Biject the reals onto the open interval @(lb,\infty)@.
fromRealLbInfty :: Double -> Double -> R
fromRealLbInfty lb r = exp r + lb

-- | Biject the reals onto the open interval @(lb,\infty)@ in log-space
fromRealLbInfty' :: Double -> Double -> Prob
fromRealLbInfty' lb r = logRToProb lb + logRToProb r

-- | Biject the open interval @(-\infty,ub)@ onto the reals.
toRealInftyUb :: Double -> Double -> R
toRealInftyUb ub x = - log (ub - x)

-- | Biject the reals onto the open interval @(-\infty,ub)@.
fromRealInftyUb :: Double -> Double -> R
fromRealInftyUb ub r = - (exp (- r) - ub)

-- | Biject the reals onto the open interval @(lb,\infty)@ in log-space
fromRealInftyUb' :: Double -> Double -> Prob
fromRealInftyUb' _ub _r =
  error "FIXME: fromRealInftyUb' is not right!"
  -- logRToProb ub + logRToProb (- r)

----------------------------------------------------------------------
-- * Example distribution types
----------------------------------------------------------------------

-- | Evaluate the density of a normal distribution at value @r@. This is
-- "unchecked" in that it assumes @sigma@ is non-negative.
normalDensityUnchecked :: Floating r => r -> r -> r -> Log.Log r
normalDensityUnchecked mu sigma x =
  Log.Exp $ - 0.5 * sq ((x - mu) / sigma) - log sigma - halfLogTwoPi
  where
    sq n = n * n
    halfLogTwoPi = 0.5 * log (2 * pi)

-- | Evaluate the density of a normal distribution at value @r@
normalDensityGen :: (Ord r, Floating r, Show r) => r -> r -> r -> Log.Log r
normalDensityGen mu sigma x =
  if sigma <= 0 then
    error $ "normalDensity: sigma is <= 0: sigma = " ++ show sigma
  else
    normalDensityUnchecked mu sigma x

-- | Evaluate the density of a normal distribution at value @r@
normalDensity :: Double -> Double -> Double -> Prob
normalDensity mu sigma x = Prob $ normalDensityGen mu sigma x

data StdNormal = StdNormal deriving Show
type instance Support StdNormal = R
instance PDFDist StdNormal where
  distDensity StdNormal x = logRToProb $ - 0.5 * sq x - halfLogTwoPi
    where
    sq n = n * n
    halfLogTwoPi = 0.5 * log (2 * pi)

instance Continuous StdNormal where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

-- | The normal distribution
data Normal = Normal R R deriving Show
type instance Support Normal = R

instance PDFDist Normal where
  distDensity (Normal mu sigma) r = normalDensity mu sigma r

instance Continuous Normal where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

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

-- | Multivariate normal distribution.
--
-- Our multivariate normal of @k@ dimensions is parameterized by the
-- mean vectur @mu = [mu1,...,muk]@ and the *Cholesky factor* @c@ of
-- the covariance matrix @Sigma@. I.e., for multivariate normal with
-- mean @mu@ and covariance matrix @Sigma@, we use
--
-- > MVNormal mu c
--
-- for
--
-- > transpose c * c = Sigma
--
-- The requirements for @c@ to be well formed are that
--
-- - @c@ is upper triangular.
--
-- - all diagonal entries of @c@ are positive.
--
-- There are several reasons to prefer the Cholesky factor @c@ to the
-- covariance matrix @Sigma@:
--
-- - computing the PDF is more efficient using @c@.
--
-- - the sampling algorithm uses @c@, so we need to compute it anyway.
--
-- - all upper triangular matrices with positive diagonal entries are
--   valid Cholesky factors, and a matrix is positive definite iff it
--   has a Cholesky factorization, and the Cholesky factorization is
--   unique. So, it's easy to uniformly sample the space of
--   multivariate normals by uniformly sampling the space upper
--   triangular matrices with positive diagonal entries. We care about
--   sampling multivariate normals e.g. when using multivariate
--   normals as a variational family in Black Box Variation Inference.
--
-- There is a more general definition of multivariate normal which
-- only requires non-negative entries on the diagonal of @c@, where
-- zero values correspond to zero variance random variables,
-- i.e. constants, but we don't support those here (the PDFs would be
-- Dirac deltas).
--
-- A covariance matrix @Sigma@ can be converted into a Cholesky factor
-- @c@ using the @Numeric.LinearAlgebra.chol@ function:
--
-- > c = chol Sigma
--
-- The covariance matrix @Sigma@ must be symmetric and positive
-- definite; you can use @Numeric.LinearAlgebra.trustsym@ to construct
-- it.
data MVNormal = MVNormal (Vector Double) (Matrix Double) deriving Show

type instance Support MVNormal = Vector Double

instance PDFDist MVNormal where
  distDensity (MVNormal mu c) x =
    mvnDensityMuC mu c x

instance (Monad m, SampleableIn m Normal) => SampleableIn m MVNormal where
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Drawing_values_from_the_distribution
  distSample (MVNormal mu c) = do
    z <- V.fromList <$> (replicateM (size mu) $ distSample (Normal 0 1))
    -- Our Cholesky factorization is @Sigma = transpose c * c@ and
    -- Wikipedia has @Sigma = c * transpose c@, so we need to
    -- transpose here.
    return $ mu + (tr c #> z)

-- | Multivariate normal density in terms of Cholesky factor.
mvnDensityMuC :: Vector Double -> Matrix Double -> Vector Double -> Prob
mvnDensityMuC mu c x =
  if not isValidCholeskyFactor
  then error "distDensity: Cholesky factor is malformed!"
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Likelihood_function.
  else logRToProb $ -0.5 * (logAbsDet + prod + k * log2pi)
  where
    log2pi = log (2 * pi)
    k = fromIntegral $ size x
    -- We have
    --
    -- > Sigma = transpose c * c
    --
    -- and so
    --
    -- > det Sigma = det c * det c .
    --
    -- We know @det c > 0@ since the diag is positive.
    logAbsDet = 2 * (log $ V.product (takeDiag c))
    -- Here @cholSolve@ solves a linear system for the matrix @Sigma@
    -- for which @c@ is the Cholesky factor. I.e. if
    --
    -- > u == cholSolve m v
    --
    -- then
    --
    -- > transpose c * c * v == u .
    prod = (x - mu) <.> flatten (cholSolve c (asColumn $ x - mu))

    isValidCholeskyFactor = all (> 0) (V.toList $ takeDiag c)

-- | Multivariate normal density in terms of covariance matrix @Sigma@.
--
-- Can use this to test the other implementation above; see
-- 'testMvDensity' below.
mvnDensityMuSigma :: Vector Double -> Matrix Double -> Vector Double -> Prob
mvnDensityMuSigma mu sigma x =
  if signDet < 0
  then error "distDensity: matrix is not positive definite!"
  -- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Likelihood_function.
  else logRToProb $ -0.5 * (logAbsDet + prod + k * log2pi)
  where
    log2pi = log (2 * pi)
    k = fromIntegral $ size x
    (sigmaInv, (logAbsDet, signDet)) = invlndet sigma
    prod = (x - mu) <.> (sigmaInv #> (x - mu))

-- | Check that the two mv normal density functions agree.
--
-- They agree if this function returns a small number.
testMvnDensity :: Vector Double -> Matrix Double -> Vector Double -> Double
testMvnDensity mu c x =
   abs $ (probToR $ mvnDensityMuSigma mu (tr c <> c) x) -
         (probToR $ mvnDensityMuC mu c x)

-- | Should return a smalllllll number.
testMvnDensityExample :: Double
testMvnDensityExample = sum [ testMvnDensity mu c x | x <- xs ]
  where
    mu = V.fromList [1,2]
    c = chol (trustSym $ (2><2) [1,0.5,0.5,1])
    xs = [ V.fromList [x0,x1] | x0 <- [-5,-4.5..5] , x1 <- [-5,-4.5..5] ]

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

-- | The uniform distribution
data Uniform = Uniform Double R deriving Show
type instance Support Uniform = R

uniformDensityGen :: (Ord r, RealFloat r, Log.Precise r) =>
                     r -> r -> r -> Log.Log r
uniformDensityGen lb ub r
  | lb >= ub =
    error "Malformed uniform distribution! Lower bound >= upper bound!"
  | r >= lb && r < ub =
    -- The expression 1 / (ub - lb) in log-space
      Log.Exp $ negate $ log $ ub - lb
  | otherwise = 0

uniformDensity :: Double -> Double -> Double -> Prob
uniformDensity lb ub r
  | lb >= ub =
    error "Malformed uniform distribution! Lower bound >= upper bound!"
  | r >= lb && r < ub = 1 / rToProb (ub - lb)
  | otherwise = 0

instance PDFDist Uniform where
  distDensity (Uniform lb ub) = uniformDensity lb ub

instance Continuous Uniform where
  toReal (Uniform lb ub) = toRealLbUb lb ub
  fromReal (Uniform lb ub) = fromRealLbUb lb ub
  --fromReal' (Uniform lb ub) = fromRealLbUb' lb ub

data StudentT = StudentT Double deriving Show
type instance Support StudentT = R

instance Continuous StudentT where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

data Cauchy = Cauchy deriving Show
type instance Support Cauchy = R
instance PDFDist Cauchy where
  distDensity Cauchy x = logRToProb $ - log (1 + sq x) - log pi
    where sq n = n * n

instance Continuous Cauchy where
  toReal _ = id
  fromReal _ = id
  --fromReal' _ = fromRealId'

data For d = forall t. For [t] (t -> d)
type instance Support (For d) = [Support d]
instance PDFDist d => PDFDist (For d) where
  distDensity (For xs f) ys =
    product [distDensity (f x) y | (x,y) <- zip xs ys]

data StdUniform = StdUniform deriving Show
type instance Support StdUniform = R

instance PDFDist StdUniform where
  distDensity _ = uniformDensity 0 1

instance Continuous StdUniform where
  toReal _ = toRealLbUb 0 1
  fromReal _ = fromRealLbUb 0 1
  --fromReal' _ = fromRealLbUb' 0 1

-- | The categorical distribution, that picks a natural number between 1 and @n@
-- from a list of probabilisties of each number. Note that the probabilities
-- need not be normalized, e.g., the list @[1,1,1]@ is treated the same as
-- @[1/3,1/3,1/3]@, i.e., that each of 0, 1, and 2 has equal probability.
data Categorical = Categorical [Prob] deriving Show
type instance Support Categorical = Int

instance PDFDist Categorical where
  distDensity (Categorical ws) n =
    let total_w = sum ws in
    if n >= length ws then 0 else ws!!n / total_w

data Exponential = Exponential R deriving Show
type instance Support Exponential = R

exponentialDensityUnchecked :: (RealFloat a, Log.Precise a) =>
                               a -> a -> Log.Log a
exponentialDensityUnchecked rate x = Log.Exp $ log rate - rate * x

exponentialDensity :: (Ord a, RealFloat a, Log.Precise a) => a -> a -> Log.Log a
exponentialDensity rate x =
  if rate > 0 then exponentialDensityUnchecked rate x else 0

instance PDFDist Exponential where
  distDensity (Exponential rate) = Prob . exponentialDensity rate

data Beta = Beta R R deriving Show
type instance Support Beta = R

betaDensityUnchecked :: (HasGamma a, Floating a) => a -> a -> a -> Log.Log a
betaDensityUnchecked alpha beta x =
  Log.Exp $
  ((alpha - 1) * log x) + ((beta - 1) * log (1 - x)) -
  (logGamma alpha + logGamma alpha - logGamma (alpha + beta))

betaDensity :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
               a -> a -> a -> Log.Log a
betaDensity alpha beta x =
  if x > 0 && x < 1 then betaDensityUnchecked alpha beta x else 0

-- | Compute the PDF of the Beta distribution where the argument is already in
-- log space
betaDensityLog :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
                  a -> a -> Log.Log a -> Log.Log a
betaDensityLog alpha beta x =
  if x > 0 && x < 1 then
    Log.Exp $
    ((alpha - 1) * Log.ln x) + ((beta - 1) * Log.ln (1 - x)) -
    (logGamma alpha + logGamma alpha - logGamma (alpha + beta))
  else 0

instance PDFDist Beta where
  distDensity (Beta alpha beta) = Prob . betaDensity alpha beta

data Gamma = Gamma R R deriving Show
type instance Support Gamma = R

-- | Calculate the PDF of the gamma distribution, given shape parameter @k@ and
-- scale parameter @theta@, but without the @1 / (gamma k * theta ** k)@
-- normalization constant
relativeGammaDensity :: Floating a => a -> a -> a -> Log.Log a
relativeGammaDensity k theta x =
  Log.Exp $ (k - 1) * log x - (x / theta)

-- | Calculate the PDF of the gamma distribution, assuming @x > 0@, given shape
-- parameter @k@ and scale parameter @theta@
gammaDensityUnchecked :: (HasGamma a, Floating a) => a -> a -> a -> Log.Log a
gammaDensityUnchecked k theta x =
  Log.Exp (Log.ln (relativeGammaDensity k theta x) - logGamma k - k * log theta)

-- | Calculate the PDF of the gamma distribution, given shape parameter @k@ and
-- scale parameter @theta@
gammaDensity :: (Ord a, HasGamma a, RealFloat a, Log.Precise a) =>
                a -> a -> a -> Log.Log a
gammaDensity k theta x =
  if x <= 0 then 0 else gammaDensityUnchecked k theta x

instance PDFDist Gamma where
  distDensity (Gamma k theta) = Prob . gammaDensity k theta

data Dirichlet = Dirichlet [R] deriving Show
type instance Support Dirichlet = [R]

-- | The log of the multivariate beta function
logMVBeta :: (HasGamma a, Floating a) => [a] -> Log.Log a
logMVBeta alphas =
  -- Gamma (a_1) * ... * Gamma (a_n) / Gamma (a_1 + ... + a_n) in log space
  Log.Exp $ sum (map logGamma alphas) - logGamma (sum alphas)

-- | Calculate the density of the Dirichlet distribution at any type that
-- supports the gamma function
dirichletDensity :: (HasGamma a, Floating a) => [a] -> [a] -> Log.Log a
dirichletDensity alphas xs =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Log.Exp $
  sum (zipWith (\alpha x -> (alpha - 1) * log x) alphas xs) -
  Log.ln (logMVBeta alphas)

-- | Calculate the density of the Dirichlet distribution over log space
dirichletDensityLog :: (HasGamma a, Floating a) => [a] -> [Log.Log a] ->
                       Log.Log a
dirichletDensityLog alphas xs =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Log.Exp $
  sum (zipWith (\alpha x -> (alpha - 1) * Log.ln x) alphas xs) -
  Log.ln (logMVBeta alphas)

-- | The log of the multivariate beta function
logMVBetaV :: RVector -> Log.Log Double
logMVBetaV (RVector alphas) =
  -- Gamma (a_1) * ... * Gamma (a_n) / Gamma (a_1 + ... + a_n) in log space
  Log.Exp $ V.foldl' (\r alpha -> r + logGamma alpha) 0 alphas

-- | Calculate the density of the Dirichlet distribution over an 'RVector'
dirichletDensityV :: RVector -> RVector -> Prob
dirichletDensityV (RVector alphas) (RVector xs) =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Prob $ Log.Exp $
  V.ifoldl' (\r i alpha -> r + (alpha - 1) * log (xs!i)) 0 alphas -
  Log.ln (logMVBetaV $ RVector alphas)

-- | Calculate the density of the Dirichlet distribution over a 'ProbVector'
dirichletDensityPV :: RVector -> ProbVector -> Prob
dirichletDensityPV (RVector alphas) (ProbVector xs) =
  -- x_1 ** (alpha_1 - 1) * ... * x_n ** (alpha_n - 1) / Beta (alphas)
  Prob $ Log.Exp $
  V.ifoldl' (\r i alpha -> r + (alpha - 1) * xs!i) 0 alphas -
  Log.ln (logMVBetaV $ RVector alphas)

instance PDFDist Dirichlet where
  distDensity (Dirichlet alphas) xs = Prob $ dirichletDensity alphas xs

data Bernoulli = Bernoulli Double deriving Show

data Binomial = Binomial Int Double deriving Show

data Wishart = Wishart (Matrix Double) Double deriving Show
type instance Support Wishart = Matrix Double
instance PDFDist Wishart where
  distDensity (Wishart _scaleChol _df) _x = undefined

-- | A dirac delta distribution, that has a 100 percent chance of returning a
-- given constant value
data Dirac a = Dirac a deriving Show
type instance Support (Dirac a) = a

instance Eq a => PDFDist (Dirac a) where
  -- FIXME: this density is not really correct for R; it should really be
  -- infinite for continuous functions and 1 for discrete ones
  distDensity (Dirac a) x = if a == x then 1 else 0

-- | A technically invalid distribution, that always returns the value @x@ but
-- always has a score of @1@
data DontCare a = DontCare a deriving Show
type instance Support (DontCare a) = a

instance PDFDist (DontCare a) where
  distDensity (DontCare _) _ = 1

----------------------------------------------------------------------
-- * Continuity
----------------------------------------------------------------------

class Continuous (d f) => ReprContinuous d f
instance Continuous (d f) => ReprContinuous d f

-- | This type class states that monad @m@ supports distribution type @d@ by
-- allowing distributions of type @d@ to be randomly sampled in @m@
class Monad m => ReprSampleableIn m d f where
  distSampleRepr :: d f -> m (f (ReprSupport d))

class DiffContinuous d f where
  toRealF   :: d f -> f (ReprSupport d) -> f R
  fromRealF :: d f -> f R -> f (ReprSupport d)

class DiffPDF d f where
  distDensityF :: d f -> f (ReprSupport d) -> f Prob

{-
instance DiffContinuous ReprNormal AD.Forward where
  toRealF _ d = d
  fromRealF _ d = d

instance DiffPDF ReprNormal AD.Forward where
  distDensityF (ReprNormal mu sigma) x =
    if sigma <= 0
      then error $ "distDensityAD Normal: sigma is <= 0: sigma = " ++ show sigma
           -- FIXME: instead of using exp below, figure out how to do an
           -- adToProb that takes a real in log space and preserves
      else adToProb $ exp $ -0.5 * sq((x - mu) / sigma) - log sigma - halfLogTwoPi
      where sq n = n * n
            halfLogTwoPi = 0.5 * log (2 * pi)

instance DiffContinuous ReprUniform AD.Forward where
  toRealF (ReprUniform lb ub) =
    toRealLbUbNum lb ub
  fromRealF (ReprUniform lb ub) =
    fromRealLbUbNum lb ub

instance DiffPDF ReprUniform AD.Forward where
  distDensityF (ReprUniform lb ub) _
    | lb >= ub = error "Malformed uniform distribution! Lower bound >= upper bound!"
    | otherwise =
      (1 / adToProb (lb - ub))
-}

----------------------------------------------------------------------
-- Entropy
----------------------------------------------------------------------

-- | FIXME: document this!
class Entropy d where
  entropy :: d -> Double

instance Entropy Normal where
  entropy (Normal _mu sigma) = k + log sigma
    where
    k = 0.5 * (1 + log (2 * pi))

----------------------------------------------------------------------
-- * The ReprDistribution stuff
----------------------------------------------------------------------

-- | This defines the support type of a distribution
type family ReprSupport (d :: (* -> *) -> *) :: *

data ReprNormal f = ReprNormal (f R) (f R)
type instance ReprSupport ReprNormal = R

data ReprUniform f = ReprUniform (f R) (f R)
type instance ReprSupport ReprUniform = R

data ReprCauchy f = ReprCauchy
type instance ReprSupport ReprCauchy = R

data ReprCategorical f = ReprCategorical [f Prob]
type instance ReprSupport ReprCategorical = Int

-- FIXME: this uses different vector and matrix types than MVNormal
data ReprMVNormal f = ReprMVNormal (f RMatrix) (f RMatrix)
type instance ReprSupport ReprMVNormal = RMatrix


-- Copyright (c) 2015-2020 Rudy Matela.
-- Distributed under the 3-Clause BSD licence (see the file LICENSE).
import Test

import System.Exit (exitFailure)
import Data.List (elemIndices)

-- import Test.LeanCheck -- already exported by Test
import Test.LeanCheck.Utils

import Data.Ratio
import Data.Complex
import Data.Int
import Data.Word

main :: IO ()
main  =  do
  max <- getMaxTestsFromArgs 200
  case elemIndices False (tests max) of
    [] -> putStrLn "Tests passed!"
    is -> do putStrLn ("Failed tests:" ++ show is)
             exitFailure

tests :: Int -> [Bool]
tests n =
  [ True

  -- interleave
  , [1,2,3] +| [0,0,0] == [1,0,2,0,3,0]
  , take 3 ([1,2] +| (0:undefined)) == [1,0,2]
  , [0,2..] +| [1,3..] =| n |= [0,1..]

  -- etc
  , tNatPairOrd n
  , tNatTripleOrd n
  , tNatQuadrupleOrd n
  , tNatQuintupleOrd n
  , tNatListOrd n
  , tListsOfNatOrd n
  , listsOf (tiers::[[Nat]]) =| 10 |= tiers

  -- tests!
  , counterExample n (\x y -> x + y /= (x::Int)) == Just ["0", "0"]
  , counterExample n (\x y -> x + y == (x::Int)) == Just ["0", "1"]
  , counterExample n (maybe True (==(0::Int))) == Just ["(Just 1)"]
  , holds n (\x -> x == (x::Int))

  -- For when NaN is in the enumeration (by default, it is not):
  --, fails 100 (\x -> x == (x::Float))  -- NaN != NaN  :-)
  --, counterExample 100 (\x -> x == (x::Float)) == Just ["NaN"]
  , counterExample n (\x y -> x + y == (x::Float))  == Just ["0.0","1.0"]
  , counterExample n (\x y -> x + y == (x::Double)) == Just ["0.0","1.0"]
  , holds          n (\x -> x + 1 /= (x::Int))
  , counterExample n (\x -> x + 1 /= (x::Float))  == Just ["Infinity"]
    || counterExample n (\x -> x + 1 /= (x::Float)) == Just ["inf"] -- bug on Hugs 2006-09?
  , counterExample n (\x -> x + 1 /= (x::Double)) == Just ["Infinity"]
    || counterExample n (\x -> x + 1 /= (x::Float)) == Just ["inf"] -- bug on Hugs 2006-09?
  , allUnique (take n list :: [Float])
  , allUnique (take n list :: [Double])

  , allUnique (take n list :: [Rational])
  , allUnique (take n list :: [Ratio Nat])
  , orderedOn (\r -> numerator r + denominator r) (take n (list :: [Ratio Nat]))
  , orderedOn (\r -> abs (numerator r) + abs(denominator r)) (take n (list :: [Rational]))

  , list == [LT, EQ, GT]
  , orderedOn length (take n (list :: [[Ordering]]))
  , orderedOn length (take n (list :: [[Bool]]))

  , strictlyOrderedOn (\xs -> (sum $ map (+1) xs, xs)) (take n (list :: [[Word]]))

  , tPairEqParams n
  , tTripleEqParams n

  , tProductsIsFilterByLength (tiers :: [[ Nat ]])   10 `all` [1..10]
  , tProductsIsFilterByLength (tiers :: [[ Bool ]])   6 `all` [1..10]
  , tProductsIsFilterByLength (tiers :: [[ [Nat] ]])  6 `all` [1..10]

  , holds n $  (\/)  ==== zipWith' (++) [] [] -:> [[uint2]]
  , holds n $  (\/)  ==== zipWith' (++) [] [] -:> [[bool]]
  , holds n $ (\\//) ==== zipWith' (+|) [] [] -:> [[uint2]]
  , holds n $ (\\//) ==== zipWith' (+|) [] [] -:> [[bool]]

  , holds n $ \x -> x == (x :: Word)
  , holds n $ \x -> x == (x :: Word8)
  , holds n $ \x -> x == (x :: Word16)
  , holds n $ \x -> x == (x :: Word32)
  , holds n $ \x -> x == (x :: Word64)

  , holds n $ \x -> x == (x :: Int)
  , holds n $ \x -> x == (x :: Int8)
  , holds n $ \x -> x == (x :: Int16)
  , holds n $ \x -> x == (x :: Int32)
  , holds n $ \x -> x == (x :: Int64)

  , holds n $ \x -> x == (x :: Complex Double)
  ]

allUnique :: Ord a => [a] -> Bool
allUnique [] = True
allUnique (x:xs) = x `notElem` xs
                && allUnique (filter (< x) xs)
                && allUnique (filter (> x) xs)

-- | 'zipwith\'' works similarly to 'zipWith', but takes neutral elements to
--   operate when one of the lists is exhausted, so, you don't loose elements.
--
-- > zipWith' f z e [x,y] [a,b,c,d] == [f x a, f y b, f z c, f z d]
--
-- > zipWith' f z e [x,y,z] [a] == [f x a, f y e, f z e]
--
-- > zipWith' (+) 0 0 [1,2,3] [1,2,3,4,5,6] == [2,4,6,4,5,6]
zipWith' :: (a->b->c) -> a -> b  -> [a] -> [b] -> [c]
zipWith' _ _  _  []     [] = []
zipWith' f _  zy xs     [] = map (`f` zy) xs
zipWith' f zx _  []     ys = map (f zx) ys
zipWith' f zx zy (x:xs) (y:ys) = f x y : zipWith' f zx zy xs ys


{-# LANGUAGE AllowAmbiguousTypes                      #-}
{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DeriveAnyClass                           #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE DerivingVia                              #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE InstanceSigs                             #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Regression (
  -- * Linear Regression
    linReg, logReg
  , LRp(..), lrAlpha, lrBeta, runLRp
  -- * Reshape
  , reshapeLRpInput, reshapeLRpOutput
  , expandLRpInput, expandLRpOutput
  , premuteLRpInput, permuteLRpOutput
  -- * ARIMA
  , arima, autoregressive, movingAverage, arma
  , ARIMAp(..), ARIMAs(..)
  , arimaPhi, arimaTheta, arimaConstant, arimaYPred, arimaYHist, arimaEHist
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Combinator
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Types
import           Backprop.Learn.Regularize
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Data.Finite
import           Data.Kind
import           Data.List
import           Data.Maybe
import           Data.Proxy
import           Data.Type.Equality
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update hiding            ((<.>))
import           Statistics.Distribution
import           Unsafe.Coerce
import qualified Data.Binary                           as Bi
import qualified Data.Vector.Generic.Sized             as SVG
import qualified Data.Vector.Sized                     as SV
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Numeric.LinearAlgebra                 as HU
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

-- | Linear Regression parameter
data LRp (i :: Nat) (o :: Nat) = LRp
    { _lrAlpha :: !(R o)
    , _lrBeta  :: !(L o i)
    }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Initialize, Backprop)

deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Num (LRp i o)
deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Fractional (LRp i o)
deriving via (GNum (LRp i o)) instance (KnownNat i, KnownNat o) => Floating (LRp i o)

makeLenses ''LRp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (LRp i o) where
    type Ref m (LRp i o) = GRef m (LRp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (PrimMonad m, KnownNat i, KnownNat o) => LinearInPlace m Double (LRp i o)

instance (KnownNat i, KnownNat o) => Regularize (LRp i o) where
    rnorm_1 = rnorm_1 . _lrBeta
    rnorm_2 = rnorm_2 . _lrBeta
    lasso r LRp{..} = LRp { _lrAlpha = 0
                          , _lrBeta  = lasso r _lrBeta
                          }
    ridge r LRp{..} = LRp { _lrAlpha = 0
                          , _lrBeta  = ridge r _lrBeta
                          }

instance (KnownNat i, KnownNat o, PrimMonad m) => Learnable m (LRp i o)

runLRp
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (LRp i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLRp lrp x = (lrp ^^. lrBeta) #> x + (lrp ^^. lrAlpha)

-- | Reshape an 'LRp' to take more or less inputs.  If more, new parameters
-- are initialized randomly according to the given distribution.
reshapeLRpInput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat i', KnownNat o)
    => d
    -> MWC.Gen (PrimState m)
    -> LRp i o
    -> m (LRp i' o)
reshapeLRpInput d g (LRp α β) =
    LRp α <$> reshapeLCols d g β

-- | Reshape an 'LRp' to return more or less outputs  If more, new
-- parameters are initialized randomly according to the given distribution.
reshapeLRpOutput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat o, KnownNat o')
    => d
    -> MWC.Gen (PrimState m)
    -> LRp i o
    -> m (LRp i o')
reshapeLRpOutput d g (LRp α β) =
    LRp <$> reshapeR d g α
        <*> reshapeLRows d g β

linReg
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LRp i o)) 'Nothing (R i) (R o)
linReg = modelStatelessD (\(PJust p) -> runLRp p)

logReg
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LRp i o)) 'Nothing (R i) (R o)
logReg = funcD logistic <~ linReg

-- | Adjust an 'LRp' to take extra inputs, initialized randomly.
--
-- Initial contributions to each output is randomized.
expandLRpInput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat j, KnownNat o)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp (i + j) o)
expandLRpInput LRp{..} d g = LRp _lrAlpha . (_lrBeta H.|||) <$> initialize d g

-- | Adjust an 'LRp' to return extra ouputs, initialized randomly
expandLRpOutput
    :: (PrimMonad m, ContGen d, KnownNat i, KnownNat o, KnownNat p)
    => LRp i o
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LRp i (o + p))
expandLRpOutput LRp{..} d g = do
    newAlpha <- initialize d g
    newBeta  <- initialize d g
    pure (LRp (_lrAlpha H.# newAlpha) (_lrBeta H.=== newBeta))

-- | Premute (or remove) inputs
--
-- Removed inputs will simply have their contributions removed from each
-- output.
premuteLRpInput
    :: (KnownNat i, KnownNat o)
    => SV.Vector i' (Finite i)
    -> LRp i o
    -> LRp i' o
premuteLRpInput is p = p { _lrBeta = colsL . fmap (β `SV.index`) $ is }
  where
    β = lCols (_lrBeta p)

-- | Premute (or remove) outputs
permuteLRpOutput
    :: (KnownNat i, KnownNat o)
    => SV.Vector o' (Finite o)
    -> LRp i o
    -> LRp i o'
permuteLRpOutput is LRp{..} =
    LRp { _lrAlpha = vecR . SVG.convert . fmap (α `SVS.index`) $ is
        , _lrBeta  = rowsL . fmap (β `SV.index`) $ is
        }
  where
    α = rVec _lrAlpha
    β = lRows _lrBeta

-- | 'ARIMA' parmaeters
data ARIMAp :: Nat -> Nat -> Type where
    ARIMAp :: { _arimaPhi      :: !(R p)
              , _arimaTheta    :: !(R q)
              , _arimaConstant :: !Double
              }
           -> ARIMAp p q
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Initialize, Backprop, Bi.Binary)

deriving via (GNum (ARIMAp p q)) instance Num (ARIMAp p q)
deriving via (GNum (ARIMAp p q)) instance Fractional (ARIMAp p q)
deriving via (GNum (ARIMAp p q)) instance Floating (ARIMAp p q)

makeLenses ''ARIMAp

-- | 'ARIMA' state
data ARIMAs :: Nat -> Nat -> Nat -> Type where
    ARIMAs :: { _arimaYPred :: !Double
              , _arimaYHist :: !(R (p + d))
              , _arimaEHist :: !(R q)
              }
          -> ARIMAs p d q
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Initialize, Backprop, Bi.Binary)

deriving via (NoRegularize (ARIMAs p d q)) instance Regularize (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Num (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Fractional (ARIMAs p d q)
deriving via (GNum (ARIMAs p d q)) instance Floating (ARIMAs p d q)

makeLenses ''ARIMAs

arima
    :: forall p d q. (KnownNat p, KnownNat d, KnownNat q)
    => Model ('Just (ARIMAp p q)) ('Just (ARIMAs p d q)) Double Double
arima = modelD $ \(PJust p) x (PJust s) ->
    let d :: L p (p + d)
        d = difference
        e  = x - (s ^^. arimaYPred)
        y  = (p ^^. arimaConstant)
           + (p ^^. arimaPhi  ) <.> (constVar d #> (s ^^. arimaYHist))
           + (p ^^. arimaTheta) <.> (s ^^. arimaEHist)
        yHist' = case Proxy @1 %<=? Proxy @(p + d) of
          LE Refl -> single y # constVar dropLast #> (s ^^. arimaYHist)
          NLE _ _ -> 0
        eHist' = case Proxy @1 %<=? Proxy @q of
          LE Refl -> single e # constVar dropLast #> (s ^^. arimaEHist)
          NLE _ _ -> 0
        s' = isoVar3 ARIMAs (\(ARIMAs pr yh eh) -> (pr,yh,eh))
                y
                yHist'
                eHist'
    in  (y, PJust s')

autoregressive
    :: KnownNat p
    => Model ('Just (ARIMAp p 0)) ('Just (ARIMAs p 0 0)) Double Double
autoregressive = arima

movingAverage
    :: KnownNat q
    => Model ('Just (ARIMAp 0 q)) ('Just (ARIMAs 0 0 q)) Double Double
movingAverage = arima

arma
    :: (KnownNat p, KnownNat q)
    => Model ('Just (ARIMAp p q)) ('Just (ARIMAs p 0 q)) Double Double
arma = arima

monosquare :: forall n. (n <=? (n ^ 2)) :~: 'True
monosquare = unsafeCoerce Refl

dropLast :: forall n. (KnownNat n, 1 <= n) => L (n - 1) n
dropLast = case monosquare @n of
    Refl -> vecL . SVS.generate $ \ij ->
      let i :: Finite n
          j :: Finite (n - 1)
          (i, j) = separateProduct ij
      in  if fromIntegral @_ @Int i == fromIntegral j
            then 1
            else 0

single :: Reifies s W => BVar s Double -> BVar s (R 1)
single = konst

difference'
    :: Int                  -- ^ initial
    -> Int                  -- ^ target
    -> HU.Matrix Double     -- ^ target x initial
difference' n m = foldl' go (HU.ident m) [m + 1 .. n]
  where
    go x k = x HU.<> d k
    d k = HU.build (k-1, k) $ \i j ->
        case round @_ @Int (j - i) of
          0 -> 1
          1 -> -1
          _ -> 0

difference :: forall n m. (KnownNat n, KnownNat m) => L n (n + m)
difference = fromJust . H.create $ difference' (n + m) n
  where
    n = fromIntegral $ natVal (Proxy @n)
    m = fromIntegral $ natVal (Proxy @m)

instance (PrimMonad m, KnownNat p, KnownNat q) => Mutable m (ARIMAp p q) where
    type Ref m (ARIMAp p q) = GRef m (ARIMAp p q)
    thawRef = gThawRef
    freezeRef = gFreezeRef
    copyRef = gCopyRef

instance (PrimMonad m, KnownNat p, KnownNat d, KnownNat q) => Mutable m (ARIMAs p d q) where
    type Ref m (ARIMAs p d q) = GRef m (ARIMAs p d q)
    thawRef = gThawRef
    freezeRef = gFreezeRef
    copyRef = gCopyRef

instance (KnownNat p, KnownNat q, PrimMonad m)  => LinearInPlace m Double (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, PrimMonad m) => LinearInPlace m Double (ARIMAs p d q)

instance (KnownNat p, KnownNat q) => Regularize (ARIMAp p q) where
    rnorm_1 ARIMAp{..} = rnorm_1 _arimaPhi + rnorm_1 _arimaTheta
    rnorm_2 ARIMAp{..} = rnorm_2 _arimaPhi + rnorm_2 _arimaTheta
    lasso r ARIMAp{..} = ARIMAp { _arimaPhi      = lasso r _arimaPhi
                                , _arimaTheta    = lasso r _arimaTheta
                                , _arimaConstant = 0
                                }
    ridge r ARIMAp{..} = ARIMAp { _arimaPhi      = ridge r _arimaPhi
                                , _arimaTheta    = ridge r _arimaTheta
                                , _arimaConstant = 0
                                }

instance (KnownNat p, KnownNat q, PrimMonad m) => Learnable m (ARIMAp p q)
instance (KnownNat p, KnownNat d, KnownNat q, PrimMonad m) => Learnable m (ARIMAs p d q)


{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE DeriveGeneric #-}
module FokkerPlanck.BrownianMotion
  ( Particle(..)
  , generatePath
  , moveParticle
  , thetaPlus
  , scalePlus
  , generateRandomNumber
  , generatePath'
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Data.DList                          as DL
import           Data.Ix
import           Data.List                           as L
import           GHC.Generics                        (Generic)
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           Statistics.Distribution.Normal
import           Statistics.Distribution.Uniform
import           System.Random.MWC
import           Text.Printf
import           Utils.Distribution

data Particle =
  Particle {-# UNPACK #-}!Double -- \phi
           {-# UNPACK #-}!Double -- \rho
           {-# UNPACK #-}!Double -- \theta
           {-# UNPACK #-}!Double -- r
           {-# UNPACK #-}!Double -- weight
  deriving (Show,Generic)

instance NFData Particle

{-# INLINE getParticleRho #-}
getParticleRho :: Particle -> Double
getParticleRho (Particle _ rho _ _ _ ) = rho

thetaCheck :: Double -> Double
thetaCheck theta =
  if theta < -pi
    then thetaCheck (theta + 2 * pi)
    else if theta >= pi
           then thetaCheck (theta - 2 * pi)
           else theta

{-# INLINE scaleCheck #-}
scaleCheck :: Double -> Double -> Double
scaleCheck maxScale scale =
  if scale >= (1 / maxScale) && scale <= maxScale
    then scale
    else error
           (printf
              "scaleCheck: %.2f is out of boundary (0,%.2f)\n"
              scale
              maxScale)

{-# INLINE thetaPlus #-}
thetaPlus :: Double -> Double -> Double
thetaPlus !x !y =
  let z = x + y
   in thetaCheck z

{-# INLINE scalePlus #-}
scalePlus :: Double -> Double -> Double
scalePlus x delta = x * exp delta

{-# INLINE scalePlusPeriodic #-}
scalePlusPeriodic :: Double -> Double -> Double -> Double
scalePlusPeriodic !logMaxScale !delta !x
  | z >= m = exp (z - 2 * m)
  | z < -m = exp (2 * m + z)
  | otherwise = exp z
  where
    !m = logMaxScale
    !z = log x + delta

{-# INLINE rhoCutoff #-}
rhoCutoff :: Double -> Double
rhoCutoff !rho =
  if rho < 0
    then 0
    else rho

{-# INLINE generateRandomNumber #-}
generateRandomNumber ::
     (Distribution d, ContGen d) => GenIO -> Maybe d -> IO Double
generateRandomNumber !gen dist =
  case dist of
    Nothing -> return 0
    Just d  -> genContVar d gen

{-# INLINE moveParticle #-}
moveParticle :: Double -> Particle -> Particle
moveParticle deltaT (Particle phi rho theta r0 v) =
  let !r = r0 * deltaT
      !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
  in Particle newPhi newRho theta r0 v

{-# INLINE diffuseParticle #-}
diffuseParticle :: Double -> Double -> Double -> Particle -> Particle
diffuseParticle !deltaTheta !deltaScale maxScale (Particle phi rho theta r v) =
  Particle
    phi
    rho
    (theta `thetaPlus` deltaTheta)
    (r `scalePlus` deltaScale)
    v

brownianMotion ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => GenIO
  -> Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion randomGen thetaDist scaleDist poissonDist deltaT maxRho maxR tao stdR2 particle xs = do
  deltaTheta <- generateRandomNumber randomGen thetaDist
  deltaScale <- generateRandomNumber randomGen scaleDist
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons
                 (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2))
                 xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao))
    then return ys
    else brownianMotion
           randomGen
           thetaDist
           scaleDist
           poissonDist
           deltaT
           maxRho
           maxR
           tao
           stdR2
           (diffuseParticle deltaTheta deltaScale maxR newParticle)
           ys

{-# INLINE generatePath #-}
generatePath ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath thetaDist scaleDist poissonDist maxRho maxR tao deltaT stdR2 randomGen =
  brownianMotion
    randomGen
    thetaDist
    scaleDist
    poissonDist
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty

{-# INLINE moveParticle' #-}
moveParticle' ::  Particle -> Particle
moveParticle' (Particle phi rho theta r v) =
  let !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
   in Particle newPhi newRho theta r v

{-# INLINE diffuseParticle' #-}
diffuseParticle' :: (Distribution d, ContGen d)
  => GenIO ->  Double -> Maybe d -> Double -> Particle -> IO Particle
diffuseParticle' randomGen thetaSigma scaleDist deltaT (Particle phi rho theta r v) = do
  deltaScale <- generateRandomNumber randomGen scaleDist
  let r1 = r `scalePlus` deltaScale
  deltaTheta <-
    genContVar (normalDistr 0 (thetaSigma * sqrt (r1 * deltaT))) randomGen
  return $ Particle phi rho (theta `thetaPlus` deltaTheta) r1 v

brownianMotion' ::
     (Distribution d, ContGen d)
  => GenIO
  -> Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion' randomGen thetaSigma scaleDist poissonLambda deltaT maxRho maxR tao stdR2 particle xs = do
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2)) xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao / (r * deltaT)))
    then return ys
    else do
      diffusedParticle <-
        diffuseParticle' randomGen thetaSigma scaleDist deltaT newParticle
      brownianMotion'
        randomGen
        thetaSigma
        scaleDist
        poissonLambda
        deltaT
        maxRho
        maxR
        tao
        stdR2
        diffusedParticle
        ys

{-# INLINE generatePath' #-}
generatePath' ::
     (Distribution d, ContGen d)
  => Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath' thetaSigma scaleDist poissonLambda maxRho maxR tao deltaT stdR2 randomGen = do
  brownianMotion'
    randomGen
    thetaSigma
    scaleDist
    poissonLambda
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty


module STCR2Z2T0S0EndModal where

import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           Image.Transform         (normalizeValueRange)
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:cutoffRadiusEndPointStr:cutoffRadiusStr:reversalFactorStr:cStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      cutoffRadiusEndPoint = read cutoffRadiusEndPointStr :: Int
      cutoffRadius = read cutoffRadiusStr :: Int
      reversalFactor = read reversalFactorStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0EndModal.noindex"
      a = 20
      b = 7
      c = read cStr :: Int
      endPointFilePath =
        folderPath </>
        (printf
           "EndPoint_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%d_%d_%d_%f.dat"
           numPoint
           (round thetaFreq :: Int)
           (round scaleFreq :: Int)
           (round maxScale :: Int)
           (round tao :: Int)
           cutoffRadiusEndPoint
           thetaSigma
           scaleSigma
           a
           b
           c
           reversalFactor)
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffRadius radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  arrR2Z2T0S0EndPoint <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff cutoffRadiusEndPoint radialArr)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  let a' = round $ (fromIntegral a) * (sqrt 2) / 2
      b' = round $ (fromIntegral b) * (sqrt 2) / 2
      c' = round $ (fromIntegral c) * (sqrt 2) / 2
      xs =
        ([R2S1RPPoint (i, i, 0, 0) | i <- [a',a' + b' .. c']] L.++
         [R2S1RPPoint (i, -i, 0, 0) | i <- [-a',-(a' + b') .. -c']] L.++
         [R2S1RPPoint (i, i, 0, 0) | i <- [-a',-(a' + b') .. -c']] L.++
         [R2S1RPPoint (i, -i, 0, 0) | i <- [a',a' + b' .. c']] L.++
         [R2S1RPPoint (i, 0, 0, 0) | i <- [a,a + b .. c]] L.++
         [R2S1RPPoint (0, i, 0, 0) | i <- [-a,-(a + b) .. -c]] L.++
         [R2S1RPPoint (i, 0, 0, 0) | i <- [-a,-(a + b) .. -c]] L.++
         [R2S1RPPoint (0, i, 0, 0) | i <- [a,a + b .. c]])
      bias = computeBiasR2T0S0' numPoint numPoint theta0Freqs scale0Freqs xs
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          thetaFreqs
          scaleFreqs
          xs
  powerMethodR2Z2T0S0EndModal
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0EndPoint
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    (printf
       "_%d_%d_%d_%d_%d_%d_%.2f_%.2f_%f"
       numPoint
       (round thetaFreq :: Int)
       (round scaleFreq :: Int)
       (round maxScale :: Int)
       (round tao :: Int)
       cutoffRadiusEndPoint
       thetaSigma
       scaleSigma
       reversalFactor)
    reversalFactor
    bias
    eigenVec


{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}

module Backprop.Learn.Model.Stochastic (
    DO(..)
  , StochFunc(..)
  , FixedStochFunc, pattern FSF, _fsfRunDeterm, _fsfRunStoch
  , rreLU
  , injectNoise, applyNoise
  , injectNoiseR, applyNoiseR
  ) where

import           Backprop.Learn.Model.Class
import           Backprop.Learn.Model.Function
import           Control.Monad.Primitive
import           Data.Bool
import           Data.Kind
import           Data.Typeable
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.LinearAlgebra.Static.Vector
import qualified Data.Vector.Storable.Sized            as SVS
import qualified Statistics.Distribution               as Stat
import qualified System.Random.MWC                     as MWC
import qualified System.Random.MWC.Distributions       as MWC

-- | Dropout layer.  Parameterized by dropout percentage (should be between
-- 0 and 1).
--
-- 0 corresponds to no dropout, 1 corresponds to complete dropout of all
-- nodes every time.
newtype DO (n :: Nat) = DO { _doRate :: Double }
  deriving (Typeable)

instance KnownNat n => Learn (R n) (R n) (DO n) where
    runLearn (DO r) _ = stateless (constVar (realToFrac (1-r)) *)
    runLearnStoch (DO r) g _ = statelessM $ \x ->
        (x *) . constVar . vecR <$> SVS.replicateM (mask g)
      where
        mask = fmap (bool 1 0) . MWC.bernoulli r

-- | Represents a random-valued function, with a possible trainable
-- parameter.
--
-- Requires both a "deterministic" and a "stochastic" mode.  The
-- deterministic mode ideally should approximate some mean of the
-- stochastic mode.
data StochFunc :: Maybe Type -> Type -> Type -> Type where
    SF :: { _sfRunDeterm :: forall s. Reifies s W => Mayb (BVar s) p -> BVar s a -> BVar s b
          , _sfRunStoch
              :: forall m s. (PrimMonad m, Reifies s W)
              => MWC.Gen (PrimState m)
              -> Mayb (BVar s) p
              -> BVar s a
              -> m (BVar s b)
          }
       -> StochFunc p a b
  deriving (Typeable)

instance Learn a b (StochFunc p a b) where
    type LParamMaybe (StochFunc p a b) = p
    type LStateMaybe (StochFunc p a b) = 'Nothing

    runLearn SF{..} = stateless . _sfRunDeterm
    runLearnStoch SF{..} g = statelessM . _sfRunStoch g

-- | Convenient alias for a 'StochFunc' (random-valued function with both
-- deterministic and stochastic modes) with no trained parameters.
type FixedStochFunc = StochFunc 'Nothing

-- | Construct a 'FixedStochFunc'
pattern FSF :: (forall s. Reifies s W => BVar s a -> BVar s b)
            -> (forall m s. (PrimMonad m, Reifies s W) => MWC.Gen (PrimState m) -> BVar s a -> m (BVar s b))
            -> FixedStochFunc a b
pattern FSF { _fsfRunDeterm, _fsfRunStoch } <- (getFSF->(getWD->_fsfRunDeterm,getWS->_fsfRunStoch))
  where
    FSF d s = SF { _sfRunDeterm = const d
                 , _sfRunStoch  = const . s
                 }
{-# COMPLETE FSF #-}

newtype WrapDeterm a b = WD { getWD :: forall s. Reifies s W => BVar s a -> BVar s b }
newtype WrapStoch  a b = WS { getWS :: forall m s. (PrimMonad m, Reifies s W) => MWC.Gen (PrimState m) -> BVar s a -> m (BVar s b) }

getFSF :: FixedStochFunc a b -> (WrapDeterm a b, WrapStoch a b)
getFSF SF{..} = ( WD (_sfRunDeterm N_)
                , WS (`_sfRunStoch` N_)
                )

-- | Random leaky rectified linear unit
rreLU
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> FixedStochFunc (R n) (R n)
rreLU d = FSF { _fsfRunDeterm = vmap' (preLU v)
              , _fsfRunStoch  = \g x -> do
                  α <- vecR <$> SVS.replicateM (Stat.genContVar d g)
                  pure (zipWithVector preLU (constVar α) x)
              }
  where
    v :: BVar s Double
    v = constVar (Stat.mean d)

-- | Inject random noise.  Usually used between neural network layers, or
-- at the very beginning to pre-process input.
--
-- In non-stochastic mode, this adds the mean of the distribution.
injectNoise
    :: (Stat.ContGen d, Stat.Mean d, Fractional a)
    => d
    -> FixedStochFunc a a
injectNoise d = FSF { _fsfRunDeterm = (realToFrac (Stat.mean d) +)
                    , _fsfRunStoch  = \g x -> do
                        e <- Stat.genContVar d g
                        pure (realToFrac e + x)
                    }

-- | 'injectNoise' lifted to 'R'
injectNoiseR
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> FixedStochFunc (R n) (R n)
injectNoiseR d = FSF { _fsfRunDeterm = (realToFrac (Stat.mean d) +)
                     , _fsfRunStoch  = \g x -> do
                         e <- vecR <$> SVS.replicateM (Stat.genContVar d g)
                         pure (constVar e + x)
                     }

-- | Multply by random noise.  Can be used to implement dropout-like
-- behavior.
--
-- In non-stochastic mode, this scales by the mean of the distribution.
applyNoise
    :: (Stat.ContGen d, Stat.Mean d, Fractional a)
    => d
    -> FixedStochFunc a a
applyNoise d = FSF { _fsfRunDeterm = (realToFrac (Stat.mean d) *)
                   , _fsfRunStoch  = \g x -> do
                       e <- Stat.genContVar d g
                       pure (realToFrac e * x)
                   }

-- | 'applyNoise' lifted to 'R'
applyNoiseR
    :: (Stat.ContGen d, Stat.Mean d, KnownNat n)
    => d
    -> FixedStochFunc (R n) (R n)
applyNoiseR d = FSF { _fsfRunDeterm = (realToFrac (Stat.mean d) *)
                   , _fsfRunStoch  = \g x -> do
                       e <- vecR <$> SVS.replicateM (Stat.genContVar d g)
                       pure (constVar e * x)
                   }


module NN.NeuralNetworkTest
  (
    testsNeuralNetwork
  ) where

import Test.HUnit
import AI.HNN.FF.Network
import Numeric.LinearAlgebra.HMatrix hiding (corr)

import NUtil
import NN.NeuralNetwork
import qualified Data.Map.Strict as Map

nnVars :: Vars
nnVars = Map.fromList [("numberInputs", 2)
                      ,("numberHidden", 2)
                      ,("numberOutputs", 1)
                      ,("timesToTrain", 3)
                      ,("learningRate", 0.8)
                      ,("sigmoidOrTanh", 0)]

samples :: Samples Double
samples = [ fromList [0, 0] --> fromList [0]
          , fromList [0, 1] --> fromList [1]
          , fromList [1, 0] --> fromList [1]
          , fromList [1, 1] --> fromList [0]
          ]

expectedOutputVals :: [Double]
expectedOutputVals = [0,1,1,0]

testsNeuralNetwork :: Test
testsNeuralNetwork = TestList [ TestLabel "label" testsLabel]

testsLabel :: Test
testsLabel = TestList [ TestLabel "test1" testCreate
                      , TestLabel "test2" testGet
                      , TestLabel "test3" testTrainV]

testCreate :: Test
testCreate = TestCase (do
                          net <- create nnVars :: IO Net
                          let smartNet = trainNTimes 1500 0.8 tanh tanh' (_network net) samples
--                              outputValues = map (get smartNet . fst) samples
                              outputValues = convert1 $ map (output smartNet tanh . fst) samples
                          assertEqual "testCreate"
                            True
                            (floatVectorCompare outputValues expectedOutputVals 0.2))

testGet :: Test
testGet = TestCase (do
                       net <- create nnVars :: IO Net
                       let smartNet = trainNTimes 1000 0.8 tanh tanh' (_network net) samples
                           outputValues = convert1 $ map (get (Net smartNet nnVars) . fst) samples
                       assertEqual "testGet"
                         True
                         (floatVectorCompare outputValues expectedOutputVals 0.2))

testTrainV :: Test
testTrainV = TestCase (do
                         net <- create nnVars :: IO Net
                         let smartNet = repeatNet trainV net samples 1000
                             outputValues = convert1 $ map (get smartNet . fst) samples
                         assertEqual "testTrainV"
                           True
                           (floatVectorCompare outputValues expectedOutputVals 0.2))

convert1 :: [Vector Double] -> [Double]
convert1 = map (! 0)

floatVectorCompare :: [Double] -> [Double] -> Double -> Bool
floatVectorCompare a b tol = foldr (&&) True $ zipWith (\x y -> abs (x - y) < tol ) a b

repeatNet :: (n -> c -> n) -> n -> c -> Int -> n
repeatNet _ net _ 0 = net
repeatNet f net c n = let net' = f net c in repeatNet f net' c (n-1)


module Parser where
import Data.Binary
import Data.Binary.Get
import qualified Data.ByteString.Lazy as BS
import Codec.Compression.GZip (decompress)
import Control.Monad
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector as V
import Control.Monad.ST

type Pixel = Word8
type Image = V.Vector (Matrix Float)
type Label = Vector Float

decodeImages :: Get [Image]
decodeImages = do
    mc <- getWord32be
    guard (mc == 0x00000803)
    [d1,d2,d3] <- many 3 getWord32be
    guard (d2 == 28 && d3 == 28)
    many d1 pic
  where
    pic :: Get Image
    pic = do
      bs <- getByteString (28*28)
      return . V.singleton . reshape 28 . toVecDouble . fromByteString $ bs
    toVecDouble :: Vector Pixel -> Vector Float
    -- mapVectorM requires the monad be strict
    -- so Identity monad shall not be used
    toVecDouble v = runST $ mapVectorM (return . (/255) . fromIntegral) v

decodeLabels :: Get [Label]
decodeLabels = do
    mc <- getWord32be
    guard (mc == 0x00000801)
    d1 <- getWord32be
    many d1 lbl
  where
    lbl :: Get Label
    lbl = do
      v <- fromIntegral <$> (get :: Get Word8)
      return $ fromList (replicate v 0 ++ [1] ++ replicate (9-v) 0)

many :: (Integral n, Monad m) => n -> m a -> m [a]
many cnt dec = sequence (replicate (fromIntegral cnt) dec)

trainingData :: IO ([Image], [Label])
trainingData = do
    s <- decompress <$> BS.readFile "tdata/train-images-idx3-ubyte.gz"
    t <- decompress <$> BS.readFile "tdata/train-labels-idx1-ubyte.gz"
    return (runGet decodeImages s, runGet decodeLabels t)

testData :: IO ([Image], [Label])
testData = do
    s <- decompress <$> BS.readFile "tdata/t10k-images-idx3-ubyte.gz"
    t <- decompress <$> BS.readFile "tdata/t10k-labels-idx1-ubyte.gz"
    return (runGet decodeImages s, runGet decodeLabels t)


{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
module Neural.Training
  ( TrainingExample(..)
  , TestingExample(..)
  , train
  , backpropagation
  )
where

import           Data.List.Split                ( chunksOf )
import           Foreign.Marshal.Utils          ( fromBool )
import           Control.Monad
import           Control.Monad.Writer
import           Control.Monad.State
import           Control.Monad.Random
import           System.Random
import           System.Random.Shuffle          ( shuffleM )
import           Numeric.LinearAlgebra
import           Text.Printf
import           GHC.Generics                   ( Generic )
import           Control.DeepSeq

import           Neural.Activation
import qualified Neural.Network                as N
import           Neural.Network          hiding ( weights
                                                , biases
                                                )
import           Neural.Layer                   ( Layer(..) )

-- | A training example with an input and desired output in vector form.
data TrainingExample a = TrainingExample (Vector a) (Vector a)
    deriving (Show, Generic, NFData)
-- | A training example with an input and desired output as an index.
data TestingExample a  = TestingExample (Vector a) Int
    deriving (Show, Generic, NFData)

-- | The Hadamar product, elementwise multiplication.
(⊙) :: Num (Vector a) => Vector a -> Vector a -> Vector a
u ⊙ v = u * v

-- | Train a neural network over some amount of epochs given some training
-- data and the size of the mini batches.
train
  :: (RandomGen g, Fractional a, Numeric a, Num (Vector a))
  => ActivationFunction a      -- ^ The activation function
  -> ActivationFunction' a     -- ^ The derivative of the activation function
  -> a                         -- ^ The learning rate
  -> [TrainingExample a]       -- ^ Training data
  -> Maybe [TestingExample a]  -- ^ Optional testing data
  -> Int                       -- ^ Number of ephocs
  -> Int                       -- ^ Mini batch size
  -> Network a                 -- ^ The untrained network
  -> StateT g (Writer [String]) (Network a) -- ^ The trained network
train σ σ' η trainingData testData epochs mbs net = do
  let tellProgress epoch net' = case testData of
        Nothing        -> tell ["Epoch " ++ show epoch ++ " complete"]
        Just testData' -> tell
          [printf "Epoch %d: %d / %d (%.2f%%)" epoch evaluation n p]
         where
          evaluation = evaluate σ testData' net'
          n          = length testData'
          p          = fromIntegral evaluation / fromIntegral n * 100 :: Double
  let trainEpoch net' epoch = do
        shuffledTrainingData <- state $ runRand $ shuffleM trainingData
        let miniBatches = chunksOf mbs shuffledTrainingData
        let tmb         = trainMiniBatch σ σ' η
        let trainedNet  = foldr tmb net' miniBatches
        lift $ tellProgress epoch trainedNet
        return trainedNet

  foldM trainEpoch net [1 .. epochs]

-- | Evaluate the network assuming the output is the index of the neuron with
-- the highest activation.
evaluate
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a     -- ^ The activation function
  -> [TestingExample a]       -- ^ Testing data
  -> Network a                -- ^ The untrained network
  -> Int                      -- ^ The amount of tests the network passed
evaluate σ testData net =
  let eval (TestingExample x y) = maxIndex (fst $ feedforward σ x net) == y
      results = map eval testData
  in  foldl (\acc r -> acc + fromBool r) 0 results

-- | Train a neural network for a mini batch of training data.
trainMiniBatch
  :: (Fractional a, Numeric a, Num (Vector a))
  => ActivationFunction a    -- ^ The activation function
  -> ActivationFunction' a   -- ^ The derivative of the activation function
  -> a                       -- ^ The learning rate
  -> [TrainingExample a]     -- ^ A mini batch of training examples
  -> Network a               -- ^ The untrained network
  -> Network a               -- ^ The trained network
trainMiniBatch σ σ' η miniBatch net =
  let bp example = backpropagation σ σ' example net
      η'     = η / fromIntegral (length miniBatch)
      ws0    = map (konst 0 . size) (N.weights net)
      bs0    = map (konst 0 . size) (N.biases net)
      deltas = map (unzip . bp) miniBatch
      f (w, b) (dw, db) = (zipWith (+) w dw, zipWith (+) b db)
      (nabla_w, nabla_b) = foldl f (ws0, bs0) deltas
      ws = zipWith (\w nw -> w - scale η' nw) (N.weights net) nabla_w
      bs = zipWith (\b nb -> b - scale η' nb) (N.biases net) nabla_b
  in  newNetwork ws bs

-- | Find the gradient of the cost function with respects to a networks
-- weights and biases for one training example.
backpropagation
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a    -- ^ The activation function
  -> ActivationFunction' a   -- ^ The derivative of the activation function
  -> TrainingExample a       -- ^ A single training example
  -> Network a               -- ^ The neural network
  -> [(Matrix a, Vector a)]  -- ^ The weight and bias gradient of the
                               -- quadratic cost function
backpropagation σ σ' (TrainingExample x y) net@(Network layers) =
  let (as, zs) = unzip $ feedforwards σ x net   -- Feed forward and save the
                                                -- results
      as_      = x : as
      nabla_aC = last as - y                    -- The gradient of the cost
                                                -- function
      calc_δ (l, z) δ = bp2 σ' (weights l) δ z  -- A function that calculates
                                                -- the error for the previous
                                                -- layer
      δ0      = bp1 σ' nabla_aC (last zs)       -- The error of the last layer
      δs      = scanr calc_δ δ0 (zip (tail layers) zs)
      nabla_b = δs                              -- The 3rd equation
      nabla_w = zipWith outer δs as_            -- The 4th equation
  in  zip nabla_w nabla_b

-- | The first equation of backpropagation.
bp1
  :: (Numeric a, Num (Vector a))
  => ActivationFunction' a    -- ^ The derivative of the activation function
  -> Vector a                 -- ^ The gradient of the cost function
                                -- with respects to the activation
  -> Vector a                 -- ^ The weighted input
  -> Vector a                 -- ^ The error in the output layer
bp1 _σ' nabla_aC z = nabla_aC ⊙ σ' z where σ' = cmap _σ'

-- | The second equation of backpropagation.
bp2
  :: (Numeric a, Num (Vector a))
  => ActivationFunction' a    -- ^ The derivative of the activation function
  -> Matrix a                 -- ^ The weights of the next layer
  -> Vector a                 -- ^ The error of the next layer
  -> Vector a                 -- ^ The weighted input of the current layer
  -> Vector a                 -- ^ The error of the current layer
bp2 _σ' w δ z = (tr' w #> δ) ⊙ σ' z where σ' = cmap _σ'