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

module Goal.Simulation.Flow (
    -- * Flows
      Flow
    -- ** Deterministic
    , autonomousODE
    , nonAutonomousODE
    , lagrangianFlow
    -- ** Stochastic
    , itoProcess
    , langevinStep
    , langevinFlow
    -- * Integral Curves
    , stepEuler
    , stepRK4
    , stepEuler'
    , stepRK4'
    ) where

--- Imports ---


-- Goal --

import Goal.Core

import Goal.Simulation.Mealy

import Goal.Geometry
import Goal.Simulation.Physics.Configuration
import Goal.Probability

-- Qualified --

import qualified Numeric.LinearAlgebra.HMatrix as M


--- Deterministic ---


-- | A 'Flow' can be used to generate a path of 'Element's through a particular
-- space 's' in time ('Double').
type Flow c m = Mealy Double (c :#: m)

autonomousODE :: Manifold m
    => (Coordinates -> Coordinates) -- ^ Differential Equation
    -> (c :#: m) -- ^ Initial State
    -> Flow c m -- ^ Differential Process
-- | Creates a process out of the ordinary differential equation described by the given
-- explicit function.
autonomousODE f' p0 =
    accumulateFunction accumulator (0,p0)
      where accumulator t' (t,p)
                | t' == t = (p,(t,p))
                | otherwise =
                      let dt = t' - t
                          p' = p <+> fromCoordinates (manifold p) (stepRK4 f' dt $ coordinates p)
                       in (p',(t',p'))

nonAutonomousODE :: Manifold m
    => (Double -> Coordinates -> Coordinates) -- ^ Differential Equation
    -> Double -- ^ Initial Time
    -> (c :#: m) -- ^ Initial State
    -> Flow c m -- ^ Differential Process
-- | Creates a process out of the ordinary differential equation described by the given
-- explicit function.
nonAutonomousODE f' t0 p0 =
    accumulateFunction accumulator (t0,p0)
      where accumulator t' (t,p)
                | t' == t = (p,(t,p))
                | otherwise =
                      let dt = t' - t
                          p' = p <+> fromCoordinates (manifold p) (stepRK4' f' t dt $ coordinates p)
                       in (p',(t',p'))

-- Stochastic --

itoProcess :: Manifold m
    => (Double -> Coordinates -> Coordinates) -- ^ The drift function
    -> (Double -> Coordinates -> M.Matrix Double) -- ^ The diffusion function
    -> Double -- ^ The initial time
    -> (c :#: m) -- ^ The initial state
    -> RandST s (Flow c m) -- ^ The Ito flow
-- | Constructs an ito process.
itoProcess mu sgma t0 p0 =
    accumulateRandomFunction (itoAccumulator mu sgma) (t0,p0)

itoAccumulator mu sgma t' (t,p)
  | t' == t = return (p,(t,p))
  | otherwise = do
        let dt = t' - t
            x = coordinates p
            x' = x + stepRK4' mu t dt x
        x'' <- generate $ muSigmaToMultivariateNormal x' (M.scale (sqrt dt) $ sgma t x)
        let p' = fromCoordinates (manifold p) x''
        return  (p',(t',p'))

-- Mechanical --

lagrangianFlow :: (Riemannian Generalized m, ForceField f m)
    => f
    -> (Partials :#: PhaseSpace m)
    -> Flow Partials (PhaseSpace m)
lagrangianFlow f p0 =
    accumulateFunction accumulator (0,p0)
      where accumulator t' (t,qdq)
                | t' == t = (qdq,(t,qdq))
                | otherwise =
                      let dt = t' - t
                          ddq xs = coordinates . vectorField f $ fromCoordinates (manifold qdq) xs
                          qdq' = qdq <+> fromCoordinates (manifold qdq) (stepRK4 ddq dt $ coordinates qdq)
                       in (qdq', (t',qdq'))

langevinStep :: (Riemannian Generalized m, ForceField f m)
    => Double
    -> f
    -> ( Partials :#: PhaseSpace m
       -> Function Differentials Differentials :#: Tensor (GeneralizedAcceleration m) (GeneralizedAcceleration m) )
    -> (Partials :#: PhaseSpace m)
    -> RandST s (Partials :#: PhaseSpace m)
langevinStep dt f sgma qdq = do
    let flx p = matrixSquareRoot $ sgma p
        dq = bundleToTangent qdq
    nrms <- replicateM (dimension $ manifold dq) . generate . chart Standard $ fromList Normal [0,1]
    let lng p = sqrt dt /> (flx p >.> fromList (Tangent $ bundleToTangent p) nrms)
        ddq xs = coordinates . vectorField (f, lng) $ fromCoordinates (manifold qdq) xs
        qdq' = qdq <+> fromCoordinates (manifold qdq) (stepRK4 ddq dt $ coordinates qdq)
    return qdq'

langevinFlow :: (Riemannian Generalized m, ForceField f m)
    => f
    -> ( Partials :#: PhaseSpace m
           -> Function Differentials Differentials :#: Tensor (GeneralizedAcceleration m) (GeneralizedAcceleration m) )
    -> (Partials :#: PhaseSpace m)
    -> RandST s (Flow Partials (PhaseSpace m))
langevinFlow f sgma p0 =
    accumulateRandomFunction (langevinAccumulator f sgma) (0,p0)

langevinAccumulator f sgma t' (t,qdq)
    | t' == t = return (qdq,(t,qdq))
    | otherwise = do
          let dt = t' - t
          qdq' <- langevinStep dt f sgma qdq
          return (qdq', (t',qdq'))


--- Integral Curves ---


stepEuler
    :: (Coordinates -> Coordinates) -- ^ The derivative of the function to simulate
    -> Double -- ^ The time step 'dt'
    -> Coordinates -- ^ The state of the system at the current time
    -> Coordinates -- ^ The difference to the state of the system at time 't' + 'dt'
-- | Returns the difference from an Euler step.
stepEuler f' dt x = realToFrac dt * f' x

stepRK4
    :: (Coordinates -> Coordinates) -- ^ The derivative of the function to simulate
    -> Double -- ^ The time step 'dt'
    -> Coordinates -- ^ The state of the system at the current time
    -> Coordinates -- ^ The difference to the state of the system at time 't' + 'dt'
-- | Returns the difference from an RK4 step.
stepRK4 f' dt x =
    let k1 = realToFrac dt * f' x
        k2 = realToFrac dt * f' (x + k1 / 2)
        k3 = realToFrac dt * f' (x + k2 / 2)
        k4 = realToFrac dt * f' (x + k3)
    in (k1 + 2 * k2 + 2 * k3 + k4) / 6

stepEuler'
    :: (Double -> Coordinates -> Coordinates) -- ^ The derivative of the function to simulate
    -> Double -- ^ The current time 't'
    -> Double -- ^ The time step 'dt'
    -> Coordinates -- ^ The state of the system at the current time
    -> Coordinates -- ^ The difference to the state of the system at time 't' + 'dt'
-- | Time inhomogenous version.
stepEuler' f' t dt x = realToFrac dt * f' t x

stepRK4'
    :: (Double -> Coordinates -> Coordinates) -- ^ The derivative of the function to simulate
    -> Double -- ^ The current time 't'
    -> Double -- ^ The time step 'dt'
    -> Coordinates -- ^ The state of the system at the current time
    -> Coordinates -- ^ The difference to the state of the system at time 't' + 'dt'
-- | Time inhomogenous version.
stepRK4' f' t dt x =
    let k1 = realToFrac dt * f' t x
        k2 = realToFrac dt * f' (t + dt / 2) (x + k1 / 2)
        k3 = realToFrac dt * f' (t + dt / 2) (x + k2 / 2)
        k4 = realToFrac dt * f' (t + dt) (x + k3)
    in (k1 + 2 * k2 + 2 * k3 + k4) / 6