-- |
-- Module     : Simulation.Aivika.Dynamics.Memo
-- 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 memo functions. The memoization creates such 'Dynamics'
-- computations, which values are cached in the integration time points. Then
-- these values are interpolated in all other time points.
--

module Simulation.Aivika.Dynamics.Memo
       (memoDynamics,
        memo0Dynamics,
        iterateDynamics,
        unzipDynamics,
        unzip0Dynamics) where

import Data.Array
import Data.Array.IO.Safe
import Data.IORef
import Control.Monad

import Simulation.Aivika.Internal.Specs
import Simulation.Aivika.Internal.Parameter
import Simulation.Aivika.Internal.Simulation
import Simulation.Aivika.Internal.Dynamics
import Simulation.Aivika.Dynamics.Extra

-- | Create a boxed array with default values.
newBoxedArray_ :: Ix i => (i, i) -> IO (IOArray i e)
newBoxedArray_ :: forall i e. Ix i => (i, i) -> IO (IOArray i e)
newBoxedArray_ = forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_

-- | Memoize and order the computation in the integration time points using 
-- the interpolation that knows of the Runge-Kutta method. The values are
-- calculated sequentially starting from 'starttime'.
memoDynamics :: Dynamics e -> Simulation (Dynamics e)
{-# INLINE memoDynamics #-}
memoDynamics :: forall e. Dynamics e -> Simulation (Dynamics e)
memoDynamics (Dynamics Point -> IO e
m) = 
  forall a. (Run -> IO a) -> Simulation a
Simulation forall a b. (a -> b) -> a -> b
$ \Run
r ->
  do let sc :: Specs
sc = Run -> Specs
runSpecs Run
r
         (Int
phl, Int
phu) = Specs -> (Int, Int)
integPhaseBnds Specs
sc
         (Int
nl, Int
nu)   = Specs -> (Int, Int)
integIterationBnds Specs
sc
     IOArray (Int, Int) e
arr   <- forall i e. Ix i => (i, i) -> IO (IOArray i e)
newBoxedArray_ ((Int
phl, Int
nl), (Int
phu, Int
nu))
     IORef Int
nref  <- forall a. a -> IO (IORef a)
newIORef Int
0
     IORef Int
phref <- forall a. a -> IO (IORef a)
newIORef Int
0
     let r :: Point -> IO e
r Point
p = 
           do let sc :: Specs
sc  = Point -> Specs
pointSpecs Point
p
                  n :: Int
n   = Point -> Int
pointIteration Point
p
                  ph :: Int
ph  = Point -> Int
pointPhase Point
p
                  phu :: Int
phu = Specs -> Int
integPhaseHiBnd Specs
sc 
                  loop :: Int -> Int -> IO e
loop Int
n' Int
ph' = 
                    if (Int
n' forall a. Ord a => a -> a -> Bool
> Int
n) Bool -> Bool -> Bool
|| ((Int
n' forall a. Eq a => a -> a -> Bool
== Int
n) Bool -> Bool -> Bool
&& (Int
ph' forall a. Ord a => a -> a -> Bool
> Int
ph)) 
                    then 
                      forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) e
arr (Int
ph, Int
n)
                    else 
                      let p' :: Point
