{-# LANGUAGE FlexibleContexts #-}

-- |
-- Module     : Simulation.Aivika.Dynamics.Memo.Unboxed
-- 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 the unboxed 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.Unboxed
       (memoDynamics,
        memo0Dynamics) 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
import Simulation.Aivika.Unboxed

-- | 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 :: Unboxed e => Dynamics e -> Simulation (Dynamics e)
{-# INLINE memoDynamics #-}
memoDynamics :: Dynamics e -> Simulation (Dynamics e)
memoDynamics (Dynamics Point -> IO e
m) = 
  (Run -> IO (Dynamics e)) -> Simulation (Dynamics e)
forall a. (Run -> IO a) -> Simulation a
Simulation ((Run -> IO (Dynamics e)) -> Simulation (Dynamics e))
-> (Run -> IO (Dynamics e)) -> Simulation (Dynamics e)
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
     IOUArray (Int, Int) e
arr   <- ((Int, Int), (Int, Int)) -> IO (IOUArray (Int, Int) e)
forall e i. (Unboxed e, Ix i) => (i, i) -> IO (IOUArray i e)
newUnboxedArray_ ((Int
phl, Int
nl), (Int
phu, Int
nu))
     IORef Int
nref  <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
     IORef Int
phref <- Int -> IO (IORef Int)
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' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n) Bool -> Bool -> Bool
|| ((Int
n' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n) Bool -> Bool -> Bool
&& (Int
ph' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ph)) 
                    then 
                      IOUArray (Int, Int) e -> (Int, Int) -> IO e
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray (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 e -> IO () -> IO ()
`seq` IOUArray (Int, Int) e -> (Int, Int) -> e -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray (Int, Int) e
arr (Int
ph', Int
n') e
a
                            if Int
ph' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
phu 
                              then do IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
phref Int
0
                                      IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
nref (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                                      Int -> Int -> IO e
loop (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0
                              else do IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
phref (Int
ph' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                                      Int -> Int -> IO e
loop Int
n' (Int
ph' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              Int
n'  <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
nref
              Int
ph' <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
phref
              Int -> Int -> IO e
loop Int
n' Int
ph'
     Dynamics e -> IO (Dynamics e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Dynamics e -> IO (Dynamics e)) -> Dynamics e -> IO (Dynamics e)
forall a b. (a -> b) -> a -> b
$ Dynamics e -> Dynamics e
forall a. Dynamics a -> Dynamics a
interpolateDynamics (Dynamics e -> Dynamics e) -> Dynamics e -> Dynamics e
forall a b. (a -> b) -> a -> b
$ (Point -> IO e) -> Dynamics e
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 :: Unboxed e => Dynamics e -> Simulation (Dynamics e)
{-# INLINE memo0Dynamics #-}
memo0Dynamics :: Dynamics e -> Simulation (Dynamics e)
memo0Dynamics (Dynamics Point -> IO e
m) = 
  (Run -> IO (Dynamics e)) -> Simulation (Dynamics e)
forall a. (Run -> IO a) -> Simulation a
Simulation ((Run -> IO (Dynamics e)) -> Simulation (Dynamics e))
-> (Run -> IO (Dynamics e)) -> Simulation (Dynamics e)
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
     IOUArray Int e
arr  <- (Int, Int) -> IO (IOUArray Int e)
forall e i. (Unboxed e, Ix i) => (i, i) -> IO (IOUArray i e)
newUnboxedArray_ (Int, Int)
bnds
     IORef Int
nref <- Int -> IO (IORef Int)
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' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n
                    then 
                      IOUArray Int e -> Int -> IO e
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray 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 e -> IO () -> IO ()
`seq` IOUArray Int e -> Int -> e -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int e
arr Int
n' e
a
                            IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
nref (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                            Int -> IO e
loop (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              Int
n' <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
nref
              Int -> IO e
loop Int
n'
     Dynamics e -> IO (Dynamics e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Dynamics e -> IO (Dynamics e)) -> Dynamics e -> IO (Dynamics e)
forall a b. (a -> b) -> a -> b
$ Dynamics e -> Dynamics e
forall a. Dynamics a -> Dynamics a
discreteDynamics (Dynamics e -> Dynamics e) -> Dynamics e -> Dynamics e
forall a b. (a -> b) -> a -> b
$ (Point -> IO e) -> Dynamics e
forall a. (Point -> IO a) -> Dynamics a
Dynamics Point -> IO e
r