-- | A general purpose library for simulating differential processes, of a deterministic or
-- stochastic nature.

module Goal.Simulation.Filter
    (
    -- * Learning
      beliefBackpropagation
    -- * Statistics
    , beliefNegativeLogLikelihoods
    -- * Structural
    , decomposeBeliefTransition
    ) where


--- Imports ---


-- Goal --

import Goal.Core

import Goal.Geometry
import Goal.Probability


--- Learning ---

beliefBackpropagation
    :: ( Manifold x, ExponentialFamily m, ExponentialFamily y, ExponentialFamily n
       , Generative Natural m, Generative Natural n
       , Riemannian Natural m, Riemannian Natural y, Legendre Natural m, Legendre Natural y)
    => [NaturalFunction :#: Harmonium m n]
    -> Int
    -> Int
    -> (Function Mixture Mixture :#: NeuralNetwork m y m)
    -> [(c, Mixture, Mixture) :#: (x, m, m)]
    -> Rand (ST s) (Differentials :#: Tangent (Function Mixture Mixture) (NeuralNetwork m y m))
-- | NB: Assumes null latent state bias on the Harmoniums at the moment.
beliefBackpropagation trnss blkn cdn nnp xnzs = do
    let zns' = decomposeBeliefTransition <$> zip xnzs (tail xnzs)
        (zs,ns') = unzip zns'
        (eys,ys,ez0s,z0s) = feedForward nnp zs
        trnss' = harmoniumTranspose <$> zipWith modulateHarmoniumBelief z0s trnss
    gbss <- zipWithM (bulkGibbsSampling0 cdn) trnss' $ replicate blkn <$> ns'
    let project trns' z0 n' = potentialMapping $ trns' >.> (z0 <+> n')
        oNss = snd . unzip . last <$> gbss
        (_,mtx,_) = splitHarmonium $ head trnss
        errs = mtx >$> [ meanPoint [ project trns' z0 oN | oN <- oNs ] <-> project trns' z0 n'
                      | (trns',z0,oNs,n') <- zip4 trnss' z0s oNss ns' ]
    return $ feedBackward nnp zs eys ys ez0s errs

beliefNegativeLogLikelihoods
    :: (Manifold x, Manifold m, AbsolutelyContinuous Natural n, ExponentialFamily n, Sample n ~ [Double])
    => NaturalFunction :#: Harmonium m n
    -> (c, Mixture, Mixture) :#: (x, m, m)
    -> (Double, Double)
beliefNegativeLogLikelihoods trns xnz =
    let (x,n,z) = splitTriple xnz
        zx = harmoniumTranspose trns >.> z
        nx = harmoniumTranspose trns >.> n
     in (negate . log . density nx $ listCoordinates x, negate . log . density zx $ listCoordinates x)



--- Internal ---


decomposeBeliefTransition
    :: (Manifold m, Manifold n)
    => ( (c, Mixture, Mixture) :#: (m, n, n)
       , (c, Mixture, Mixture) :#: (m, n, n) )
    -> (Mixture :#: n, Mixture :#: n)
decomposeBeliefTransition (xnz,xnz') =
    let (_,_,z) = splitTriple xnz
        (_,n',_) = splitTriple xnz'
     in (z, n')