p' = Point
p { pointIteration :: Int
pointIteration = Int
n', pointPhase :: Int
pointPhase = Int
ph',
                                   pointTime :: Double
pointTime = Specs -> Int -> Int -> Double
basicTime Specs
sc Int
n' Int
ph' }
                      in do e
a <- Point -> IO e
m Point
p'
                            e
a seq :: forall a b. a -> b -> b
`seq` forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray (Int, Int) e
arr (Int
ph', Int
n') e
a
                            if Int
ph' forall a. Ord a => a -> a -> Bool
>= Int
phu 
                              then do forall a. IORef a -> a -> IO ()
writeIORef IORef Int
phref Int
0
                                      forall a. IORef a -> a -> IO ()
writeIORef IORef Int
nref (Int
n' forall a. Num a => a -> a -> a
+ Int
1)
                                      Int -> Int -> IO e
loop (Int
n' forall a. Num a => a -> a -> a
+ Int
1) Int
0
                              else do forall a. IORef a -> a -> IO ()
writeIORef IORef Int
phref (Int
ph' forall a. Num a => a -> a -> a
+ Int
1)
                                      Int -> Int -> IO e
loop Int
n' (Int
ph' forall a. Num a => a -> a -> a
+ Int
1)
              Int
n'  <- forall a. IORef a -> IO a
readIORef IORef Int
nref
              Int
ph' <- forall a. IORef a -> IO a
readIORef IORef Int
phref
              Int -> Int -> IO e
loop Int
n' Int
ph'
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Dynamics a -> Dynamics a
interpolateDynamics forall a b. (a -> b) -> a -> b
$ forall a. (Point -> IO a) -> Dynamics a
Dynamics Point -> IO e
r

-- | Memoize and order the computation in the integration time points using 
-- the 'discreteDynamics' interpolation. It consumes less memory than the 'memoDynamics'
-- function but it is not aware of the Runge-Kutta method. There is a subtle
-- difference when we request for values in the intermediate time points
-- that are used by this method to integrate. In general case you should 
-- prefer the 'memo0Dynamics' function above 'memoDynamics'.
memo0Dynamics :: Dynamics e -> Simulation (Dynamics e)
{-# INLINE memo0Dynamics #-}
memo0Dynamics :: forall e. Dynamics e -> Simulation (Dynamics e)
memo0Dynamics (Dynamics Point -> IO e
m) = 
  forall a. (Run -> IO a) -> Simulation a
Simulation forall a b. (a -> b) -> a -> b
$ \Run
r ->
  do let sc :: Specs
sc   = Run -> Specs
runSpecs Run
r
         bnds :: (Int, Int)
bnds = Specs -> (Int, Int)
integIterationBnds Specs
sc
     IOArray Int e
arr  <- forall i e. Ix i => (i, i) -> IO (IOArray i e)
newBoxedArray_ (Int, Int)
bnds
     IORef Int
nref <- forall a. a -> IO (IORef a)
newIORef Int
0
     let r :: Point -> IO e
r Point
p =
           do let sc :: Specs
sc = Point -> Specs
pointSpecs Point
p
                  n :: Int
n  = Point -> Int
pointIteration Point
p
                  loop :: Int -> IO e
loop Int
n' = 
                    if Int
n' forall a. Ord a => a -> a -> Bool
> Int
n
                    then 
                      forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray Int e
arr Int
n
                    else 
                      let p' :: Point
p' = Point
p { pointIteration :: Int
pointIteration = Int
n', pointPhase :: Int
pointPhase = Int
0,
                                   pointTime :: Double
pointTime = Specs -> Int -> Int -> Double
basicTime Specs
sc Int
n' Int
0 }
                      in do e
a <- Point -> IO e
m Point
p'
                            e
a seq :: forall a b. a -> b -> b
`seq` forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray Int e
arr Int
n' e
a
                            forall a. IORef a -> a -> IO ()
writeIORef IORef Int
nref (Int
n' forall a. Num a => a -> a -> a
+ Int
1)
                            Int -> IO e
loop (Int
n' forall a. Num a => a -> a -> a
+ Int
1)
              Int
n' <- forall a. IORef a -> IO a
readIORef IORef Int
nref
              Int -> IO e
loop Int
n'
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Dynamics a -> Dynamics a
discreteDynamics forall a b. (a -> b) -> a -> b
$ forall a. (Point -> IO a) -> Dynamics a
Dynamics Point -> IO e
r

-- | Iterate sequentially the dynamic process with side effects in 
-- the integration time points. It is equivalent to a call of the
-- 'memo0Dynamics' function but significantly more efficient, for the array 
-- is not created.
iterateDynamics :: Dynamics () -> Simulation (Dynamics ())
{-# INLINE iterateDynamics #-}
iterateDynamics :: Dynamics () -> Simulation (Dynamics ())
iterateDynamics (Dynamics Point -> IO ()
m) = 
  forall a. (Run -> IO a) -> Simulation a
Simulation forall a b. (a -> b) -> a -> b
$ \Run
r ->
  do let sc :: Specs
sc = Run -> Specs
runSpecs Run
r
     IORef Int
nref <- forall a. a -> IO (IORef a)
newIORef Int
0
     let r :: Point -> IO ()
r Point
p =
           do let sc :: Specs
sc = Point -> Specs
pointSpecs Point
p
                  n :: Int
n  = Point -> Int
pointIteration Point
p
                  loop :: Int -> IO ()
loop Int
n' = 
                    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
n' forall a. Ord a => a -> a -> Bool
> Int
n) forall a b. (a -> b) -> a -> b
$
                    let p' :: Point
p' = Point
p { pointIteration :: Int
pointIteration = Int
n', pointPhase :: Int
pointPhase = Int
0,
                                 pointTime :: Double
pointTime = Specs -> Int -> Int -> Double
basicTime Specs
sc Int
n' Int
0 }
                    in do ()
a <- Point -> IO ()
m Point
p'
                          ()
a seq :: forall a b. a -> b -> b
`seq` forall a. IORef a -> a -> IO ()
writeIORef IORef Int
nref (Int
n' forall a. Num a => a -> a -> a
+ Int
1)
                          Int -> IO ()
loop (Int
n' forall a. Num a => a -> a -> a
+ Int
1)
              Int
n' <- forall a. IORef a -> IO a
readIORef IORef Int
nref
              Int -> IO ()
loop Int
n'
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Dynamics a -> Dynamics a
discreteDynamics forall a b. (a -> b) -> a -> b
$ forall a. (Point -> IO a) -> Dynamics a
Dynamics Point -> IO ()
r

-- | Memoize and unzip the computation of pairs, applying the 'memoDynamics' function.
unzipDynamics :: Dynamics (a, b) -> Simulation (Dynamics a, Dynamics b)
unzipDynamics :: forall a b. Dynamics (a, b) -> Simulation (Dynamics a, Dynamics b)
unzipDynamics Dynamics (a, b)
m =
  forall a. (Run -> IO a) -> Simulation a
Simulation forall a b. (a -> b) -> a -> b
$ \Run
r ->
  do Dynamics (a, b)
m' <- forall a. Run -> Simulation a -> IO a
invokeSimulation Run
r (forall e. Dynamics e -> Simulation (Dynamics e)
memoDynamics Dynamics (a, b)
m)
     let ma :: Dynamics a
ma =
           forall a. (Point -> IO a) -> Dynamics a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point
p ->
           do (a
a, b
_) <- forall a. Point -> Dynamics a -> IO a
invokeDynamics Point
p Dynamics (a, b)
m'
              forall (m :: * -> *) a. Monad m => a -> m a
return a
a
         mb :: Dynamics b
mb =
           forall a. (Point -> IO a) -> Dynamics a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point
p ->
           do (a
_, b
b) <- forall a. Point -> Dynamics a -> IO a
invokeDynamics Point
p Dynamics (a, b)
m'
              forall (m :: * -> *) a. Monad m => a -> m a
return b
b
     forall (m :: * -> *) a. Monad m => a -> m a
return (Dynamics a
ma, Dynamics b
mb)

-- | Memoize and unzip the computation of pairs, applying the 'memo0Dynamics' function.
unzip0Dynamics :: Dynamics (a, b) -> Simulation (Dynamics a, Dynamics b)
unzip0Dynamics :: forall a b. Dynamics (a, b) -> Simulation (Dynamics a, Dynamics b)
unzip0Dynamics Dynamics (a, b)
m =
  forall a. (Run -> IO a) -> Simulation a
Simulation forall a b. (a -> b) -> a -> b
$ \Run
r ->
  do Dynamics (a, b)
m' <- forall a. Run -> Simulation a -> IO a
invokeSimulation Run
r (forall e. Dynamics e -> Simulation (Dynamics e)
memo0Dynamics Dynamics (a, b)
m)
     let ma :: Dynamics a
ma =
           forall a. (Point -> IO a) -> Dynamics a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point
p ->
           do (a
a, b
_) <- forall a. Point -> Dynamics a -> IO a
invokeDynamics Point
p Dynamics (a, b)
m'
              forall (m :: * -> *) a. Monad m => a -> m a
return a
a
         mb :: Dynamics b
mb =
           forall a. (Point -> IO a) -> Dynamics a
Dynamics forall a b. (a -> b) -> a -> b
$ \Point
p ->
           do (a
_, b
b) <- forall a. Point -> Dynamics a -> IO a
invokeDynamics Point
p Dynamics (a, b)
m'
              forall (m :: * -> *) a. Monad m => a -> m a
return b
b
     forall (m :: * -> *) a. Monad m => a -> m a
return (Dynamics a
ma, Dynamics b
mb)