{-# LANGUAGE RecursiveDo #-}

-- |
-- Module     : Simulation.Aivika.Trans.Dynamics.Extra
-- Copyright  : Copyright (c) 2009-2017, David Sorokin <david.sorokin@gmail.com>
-- License    : BSD3
-- Maintainer : David Sorokin <david.sorokin@gmail.com>
-- Stability  : experimental
-- Tested with: GHC 8.0.1
--
-- This module defines auxiliary functions such as interpolation ones
-- that complement the memoization, for example. There are scan functions too.
--

module Simulation.Aivika.Trans.Dynamics.Extra
       (-- * Interpolation
        initDynamics,
        discreteDynamics,
        interpolateDynamics,
        -- * Scans
        scanDynamics,
        scan1Dynamics) where

import Control.Monad.Fix

import Simulation.Aivika.Trans.Comp
import Simulation.Aivika.Trans.Internal.Specs
import Simulation.Aivika.Trans.Internal.Simulation
import Simulation.Aivika.Trans.Internal.Dynamics

-- | Return the initial value.
initDynamics :: Dynamics m a -> Dynamics m a
{-# INLINE initDynamics #-}
initDynamics :: forall (m :: * -> *) a. Dynamics m a -> Dynamics m a
initDynamics (Dynamics Point m -> m a
m) =
  forall (m :: * -> *) a. (Point m -> m a) -> Dynamics m a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point m
p ->
  let sc :: Specs m
sc = forall (m :: * -> *). Point m -> Specs m
pointSpecs Point m
p
  in Point m -> m a
m forall a b. (a -> b) -> a -> b
$ Point m
p { pointTime :: Double
pointTime = forall (m :: * -> *). Specs m -> Int -> Int -> Double
basicTime Specs m
sc Int
0 Int
0,
             pointIteration :: Int
pointIteration = Int
0,
             pointPhase :: Int
pointPhase = Int
0 }

-- | Discretize the computation in the integration time points.
discreteDynamics :: Dynamics m a -> Dynamics m a
{-# INLINE discreteDynamics #-}
discreteDynamics :: forall (m :: * -> *) a. Dynamics m a -> Dynamics m a
discreteDynamics (Dynamics Point m -> m a
m) =
  forall (m :: * -> *) a. (Point m -> m a) -> Dynamics m a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point m
p ->
  if forall (m :: * -> *). Point m -> Int
pointPhase Point m
p forall a. Eq a => a -> a -> Bool
== Int
0 then
    Point m -> m a
m Point m
p
  else
    let sc :: Specs m
sc = forall (m :: * -> *). Point m -> Specs m
pointSpecs Point m
p
        n :: Int
n  = forall (m :: * -> *). Point m -> Int
pointIteration Point m
p
    in Point m -> m a
m forall a b. (a -> b) -> a -> b
$ Point m
p { pointTime :: Double
pointTime = forall (m :: * -> *). Specs m -> Int -> Int -> Double
basicTime Specs m
sc Int
n Int
0,
               pointPhase :: Int
pointPhase = Int
0 }

-- | Interpolate the computation based on the integration time points only.
-- Unlike the 'discreteDynamics' function it knows about the intermediate 
-- time points that are used in the Runge-Kutta method.
interpolateDynamics :: Dynamics m a -> Dynamics m a
{-# INLINE interpolateDynamics #-}
interpolateDynamics :: forall (m :: * -> *) a. Dynamics m a -> Dynamics m a
interpolateDynamics (Dynamics Point m -> m a
m) = 
  forall (m :: * -> *) a. (Point m -> m a) -> Dynamics m a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point m
p -> 
  if forall (m :: * -> *). Point m -> Int
pointPhase Point m
p forall a. Ord a => a -> a -> Bool
>= Int
0 then 
    Point m -> m a
m Point m
p
  else 
    let sc :: Specs m
sc = forall (m :: * -> *). Point m -> Specs m
pointSpecs Point m
p
        n :: Int
n  = forall (m :: * -> *). Point m -> Int
pointIteration Point m
p
    in Point m -> m a
m forall a b. (a -> b) -> a -> b
$ Point m
p { pointTime :: Double
pointTime = forall (m :: * -> *). Specs m -> Int -> Int -> Double
basicTime Specs m
sc Int
n Int
0,
               pointPhase :: Int
pointPhase = Int
0 }

-- | Like the standard 'scanl1' function but applied to values in 
-- the integration time points. The accumulator values are transformed
-- according to the second argument, which should be either function 
-- 'memo0Dynamics' or its unboxed version.
scan1Dynamics :: MonadFix m
                 => (a -> a -> a)
                 -> (Dynamics m a -> Simulation m (Dynamics m a))
                 -> (Dynamics m a -> Simulation m (Dynamics m a))
{-# INLINABLE scan1Dynamics #-}
scan1Dynamics :: forall (m :: * -> *) a.
MonadFix m =>
(a -> a -> a)
-> (Dynamics m a -> Simulation m (Dynamics m a))
-> Dynamics m a
-> Simulation m (Dynamics m a)
scan1Dynamics a -> a -> a
f Dynamics m a -> Simulation m (Dynamics m a)
tr Dynamics m a
m =
  mdo Dynamics m a
y <- Dynamics m a -> Simulation m (Dynamics m a)
tr forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (Point m -> m a) -> Dynamics m a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point m
p ->
        case forall (m :: * -> *). Point m -> Int
pointIteration Point m
p of
          Int
0 -> 
            forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
p Dynamics m a
m
          Int
n -> do 
            let sc :: Specs m
sc = forall (m :: * -> *). Point m -> Specs m
pointSpecs Point m
p
                ty :: Double
ty = forall (m :: * -> *). Specs m -> Int -> Int -> Double
basicTime Specs m
sc (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
0
                py :: Point m
py = Point m
p { pointTime :: Double
pointTime = Double
ty, pointIteration :: Int
pointIteration = Int
n forall a. Num a => a -> a -> a
- Int
1, pointPhase :: Int
pointPhase = Int
0 }
            a
s <- forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
py Dynamics m a
y
            a
x <- forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
p Dynamics m a
m
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! a -> a -> a
f a
s a
x
      forall (m :: * -> *) a. Monad m => a -> m a
return Dynamics m a
y

-- | Like the standard 'scanl' function but applied to values in 
-- the integration time points. The accumulator values are transformed
-- according to the third argument, which should be either function
-- 'memo0Dynamics' or its unboxed version.
scanDynamics :: MonadFix m
                => (a -> b -> a)
                -> a
                -> (Dynamics m a -> Simulation m (Dynamics m a))
                -> (Dynamics m b -> Simulation m (Dynamics m a))
{-# INLINABLE scanDynamics #-}
scanDynamics :: forall (m :: * -> *) a b.
MonadFix m =>
(a -> b -> a)
-> a
-> (Dynamics m a -> Simulation m (Dynamics m a))
-> Dynamics m b
-> Simulation m (Dynamics m a)
scanDynamics a -> b -> a
f a
acc Dynamics m a -> Simulation m (Dynamics m a)
tr Dynamics m b
m =
  mdo Dynamics m a
y <- Dynamics m a -> Simulation m (Dynamics m a)
tr forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (Point m -> m a) -> Dynamics m a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point m
p ->
        case forall (m :: * -> *). Point m -> Int
pointIteration Point m
p of
          Int
0 -> do
            b
x <- forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
p Dynamics m b
m
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! a -> b -> a
f a
acc b
x
          Int
n -> do 
            let sc :: Specs m
sc = forall (m :: * -> *). Point m -> Specs m
pointSpecs Point m
p
                ty :: Double
ty = forall (m :: * -> *). Specs m -> Int -> Int -> Double
basicTime Specs m
sc (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
0
                py :: Point m
py = Point m
p { pointTime :: Double
pointTime = Double
ty, pointIteration :: Int
pointIteration = Int
n forall a. Num a => a -> a -> a
- Int
1, pointPhase :: Int
pointPhase = Int
0 }
            a
s <- forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
py Dynamics m a
y
            b
x <- forall (m :: * -> *) a. Point m -> Dynamics m a -> m a
invokeDynamics Point m
p Dynamics m b
m
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! a -> b -> a
f a
s b
x
      forall (m :: * -> *) a. Monad m => a -> m a
return Dynamics m a
y