{-# LANGUAGE CPP                                      #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE PolyKinds                                #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType #-}
#endif

-- |
-- Module      : Numeric.LinearAlgebra.Static.Backprop
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- A wrapper over "Numeric.LinearAlgebra.Static" (type-safe vector and
-- matrix operations based on blas/lapack) that allows its operations to
-- work with <https://hackage.haskell.org/package/backprop backprop>.  Also
-- provides orphan instances of 'Backprop' for types in
-- "Numeric.LinearAlgebra.Static".
--
-- In short, these functions are "lifted" to work with 'BVar's.
--
-- Using 'evalBP' will run the original operation:
--
-- @
-- 'evalBP' :: (forall s. 'Reifies' s 'W'. 'BVar' s a -> 'BVar' s b) -> a -> b
-- @
--
-- But using 'gradBP' or 'backprop' will give you the gradient:
--
-- @
-- 'gradBP' :: (forall s. 'Reifies' s 'W'. 'BVar' s a -> 'BVar' s b) -> a -> a
-- @
--
-- These can act as a drop-in replacement to the API of
-- "Numeric.LinearAlgebra.Static".  Just change your imports, and your
-- functions are automatically backpropagatable.  Useful types are all
-- re-exported.
--
-- Also contains 'sumElements' 'BVar' operation.
--
-- Formulas for gradients come from the following papers:
--
--     * https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
--     * http://www.dtic.mil/dtic/tr/fulltext/u2/624426.pdf
--     * http://www.cs.cmu.edu/~zkolter/course/15-884/linalg-review.pdf
--     * https://arxiv.org/abs/1602.07527
--
-- Some functions are notably unlifted:
--
--     * 'H.svd': I can't find any resources that allow you to backpropagate
--     if the U and V matrices are used!  If you find one, let me know, or
--     feel free to submit a PR!  Because of this, Currently only a version
--     that exports only the singular values is exported.
--     * 'H.svdTall', 'H.svdFlat': Not sure where to start for these
--     * 'qr': Same story.
--     https://github.com/tensorflow/tensorflow/issues/6504 might yield
--     a clue?
--     * 'H.her': No 'Num' instance for 'H.Her' makes this impossible at
--     the moment with the current backprop API
--     * 'H.exmp': Definitely possible, but I haven't dug deep enough to
--     figure it out yet!  There is a description here
--     https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf but it
--     requires some things I am not familiar with yet.  Feel free to
--     submit a PR!
--     * 'H.sqrtm': Also likely possible.  Maybe try to translate
--     http://people.cs.umass.edu/~smaji/projects/matrix-sqrt/ ?  PRs
--     welcomed!
--     * 'H.linSolve': Haven't figured out where to start!
--     * 'H.</>': Same story
--     * Functions returning existential types, like 'H.withNullSpace',
--     'H.withOrth', 'H.withRows', etc.; not quite sure what the best way
--     to handle these are at the moment.
--     * 'H.withRows' and 'H.withColumns' made "type-safe", without
--     existential types, with 'fromRows' and 'fromColumns'.

module Numeric.LinearAlgebra.Static.Backprop (
  -- * Vector
    H.R
  , H.ℝ
  , vec2
  , vec3
  , vec4
  , (&)
  , (#)
  , split
  , headTail
  , vector
  , linspace
  , H.range
  , H.dim
  -- * Matrix
  , H.L
  , H.Sq
  , row
  , col
  , (|||)
  , (===)
  , splitRows
  , splitCols
  , unrow
  , uncol
  , tr
  , H.eye
  , diag
  , matrix
  -- * Complex
  , H.ℂ
  , H.C
  , H.M
  , H.𝑖
  , toComplex
  , fromComplex
  , complex
  , real
  , imag
  , sqMagnitude
  , magnitude
  -- * Products
  , (<>)
  , (#>)
  , (<.>)
  -- * Factorizations
  , svd
  , svd_
  , H.Eigen
  , eigensystem
  , eigenvalues
  , chol
  -- * Norms
  , H.Normed
  , norm_0
  , norm_1V
  , norm_1M
  , norm_2V
  , norm_2M
  , norm_InfV
  , norm_InfM
  -- * Misc
  , mean
  , meanCov
  , meanL
  , cov
  , H.Disp(..)
  -- ** Domain
  , H.Domain
  , mul
  , app
  , dot
  , cross
  , diagR
  , vmap
  , vmap'
  , dvmap
  , mmap
  , mmap'
  , dmmap
  , outer
  , zipWithVector
  , zipWithVector'
  , dzipWithVector
  , det
  , invlndet
  , lndet
  , inv
  -- ** Conversions
  , toRows
  , toColumns
  , fromRows
  , fromColumns
  -- ** Misc Operations
  , konst
  , sumElements
  , extractV
  , extractM
  , create
  , H.Diag
  , takeDiag
  , H.Sym
  , sym
  , mTm
  , unSym
  , (<·>)
  -- * Backprop types re-exported
  -- | Re-exported for convenience.
  --
  -- @since 0.1.1.0
  , BVar
  , Backprop
  , Reifies
  , W
  ) where

import           Data.Bifunctor
import           Data.Coerce
import           Data.Functor.Identity
import           Data.Maybe
import           Data.Proxy
import           Data.Vinyl                          (Rec(..))
import           Foreign.Storable
import           GHC.TypeLits
import           Lens.Micro hiding                   ((&))
import           Numeric.Backprop
import           Numeric.Backprop.Class
import           Unsafe.Coerce
import qualified Data.Vector                         as V
import qualified Data.Vector.Generic                 as VG
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.Backprop.Explicit           as BE
import qualified Numeric.LinearAlgebra               as HU
import qualified Numeric.LinearAlgebra.Static        as H
import qualified Numeric.LinearAlgebra.Static.Vector as H
import qualified Prelude.Backprop                    as B

#if MIN_VERSION_base(4,11,0)
import           Prelude hiding               ((<>))
#endif

instance Backprop (H.R n) where
    zero = zeroNum
    add = addNum
    one = oneNum

instance Backprop (H.C n) where
    zero = zeroNum
    add = addNum
    one = oneNum

instance (KnownNat n, KnownNat m) => Backprop (H.L n m) where
    zero = zeroNum
    add = addNum
    one = oneNum

instance (KnownNat n, KnownNat m) => Backprop (H.M n m) where
    zero = zeroNum
    add = addNum
    one = oneNum

instance KnownNat n => Backprop (H.Sym n) where
    zero = zeroNum
    add = addNum
    one = oneNum

vec2
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 2)
vec2 = isoVar2 H.vec2 (\(H.rVec->v) -> (SVS.index v 0, SVS.index v 1))
{-# INLINE vec2 #-}

vec3
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 3)
vec3 = isoVar3 H.vec3 (\(H.rVec->v) -> (SVS.index v 0, SVS.index v 1, SVS.index v 2))
{-# INLINE vec3 #-}

vec4
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 4)
vec4 vX vY vZ vW = isoVarN
    (\(Identity x :& Identity y :& Identity z :& Identity w :& RNil) -> H.vec4 x y z w)
    (\(H.rVec->v) -> Identity (SVS.index v 0)
                  :& Identity (SVS.index v 1)
                  :& Identity (SVS.index v 2)
                  :& Identity (SVS.index v 3)
                  :& RNil
    )
    (vX :& vY :& vZ :& vW :& RNil)
{-# INLINE vec4 #-}

(&) :: (KnownNat n, 1 <= n, KnownNat (n + 1), Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
    -> BVar s (H.R (n + 1))
(&) = isoVar2 (H.&) (\(H.split->(dxs,dy)) -> (dxs, fst (H.headTail dy)))
infixl 4 &
{-# INLINE (&) #-}

(#) :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R m)
    -> BVar s (H.R (n + m))
(#) = isoVar2 (H.#) H.split
infixl 4 #
{-# INLINE (#) #-}

split
    :: forall p n s. (KnownNat p, KnownNat n, p <= n, Reifies s W)
    => BVar s (H.R n)
    -> (BVar s (H.R p), BVar s (H.R (n - p)))
split v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.split (uncurry (H.#)) v
    {-# NOINLINE t #-}
{-# INLINE split #-}

headTail
    :: (Reifies s W, KnownNat n, 1 <= n)
    => BVar s (H.R n)
    -> (BVar s H.ℝ, BVar s (H.R (n - 1)))
headTail v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.headTail
                   (\(d, dx) -> (H.konst d :: H.R 1) H.# dx)
                   v
    {-# NOINLINE t #-}
{-# INLINE headTail #-}

-- | Potentially extremely bad for anything but short lists!!!
vector
    :: forall n s. (KnownNat n, Reifies s W)
    => SV.Vector n (BVar s H.ℝ)
    -> BVar s (H.R n)
vector = BE.isoVar afSV
            (H.vecR . SVG.convert) (SVG.convert . H.rVec)
       . collectVar
{-# INLINE vector #-}

linspace
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R n)
linspace = liftOp2 . op2 $ \l u ->
    ( H.linspace (l, u)
    , \d -> let n1 = fromInteger $ natVal (Proxy @n) - 1
                dDot = ((H.range - 1) H.<.> d) / n1
                dSum = HU.sumElements . H.extract $ d
            in  (dSum - dDot, dDot)
    )
{-# INLINE linspace #-}

row :: Reifies s W
    => BVar s (H.R n)
    -> BVar s (H.L 1 n)
row = isoVar H.row H.unrow
{-# INLINE row #-}

col :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.L n 1)
col = isoVar H.col H.uncol
{-# INLINE col #-}

(|||) :: (KnownNat c, KnownNat r1, KnownNat (r1 + r2), Reifies s W)
      => BVar s (H.L c r1)
      -> BVar s (H.L c r2)
      -> BVar s (H.L c (r1 + r2))
(|||) = isoVar2 (H.|||) H.splitCols
infixl 3 |||
{-# INLINE (|||) #-}

(===) :: (KnownNat c, KnownNat r1, KnownNat (r1 + r2), Reifies s W)
      => BVar s (H.L r1        c)
      -> BVar s (H.L r2        c)
      -> BVar s (H.L (r1 + r2) c)
(===) = isoVar2 (H.===) H.splitRows
infixl 2 ===
{-# INLINE (===) #-}

splitRows
    :: forall p m n s. (KnownNat p, KnownNat m, KnownNat n, p <= m, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L p n), BVar s (H.L (m - p) n))
splitRows v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.splitRows (uncurry (H.===)) v
    {-# NOINLINE t #-}
{-# INLINE splitRows #-}

splitCols
    :: forall p m n s. (KnownNat p, KnownNat m, KnownNat n, KnownNat (n - p), p <= n, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L m p), BVar s (H.L m (n - p)))
splitCols v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.splitCols (uncurry (H.|||)) v
    {-# NOINLINE t #-}
{-# INLINE splitCols #-}

unrow
    :: (KnownNat n, Reifies s W)
    => BVar s (H.L 1 n)
    -> BVar s (H.R n)
unrow = isoVar H.unrow H.row
{-# INLINE unrow #-}

uncol
    :: (KnownNat n, Reifies s W)
    => BVar s (H.L n 1)
    -> BVar s (H.R n)
uncol = isoVar H.uncol H.col
{-# INLINE uncol #-}

tr  :: (HU.Transposable m mt, HU.Transposable mt m, Backprop m, Reifies s W)
    => BVar s m
    -> BVar s mt
tr = isoVar H.tr H.tr
{-# INLINE tr #-}

diag
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.Sq n)
diag = liftOp1 . op1 $ \x -> (H.diag x, H.takeDiag)
{-# INLINE diag #-}

-- | Potentially extremely bad for anything but short lists!!!
matrix
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => [BVar s H.ℝ]
    -> BVar s (H.L m n)
matrix = maybe (error "matrix: invalid number of elements")
               ( isoVar (H.vecL . SVG.convert . runABP) (ABP . SVG.convert . H.lVec)
               . collectVar
               . ABP
               )
       . SV.fromList @(m * n)
{-# INLINE matrix #-}

-- | Matrix product
(<>)
    :: (KnownNat m, KnownNat k, KnownNat n, Reifies s W)
    => BVar s (H.L m k)
    -> BVar s (H.L k n)
    -> BVar s (H.L m n)
(<>) = mul
infixr 8 <>
{-# INLINE (<>) #-}

-- | Matrix-vector product
(#>)
    :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
    -> BVar s (H.R m)
(#>) = app
infixr 8 #>
{-# INLINE (#>) #-}

-- | Dot product
(<.>)
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s H.ℝ
(<.>) = dot
infixr 8 <.>
{-# INLINE (<.>) #-}

-- | Can only get the singular values, for now.  Let me know if you find an
-- algorithm that can compute the gradients based on differentials for the
-- other matricies!
--
svd :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
svd = liftOp1 . op1 $ \x ->
    let (u, σ, v) = H.svd x
    in  ( σ
        , \(dΣ :: H.R n) -> (u H.<> H.diagR 0 dΣ) H.<> H.tr v
                -- must manually associate because of bug in diagR in
                -- hmatrix-0.18.2.0
        )
{-# INLINE svd #-}

-- | Version of 'svd' that returns the full SVD, but if you attempt to find
-- the gradient, it will fail at runtime if you ever use U or V.
svd_
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L m m), BVar s (H.R n), BVar s (H.L n n))
svd_ r = (t ^^. _1, t ^^. _2, t ^^. _3)
  where
    o :: Op '[H.L m n] (H.L m m, H.R n, H.L n n)
    o = op1 $ \x ->
        let msv@(u, _, v) = H.svd x
        in  ( msv
            , \(dU, dΣ, dV) ->
                    if H.norm_0 dU == 0 && H.norm_0 dV == 0
                      then (u H.<> H.diagR 0 dΣ) H.<> H.tr v
                      else error "svd_: Cannot backprop if U and V are used."
            )
    {-# INLINE o #-}
    t = liftOp1 o r
    {-# NOINLINE t #-}
{-# INLINE svd_ #-}

helpEigen :: KnownNat n => H.Sym n -> (H.R n, H.L n n, H.L n n, H.L n n)
helpEigen x = (l, v, H.inv v, H.tr v)
  where
    (l, v) = H.eigensystem x
{-# INLINE helpEigen #-}

-- | /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
eigensystem
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> (BVar s (H.R n), BVar s (H.L n n))
eigensystem u = (t ^^. _1, t ^^. _2)
  where
    o :: Op '[H.Sym n] (H.R n, H.L n n)
    o = op1 $ \x ->
        let (l, v, vInv, vTr) = helpEigen x
            lRep = H.rowsL . SV.replicate $ l
            fMat = (1 - H.eye) * (lRep - H.tr lRep)
        in  ( (l, v)
            , \(dL, dV) -> unsafeCoerce $
                       H.tr vInv
                  H.<> (H.diag dL + fMat * (vTr H.<> dV))
                  H.<> vTr
            )
    {-# INLINE o #-}
    t = liftOp1 o u
    {-# NOINLINE t #-}
{-# INLINE eigensystem #-}

-- | /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
eigenvalues
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.R n)
eigenvalues = liftOp1 . op1 $ \x ->
    let (l, _, vInv, vTr) = helpEigen x
    in  ( l
        , \dL -> unsafeCoerce $
                 H.tr vInv H.<> H.diag dL H.<> vTr
        )
{-# INLINE eigenvalues #-}

-- | Algorithm from https://arxiv.org/abs/1602.07527
--
-- The paper also suggests a potential imperative algorithm that might
-- help.  Need to benchmark to see what is best.
--
-- /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
chol
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.Sq n)
chol = liftOp1 . op1 $ \x ->
    let l = H.chol x
        lInv = H.inv l
        phi :: H.Sq n
        phi = H.build $ \i j -> case compare i j of
                                  LT -> 1
                                  EQ -> 0.5
                                  GT -> 0
    in  ( l
        , \dL -> let s = H.tr lInv H.<> (phi * (H.tr l H.<> dL)) H.<> lInv
                 in  unsafeCoerce $ s + H.tr s - H.eye * s
        )
{-# INLINE chol #-}

-- | Number of non-zero items
norm_0
    :: (H.Normed a, Backprop a, Reifies s W)
    => BVar s a
    -> BVar s H.ℝ
norm_0 = liftOp1 . op1 $ \x -> (H.norm_0 x, const (zero x))
{-# INLINE norm_0 #-}

-- | Sum of absolute values
norm_1V
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_1V = liftOp1 . op1 $ \x -> (H.norm_1 x, (* signum x) . H.konst)
{-# INLINE norm_1V #-}

-- | Maximum 'H.norm_1' of columns
norm_1M
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_1M = liftOp1 . op1 $ \x ->
    let n = H.norm_1 x
    in  (n, \d -> let d' = H.konst d
                  in  H.colsL
                    . SV.map (\c -> if H.norm_1 c == n
                                      then d' * signum c
                                      else 0
                             )
                    . H.lCols
                    $ x
        )
{-# INLINE norm_1M #-}

-- | Square root of sum of squares
--
-- Be aware that gradient diverges when the norm is zero
norm_2V
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_2V = liftOp1 . op1 $ \x ->
    let n = H.norm_2 x
    in (n, \d -> x * H.konst (d / n))
{-# INLINE norm_2V #-}

-- | Maximum singular value
norm_2M
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_2M = liftOp1 . op1 $ \x ->
    let n = H.norm_2 x
        (head.H.toColumns->u1,_,head.H.toColumns->v1) = H.svd x
    in (n, \d -> H.konst d * (u1 `H.outer` v1))
{-# INLINE norm_2M #-}

-- | Maximum absolute value
norm_InfV
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_InfV = liftOp1 . op1 $ \x ->
    let n :: H.ℝ
        n = H.norm_Inf x
    in  (n, \d -> H.vecR
                . SVS.map (\e -> if abs e == n
                                   then signum e * d
                                   else 0
                          )
                . H.rVec
                $ x
        )
{-# ANN norm_InfV "HLint: ignore Use camelCase" #-}
{-# INLINE norm_InfV #-}

-- | Maximum 'H.norm_1' of rows
norm_InfM
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_InfM = liftOp1 . op1 $ \x ->
    let n = H.norm_Inf x
    in  (n, \d -> let d' = H.konst d
                  in  H.rowsL
                    . SV.map (\c -> if H.norm_1 c == n
                                      then d' * signum c
                                      else 0
                             )
                    . H.lRows
                    $ x
        )
{-# ANN norm_InfM "HLint: ignore Use camelCase" #-}
{-# INLINE norm_InfM #-}

mean
    :: (KnownNat n, 1 <= n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
mean = liftOp1 . op1 $ \x -> (H.mean x, H.konst . (/ H.norm_0 x))
{-# INLINE mean #-}

gradCov
    :: forall m n. (KnownNat m, KnownNat n)
    => H.L m n
    -> H.R n
    -> H.Sym n
    -> H.L m n
gradCov x μ dσ = H.rowsL
               . SV.map (subtract (dDiffsSum / m))
               . H.lRows
               $ dDiffs
  where
    diffs = H.rowsL . SV.map (subtract μ) . H.lRows $ x
    dDiffs = H.konst (2/n) * (diffs H.<> H.tr (H.unSym dσ))
    dDiffsSum = sum . H.toRows $ dDiffs
    m = fromIntegral $ natVal (Proxy @m)
    n = fromIntegral $ natVal (Proxy @n)
{-# INLINE gradCov #-}

-- | Mean and covariance.  If you know you only want to use one or the
-- other, use 'meanL' or 'cov'.
meanCov
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.R n), BVar s (H.Sym n))
meanCov v = (t ^^. _1, t ^^. _2)
  where
    m = fromInteger $ natVal (Proxy @m)
    t = ($ v) . liftOp1 . op1 $ \x ->
        let ms@(μ, _) = H.meanCov x
        in  ( ms
            , \(dμ, dσ) ->
                let gradMean = H.rowsL
                             . SV.replicate
                             $ (dμ / H.konst m)
                in  gradMean + gradCov x μ dσ
            )
    {-# NOINLINE t #-}
{-# INLINE meanCov #-}

-- | 'meanCov', but if you know you won't use the covariance.
meanL
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
meanL = liftOp1 . op1 $ \x ->
    ( fst (H.meanCov x)
    , H.rowsL . SV.replicate . (/ H.konst m)
    )
  where
    m = fromInteger $ natVal (Proxy @m)
{-# INLINE meanL #-}

-- | 'cov', but if you know you won't use the covariance.
cov
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.Sym n)
cov = liftOp1 . op1 $ \x ->
    let (μ, σ) = H.meanCov x
    in  (σ, gradCov x μ)
{-# INLINE cov #-}

mul :: ( KnownNat m
       , KnownNat k
       , KnownNat n
       , H.Domain field vec mat
       , Backprop (mat m k)
       , Backprop (mat k n)
       , HU.Transposable (mat m k) (mat k m)
       , HU.Transposable (mat k n) (mat n k)
       , Reifies s W
       )
    => BVar s (mat m k)
    -> BVar s (mat k n)
    -> BVar s (mat m n)
mul = liftOp2 . op2 $ \x y ->
    ( x `H.mul` y
    , \d -> (d `H.mul` H.tr y, H.tr x `H.mul` d)
    )
{-# INLINE mul #-}

app :: ( KnownNat m
       , KnownNat n
       , H.Domain field vec mat
       , HU.Transposable (mat m n) (mat n m)
       , Backprop (mat m n)
       , Backprop (vec n)
       , Reifies s W
       )
    => BVar s (mat m n)
    -> BVar s (vec n)
    -> BVar s (vec m)
app = liftOp2 . op2 $ \xs y ->
    ( xs `H.app` y
    , \d -> (d `H.outer` y, H.tr xs `H.app` d)
    )
{-# INLINE app #-}

dot :: ( KnownNat n
       , H.Domain field vec mat
       , H.Sized field (vec n) d
       , Num (vec n)
       , Backprop (vec n)
       , Reifies s W
       )
    => BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s field
dot = liftOp2 . op2 $ \x y ->
    ( x `H.dot` y
    , \d -> let d' = H.konst d
            in  (d' * y, x * d')
    )
{-# INLINE dot #-}

cross
    :: ( H.Domain field vec mat
       , Reifies s W
       , Backprop (vec 3)
       )
    => BVar s (vec 3)
    -> BVar s (vec 3)
    -> BVar s (vec 3)
cross = liftOp2 . op2 $ \x y ->
    ( x `H.cross` y
    , \d -> (y `H.cross` d, d `H.cross` x)
    )
{-# INLINE cross #-}

-- | Create matrix with diagonal, and fill with default entries
diagR
    :: forall m n k field vec mat s.
       ( H.Domain field vec mat
       , Num (vec k)
       , Num (mat m n)
       , KnownNat m
       , KnownNat n
       , KnownNat k
       , HU.Container HU.Vector field
       , H.Sized field (mat m n) HU.Matrix
       , H.Sized field (vec k) HU.Vector
       , Backprop field
       , Backprop (vec k)
       , Reifies s W
       )
    => BVar s field             -- ^ default value
    -> BVar s (vec k)           -- ^ diagonal
    -> BVar s (mat m n)
diagR = liftOp2 . op2 $ \c x ->
    ( H.diagR c x
    , \d -> ( HU.sumElements . H.extract $ H.diagR 1 (0 :: vec k) * d
            , fromJust . H.create . HU.takeDiag . H.extract $ d
            )
    )
{-# INLINE diagR #-}

-- | Note: if possible, use the potentially much more performant 'vmap''.
vmap
    :: (KnownNat n, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
vmap f = isoVar (H.vecR . SVG.convert @V.Vector . runABP)
                (ABP . SVG.convert . H.rVec)
       . B.fmap f
       . isoVar (ABP . SVG.convert . H.rVec) (H.vecR . SVG.convert . runABP)
{-# INLINE vmap #-}

-- | 'vmap', but potentially more performant.  Only usable if the mapped
-- function does not depend on any external 'BVar's.
vmap'
    :: ( Num (vec n)
       , Storable field
       , H.Sized field (vec n) HU.Vector
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
vmap' f = liftOp1 . op1 $ bimap (fromJust . H.create . VG.convert)
                                ((*) . fromJust . H.create . VG.convert)
                        . V.unzip
                        . V.map (backprop f)
                        . VG.convert
                        . H.extract
{-# INLINE vmap' #-}

-- TODO: Can be made more efficient if backprop exports
-- a custom-total-derivative version

-- | Note: Potentially less performant than 'vmap''.
dvmap
    :: ( KnownNat n
       , H.Domain field vec mat
       , Num (vec n)
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
dvmap f = liftOp1 . op1 $ \x ->
    ( H.dvmap (evalBP f) x
    , (H.dvmap (gradBP f) x *)
    )
{-# INLINE dvmap #-}

-- | Note: if possible, use the potentially much more performant 'mmap''.
mmap
    :: (KnownNat n, KnownNat m, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.L n m)
    -> BVar s (H.L n m)
mmap f = isoVar (H.vecL . SVG.convert @V.Vector . runABP)
                (ABP . SVG.convert . H.lVec)
       . B.fmap f
       . isoVar (ABP . SVG.convert . H.lVec) (H.vecL . SVG.convert . runABP)
{-# INLINE mmap #-}

-- | 'mmap', but potentially more performant.  Only usable if the mapped
-- function does not depend on any external 'BVar's.
mmap'
    :: forall n m mat field s.
       ( KnownNat m
       , Num (mat n m)
       , Backprop (mat n m)
       , Backprop field
       , H.Sized field (mat n m) HU.Matrix
       , HU.Element field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (mat n m)
    -> BVar s (mat n m)
mmap' f = liftOp1 . op1 $ bimap (fromJust . H.create . HU.reshape m . VG.convert)
                                ((*) . fromJust . H.create . HU.reshape m . VG.convert)
                        . V.unzip
                        . V.map (backprop f)
                        . VG.convert
                        . HU.flatten
                        . H.extract
  where
    m :: Int
    m = fromInteger $ natVal (Proxy @m)
{-# INLINE mmap' #-}

-- | Note: Potentially less performant than 'mmap''.
dmmap
    :: ( KnownNat n
       , KnownNat m
       , H.Domain field vec mat
       , Num (mat n m)
       , Backprop (mat n m)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (mat n m)
    -> BVar s (mat n m)
dmmap f = liftOp1 . op1 $ \x ->
    ( H.dmmap (evalBP f) x
    , (H.dmmap (gradBP f) x *)
    )
{-# INLINE dmmap #-}

outer
    :: ( KnownNat m
       , KnownNat n
       , H.Domain field vec mat
       , HU.Transposable (mat n m) (mat m n)
       , Backprop (vec n)
       , Backprop (vec m)
       , Reifies s W
       )
    => BVar s (vec n)
    -> BVar s (vec m)
    -> BVar s (mat n m)
outer = liftOp2 . op2 $ \x y ->
    ( x `H.outer` y
    , \d -> ( d `H.app` y
            , H.tr d `H.app` x)
    )
{-# INLINE outer #-}

-- | Note: if possible, use the potentially much more performant
-- 'zipWithVector''.
zipWithVector
    :: (KnownNat n, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
zipWithVector f x y = isoVar (H.vecR . SVG.convert . runABP)
                             (ABP . SVG.convert . H.rVec)
                    $ B.liftA2 @(ABP (SV.Vector _)) f (iv x) (iv y)
  where
    iv = isoVar (ABP . SVG.convert . H.rVec) (H.vecR . SVG.convert . runABP)
{-# INLINE zipWithVector #-}

zipWithVector'
    :: ( Num (vec n)
       , Backprop (vec n)
       , Storable field
       , Backprop field
       , H.Sized field (vec n) HU.Vector
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s (vec n)
zipWithVector' f = liftOp2 . op2 $ \(VG.convert.H.extract->x) (VG.convert.H.extract->y) ->
    let (z, dx, dy) = V.unzip3 $ V.zipWith (\x' -> retup . backprop2 f x') x y
    in  ( fromJust (H.create (VG.convert z))
        , \d -> ( d * fromJust (H.create (VG.convert dx))
                , d * fromJust (H.create (VG.convert dy))
                )
        )
  where
    retup (x, (y, z)) = (x, y, z)
{-# INLINE zipWithVector' #-}

-- | A version of 'zipWithVector'' that is potentially less performant but
-- is based on 'H.zipWithVector' from 'H.Domain'.
dzipWithVector
    :: ( KnownNat n
       , H.Domain field vec mat
       , Num (vec n)
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s (vec n)
dzipWithVector f = liftOp2 . op2 $ \x y ->
    ( H.zipWithVector (evalBP2 f) x y
    , \d -> let dx = H.zipWithVector (\x' -> fst . gradBP2 f x') x y
                dy = H.zipWithVector (\x' -> snd . gradBP2 f x') x y
            in  (d * dx, d * dy)
    )
{-# INLINE dzipWithVector #-}

det :: ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s field
det = liftOp1 . op1 $ \x ->
    let xDet = H.det x
        xInv = H.inv x
    in  ( xDet, \d -> H.konst (d * xDet) * H.tr xInv )
{-# INLINE det #-}

-- | The inverse and the natural log of the determinant together.  If you
-- know you don't need the inverse, it is best to use 'lndet'.
invlndet
    :: forall n mat field vec d s.
       ( KnownNat n
       , Num (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Backprop field
       , Backprop (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> (BVar s (mat n n), (BVar s field, BVar s field))
invlndet v = (t ^^. _1, (t ^^. _2, t ^^. _3))
  where
    o :: Op '[mat n n] (mat n n, field, field)
    o = op1 $ \x ->
      let (i,(ldet, s)) = H.invlndet x
          iTr           = H.tr i
      in  ( (i, ldet, s)
          , \(dI, dLDet, _) ->
                let gradI    = - iTr `H.mul` dI `H.mul` iTr
                    gradLDet = H.konst dLDet * H.tr i
                in  gradI + gradLDet
          )
    {-# INLINE o #-}
    t = liftOp1 o v
    {-# NOINLINE t #-}
{-# INLINE invlndet #-}

-- | The natural log of the determinant.
lndet
    :: forall n mat field vec d s.
       ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s field
lndet = liftOp1 . op1 $ \x ->
          let (i,(ldet,_)) = H.invlndet x
          in  (ldet, (* H.tr i) . H.konst)
{-# INLINE lndet #-}

inv :: ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s (mat n n)
inv = liftOp1 . op1 $ \x ->
    let xInv   = H.inv x
        xInvTr = H.tr xInv
    in  ( xInv, \d -> - xInvTr `H.mul` d `H.mul` xInvTr )
{-# INLINE inv #-}

toRows
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> SV.Vector m (BVar s (H.R n))
toRows = runABP . sequenceVar
       . isoVar (coerce (H.lRows @m)) (coerce H.rowsL)
{-# INLINE toRows #-}

toColumns
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> SV.Vector n (BVar s (H.R m))
toColumns = runABP . sequenceVar
          . isoVar (coerce (H.lCols @_ @n)) (coerce H.colsL)
{-# INLINE toColumns #-}

fromRows
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => SV.Vector m (BVar s (H.R n))
    -> BVar s (H.L m n)
fromRows = isoVar (coerce H.rowsL) (coerce (H.lRows @m))
         . collectVar . ABP
{-# INLINE fromRows #-}

fromColumns
    :: forall m n s. (KnownNat n, KnownNat m, Reifies s W)
    => SV.Vector n (BVar s (H.R m))
    -> BVar s (H.L m n)
fromColumns = isoVar (coerce H.colsL) (coerce (H.lCols @_ @n))
            . collectVar . ABP
{-# INLINE fromColumns #-}

konst
    :: forall t s d q.
     ( H.Sized t s d
     , HU.Container d t
     , Backprop t
     , Reifies q W
     )
    => BVar q t
    -> BVar q s
konst = liftOp1 . op1 $ \x ->
    ( H.konst x
    , HU.sumElements . H.extract
    )
{-# INLINE konst #-}

sumElements
    :: forall t s d q.
     ( H.Sized t s d
     , HU.Container d t
     , Backprop s
     , Reifies q W
     )
    => BVar q s
    -> BVar q t
sumElements = liftOp1 . op1 $ \x ->
    ( HU.sumElements . H.extract $ x
    , H.konst
    )
{-# INLINE sumElements #-}

-- | If there are extra items in the total derivative, they are dropped.
-- If there are missing items, they are treated as zero.
extractV
    :: forall t s q.
       ( H.Sized t s HU.Vector
       , HU.Konst t Int HU.Vector
       , HU.Container HU.Vector t
       , Backprop s
       , Reifies q W
       )
    => BVar q s
    -> BVar q (HU.Vector t)
extractV = liftOp1 . op1 $ \x ->
    let n = H.size x
    in  ( H.extract x
        , \d -> let m  = HU.size d
                    m' = case compare n m of
                            LT -> HU.subVector 0 n d
                            EQ -> d
                            GT -> HU.vjoin [d, HU.konst 0 (n - m)]
                in  fromJust . H.create $ m'
        )
{-# INLINE extractV #-}

-- | If there are extra items in the total derivative, they are dropped.
-- If there are missing items, they are treated as zero.
extractM
    :: forall t s q.
       ( H.Sized t s HU.Matrix
       , Backprop s
       , HU.Konst t (Int, Int) HU.Matrix
       , HU.Container HU.Matrix t
       , Num (HU.Matrix t)
       , Reifies q W
       )
    => BVar q s
    -> BVar q (HU.Matrix t)
extractM = liftOp1 . op1 $ \x ->
    let (xI,xJ) = H.size x
    in  ( H.extract x
        , \d -> let (dI,dJ) = HU.size d
                    m' = case (compare xI dI, compare xJ dJ) of
                           (LT, LT) -> d HU.?? (HU.Take xI, HU.Take xJ)
                           (LT, EQ) -> d HU.?? (HU.Take xI, HU.All)
                           (LT, GT) -> d HU.?? (HU.Take xI, HU.All)
                                HU.||| HU.konst 0 (xI, xJ - dJ)
                           (EQ, LT) -> d HU.?? (HU.All    , HU.Take xJ)
                           (EQ, EQ) -> d
                           (EQ, GT) -> d HU.?? (HU.All, HU.All)
                                HU.||| HU.konst 0 (xI, xJ - dJ)
                           (GT, LT) -> d HU.?? (HU.All, HU.Take xJ)
                                HU.=== HU.konst 0 (xI - dI, xJ)
                           (GT, EQ) -> d HU.?? (HU.All, HU.All)
                                HU.=== HU.konst 0 (xI - dI, xJ)
                           (GT, GT) -> HU.fromBlocks
                              [[d,0                            ]
                              ,[0,HU.konst 0 (xI - dI, xJ - dJ)]
                              ]
                in  fromJust . H.create $ m'
        )
{-# INLINE extractM #-}

create
    :: (H.Sized t s d, Backprop s, Num (d t), Backprop (d t), Reifies q W)
    => BVar q (d t)
    -> Maybe (BVar q s)
create = sequenceVar . isoVar H.create (maybe 0 H.extract)
{-# INLINE create #-}

takeDiag
    :: ( KnownNat n
       , H.Diag (mat n n) (vec n)
       , H.Domain field vec mat
       , Num field
       , Backprop (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s (vec n)
takeDiag = liftOp1 . op1 $ \x ->
    ( H.takeDiag x
    , H.diagR 0
    )
{-# INLINE takeDiag #-}

-- |
-- \[
-- \frac{1}{2} (M + M^T)
-- \]
sym :: (KnownNat n, Reifies s W)
    => BVar s (H.Sq n)
    -> BVar s (H.Sym n)
sym = liftOp1 . op1 $ \x ->
    ( H.sym x
    , H.unSym . H.sym . H.unSym
    )
{-# INLINE sym #-}

-- |
-- \[
-- M^T M
-- \]
mTm :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.Sym n)
mTm = liftOp1 . op1 $ \x ->
    ( H.mTm x
    , \d -> 2 * (x H.<> H.unSym d)
    )
{-# INLINE mTm #-}

-- | Warning: the gradient is going necessarily symmetric, and so is /not/
-- meant to be used directly.  Rather, it is meant to be used in the middle
-- (or at the end) of a longer computation.
unSym
    :: (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.Sq n)
unSym = isoVar H.unSym unsafeCoerce
{-# INLINE unSym #-}

-- | Unicode synonym for '<.>>'
(<·>)
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s H.ℝ
(<·>) = dot
infixr 8 <·>
{-# INLINE (<·>) #-}

afSV :: Backprop a => BE.AddFunc (SV.Vector n a)
afSV = BE.AF (SV.zipWith add)
{-# INLINE afSV #-}

-- support for complex types
toComplex :: (KnownNat n, Reifies s W)
          => BVar s (H.R n) -> BVar s (H.R n) -> BVar s (H.C n)
toComplex = isoVar2 (curry H.toComplex) H.fromComplex
{-# INLINE toComplex #-}

fromComplex :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> (BVar s (H.R n), BVar s (H.R n))
fromComplex c = let ri = isoVar H.fromComplex H.toComplex c in (ri ^^. _1, ri ^^. _2)
{-# INLINE fromComplex #-}

complex :: (KnownNat n, Reifies s W)
       => BVar s (H.R n) -> BVar s (H.C n)
complex = isoVar H.complex H.real
{-# INLINE complex #-}

real :: (KnownNat n, Reifies s W)
     => BVar s (H.C n) -> BVar s (H.R n)
real = isoVar H.real (\r -> H.toComplex (r, H.konst 0))
{-# INLINE real #-}

imag :: (KnownNat n, Reifies s W)
     => BVar s (H.C n) -> BVar s (H.R n)
imag = isoVar H.imag (\r -> H.toComplex (H.konst 0, r))
{-# INLINE imag #-}

sqMagnitude :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> BVar s (H.R n)
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2
{-# INLINE sqMagnitude #-}

magnitude :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> BVar s (H.R n)
magnitude c = sqrt $ sqMagnitude c
{-# INLINE magnitude #-}


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
module DataRep where

import Prelude hiding ((.), id, curry, uncurry)
import Control.Exception (assert)
import Data.List
import Data.Proxy
import qualified Data.Vector.Generic as VG
import Foreign.Storable
import qualified Numeric.LinearAlgebra as HM

import qualified Base
import Base hiding (Scalable (..))

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

class (Fractional (Scalar a), Storable (Scalar a), HM.Numeric (Scalar a), Additive a) => VectorSpace a where
  type Scalar a
  dim :: Proxy a -> Int
  toVector :: a -> HM.Vector (Scalar a)
  fromVector :: HM.Vector (Scalar a) -> a

  scale :: Scalar a -> a -> a
  scale a = fromVector . HM.scale a . toVector

linComb :: VectorSpace a => [(a, Scalar a)] -> a
linComb xs = foldl' (.+.) zero [scale s a | (a,s) <- xs]

instance VectorSpace Double where
  type Scalar Double = Double
  scale = (*)
  dim _ = 1
  toVector = VG.singleton
  fromVector = (VG.! 0)

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => VectorSpace (a, b) where
  type Scalar (a, b) = Scalar a
  scale s (a,b) = (scale s a, scale s b)
  dim _ = dim (Proxy :: Proxy a) + dim (Proxy :: Proxy b)
  toVector (a,b) = toVector a <> toVector b
  fromVector v =
    case VG.splitAt (dim (Proxy :: Proxy a)) v of
      (v1, v2) -> (fromVector v1, fromVector v2)

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

newtype LinMap s a b = LinMap (HM.Matrix s)

{-
infixr 0 ⊸

data a ⊸ b where
  LinMap :: (Scalar a ~ Scalar b) => !(HM.Matrix (Scalar a)) -> a ⊸ b

と書けると格好良いが、スカラー型の異なるベクトル空間の直積とかで困ったことになる。
-}

instance Category (LinMap s) where
  type Obj (LinMap s) a = (VectorSpace a, Scalar a ~ s)

  id :: forall a. Obj (LinMap s) a => LinMap s a a
  id = LinMap $ HM.ident $ dim (Proxy :: Proxy a)

  LinMap m1 . LinMap m2 = LinMap (m1 HM.<> m2)

instance ToFun (LinMap s) where
  toFun (LinMap m) a = fromVector (m HM.#> toVector a)

instance (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => Additive (LinMap s a b) where
  zero = LinMap $ HM.konst 0 (m, n)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)
  LinMap m1 .+. LinMap m2 = LinMap $ HM.add m1 m2

instance (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => VectorSpace (LinMap s a b) where
  type Scalar (LinMap s a b) = s
  scale s (LinMap m) = LinMap (HM.scale s m)
  dim _ = dim (Proxy :: Proxy a) * dim (Proxy :: Proxy b)
  toVector (LinMap m) = HM.flatten m
  fromVector v = LinMap (HM.reshape (dim (Proxy :: Proxy b)) v)

instance Monoidal (LinMap s) where
  LinMap m1 >< LinMap m2 = LinMap $ HM.diagBlock [m1, m2]
  monObj = ObjR

instance Cartesian (LinMap s) where
  exl :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s (a,b) a
  exl = LinMap $ HM.ident m HM.||| HM.konst 0 (m, n)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  exr :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s (a,b) b
  exr = LinMap $ HM.konst 0 (n, m) HM.||| HM.ident n
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  dup :: forall a. Obj (LinMap s) a => LinMap s a (a,a)
  dup = LinMap $
      HM.ident n
      HM.===
      HM.ident n
    where
      n = dim (Proxy :: Proxy a)

instance Cocartesian (LinMap s) where
  inl :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s a (a,b)
  inl = LinMap $
      HM.ident m
      HM.===
      HM.konst 0 (n, m)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  inr :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s b (a,b)
  inr = LinMap $
      HM.konst 0 (n, m)
      HM.===
      HM.ident m
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  jam :: forall a. (Obj (LinMap s) a) => LinMap s (a,a) a
  jam = LinMap $ HM.ident n HM.||| HM.ident n
    where
      n = dim (Proxy :: Proxy a)

instance (VectorSpace s, Scalar s ~ s) => Base.Scalable (LinMap s) s where
  scale s = assert (dim (Proxy :: Proxy s) == 1) $ LinMap ((1 HM.>< 1) [s])

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

-- | Dual vector space
newtype Dual a = Dual (LinMap (Scalar a) a (Scalar a))

instance VectorSpace a => Additive (Dual a) where
  zero = Dual $ LinMap $ HM.konst 0 (n, 1)
    where
      n = dim (Proxy :: Proxy a)
  Dual (LinMap m1) .+. Dual (LinMap m2) = Dual $ LinMap $ HM.add m1 m2

instance VectorSpace a => VectorSpace (Dual a) where
  type Scalar (Dual a) = Scalar a
  scale s (Dual (LinMap m)) = Dual (LinMap (HM.scale s m))
  dim _ = dim (Proxy :: Proxy a)
  toVector (Dual (LinMap m)) = HM.flatten m
  fromVector v = Dual (LinMap (HM.reshape 1 v))

toDual :: VectorSpace a => a -> Dual a
toDual a = Dual $ LinMap $ HM.asRow (toVector a)

toDualMap :: forall a b s. (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => LinMap s a b -> LinMap s (Dual b) (Dual a)
toDualMap (LinMap m) = LinMap (HM.tr' m)

fromDual :: forall a. VectorSpace a => Dual a -> a
fromDual (Dual (LinMap m)) = fromVector $ HM.flatten m

fromDualMap :: forall a b s. (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => LinMap s (Dual b) (Dual a) -> LinMap s a b
fromDualMap (LinMap m) = LinMap (HM.tr' m)

instance (VectorSpace u, VectorSpace s, Scalar u ~ s, Scalar s ~ s) => Base.HasDot (LinMap s) s u where
  dot x = case toDual x of Dual f -> f
  undot f = fromDual (Dual f)

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

testDual = (toFun f (2,1) == 7, toFun f2 (2,1) == 7)
  where
    f :: LinMap Double (Double, Double) Double
    f = LinMap $ HM.asRow (VG.fromList [2,3])
    f2 = fromDualMap (toDualMap f)

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

data a :⊗ b where
  TensorProd :: (Scalar a ~ s, Scalar b ~ s) => HM.Matrix s -> a :⊗ b

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => Additive (a :⊗ b) where
  zero = TensorProd (HM.konst 0 (m, n))
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)
  TensorProd m1 .+. TensorProd m2 = TensorProd $ m1 `HM.add` m2

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => VectorSpace (a :⊗ b) where
  type Scalar (a :⊗ b) = Scalar a
  scale s (TensorProd m) = TensorProd $ HM.scale s m
  dim _ = dim (Proxy :: Proxy a) * dim (Proxy :: Proxy b)
  toVector (TensorProd m) = HM.flatten m
  fromVector v = TensorProd (HM.reshape (dim (Proxy :: Proxy b)) v)

curry
  :: forall a b c s. (VectorSpace a, VectorSpace b, VectorSpace c, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s)
  => LinMap s (a :⊗ b) c -> LinMap s a (LinMap s b c)
curry (LinMap m) = LinMap $ HM.reshape (nb*nc) $ HM.flatten m
  where
    nb = dim (Proxy :: Proxy b)
    nc = dim (Proxy :: Proxy c)

uncurry
  :: forall a b c s. (VectorSpace a, VectorSpace b, VectorSpace c, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s)
  => LinMap s a (LinMap s b c) -> LinMap s (a :⊗ b) c
uncurry (LinMap m) = LinMap $ HM.reshape nc $ HM.flatten m
  where
    nc = dim (Proxy :: Proxy c)

mapTensor
  :: forall a b c d s.
     (VectorSpace a, VectorSpace b, VectorSpace c, VectorSpace d, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s, Scalar d ~ s)
  => LinMap s a b -> LinMap s c d -> LinMap s (a :⊗ c) (b :⊗ d)
mapTensor (LinMap m1) (LinMap m2) = LinMap $ ((na*nc) HM.>< (nb*nd)) [(m1 `HM.atIndex` (i, k)) * (m2 `HM.atIndex` (j, l)) | i <- [0..na-1], j <- [0..nc-1], k <- [0..nb-1], l <- [0..nd-1]]
  where
    na = dim (Proxy :: Proxy a)
    nb = dim (Proxy :: Proxy b)
    nc = dim (Proxy :: Proxy c)
    nd = dim (Proxy :: Proxy d)

TensorProd test_Tensor = toFun h t
  where
    f, g :: LinMap Double Double Double
    f = Base.scale 2
    g = Base.scale 3
    h :: LinMap Double (Double :⊗ Double) (Double :⊗ Double)
    h = mapTensor f g

    t :: Double :⊗ Double
    t = TensorProd $ (1 HM.>< 1) [4]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

-- |
-- Module      : Swarm
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Swarm
  ( -- * Idea
    --
    -- | This module implements different variants of Particle Swarm
    -- Optimisation (PSO), so if you're completely unfamiliar with PSO, go read
    -- https://en.wikipedia.org/wiki/Particle_swarm_optimization first, it's a
    -- nice introduction.
    --
    -- Let there be \(M\in\mathbb{N}\) 'Bee's in the swarm. Bees live in a phase
    -- space (which is a vector space). What this means is that each bee is
    -- described by a phase \(\psi\). The two most common choices of the phase
    -- are:
    --
    -- * \((q, p)\in\mathcal{H}^2\) with \(q\) being the canonical coordinate
    -- and \(p\) being the canonical momentum.
    -- * \(|\psi\rangle\in\mathcal{H}^2\), i.e. a vector in Hilbert space.
    --
    -- While flying around, bees may pick up some information about the
    -- surrounding environment. We call the part of it that bees remember a
    -- local (or private) /guide/. It's a guide because it helps bees decide
    -- where to move next. The simplest and most common guide is 'BeeGuide'.
    --
    -- The whole point of having a swarm of bees rather than one bee is
    -- collective behavior. In case of PSO this is achieved by having a /global
    -- guide/. Each bee then uses this global guide, its own local guide, and
    -- its current phase to compute the next phase. This is represented by
    -- 'PhaseUpdater'.
    --
    -- The whole optimisation process then looks more or less like this:
    --
    --    1. __Initialisation__ (see 'mkSwarm' function).
    --
    --        * Generate initial phases \((\psi^{(0)}_i)_{i\in\{1\dots M\}}\).
    --        * Evaluate fitness in these phases, i.e. compute
    --          \((f(\psi^{(0)}_i))_{i\in\{1\dots M\}}\) (see 'mkBee' function).
    --        * Initialise local guides \((g^{(0)}_{l,i})_{i\in\{1\dots M\}}\)
    --          (see 'mkLocalG' function).
    --        * Construct bees using phases, fitness values, and guides (see
    --          'mkBee' function).
    --        * Initialise global guide \(g^{(0)}_g\) (see 'mkGlobalG'
    --          function).
    --        * Construct swarm using bees and global guide.
    --
    --    2. __Main Loop__. While the termination condition is not met, do the
    --    following (see 'updateSwarm' function):
    --
    --        * Update each bee (see 'updateBee' function), i.e.
    --
    --            * Using global guide \(g^{(n)}_g\), local guide
    --              \(g^{(n)}_{l,i}\), current phase \(\psi^{(n)}_i\), and
    --              \(f(\psi^{(n)}_i)\), come up with a new phase
    --              \(\psi^{(n+1)}_i\) (see 'PhaseUpdater').
    --            * Evaluate the fitness in this new phase, i.e. compute
    --              \(f(\psi^{(n+1)}_i)\).
    --            * Using \(\psi^{(n+1)}_i\) and \(f(\psi^{(n+1)}_i)\), compute
    --              the local guide in the next iteration \(g^{(n+1)}_i\) (see
    --              'updateLocalG' function).
    --
    --        * Update the global guide (see 'updateGlobalG' function).
    --

    -- * Basic types
    Bee(..)
  , CMState(..)
  , QMState(..)
  , LocalGuide(..)
  , BeeGuide(..)
  , GlobalGuide(..)
  , SwarmGuide(..)
  , Swarm(..)
  , PhaseUpdater(..)

    -- * Important functions
  , wpgUpdater
  -- , kickUpdater
  , deltaUpdater
  , absorbUpdater
  , mkCMState
  , mkQMState
  , optimiseND
  , optimiseNDFromList

    -- * Details
  , VelocityUpdater(..)
  , upScale
  , upLocal
  , upGlobal
  , upStep
  , updateBee
  , updateSwarm
  , mkBee
  , mkSwarm
  , fromPure
  , iterateNM
  , iterateWhileM
  , Absorbable(..)
  , DeltaWell(..)
  , Scalable(..)

    -- * Lenses
  , HasPos(..)
  , HasVal(..)
  , HasVar(..)
  , HasGuide(..)
  , HasState(..)
  , HasRun(..)
  , HasBees(..)
  , HasUpdater(..)
  , HasIteration(..)
  ) where

import Data.Complex
import Data.Ord(comparing)
import qualified Data.List as List
-- import Data.Monoid
import Data.Semigroup
import qualified Data.Vector.Generic as GV
import qualified Data.Vector.Storable as V

import Control.Arrow
-- import Control.Lens
import Lens.Micro
import Lens.Micro.TH
import Lens.Micro.Extras
import Control.DeepSeq
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Reader

import qualified Foreign
import Foreign.Storable.Tuple

import GHC.Float (float2Double)
import GHC.Generics(Generic)

-- import qualified Numeric.LinearAlgebra as LA
-- import Numeric.LinearAlgebra.Devel(matrixFromVector, orderOf, MatrixOrder(..))
-- import Numeric.LinearAlgebra.Data(flatten, tr)
-- import qualified Numeric.LinearAlgebra.Devel as LADevel

import PSO.VectorSpace
import PSO.Random

-- | A bee in our swarm.
--
-- * @α@ is the type of a point in our phase space (i.e. position-velocity pair
-- in the classical case and wave function in the quantum mechanical)
-- * @β@ is the type of the local guide, i.e. \"memory\" of the bee.
-- * @r@ return type of the function we're minimising.
--
-- /Note:/ use 'state', 'guide', and 'val' lenses to access the fields.
data Bee α β r = Bee
  { _beeState :: !α
  , _beeGuide :: !β
  , _beeVal   :: !r
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''Bee

-- | Point in the classical phase space.
--
-- @χ@ -- vector in coordinate space. We always use time steps of 1, so physical
-- dimension of velocity is equal to the one of coordinate. The type of velocity
-- is thus also @χ@.
--
-- /Note:/ use 'pos' and 'vel' lenses to access the fields.
data CMState χ = CMState
  { _cmstatePos :: !χ
  , _cmstateVel :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''CMState

-- | Point in our quantum mechanical phase space (i.e. Hilbert space).
--
-- /Note:/ use 'pos' lens to access the field.
data QMState χ = QMState
  { _qmstatePos :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''QMState

-- | Rather than limiting ourselves to a single type of local guide (because who
-- knows which info we might want to keep track of), we define a type class.
-- @LocalGuide β α r@ means that @β@ is a local guide for a bee with phase of
-- type @α@ and fitness function returning @r@.
class LocalGuide β α r where
  -- | Initialisation of a guide. Given a phase \(\psi^{(0)}\) and its fitness
  -- \(f(\psi^{(0)})\), @mkLocalG@ creates local guide \(g^{(0)}_l\).
  mkLocalG     :: α -> r -> β
  -- | Update step. After computing the phase \(\psi^{(n+1)}\) (of type @α@) and
  -- its fitness \(f(\psi^{(n+1)})\) (of type @r@) in the @(n+1)@'st iteration,
  -- we use @updateLocalG@ to compute \(g^{(n+1)}_l\) given \(g^{(n)}_l\).
  updateLocalG :: α -> r -> β -> β

-- | Similar to 'LocalGuide' we define a class of global guides rather than a
-- single one. Here, @GlobalGuide γ β@ means that @γ@ can serve as a global
-- guide for local guides of type @β@.
class (LocalGuide β α r) => GlobalGuide γ β α r where
  -- | Initialisation of a guide. Given a collection of bees, @mkLocalG@
  -- initialises global guide \(g^{(0)}_g\).
  mkGlobalG     :: (Foldable t, Functor t) => t (Bee α β r) -> γ
  -- | Update step. After computing phases, values, and local guides of bees in
  -- the @(n+1)@'st iteration, we use @updateLocalG@ to compute \(g^{(n+1)}_g\)
  -- given \(g^{(n)}_g\).
  updateGlobalG :: (Foldable t, Functor t) => t (Bee α β r) -> γ -> γ

-- | A guide that remembers position and its corresponding fitness.
--
-- /Note:/ use 'pos' and 'val' lenses to access the fields.
data BeeGuide χ r = BeeGuide
  { _beeguidePos :: !χ -- ^ Position.
  , _beeguideVal :: !r -- ^ Fitness value.
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''BeeGuide

instance (NFData χ, NFData r) => NFData (BeeGuide χ r)

-- | We compare guides by their fitness values.
instance Eq r => Eq (BeeGuide χ r) where
  x == y = (x ^. val) == (y ^. val)

-- | We compare guides by their fitness values.
instance Ord r => Ord (BeeGuide χ r) where
  compare = comparing (view val)

-- | 'BeeGuide' is the simplest and, probably, the most useful example of a
-- 'LocalGuide'. It keeps track of the best position and its fitness. That is,
-- if we've run @n@ iterations (calls to 'updateLocalG'), then the best position
-- is given by \(\operatorname{argmin}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\), and
-- its fitness is just \(\operatorname{min}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\).
instance (HasPos α χ, Ord r) => LocalGuide (BeeGuide χ r) α r where
  mkLocalG ψ fψ = BeeGuide (ψ ^. pos) fψ
  updateLocalG ψ fψ g
    | fψ < g ^. val = g & (pos .~ (ψ ^. pos))
                        & (val .~ fψ)
    | otherwise = g

-- | 'BeeGuide' can also serve as a 'GlobalGuide'. It keeps track of the overall
-- best position among all bees and iterations and its fitness. So after @n@
-- iterations, the best position is given by
-- \(\operatorname{argmin}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\),
-- and its fitness is
-- \(\operatorname{min}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\).
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (BeeGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG = minimum . fmap (view guide)
    updateGlobalG xs _ = minimum . fmap (view guide) $ xs

-- | A guide that remembers
--
--    * Best position throughout the whole swarm.
--    * Fitness value at this position.
--    * Current iteration.
--    * Variance in the fitness values of the local guides.
--
-- /Note:/ use 'pos', 'val', 'iteration', and 'var' lenses to access the fields.
data SwarmGuide χ r = SwarmGuide
  { _swarmguidePos       :: !χ
  , _swarmguideVal       :: !r
  , _swarmguideIteration :: !Int
  , _swarmguideVar       :: () -- !r
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''SwarmGuide

instance (NFData χ, NFData r) => NFData (SwarmGuide χ r)

-- | Calculates the mean value of a sequence.
mean :: (Fractional a, Foldable v) => v a -> a
mean xs = sum xs / fromIntegral (length xs)

-- | Calculates the variance in a sequence.
variance :: (Fractional a, Foldable v, Functor v) => v a -> a
variance xs = let m = mean xs in mean . fmap (\x -> (x - m)^^2) $ xs

-- | 'SwarmGuide' is another example of a 'GlobalGuide'. On top of keeping track
-- of the best position and its fitness value (what 'BeeGuide' does), we
-- remember the current iteration ('mkGlobalG' sets it to @0@, and each call to
-- 'updateGlobalG' increases the iteration by @1@). And finally, each iteration
-- we compute the variance \(\operatorname{Var}\{f(\psi^{(n)}_i)\}_i\) in the
-- fitness values of bees.
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (SwarmGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG xs = SwarmGuide (x ^. pos) (x ^. val) 0 var
      where x   = minimum . fmap (view guide) $ xs
            var = () -- variance . fmap (view val) $ xs
    updateGlobalG xs g = mkGlobalG xs & iteration .~ (g ^. iteration + 1)

-- | Given global guide and bee in the @n@'th iteration, calculates the phase of
-- the bee in the @(n+1)@'st iteration. This calculation is done in some 'Monad'
-- @m@ which allows us to, for example, use random numbers when updating the
-- phase.
--
-- /Note:/ use 'run' lens to unpack the function.
newtype PhaseUpdater m γ β α r = PhaseUpdater
  { _phaseupdaterRun :: γ -> Bee α β r -> m α
  }

makeLensesWith abbreviatedFields ''PhaseUpdater

-- | 'PhaseUpdater's can be combined to build more difficult ones.
instance (Monad m) => Semigroup (PhaseUpdater m γ β α r) where
  (PhaseUpdater a) <> (PhaseUpdater b) = PhaseUpdater update
    where update guide bee = a guide bee >>= \ψ ->
            b guide (bee & state .~ ψ)

-- | Our swarm.
data Swarm m γ β α r = Swarm
  { _swarmBees      :: ![Bee α β r] -- ^ Collection of bees.
  , _swarmGuide     :: !γ -- ^ Global guide.
  , _swarmFunc      :: !(α -> m r) -- ^ Fitness function.
  , _swarmUpdater   :: !(PhaseUpdater m γ β α r) -- ^ Updater.
  }

makeLensesWith abbreviatedFields ''Swarm

-- | Updates the canonical momentum in the classical phase. Given a global guide
-- and a bee described by classical mechanics, returns the momentum of the bee
-- in the next iteration. 'VelocityUpdater's can be combined using '<>' operator
-- and afterwards converted to 'PhaseUpdater's using 'upStep' function.
newtype VelocityUpdater m γ χ r = VelocityUpdater
  { _velocityupdaterRun :: γ -> Bee (CMState χ) (BeeGuide χ r) r -> m χ
  }

makeLensesWith abbreviatedFields ''VelocityUpdater

-- | 'VelocityUpdater's can be combined to form more complex ones.
instance (Monad m) => Semigroup (VelocityUpdater m γ χ r) where
  (VelocityUpdater a) <> (VelocityUpdater b) = VelocityUpdater $ update
    where update guide bee =
            a guide bee >>= \v -> b guide (bee & (state . vel) .~ v)

-- | Given a constant \(\varphi_l\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_l U(0, 1) \circ (g^{(n)}_l - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_l\) is local best position in the @n@'th iteration.
upLocal :: (RandomScalable m χ, VectorSpace χ λ) => λ -> VelocityUpdater m γ χ r
upLocal φl = VelocityUpdater $ upLocalImpl φl

upLocalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos β χ)
  => λ -> γ -> Bee α β r -> m χ
upLocalImpl φ _ bee =
  let δq = bee ^. guide . pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\varphi_g\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_g U(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_g\) is global best position in the @n@'th iteration.
upGlobal ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => λ -> VelocityUpdater m γ χ r
upGlobal φg = VelocityUpdater $ upGlobalImpl φg

upGlobalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos γ χ)
  => λ -> γ -> Bee α β r -> m χ
upGlobalImpl φ g bee = do
  let δq = g ^. pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\omega\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = \omega p^{(n)} \;.
-- \]
upScale :: (Monad m, Scalable λ χ) => λ -> VelocityUpdater m γ χ r
upScale ω = VelocityUpdater $ upScaleImpl ω

upScaleImpl :: (Monad m, Scalable λ χ, HasVel α χ) => λ -> γ -> Bee α β r -> m χ
upScaleImpl ω _ = return . scale ω . view (state . vel)

-- | Converts a 'VelocityUpdater' into a 'PhaseUpdater'. This is done by first
-- running the velocity updater to compute \(p^{(n+1)}\). \(q^{(n+1)}\) is then
-- calculated as \(q^{(n)} + p^{(n+1)}\).
upStep ::
     (Monad m, Num χ)
  => VelocityUpdater m γ χ r -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
upStep (VelocityUpdater upVel) = PhaseUpdater update
  where update guide bee = do
          δq <- upVel guide bee
          let x = bee ^. state
              q = x ^. pos
          return $ x & (vel .~ δq)
                     & (pos .~ q + δq)

-- | Creates the standard \"WPG\" updater. Given constants
-- \(\omega,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \omega p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
wpgUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
wpgUpdater (ω, φl, φg) =
  upStep $ (upScale ω) <> (upLocal φl) <> (upGlobal φg)

-- | Creates the standard updater which uses constriction rather than inertia
-- parameter. Given constants
-- \(\kappa,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \kappa \left( p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \right) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
constrictedUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
constrictedUpdater (κ, φl, φg) =
  upStep $ (upLocal φl) <> (upGlobal φg) <> (upScale κ)

-- randomWalkUpdater :: (UniformDist p gen m)
--                   => (p, p) -> Updater m gen s g p r
-- randomWalkUpdater bounds = Updater $ f
--   where f _ _ = do
--           randGen <- ask
--           v' <- lift $ uniform bounds randGen
--           return v'
--
-- kickUpdater ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      , s ~ SwarmGuide p r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> Updater m gen s g p r
-- kickUpdater cutoff bounds wpg = Updater $ kickUpdaterImpl cutoff bounds wpg
--
-- kickUpdaterImpl ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> s -> Bee p r g -> ReaderT gen m p
-- kickUpdaterImpl (cutoff, c) bounds wpg stats x = runUpdater updater stats x
--   where updater = if (stats ^. var) < cutoff
--                      then randomWalkUpdater bounds
--                           <> upScale (c * sqrt (stats ^.var))
--                      else standardUpdater wpg

-- | Update the position according to QDPSO (Quantum Delta well Particle Swarm
-- Optimisation).
class (Monad m) => DeltaWell m λ χ where
  upDeltaWell :: λ -> χ -> χ -> m χ

-- | 1D QDPSO updater. Given a constant \(\kappa\in\mathbb{R}\), center
-- \(g^{(n)}\in\mathbb{R}\) of the density distribution, and bees current
-- position \(q^{(n)}\in\mathbb{R}\), its position in the next iteration is
-- computed as
-- \[
--  q^{(n+1)} =
--    q^{(n)} + \sgn(2\cdot\operatorname{rand}(0,1) - 1) \kappa
--                \log(\frac{1}{\operatorname{rand}(0, 1)}) \cdot
--                  |q^{(n)} - g^{(n)}| \;.
-- \]
instance (Monad m, RealFloat λ, Ord λ, Randomisable m λ)
  => DeltaWell m λ λ where
    upDeltaWell κ p x = do
      u    <- random
      sign <- (\c -> signum (2 * c - 1)) <$> random
      return $ p + (sign * (log . recip $ u) / κ) * abs (x - p)

-- | Extending QDPSO to \(\mathbb{C}\). 1D real version is simply applied to
-- both real and imaginary parts.
instance (RealFloat λ, DeltaWell m λ λ)
  => DeltaWell m λ (Complex λ) where
    upDeltaWell κ p x =
      let f getter = uncurry (upDeltaWell κ) . over each getter
       in return (:+) `ap` f realPart (p, x)
                      `ap` f imagPart (p, x)

-- | Extending QDPSO to multidimensional case. 1D version is simply applied to
-- each component.
instance (Monad m, Foreign.Storable ξ, DeltaWell m λ ξ)
  => DeltaWell m λ (V.Vector ξ) where
    upDeltaWell κ p x = V.zipWithM (upDeltaWell κ) p x

{-
liftMatrix2M ::
     (Monad m, LA.Element α, LA.Element β, LA.Element γ,
      LA.Container V.Vector α, Num α,
      LA.Container V.Vector β, Num β,
      LA.Transposable (LA.Matrix α) (LA.Matrix α),
      LA.Transposable (LA.Matrix β) (LA.Matrix β))
  => (LA.Vector α -> LA.Vector β -> m (LA.Vector γ))
  -> LA.Matrix α -> LA.Matrix β -> m (LA.Matrix γ)
liftMatrix2M f m1@(LA.size -> (r,c)) m2
  | LA.size m1 /= LA.size m2 = error "nonconformant matrices in liftMatrix2M"
  | orderOf m1 == RowMajor =
      return (matrixFromVector RowMajor r c) `ap` f (flatten m1) (flatten m2)
  | otherwise =
      return (matrixFromVector ColumnMajor r c) `ap`
        f (flatten . tr $ m1) (flatten . tr $ m2)

instance (Monad m, Num ξ, LA.Container V.Vector ξ, LA.Transposable (LA.Matrix ξ)(LA.Matrix ξ), DeltaWell m λ ξ)
  => DeltaWell m λ (LA.Matrix ξ) where
    upDeltaWell κ = liftMatrix2M (V.zipWithM (upDeltaWell κ))
-}

-- | Creates a QDPSO updater.
deltaUpdater :: forall m g γ β χ λ r.
     ( Randomisable m λ
     , RealFloat λ
     , Ord λ
     , VectorSpace χ λ
     , DeltaWell m λ χ
     , HasPos β χ
     , HasPos γ χ
     )
  => λ -> γ -> Bee (QMState χ) β r -> m (QMState χ)
deltaUpdater κ g bee = do
  φ₁ <- random :: m λ
  φ₂ <- random :: m λ
  -- lift $ putStrLn $ "φ₁ = " ++ show φ₁ ++ ", φ₂ = " ++ show φ₂
  let center = (φ₁ / (φ₁ + φ₂)) `scale` (g ^. pos)
             + (φ₂ / (φ₁ + φ₂)) `scale` (bee ^. guide . pos)
  x <- upDeltaWell κ center (bee ^. state . pos)
  -- lift $ putStrLn $ "p - x = " ++ show (bee ^. state . pos - p)
  return $ (bee ^. state) & (pos .~ x)

data RelativePosition = Inside | LeftOf | RightOf

class RelativelyPositioned α β where
  getRelPos :: α -> β -> RelativePosition

instance (Ord λ) => RelativelyPositioned (λ, λ) λ where
  getRelPos (l, h) x
    | x < l     = LeftOf
    | x > h     = RightOf
    | otherwise = Inside

class Absorbable λ α where
  absorb :: (λ, λ) -> α -> α

class Reflectable λ α where
  reflect :: (λ, λ) -> α -> α

instance (Ord χ) => Absorbable χ χ where
  absorb (l, h) x
    | x < l     = l
    | x > h     = h
    | otherwise = x

instance (Ord λ) => Absorbable λ (Complex λ) where
  absorb bs = over each (absorb bs)

instance (Ord λ, Num λ)
  => Absorbable λ (λ, λ) where
    absorb (l, h) (q, p) = case getRelPos (l, h) q of
                             Inside  -> (q, p)
                             LeftOf  -> (l, 0)
                             RightOf -> (h, 0)

zipWithTuple :: (a -> b -> c) -> (a, a) -> (b, b) -> (c, c)
zipWithTuple f (x₁, x₂) (y₁, y₂) = (f x₁ y₁, f x₂ y₂)

instance (Ord λ, Num λ)
  => Absorbable λ (Complex λ, Complex λ) where
    absorb bounds ψ =
      let f getter = absorb bounds . over each getter
       in zipWithTuple (:+) (f realPart ψ) (f imagPart ψ)

instance (Absorbable λ (ξ, ξ), GV.Vector χ ξ, GV.Vector χ (ξ, ξ))
  => Absorbable λ (CMState (χ ξ)) where
    absorb bounds ψ =
      let (q, p) = GV.unzip . GV.map (absorb bounds)
                 $ GV.zip (ψ ^. pos) (ψ ^. vel)
       in ψ & (pos .~ q) & (vel .~ p)

instance (Absorbable λ ξ, GV.Vector χ ξ)
  => Absorbable λ (QMState (χ ξ)) where
    absorb bounds = pos %~ GV.map (absorb bounds)

absorbUpdater :: (Monad m, Absorbable λ α) => (λ, λ) -> PhaseUpdater m γ β α r
absorbUpdater bounds = PhaseUpdater $ upAbsorbImpl bounds

upAbsorbImpl :: (Monad m, Absorbable λ α) => (λ, λ) -> γ -> Bee α β r -> m α
upAbsorbImpl bounds _ = return . absorb bounds . view state

-- | Given a swarm and bee in the @n@'th iteration, returns a computation that
-- returns the bee in the @(n+1)@'st iteration.
updateBee ::
     (Monad m, LocalGuide β α r)
  => Swarm m γ β α r -> Bee α β r -> m (Bee α β r)
updateBee xs x = do
  ψ'  <- (xs ^. updater . run) (xs ^. guide) x
  fψ' <- (xs ^. func) ψ'
  return $ x & (state .~ ψ')
             & (val   .~ fψ')
             & (guide %~ updateLocalG ψ' fψ')

-- | Given a swarm in the @n@'th iteration, returns a computation that returns
-- the swarm in the @(n+1)@'st iteration.
updateSwarm :: (Monad m, GlobalGuide γ β α r)
            => Swarm m γ β α r -> m (Swarm m γ β α r)
updateSwarm swarm = do
  xs <- mapM (updateBee swarm) (swarm ^. bees)
  return $ swarm & (bees .~ xs)
                 & (guide %~ updateGlobalG xs)

-- | Applies \(f\) to \(x\) \(\max(0, n - 1)\) times and accumulates the
-- results. Returned list will have \(n\) elements.
iterateNM ::
     (Monad m)
  => Int -- ^ \(n\), number of times to apply \(f\)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateNM n f x
  | n == 0 = return [x]
  | otherwise = do
    y <- f x
    ys <- iterateNM (n - 1) f y
    return $ x : ys

iterateWhileM ::
     (Monad m)
  => (a -> Bool)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateWhileM predicate func x
  | predicate x = do
      y  <- func x
      ys <- iterateWhileM predicate func y
      return $ x : ys
  | otherwise   = return [x]

mkCMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (CMState χ)
mkCMState bounds = liftM (\q -> CMState q 0) $ uniform bounds

mkQMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (QMState χ)
mkQMState bounds = liftM (\q -> QMState q)   $ uniform bounds

-- | @mkSwarm updater initialiser inside func n@ creates a new 'Swarm' where
--
-- * @updater@ is the update strategy used to obtain velocities of bees in the
-- next iteration.
-- * @initialiser@ creates a new random 'Bee'. It is assumed that position of a
-- bee created using @initialiser@ will satisfy @inside@.
-- * Given a position \(x\in\mathcal{H}\) returns whether \(x\in\mathcal{V}\)
-- holds.
-- * @func@ is the function \(f\) we're trying to minimise.
-- * @n@ is number of bees in the newly created swarm.
mkSwarm ::
     (Monad m, GlobalGuide γ β α r)
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> m (Swarm m γ β α r)
mkSwarm newState updater func n = (sequence . replicate n $ newState) >>=
  mapM (mkBee func) >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

mkSwarmFromList ::
     (Monad m, GlobalGuide γ β α r)
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> m (Swarm m γ β α r)
mkSwarmFromList states updater func =
  mapM (mkBee func) states >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

fromPure :: (Monad m, HasPos α χ) => (χ -> r) -> α -> m r
fromPure f ψ = return $ f (ψ ^. pos)
--
-- class Between a b where
--   isBetween :: (a, a) -> b -> Bool
--
-- instance (Ord a) => Between a a where
--   isBetween (l, h) x = l <= x && x <= h
--
-- instance {-# OVERLAPS #-} (Between a a) => Between (Complex a) (Complex a) where
--   isBetween (rl :+ il, rh :+ ih) (rx :+ ix) =
--     isBetween (rl, rh) rx && isBetween (il, ih) ix
--
-- instance {-# OVERLAPS #-} (Foreign.Storable a, Between a a)
--   => Between (V.Vector a) (V.Vector a) where
--     isBetween (low, high) x =
--       V.and $ V.zipWith3 (\l h x -> isBetween (l, h) x) low high x
--
-- -- | Creates a new 'Bee' with a random position inside a given 1D interval and
-- -- zero velocity.
-- randomBee1D ::
--      (Num p, UniformDist p gen m)
--   => (p, p) -- ^ @[low, high]@ interval. It is assumed that @low <= high@.
--   -> (p -> ReaderT gen m r) -- ^ Function we're minimising.
--   -> ReaderT gen m (Bee p r (BeeGuide p r))
-- randomBee1D bounds func = do
--   g <- ask
--   x <- lift $ uniform bounds g
--   fx <- func x
--   return $ Bee x 0 fx (BeeGuide x fx)

-- | Given a fitness function \(f\) and phase \(\psi^{(0)}\), creates a new
-- 'Bee'.
mkBee :: (Monad m, LocalGuide β α r)
  => (α -> m r) -> α -> m (Bee α β r)
mkBee f ψ = do
  fψ <- f ψ
  return $ Bee ψ (mkLocalG ψ fψ) fψ

-- -- | Simplified @mkSwarm@ for 1D pure functions.
-- -- simpleInit1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a) -- ^ WPG parameters \((\omega,\varphi_l,\varphi_g)\).
-- --   -> (a, a) -- ^ Boundaries \([a, b]\).
-- --   -> (a -> r) -- ^ Function \(f\) we're minimising.
-- --   -> Int -- ^ Number of bees.
-- --   -> ReaderT gen m (Swarm m gen s g a r)
-- -- simpleInit1D wpg (low, high) func =
-- --   let updater = standardUpdater wpg
-- --       initialiser = randomBee1D (low, high)
-- --       boundaries = isBetween (low, high)
-- --   in mkSwarm updater initialiser boundaries (fromPure func)
--
--
-- unpackToV :: (Foreign.Storable a) => [(a, a)] -> (V.Vector a, V.Vector a)
-- unpackToV xs = let a = V.fromList . map fst $ xs
--                    b = V.fromList . map snd $ xs
--                 in (a, b)
--
-- -- | Simplified @mkSwarm@ for N-dimensional functions.
-- simpleInitND ::
--      ( Foreign.Storable a
--      , VectorSpace p a
--      , RandomScalable p gen m
--      , UniformDist p gen m
--      , Between p p
--      , Ord r
--      , p ~ V.Vector a
--      , GlobalGuide s g p r
--      )
--   => Updater m gen s g p r -- ^ PWG parameters \((\omega,\varphi_l,\varphi_g)\)
--   -> [(a, a)]
--   -> [(a, a)] -- ^ List of intervals \([a_i,b_i]\). \(\mathcal{V}\) is then
--   -- \(\prod_i [a_i,b_i]\).
--   -> (p -> r) -- ^ Function \(f\) we're minimising.
--   -> Int -- ^ Number of bees.
--   -> ReaderT gen m (Swarm m gen s g p r)
-- simpleInitND updater initBounds bounds func =
--   let initialiser = randomBeeND (unpackToV initBounds)
--       boundaries = isBetween (unpackToV bounds)
--   in mkSwarm updater initialiser boundaries (fromPure func)
--
-- -- optimise1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a)
-- --   -> (a, a)
-- --   -> (a -> r)
-- --   -> Int
-- --   -> (Swarm m gen s g a r -> Bool)
-- --   -> ReaderT gen m [Swarm m gen s g a r]
-- -- optimise1D wpg bounds func n predicate = do
-- --   swarm <- simpleInit1D wpg bounds func n
-- --   iterateWhileM (not . predicate) updateSwarm swarm

optimiseND ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> (Swarm m γ β α r -> Bool)
  -> m [Swarm m γ β α r]
optimiseND newState updater func n predicate = do
  swarm <- mkSwarm newState updater func n
  iterateWhileM (not . predicate) updateSwarm swarm

optimiseNDFromList ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> (Swarm m γ β α r -> Bool)
  -> (Swarm m γ β α r -> m ())
  -> m (Swarm m γ β α r)
optimiseNDFromList states updater func predicate process =
  do
    swarm <- mkSwarmFromList states updater func
    go swarm
  where
        go x
          | predicate x = return x
          | otherwise   = process x >> updateSwarm x >>= go


module Gosadenban
( main3

) where

-------------------------------------------------------------
--module
-------------------------------------------------------------

import Numeric.LinearAlgebra
import System.Random
import qualified Data.ByteString.Lazy as B
import Data.List
import Data.Time

import Tools
import Durstenfeld
import ForMNIST

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

data ActivationFunction = ActivationFunction
    { func  :: Double -> Double ,
      func' :: Double -> Double ,
      description :: String
    }

data Params = Params
    { weight :: Matrix Double ,
      bias :: Matrix Double
    } deriving (Show)

data DParams = DParams
    { dweight :: Matrix Double ,
      dbias :: Matrix Double
    } deriving (Show)

data LayerProp = LayerProp
    { params :: Params ,
      af :: ActivationFunction
    }

data HyperProp = HyperProp
    { vw0 :: Matrix Double ,
      vb0 :: Matrix Double
    }

data Hyperparameter = Hyperparameter
    { learningrate :: Double ,
      alpha0 :: Double ,
      beta0  :: Double ,
      beta1  :: Double
    }

data Network = Network
    { layers :: [LayerProp] ,
      hprop  :: [HyperProp] ,
      hpara  :: Hyperparameter
    }

data PropagateLayer
    = PropagateLayer
        { inData  :: Matrix Double ,
          outData :: Matrix Double ,
          pAf :: ActivationFunction ,
          pParams :: Params
        }
    | PropagateSensorLayer
        { outData :: Matrix Double
        }

data BackPropagateLayer
    = BackPropagateLayer
        { inGrad  :: Matrix Double ,
          outGrad :: Matrix Double ,
          bpAf :: ActivationFunction ,
          bpParams :: Params ,
          dParams :: DParams
        }
    | BackPropagateLastLayer
        { outGrad :: Matrix Double
        }

data ErrorLayer
    = ErrorLayer
        { lastData :: Matrix Double ,
          teachData :: Matrix Double ,
          errorValue :: Double ,
          grad :: Matrix Double
        } deriving (Show)

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

propagate :: PropagateLayer -> LayerProp -> PropagateLayer
propagate dataJ layerK = PropagateLayer
    { inData  = x ,
      outData = y ,
      pAf     = functions ,
      pParams = p
    }
    where x = outData dataJ
          p = params layerK
          w = weight p
          b = bias p
          a = add b $ x Numeric.LinearAlgebra.<> w
          functions =  af layerK
          h = func functions
          y = cmap h a

propagateNetwork :: Matrix Double -> Network -> [PropagateLayer]
propagateNetwork input net = calcs
    where calcs = scanl propagate layer0 (layers net)
          layer0 = PropagateSensorLayer { outData = input}

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

backpropagate :: PropagateLayer -> BackPropagateLayer -> BackPropagateLayer
backpropagate pLayerJ bpLayerK = BackPropagateLayer
    { inGrad = indx ,
      outGrad = dx ,
      bpAf = af ,
      bpParams  = p ,
      dParams = dP
    }
    where af = pAf pLayerJ
          p = pParams pLayerJ
          h' = func' af
          indx = outGrad bpLayerK
          y = outData pLayerJ
          dL' = cmap h' y
          dL = indx * dL'
          db = dL
          x = inData pLayerJ
          w = weight p
          dw = (tr x) Numeric.LinearAlgebra.<> dL
          dx = dL Numeric.LinearAlgebra.<> (tr w)
          dP = DParams { dweight = dw, dbias = db}

backpropagateNetwork :: ErrorLayer -> [PropagateLayer] ->[BackPropagateLayer]
backpropagateNetwork el prL = tail ans1
    where ans1 = scanr backpropagate layerFin prL
          layerFin = BackPropagateLastLayer {outGrad = grad0}
          grad0 = grad el

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

--Softmax & cross entropy error
errorNetwork :: Matrix Double -> Matrix Double -> ErrorLayer
errorNetwork ld td = ErrorLayer
    { lastData = dataFromSoftmaxFunction ,
      teachData = td ,
      errorValue = e / len ,
      grad = g
    }
    where dataFromSoftmaxFunction = softmaxBatch' ld
          g = add dataFromSoftmaxFunction $ scale (-1) td
          y = map asRow $ toRows dataFromSoftmaxFunction
          t = map asRow $ toRows td
          e' = zipWith cee t y
          e = sum e'
          len = length' t

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

inference :: Matrix Double -> Int -> Double
inference x t = ans
    where (r, c) = maxIndex x
          ans = if c == t then 1.0 else 0.0

accuracyNet :: Network -> Matrix Double -> Int -> Double
accuracyNet net x t = a
    where ans' = propagateNetwork x net
          ans  = last $ map outData ans'
          a    = inference ans t

accuracyNetwork :: Network -> [Matrix Double] -> [Int] -> Double
accuracyNetwork net xs ts = sumans / lenans
    where f x t = accuracyNet net x t
          ans = zipWith f xs ts
          sumans = sum ans
          lenans = length' xs

----------------------------------------------------------------------
sgd :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double
sgd x dx hp = x - scale eta dx
    where eta = learningrate hp

update :: Hyperparameter -> Params -> DParams -> Params
update hp p dp = newParam
    where b  = bias p
          w  = weight p
          db = dbias dp
          dw = dweight dp
          new_b = sgd b db hp
          new_w = sgd w dw hp
          newParam = Params {bias = new_b, weight = new_w}

updateNetowrk :: Network -> [DParams] -> Network
updateNetowrk net dp = new_Network
    where layerList = layers net
          hp = hpara net
          hpp = hprop net
          paramList = map params layerList
          afList = map af layerList
          f = update hp
          new_Para = zipWith f paramList dp
          new_Layer = zipWith makeLayerProp new_Para afList
          new_Network = Network{layers = new_Layer, hpara = hp, hprop = hpp  }

momentum :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double -> ( Matrix Double, Matrix Double)
momentum x dx hp v = (x + new_v, new_v)
    where eta = learningrate hp
          alpha = alpha0 hp
          new_v = (scale alpha v) - scale eta dx

adagrad :: Matrix Double -> Matrix Double -> Hyperparameter -> Matrix Double -> ( Matrix Double, Matrix Double)
adagrad x dx hp v = (new_x, new_v')
    where eta = learningrate hp
          alpha = alpha0 hp
          new_v' = v + dx*dx
          new_v = cmap sqrt new_v'
          new_v'' = cmap devdev new_v
          new_x = x - scale eta ( dx/new_v'' )

devdev :: Double -> Double
devdev x =
    if x == 0
        then 1.0
        else x

update' :: Hyperparameter -> (HyperProp , Params) -> DParams -> (Params, HyperProp)
update' hp (hpp, p) dp = (newParam, newHpp)
    where b  = bias p
          w  = weight p
          db = dbias dp
          dw = dweight dp
          vw = vw0 hpp
          vb = vb0 hpp
          (new_w, new_vw) = momentum w dw hp vw
          (new_b, new_vb) = momentum b db hp vb
          newParam = Params {bias = new_b, weight = new_w}
          newHpp = HyperProp {vw0 = new_vw, vb0 = new_vb}

updateNetowrk' :: Network -> [DParams] -> Network
updateNetowrk' net dp = new_Network
    where layerList = layers net
          hp = hpara net
          hpp = hprop net
          paramList = map params layerList
          afList = map af layerList
          f = update' hp
          hppp = zip hpp paramList
          new_Para_And_Hpp = zipWith f hppp dp
          new_Para = map fst new_Para_And_Hpp
          new_Hpp  = map snd new_Para_And_Hpp
          new_Layer = zipWith makeLayerProp new_Para afList
          new_Network = Network{layers = new_Layer, hpara = hp, hprop = new_Hpp  }

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

makeLayerProp :: Params -> ActivationFunction -> LayerProp
makeLayerProp p aF = LayerProp {params = p, af = aF }

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

loadL :: B.ByteString -> [Matrix Double]
loadL xs = tL4
    where tLabel  = loadLabel xs
          f x = pickUpNum x tLabel
          tL1 = map f batch_List
          tL2 = map ( map oneHotEncoder ) tL1
          tL3 = map ( map flatten ) tL2
          tL4 = map fromRows tL3

loadI :: B.ByteString -> [Matrix Double]
loadI xs = tI4
    where tImage  = loadImage xs
          g x = pickUpNum x tImage
          tI1 = map g batch_List
          tI2 = map ( map (scale (1/255.0) )) tI1
          tI3 = map ( map flatten ) tI2
          tI4 = map fromRows tI3

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

finPara :: Matrix Double -> Matrix Double -> Params
finPara b w = Params {bias = b, weight = w}

finLayer :: Params -> ActivationFunction -> LayerProp
finLayer p a = LayerProp {params = p, af = a }

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

sigmoidFunction :: Floating a => a -> a
sigmoidFunction x = 1 / (1 + exp (-x) )

sigmoidFunctionBack :: Floating a => a -> a
sigmoidFunctionBack  x = x*(1-x)

sigmoidFunc :: ActivationFunction
sigmoidFunc = ActivationFunction
    { func  = sigmoidFunction,
      func' = sigmoidFunctionBack,
      description = "Sigmoid Function"
    }

reluFunction :: Double -> Double
reluFunction x
    | x <= 0 = 0.0
    | x > 0  = x

reluFunctionBack :: Double -> Double
reluFunctionBack x
    | x <= 0 = 0.0
    | x > 0  = 1.0

reluFunc :: ActivationFunction
reluFunc = ActivationFunction
    { func  = reluFunction,
      func' = reluFunctionBack,
      description = "Relu Function"
    }

identity :: ActivationFunction
identity = ActivationFunction
    { func  = id,
      func' = backOne,
      description = "id"
    }

backOne :: Double -> Double
backOne x = 1.0

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

softmax :: (Linear t c, Floating t, Container c t) => c t -> c t
softmax m = scale  (1/bunbo) newM
    where maxValue = maxElement m
          newM     = cmap  (\x -> exp $ x - maxValue ) m
          bunbo    = sumElements newM

softmaxBatch' :: Matrix Double -> Matrix Double
softmaxBatch' m  =  m1
    where list  = toRows m
          list1 = fmap softmax list
          m1 = fromRows list1

cee :: Matrix Double -> Matrix Double -> Double
cee t y = (-1.0)*e
    where delta = 0.0000001
          y1 = cmap (\x -> log (x + delta) ) y
          y2 = flatten y1
          y3 = diag y2
          e = sumElements (t Numeric.LinearAlgebra.<> y3)

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

checkIt (e0, net) (x ,t) = (e, newNetwork)
    where ans' = propagateNetwork x net
          ans  = map outData ans'
          answer  = last ans
          answer2 = errorNetwork answer t
          e = errorValue answer2
          finalAnswer' = backpropagateNetwork answer2 ans'
          finalA = map dParams $ init finalAnswer'
          newNetwork = updateNetowrk' net finalA

checking net0 lt = ans
    where ans = scanl checkIt net0 lt

checking' net0 lt = ans
    where ans = foldl checkIt net0 lt

-------------------------------------------------------------
--Parameter
-------------------------------------------------------------

data_size = 10000
batch_size = 100
trials_num = 2000

input_size  = 784
hidden_size = 50
output_size = 10

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

-------------------------------------------------------------
--Initial Condition (?)
-------------------------------------------------------------

nw1 = randomMatrix input_size  hidden_size (mkStdGen 101) :: (Matrix Double)
nw2 = randomMatrix hidden_size output_size (mkStdGen 102) :: (Matrix Double)

nb1 = matrix hidden_size $ take (batch_size * hidden_size) $ repeat 0.0
nb2 = matrix output_size $ take (batch_size * output_size) $ repeat 0.0

nv_w1 =  matrix hidden_size $ take (input_size  * hidden_size) $ repeat 0.0
nv_w2 =  matrix output_size $ take (hidden_size * output_size) $ repeat 0.0

nv_b1 =  matrix hidden_size $ take (batch_size * hidden_size) $ repeat 0.0
nv_b2 =  matrix output_size $ take (batch_size * output_size) $ repeat 0.0

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

param1 = Params {bias = nb1, weight = nw1}
layer1 = LayerProp {params = param1, af = reluFunc}
hprop1 = HyperProp { vw0 = nv_w1, vb0 = nv_b1 }

param2 = Params {bias = nb2, weight = nw2}
layer2 = LayerProp {params = param2, af = identity }
hprop2  = HyperProp { vw0 = nv_w2, vb0 = nv_b2}

hyperParam = Hyperparameter {learningrate = 0.0001, alpha0 = 0.9, beta0 = 0.0, beta1 = 0.0}

networkA = Network {layers = [layer1,layer2], hpara = hyperParam, hprop = [hprop1,hprop2]}

--makeBatchList gen num_of_trials data_size batch_size
batch_List = makeBatchList 100000 trials_num data_size batch_size

main3 = do

    z <- getZonedTime
    print z

    x <- getCurrentTime

    contents1 <- B.readFile "MNIST/train-labels.idx1-ubyte"
    contents2 <- B.readFile "MNIST/train-images.idx3-ubyte"

    let tL4 = loadL contents1
        --tL4 = [[Batch X 10],[Batch X 10],...] :: [Matrix Double]

        tI4 = loadI contents2
        --tI4 = [[Batch X data_size],[Batch X data_size],...] :: [Matrix Double]

        test_Image_and_Label = zip tI4 tL4

        --e1 = checking' (0.0,networkA) test_Image_and_Label
        --e2 =  fst e1

        e1' =  checking (0.0,networkA) test_Image_and_Label
        e2' =  tail $ map fst e1'
        net1' = last $ map snd e1'

    --print e2
    writeFile "logs/e_mmt-0.0001.dat" $ encorderForDat2 e2'

    contents3 <- B.readFile "MNIST/t10k-labels.idx1-ubyte"
    contents4 <- B.readFile "MNIST/t10k-images.idx3-ubyte"

    let xs = loadImageTest contents4
        ts = loadLabelTest contents3
        --a = accuracyNet net1' xs ts
        --acc = accuracyNetwork net1' xs ts
        bs = map toRows $ map bias $ map params $ layers net1'
        bs' = map (foldl1 add) bs
        bsl = map recip $ map length' bs
        newbs = map asRow $ zipWith scale bsl bs'
        w = map weight $ map params $ layers net1'
        para = zipWith finPara newbs w
        acfunc = map af $ layers net1'
        newLayer = zipWith finLayer para acfunc
        finNetwork = Network {layers = newLayer, hpara = Hyperparameter {learningrate = 1.0, alpha0 = 0.0, beta0 = 0.0, beta1 = 0.0},hprop = [hprop1,hprop2]}

        acc = accuracyNetwork finNetwork xs ts

    print acc

    y <- getCurrentTime
    print$ diffUTCTime y x

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

loopCheck x net t = (e, finalA, newNetwork)
    where ans' = propagateNetwork x net
          ans  = map outData ans'
          answer  = last ans
          answer2 = errorNetwork answer t
          e = errorValue answer2
          finalAnswer' = backpropagateNetwork answer2 ans'
          finalA = map dParams $ init finalAnswer'
          newNetwork = updateNetowrk net finalA

loop _  _  _ 0 = []
loop x network t n = e : loop x newNetwork t (n-1)
    where  (e, finalA, newNetwork) = loopCheck x network t

randomList = randomChoice 10 10 (mkStdGen 10 )

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


{-# LANGUAGE FunctionalDependencies #-}

module Distribution (
    Dist, pdf, logpdf, samplefrom,
    Uniform(..), Normal(..), Bernoulli(..),
    Exponential(..), Gamma(..), Poisson(..),
    Binomial(..), Beta(..),
    Prob, Weight, Weighted
) where

import Randomness ( Random(..), uniform, Prob, Weight, Weighted )
import Numeric.Log ( Log(..) )
import qualified Data.Random.Distribution.Uniform as R ( uniform )
import qualified Data.Random.Distribution.Normal as R ( normal )
import qualified Data.Random.Distribution.Bernoulli as R ( bernoulli )
import qualified Data.Random.Distribution.Binomial as R ( binomial )
import qualified Data.Random.Distribution.Poisson as R ( poisson )
import qualified Statistics.Distribution as S ( ContDistr(density, logDensity, quantile), DiscreteDistr(probability, logProbability) )
import qualified Statistics.Distribution.Uniform as S ( uniformDistr )
import qualified Statistics.Distribution.Normal as S ( normalDistr )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Poisson as S ( poisson )
import qualified Statistics.Distribution.Gamma as S ( gammaDistr )
import qualified Statistics.Distribution.Exponential as S ( exponential )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Beta as S ( betaDistr )
import qualified Statistics.Distribution.DiscreteUniform as S ( discreteUniformAB )

class Dist d a | d -> a where
    samplefrom :: d -> Random a

    pdf :: d -> a -> Double
    pdf d = exp . ln . logpdf d

    logpdf :: d -> a -> Weight
    logpdf d = Exp . log . pdf d

data Uniform = Uniform Double Double
-- data DiscreteUniform = DiscreteUniform Int Int
data Normal = Normal Double Double
data Bernoulli = Bernoulli Prob
data Exponential = Exponential Double
data Gamma = Gamma Double Double
data Poisson = Poisson Double
data Binomial = Binomial Int Prob
data Beta = Beta Double Double

instance Dist Uniform Double where
    pdf (Uniform a b) = S.density $ S.uniformDistr a b
    logpdf (Uniform a b) = Exp . S.logDensity (S.uniformDistr a b)
    samplefrom (Uniform a b) = do
        u <- uniform
        return $ a + u * (b - a)

{-
instance Dist DiscreteUniform Int where
    pdf (DiscreteUniform a b) = S.probability $ S.discreteUniformAB a b
    logpdf (DiscreteUniform a b) = Exp . S.logProbability (S.discreteUniformAB a b)
    samplefrom (DiscreteUniform a b) = uniform (a, b)
-}

instance Dist Normal Double where
    pdf (Normal mu sig) = S.density $ S.normalDistr mu sig
    logpdf (Normal mu sig) = Exp . S.logDensity (S.normalDistr mu sig)
    samplefrom (Normal mu sig) = S.quantile (S.normalDistr mu sig) <$> uniform

instance Dist Bernoulli Bool where
    pdf (Bernoulli p) b = if b then p else 1-p
    logpdf (Bernoulli p) b = if b then Exp $ log p else Exp $ log (1-p)
    samplefrom (Bernoulli p) = (<= p) <$> uniform

instance Dist Gamma Double where
    pdf (Gamma k theta) = S.density $ S.gammaDistr k theta
    logpdf (Gamma k theta) = Exp . S.logDensity (S.gammaDistr k theta)
    samplefrom (Gamma k theta) = S.quantile (S.gammaDistr k theta) <$> uniform

instance Dist Exponential Double where
    pdf (Exponential lambda) = S.density $ S.exponential lambda
    logpdf (Exponential lambda) = Exp . S.logDensity (S.exponential lambda)
    samplefrom (Exponential lambda) = S.quantile (S.exponential lambda) <$> uniform

instance Dist Beta Double where
    pdf (Beta a b) = S.density $ S.betaDistr a b
    logpdf (Beta a b) = Exp . S.logDensity (S.betaDistr a b)
    samplefrom (Beta a b) = do
        x <- samplefrom (Gamma a 1)
        y <- samplefrom (Gamma b 1)
        return $ x / (x+y)

instance Dist Poisson Int where
    pdf (Poisson lambda) = S.probability $ S.poisson lambda
    logpdf (Poisson lambda) = Exp . S.logProbability (S.poisson lambda)
    samplefrom (Poisson lambda) = step 0 1.0
        where
            l = exp (-lambda)
            step k p = do
                u <- uniform
                let p' = p*u
                if p' <= l then return k
                           else step (k+1) p'

instance Dist Binomial Int where
    pdf (Binomial n p) = S.probability $ S.binomial n p
    logpdf (Binomial n p) = Exp . S.logProbability (S.binomial n p)
    samplefrom d@(Binomial n p) | n <= 10 = sumBernoullis n 0
                                | otherwise = do
        let a = 1 + n `div` 2
        let b = 1 + n - a
        x <- samplefrom $ Beta (fromIntegral a) (fromIntegral b)
        if x >= p
            then samplefrom $ Binomial (a-1) (p/x)
            else do
                s <- samplefrom (Binomial (b-1) ((p-x)/(1-x)))
                return (s + a)
        where
            sumBernoullis 0 s = return s
            sumBernoullis n s = do
                u <- uniform
                if u <= p then sumBernoullis (n-1) (s+1) else sumBernoullis (n-1) s