{-# LANGUAGE RecursiveDo #-}

-- |
-- Module     : Simulation.Aivika.Lattice.Internal.LIO
-- Copyright  : Copyright (c) 2016-2017, David Sorokin <david.sorokin@gmail.com>
-- License    : BSD3
-- Maintainer : David Sorokin <david.sorokin@gmail.com>
-- Stability  : experimental
-- Tested with: GHC 7.10.3
--
-- This module defines the 'LIO' computation.
--
module Simulation.Aivika.Lattice.Internal.LIO
       (LIOParams(..),
        LIO(..),
        LIOLattice(..),
        lattice,
        newRandomLattice,
        newRandomLatticeWithProb,
        invokeLIO,
        runLIO,
        lioParams,
        rootLIOParams,
        parentLIOParams,
        upSideLIOParams,
        downSideLIOParams,
        shiftLIOParams,
        lioParamsAt,
        latticeTimeIndex,
        latticeMemberIndex,
        latticeParentMemberIndex,
        latticeTime,
        latticeTimes,
        latticeTimeStep,
        latticePoint,
        latticeSize,
        findLatticeTimeIndex) where

import Data.IORef
import Data.Maybe

import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import Control.Monad.Fix
import Control.Exception (throw, catch, finally)

import Simulation.Aivika.Trans.Exception
import Simulation.Aivika.Trans.Internal.Types
import Simulation.Aivika.Trans.Internal.Specs
import Simulation.Aivika.Trans.Parameter

import Simulation.Aivika.Lattice.Internal.Lattice

-- | The 'LIO' computation that can be run as nested one on the lattice node.
newtype LIO a = LIO { forall a. LIO a -> LIOParams -> IO a
unLIO :: LIOParams -> IO a
                      -- ^ Unwrap the computation.
                    }

-- | The parameters of the 'LIO' computation.
data LIOParams =
  LIOParams { LIOParams -> LIOLattice
lioLattice :: LIOLattice,
              -- ^ The lattice.
              LIOParams -> Int
lioTimeIndex :: !Int,
              -- ^ The time index.
              LIOParams -> Int
lioMemberIndex :: !Int
              -- ^ The member index.
            }

instance Monad LIO where

  {-# INLINE (>>=) #-}
  (LIO LIOParams -> IO a
m) >>= :: forall a b. LIO a -> (a -> LIO b) -> LIO b
>>= a -> LIO b
k = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
    LIOParams -> IO a
m LIOParams
ps forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a ->
    let m' :: LIOParams -> IO b
m' = forall a. LIO a -> LIOParams -> IO a
unLIO (a -> LIO b
k a
a) in LIOParams -> IO b
m' LIOParams
ps

instance Applicative LIO where

  {-# INLINE pure #-}
  pure :: forall a. a -> LIO a
pure = forall a. (LIOParams -> IO a) -> LIO a
LIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure

  {-# INLINE (<*>) #-}
  <*> :: forall a b. LIO (a -> b) -> LIO a -> LIO b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Functor LIO where

  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> LIO a -> LIO b
fmap a -> b
f (LIO LIOParams -> IO a
m) = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOParams -> IO a
m 

instance MonadIO LIO where

  {-# INLINE liftIO #-}
  liftIO :: forall a. IO a -> LIO a
liftIO = forall a. (LIOParams -> IO a) -> LIO a
LIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

instance MonadFix LIO where

  mfix :: forall a. (a -> LIO a) -> LIO a
mfix a -> LIO a
f = 
    forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
    do { rec { a
a <- forall a. LIOParams -> LIO a -> IO a
invokeLIO LIOParams
ps (a -> LIO a
f a
a) }; forall (m :: * -> *) a. Monad m => a -> m a
return a
a }

instance MonadException LIO where

  catchComp :: forall e a. Exception e => LIO a -> (e -> LIO a) -> LIO a
catchComp (LIO LIOParams -> IO a
m) e -> LIO a
h = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
    forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (LIOParams -> IO a
m LIOParams
ps) (\e
e -> forall a. LIO a -> LIOParams -> IO a
unLIO (e -> LIO a
h e
e) LIOParams
ps)

  finallyComp :: forall a b. LIO a -> LIO b -> LIO a
finallyComp (LIO LIOParams -> IO a
m1) (LIO LIOParams -> IO b
m2) = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
    forall a b. IO a -> IO b -> IO a
finally (LIOParams -> IO a
m1 LIOParams
ps) (LIOParams -> IO b
m2 LIOParams
ps)
  
  throwComp :: forall e a. Exception e => e -> LIO a
throwComp e
e = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
    forall a e. Exception e => e -> a
throw e
e

-- | Invoke the computation.
invokeLIO :: LIOParams -> LIO a -> IO a
{-# INLINE invokeLIO #-}
invokeLIO :: forall a. LIOParams -> LIO a -> IO a
invokeLIO LIOParams
ps (LIO LIOParams -> IO a
m) = LIOParams -> IO a
m LIOParams
ps

-- | Run the 'LIO' computation using the specified lattice.
runLIO :: LIOLattice -> LIO a -> IO a
runLIO :: forall a. LIOLattice -> LIO a -> IO a
runLIO LIOLattice
lattice LIO a
m = forall a. LIO a -> LIOParams -> IO a
unLIO LIO a
m forall a b. (a -> b) -> a -> b
$ LIOLattice -> LIOParams
rootLIOParams LIOLattice
lattice

-- | Return the parameters of the computation.
lioParams :: LIO LIOParams
lioParams :: LIO LIOParams
lioParams = forall a. (LIOParams -> IO a) -> LIO a
LIO forall (m :: * -> *) a. Monad m => a -> m a
return

-- | Return the root node parameters.
rootLIOParams :: LIOLattice -> LIOParams
rootLIOParams :: LIOLattice -> LIOParams
rootLIOParams LIOLattice
lattice =
  LIOParams { lioLattice :: LIOLattice
lioLattice = LIOLattice
lattice,
              lioTimeIndex :: Int
lioTimeIndex = Int
0,
              lioMemberIndex :: Int
lioMemberIndex = Int
0 }

-- | Return the parent parameters.
parentLIOParams :: LIOParams -> Maybe LIOParams
parentLIOParams :: LIOParams -> Maybe LIOParams
parentLIOParams LIOParams
ps
  | Int
i forall a. Eq a => a -> a -> Bool
== Int
0    = forall a. Maybe a
Nothing
  | Bool
otherwise = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ LIOParams
ps { lioTimeIndex :: Int
lioTimeIndex = Int
i forall a. Num a => a -> a -> a
- Int
1, lioMemberIndex :: Int
lioMemberIndex = Int
k' }
  where i :: Int
i  = LIOParams -> Int
lioTimeIndex LIOParams
ps
        k :: Int
k  = LIOParams -> Int
lioMemberIndex LIOParams
ps
        k' :: Int
k' = LIOLattice -> Int -> Int -> Int
lioParentMemberIndex (LIOParams -> LIOLattice
lioLattice LIOParams
ps) Int
i Int
k

-- | Return the next up side parameters.
upSideLIOParams :: LIOParams -> LIOParams
upSideLIOParams :: LIOParams -> LIOParams
upSideLIOParams LIOParams
ps = LIOParams
ps { lioTimeIndex :: Int
lioTimeIndex = Int
1 forall a. Num a => a -> a -> a
+ Int
i }
  where i :: Int
i = LIOParams -> Int
lioTimeIndex LIOParams
ps

-- | Return the next down side parameters.
downSideLIOParams :: LIOParams -> LIOParams
downSideLIOParams :: LIOParams -> LIOParams
downSideLIOParams LIOParams
ps = LIOParams
ps { lioTimeIndex :: Int
lioTimeIndex = Int
1 forall a. Num a => a -> a -> a
+ Int
i, lioMemberIndex :: Int
lioMemberIndex = Int
1 forall a. Num a => a -> a -> a
+ Int
k }
  where i :: Int
i = LIOParams -> Int
lioTimeIndex LIOParams
ps
        k :: Int
k = LIOParams -> Int
lioMemberIndex LIOParams
ps

-- | Return the derived parameters with the specified shift in 'latticeTimeIndex' and
-- 'latticeMemberIndex' respectively, where the first parameter can be positive only.
shiftLIOParams :: Int
                  -- ^ a positive shift the lattice time index
                  -> Int
                  -- ^ a shift of the lattice member index
                  -> LIOParams
                  -- ^ the source parameters
                  -> LIOParams
shiftLIOParams :: Int -> Int -> LIOParams -> LIOParams
shiftLIOParams Int
di Int
dk LIOParams
ps
  | Int
i' forall a. Ord a => a -> a -> Bool
< Int
0    = forall a. HasCallStack => [Char] -> a
error [Char]
"The time index cannot be negative: shiftLIOParams"
  | Int
k' forall a. Ord a => a -> a -> Bool
< Int
0    = forall a. HasCallStack => [Char] -> a
error [Char]
"The member index cannot be negative: shiftLIOParams"
  | Int
k' forall a. Ord a => a -> a -> Bool
> Int
i'   = forall a. HasCallStack => [Char] -> a
error [Char]
"The member index cannot be greater than the time index: shiftLIOParams"
  | Bool
otherwise = LIOParams
ps { lioTimeIndex :: Int
lioTimeIndex = Int
i', lioMemberIndex :: Int
lioMemberIndex = Int
k' }
  where i :: Int
i  = LIOParams -> Int
lioTimeIndex LIOParams
ps
        i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
+ Int
di
        k :: Int
k  = LIOParams -> Int
lioMemberIndex LIOParams
ps
        k' :: Int
k' = Int
k forall a. Num a => a -> a -> a
+ Int
dk

-- | Return the parameters at the specified 'latticeTimeIndex' and 'latticeMemberIndex'.
lioParamsAt :: Int
               -- ^ the lattice time index
               -> Int
               -- ^ the lattice member index
               -> LIOParams
               -- ^ the source parameters
               -> LIOParams
lioParamsAt :: Int -> Int -> LIOParams -> LIOParams
lioParamsAt Int
i Int
k LIOParams
ps
  | Int
i forall a. Ord a => a -> a -> Bool
< Int
0     = forall a. HasCallStack => [Char] -> a
error [Char]
"The time index cannot be negative: lioParamsAt"
  | Int
k forall a. Ord a => a -> a -> Bool
< Int
0     = forall a. HasCallStack => [Char] -> a
error [Char]
"The member index cannot be negative: lioParamsAt"
  | Int
k forall a. Ord a => a -> a -> Bool
> Int
i     = forall a. HasCallStack => [Char] -> a
error [Char]
"The member index cannot be greater than the time index: lioParamsAt"
  | Bool
otherwise = LIOParams
ps { lioTimeIndex :: Int
lioTimeIndex = Int
i, lioMemberIndex :: Int
lioMemberIndex = Int
k }

-- | Return the lattice time index starting from 0. The index should be less than or equaled to 'latticeSize'. 
latticeTimeIndex :: LIO Int
latticeTimeIndex :: LIO Int
latticeTimeIndex = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOParams -> Int
lioTimeIndex

-- | Return the lattice member index starting from 0. It is always less than or equaled to 'latticeTimeIndex'.
latticeMemberIndex :: LIO Int
latticeMemberIndex :: LIO Int
latticeMemberIndex = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOParams -> Int
lioMemberIndex

-- | Return the parent member index starting from 0 for non-root lattice nodes.
latticeParentMemberIndex :: LIO (Maybe Int)
latticeParentMemberIndex :: LIO (Maybe Int)
latticeParentMemberIndex = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LIOParams -> Int
lioMemberIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOParams -> Maybe LIOParams
parentLIOParams

-- | Return the time for the current lattice node.
latticeTime :: Parameter LIO Double
latticeTime :: Parameter LIO Double
latticeTime =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
  let i :: Int
i = LIOParams -> Int
lioTimeIndex LIOParams
ps
  in forall a. LIOParams -> LIO a -> IO a
invokeLIO LIOParams
ps forall a b. (a -> b) -> a -> b
$
     forall (m :: * -> *) a. Run m -> Parameter m a -> m a
invokeParameter Run LIO
r forall a b. (a -> b) -> a -> b
$
     Int -> Parameter LIO Double
getLatticeTimeByIndex Int
i

-- | Return the time values in the lattice nodes.
latticeTimes :: Parameter LIO [Double]
latticeTimes :: Parameter LIO [Double]
latticeTimes =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
  let m :: Int
m  = LIOLattice -> Int
lioSize forall a b. (a -> b) -> a -> b
$ LIOParams -> LIOLattice
lioLattice LIOParams
ps
  in forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int
m] forall a b. (a -> b) -> a -> b
$ \Int
i ->
     forall a. LIOParams -> LIO a -> IO a
invokeLIO LIOParams
ps forall a b. (a -> b) -> a -> b
$
     forall (m :: * -> *) a. Run m -> Parameter m a -> m a
invokeParameter Run LIO
r forall a b. (a -> b) -> a -> b
$
     Int -> Parameter LIO Double
getLatticeTimeByIndex Int
i

-- | Return the point in the corresponding lattice node.
latticePoint :: Parameter LIO (Point LIO)
latticePoint :: Parameter LIO (Point LIO)
latticePoint =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  do Double
t <- forall (m :: * -> *) a. Run m -> Parameter m a -> m a
invokeParameter Run LIO
r Parameter LIO Double
latticeTime
     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Run m -> Double -> Int -> Point m
pointAt Run LIO
r Double
t Int
0

-- | Return the lattice time step.
latticeTimeStep :: Parameter LIO Double
latticeTimeStep :: Parameter LIO Double
latticeTimeStep =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
  do let sc :: Specs LIO
sc = forall (m :: * -> *). Run m -> Specs m
runSpecs Run LIO
r
         t0 :: Double
t0 = forall (m :: * -> *). Specs m -> Double
spcStartTime Specs LIO
sc
         t2 :: Double
t2 = forall (m :: * -> *). Specs m -> Double
spcStopTime Specs LIO
sc
         m :: Int
m  = LIOLattice -> Int
lioSize forall a b. (a -> b) -> a -> b
$ LIOParams -> LIOLattice
lioLattice LIOParams
ps
         dt :: Double
dt = (Double
t2 forall a. Num a => a -> a -> a
- Double
t0) forall a. Fractional a => a -> a -> a
/ (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m)
     forall (m :: * -> *) a. Monad m => a -> m a
return Double
dt

-- | Return the lattice size.
latticeSize :: LIO Int
latticeSize :: LIO Int
latticeSize = forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOLattice -> Int
lioSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. LIOParams -> LIOLattice
lioLattice

-- | Find the lattice time index by the specified modeling time.
findLatticeTimeIndex :: Double -> Parameter LIO Int
findLatticeTimeIndex :: Double -> Parameter LIO Int
findLatticeTimeIndex Double
t =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
  do let sc :: Specs LIO
sc = forall (m :: * -> *). Run m -> Specs m
runSpecs Run LIO
r
         t0 :: Double
t0 = forall (m :: * -> *). Specs m -> Double
spcStartTime Specs LIO
sc
         t2 :: Double
t2 = forall (m :: * -> *). Specs m -> Double
spcStopTime Specs LIO
sc
         m :: Int
m  = LIOLattice -> Int
lioSize forall a b. (a -> b) -> a -> b
$ LIOParams -> LIOLattice
lioLattice LIOParams
ps
         i :: Int
i | Double
t forall a. Eq a => a -> a -> Bool
== Double
t0   = Int
0
           | Double
t forall a. Eq a => a -> a -> Bool
== Double
t2   = Int
m
           | Bool
otherwise = forall a b. (RealFrac a, Integral b) => a -> b
floor (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m forall a. Num a => a -> a -> a
* ((Double
t forall a. Num a => a -> a -> a
- Double
t0) forall a. Fractional a => a -> a -> a
/ (Double
t2 forall a. Num a => a -> a -> a
- Double
t0)))
     forall (m :: * -> *) a. Monad m => a -> m a
return Int
i

-- | Get the modeling time in the lattice node by the specified time index. 
getLatticeTimeByIndex :: Int -> Parameter LIO Double
getLatticeTimeByIndex :: Int -> Parameter LIO Double
getLatticeTimeByIndex Int
i =
  forall (m :: * -> *) a. (Run m -> m a) -> Parameter m a
Parameter forall a b. (a -> b) -> a -> b
$ \Run LIO
r ->
  forall a. (LIOParams -> IO a) -> LIO a
LIO forall a b. (a -> b) -> a -> b
$ \LIOParams
ps ->
  let sc :: Specs LIO
sc = forall (m :: * -> *). Run m -> Specs m
runSpecs Run LIO
r
      t0 :: Double
t0 = forall (m :: * -> *). Specs m -> Double
spcStartTime Specs LIO
sc
      t2 :: Double
t2 = forall (m :: * -> *). Specs m -> Double
spcStopTime Specs LIO
sc
      m :: Int
m  = LIOLattice -> Int
lioSize forall a b. (a -> b) -> a -> b
$ LIOParams -> LIOLattice
lioLattice LIOParams
ps
      dt :: Double
dt = (Double
t2 forall a. Num a => a -> a -> a
- Double
t0) forall a. Fractional a => a -> a -> a
/ (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m)
      t :: Double
t | Int
i forall a. Eq a => a -> a -> Bool
== Int
0    = Double
t0
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
m    = Double
t2
        | Bool
otherwise = Double
t0 forall a. Num a => a -> a -> a
+ (forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int
i) forall a. Num a => a -> a -> a
* Double
dt
  in forall (m :: * -> *) a. Monad m => a -> m a
return Double
t