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

module Goal.Simulation.Filter.Flow
    (
    -- * Belief Processes
      coxProcess
    , filteringProcess
    , transducerProcess
    -- ** Statistics
    , locationBeliefField
    ) where


--- Imports ---


-- Goal --

import Goal.Core

import Goal.Simulation.Mealy

import Goal.Geometry
import Goal.Probability
import Goal.Simulation.Flow


--- Processes ---


coxProcess :: (ExponentialFamily m, Generative Natural m, Sample m ~ [Double])
    => (Double -> Double) -- ^ Gain over time
    -> (NaturalFunction :#: Harmonium (Replicated Poisson) m)
    -> RandST s (Mealy (Double, [Double]) [Int])
-- | Constructs a inhomogeneous poisson process. Since we have yet to
-- develop tools for describing distributions on manifolds directly, the
-- dynamic response works simply on sample spaces.
coxProcess gnf trns =
    accumulateRandomFunction (coxAccumulator gnf trns) 0

coxAccumulator gnf trns (t',cs) t = do
    let dt = t' - t
        trns' = modulateTransducerGain (dt*gnf t') trns
    n' <- generate $ conditionalLatentDistribution trns' cs
    return (n',t')

transducerProcess :: Manifold m
    => [Double]
    -> (Double -> Double)
    -> NaturalFunction :#: Harmonium (Replicated Poisson) m
    -> [NaturalFunction :#: Harmonium (Replicated Poisson) m]
-- | Right now assumes that dt is constant. Watch out for that.
transducerProcess ts@(t0:t1:_) gnf trns =
    let dt = t1 - t0
     in  [ modulateTransducerGain (dt * gnf t) trns | t <- ts ]
transducerProcess _ _ _ = error "Try Harder"


filteringProcess
    :: (ExponentialFamily m, Apply Mixture Mixture f, Domain f ~ m, Codomain f ~ m, Manifold x)
    => Flow c x -- ^ Underlying Flow
    -> Mealy (Double, [Double]) (Sample m) -- ^ Emission Process
    -> (Function Mixture Mixture :#: f) -- ^ Belief Dynamics
    -> (Mixture :#: m) -- ^ Initial Belief
    -> Flow (c, Mixture, Mixture) (x, m, m) -- ^ The belief process
filteringProcess flow cox nnp z0 =
    accumulateMealy z0 $ proc (t',z) -> do
        cm' <- flow -< t'
        n' <- cox -< (t', listCoordinates cm')
        let mn' = sufficientStatistic (manifold z) n'
            z' = nnp >.> z <+> mn'
            p' = joinTriple cm' mn' z'
        returnA -< (p', z')


--- Plot Oriented ---

locationBeliefField
    :: ( ExponentialFamily m, ExponentialFamily n, Generative Natural n
       , Apply Mixture Mixture f, Domain f ~ m, Codomain f ~ m
       , SampleSpace n ~ Continuum, Transition Natural Standard n )
    => NaturalFunction :#: Harmonium m (Replicated n)
    -> Double
    -> Double
    -> Int
    -> Int
    -> [Double]
    -> Function Mixture Mixture :#: f
    -> (Double,Double)
    -> (Double,Double)
locationBeliefField trns dt scl ix iy xs nnp (x,y) =
    let xs' = take ix xs ++ x : drop (ix+1) xs
        xs'' = take iy xs' ++ y : drop (iy+1) xs'
        sp = chart Standard . transition . (harmoniumTranspose trns >.>)
            . (nnp >.>) . potentialMapping $ conditionalLatentDistribution trns xs''
     in (scl/dt * (coordinate (2*ix) sp - (xs'' !! ix)), scl/dt * (coordinate (2*iy) sp - (xs'' !! iy)))