module Simulation.Aivika.Trans.SystemDynamics
       (
        (.==.),
        (./=.),
        (.<.),
        (.>=.),
        (.>.),
        (.<=.),
        maxDynamics,
        minDynamics,
        ifDynamics,
        
        integ,
        integEither,
        smoothI,
        smooth,
        smooth3I,
        smooth3,
        smoothNI,
        smoothN,
        delay1I,
        delay1,
        delay3I,
        delay3,
        delayNI,
        delayN,
        forecast,
        trend,
        
        diffsum,
        diffsumEither,
        
        lookupDynamics,
        lookupStepwiseDynamics,
        
        delay,
        delayI,
        delayByDT,
        delayIByDT,
        step,
        pulse,
        pulseP,
        ramp,
        
        npv,
        npve) where
import Data.Array
import Control.Monad
import Control.Monad.Trans
import Control.Monad.Fix
import Simulation.Aivika.Trans.Internal.Specs
import Simulation.Aivika.Trans.Internal.Parameter
import Simulation.Aivika.Trans.Internal.Simulation
import Simulation.Aivika.Trans.Internal.Dynamics
import Simulation.Aivika.Trans.Dynamics.Extra
import Simulation.Aivika.Trans.Table
import Simulation.Aivika.Trans.SD
import qualified Simulation.Aivika.Trans.Dynamics.Memo as M
import qualified Simulation.Aivika.Trans.Dynamics.Memo.Unboxed as MU
(.==.) :: (Monad m, Eq a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(.==.) = liftM2 (==)
(./=.) :: (Monad m, Eq a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(./=.) = liftM2 (/=)
(.<.) :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(.<.) = liftM2 (<)
(.>=.) :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(.>=.) = liftM2 (>=)
(.>.) :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(.>.) = liftM2 (>)
(.<=.) :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m Bool
(.<=.) = liftM2 (<=)
maxDynamics :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m a
maxDynamics = liftM2 max
minDynamics :: (Monad m, Ord a) => Dynamics m a -> Dynamics m a -> Dynamics m a
minDynamics = liftM2 min
ifDynamics :: Monad m => Dynamics m Bool -> Dynamics m a -> Dynamics m a -> Dynamics m a
ifDynamics cond x y =
  do a <- cond
     if a then x else y
integEuler :: Monad m
              => Dynamics m Double
              -> Dynamics m Double 
              -> Dynamics m Double 
              -> Point m
              -> m Double
integEuler (Dynamics f) (Dynamics i) (Dynamics y) p = 
  case pointIteration p of
    0 -> 
      i p
    n -> do 
      let sc = pointSpecs p
          ty = basicTime sc (n  1) 0
          py = p { pointTime = ty, pointIteration = n  1, pointPhase = 0 }
      a <- y py
      b <- f py
      let !v = a + spcDT (pointSpecs p) * b
      return v
integRK2 :: Monad m
            => Dynamics m Double
            -> Dynamics m Double
            -> Dynamics m Double
            -> Point m
            -> m Double
integRK2 (Dynamics f) (Dynamics i) (Dynamics y) p =
  case pointPhase p of
    0 -> case pointIteration p of
      0 ->
        i p
      n -> do
        let sc = pointSpecs p
            ty = basicTime sc (n  1) 0
            t1 = ty
            t2 = basicTime sc (n  1) 1
            py = p { pointTime = ty, pointIteration = n  1, pointPhase = 0 }
            p1 = py
            p2 = p { pointTime = t2, pointIteration = n  1, pointPhase = 1 }
        vy <- y py
        k1 <- f p1
        k2 <- f p2
        let !v = vy + spcDT sc / 2.0 * (k1 + k2)
        return v
    1 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t1 = ty
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p1 = py
      vy <- y py
      k1 <- f p1
      let !v = vy + spcDT sc * k1
      return v
    _ -> 
      error "Incorrect phase: integRK2"
integRK4 :: Monad m
            => Dynamics m Double
            -> Dynamics m Double
            -> Dynamics m Double
            -> Point m
            -> m Double
integRK4 (Dynamics f) (Dynamics i) (Dynamics y) p =
  case pointPhase p of
    0 -> case pointIteration p of
      0 -> 
        i p
      n -> do
        let sc = pointSpecs p
            ty = basicTime sc (n  1) 0
            t1 = ty
            t2 = basicTime sc (n  1) 1
            t3 = basicTime sc (n  1) 2
            t4 = basicTime sc (n  1) 3
            py = p { pointTime = ty, pointIteration = n  1, pointPhase = 0 }
            p1 = py
            p2 = p { pointTime = t2, pointIteration = n  1, pointPhase = 1 }
            p3 = p { pointTime = t3, pointIteration = n  1, pointPhase = 2 }
            p4 = p { pointTime = t4, pointIteration = n  1, pointPhase = 3 }
        vy <- y py
        k1 <- f p1
        k2 <- f p2
        k3 <- f p3
        k4 <- f p4
        let !v = vy + spcDT sc / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
        return v
    1 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t1 = ty
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p1 = py
      vy <- y py
      k1 <- f p1
      let !v = vy + spcDT sc / 2.0 * k1
      return v
    2 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t2 = basicTime sc n 1
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p2 = p { pointTime = t2, pointIteration = n, pointPhase = 1 }
      vy <- y py
      k2 <- f p2
      let !v = vy + spcDT sc / 2.0 * k2
      return v
    3 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t3 = basicTime sc n 2
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p3 = p { pointTime = t3, pointIteration = n, pointPhase = 2 }
      vy <- y py
      k3 <- f p3
      let !v = vy + spcDT sc * k3
      return v
    _ -> 
      error "Incorrect phase: integRK4"
integRK4b :: Monad m
             => Dynamics m Double
             -> Dynamics m Double
             -> Dynamics m Double
             -> Point m
             -> m Double
integRK4b (Dynamics f) (Dynamics i) (Dynamics y) p =
  case pointPhase p of
    0 -> case pointIteration p of
      0 -> 
        i p
      n -> do
        let sc = pointSpecs p
            ty = basicTime sc (n  1) 0
            t1 = ty
            t2 = basicTime sc (n  1) 1
            t3 = basicTime sc (n  1) 2
            t4 = basicTime sc (n  1) 3
            py = p { pointTime = ty, pointIteration = n  1, pointPhase = 0 }
            p1 = py
            p2 = p { pointTime = t2, pointIteration = n  1, pointPhase = 1 }
            p3 = p { pointTime = t3, pointIteration = n  1, pointPhase = 2 }
            p4 = p { pointTime = t4, pointIteration = n  1, pointPhase = 3 }
        vy <- y py
        k1 <- f p1
        k2 <- f p2
        k3 <- f p3
        k4 <- f p4
        let !v = vy + spcDT sc / 8.0 * (k1 + 3.0 * (k2 + k3) + k4)
        return v
    1 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t1 = ty
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p1 = py
      vy <- y py
      k1 <- f p1
      let !v = vy + spcDT sc / 3.0 * k1
      return v
    2 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t1 = ty
          t2 = basicTime sc n 1
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p1 = py
          p2 = p { pointTime = t2, pointIteration = n, pointPhase = 1 }
      vy <- y py
      k1 <- f p1
      k2 <- f p2
      let !v = vy + spcDT sc * ( k1 / 3.0 + k2)
      return v
    3 -> do
      let sc = pointSpecs p
          n  = pointIteration p
          ty = basicTime sc n 0
          t1 = ty
          t2 = basicTime sc n 1
          t3 = basicTime sc n 2
          py = p { pointTime = ty, pointIteration = n, pointPhase = 0 }
          p1 = py
          p2 = p { pointTime = t2, pointIteration = n, pointPhase = 1 }
          p3 = p { pointTime = t3, pointIteration = n, pointPhase = 2 }
      vy <- y py
      k1 <- f p1
      k2 <- f p2
      k3 <- f p3
      let !v = vy + spcDT sc * (k1  k2 + k3)
      return v
    _ -> 
      error "Incorrect phase: integRK4b"
integ :: (MonadSD m, MonadFix m)
         => Dynamics m Double                  
         -> Dynamics m Double                  
         -> Simulation m (Dynamics m Double)   
integ diff i =
  mdo y <- MU.memoDynamics z
      z <- Simulation $ \r ->
        case spcMethod (runSpecs r) of
          Euler -> return $ Dynamics $ integEuler diff i y
          RungeKutta2 -> return $ Dynamics $ integRK2 diff i y
          RungeKutta4 -> return $ Dynamics $ integRK4 diff i y
          RungeKutta4b -> return $ Dynamics $ integRK4b diff i y
      return y
integEulerEither :: Monad m
                    => Dynamics m (Either Double Double)
                    -> Dynamics m Double 
                    -> Dynamics m Double 
                    -> Point m
                    -> m Double
integEulerEither (Dynamics f) (Dynamics i) (Dynamics y) p = 
  case pointIteration p of
    0 -> 
      i p
    n -> do 
      let sc = pointSpecs p
          ty = basicTime sc (n  1) 0
          py = p { pointTime = ty, pointIteration = n  1, pointPhase = 0 }
      b <- f py
      case b of
        Left v ->
          return v
        Right b -> do
          a <- y py
          let !v = a + spcDT (pointSpecs p) * b
          return v
integEither :: (MonadSD m, MonadFix m)
               => Dynamics m (Either Double Double)
               
               -> Dynamics m Double
               
               -> Simulation m (Dynamics m Double)
integEither diff i =
  mdo y <- MU.memoDynamics z
      z <- Simulation $ \r ->
        return $ Dynamics $ integEulerEither diff i y
      return y
smoothI :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Dynamics m Double                  
           -> Simulation m (Dynamics m Double)   
smoothI x t i =
  mdo y <- integ ((x  y) / t) i
      return y
smooth :: (MonadSD m, MonadFix m)
          => Dynamics m Double                  
          -> Dynamics m Double                  
          -> Simulation m (Dynamics m Double)   
smooth x t = smoothI x t x
smooth3I :: (MonadSD m, MonadFix m)
            => Dynamics m Double                  
            -> Dynamics m Double                  
            -> Dynamics m Double                  
            -> Simulation m (Dynamics m Double)   
smooth3I x t i =
  mdo y  <- integ ((s2  y) / t') i
      s2 <- integ ((s1  s2) / t') i
      s1 <- integ ((x  s1) / t') i
      let t' = t / 3.0
      return y
smooth3 :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Simulation m (Dynamics m Double)   
smooth3 x t = smooth3I x t x
smoothNI :: (MonadSD m, MonadFix m)
            => Dynamics m Double                  
            -> Dynamics m Double                  
            -> Int                                
            -> Dynamics m Double                  
            -> Simulation m (Dynamics m Double)   
smoothNI x t n i =
  mdo s <- forM [1 .. n] $ \k ->
        if k == 1
        then integ ((x  a ! 1) / t') i
        else integ ((a ! (k  1)  a ! k) / t') i
      let a  = listArray (1, n) s 
          t' = t / fromIntegral n
      return $ a ! n
smoothN :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Int                                
           -> Simulation m (Dynamics m Double)   
smoothN x t n = smoothNI x t n x
delay1I :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Dynamics m Double                  
           -> Simulation m (Dynamics m Double)   
delay1I x t i =
  mdo y <- integ (x  y / t) (i * t)
      return $ y / t
delay1 :: (MonadSD m, MonadFix m)
          => Dynamics m Double                  
          -> Dynamics m Double                  
          -> Simulation m (Dynamics m Double)   
delay1 x t = delay1I x t x
delay3I :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Dynamics m Double                  
           -> Simulation m (Dynamics m Double)   
delay3I x t i =
  mdo y  <- integ (s2 / t'  y / t') (i * t')
      s2 <- integ (s1 / t'  s2 / t') (i * t')
      s1 <- integ (x  s1 / t') (i * t')
      let t' = t / 3.0
      return $ y / t'         
delay3 :: (MonadSD m, MonadFix m)
          => Dynamics m Double                  
          -> Dynamics m Double                  
          -> Simulation m (Dynamics m Double)   
delay3 x t = delay3I x t x
delayNI :: (MonadSD m, MonadFix m)
           => Dynamics m Double                  
           -> Dynamics m Double                  
           -> Int                                
           -> Dynamics m Double                  
           -> Simulation m (Dynamics m Double)   
delayNI x t n i =
  mdo s <- forM [1 .. n] $ \k ->
        if k == 1
        then integ (x  (a ! 1) / t') (i * t')
        else integ ((a ! (k  1)) / t'  (a ! k) / t') (i * t')
      let a  = listArray (1, n) s
          t' = t / fromIntegral n
      return $ (a ! n) / t'
delayN :: (MonadSD m, MonadFix m)
          => Dynamics m Double                  
          -> Dynamics m Double                  
          -> Int                                
          -> Simulation m (Dynamics m Double)   
delayN x t n = delayNI x t n x
forecast :: (MonadSD m, MonadFix m)
            => Dynamics m Double                  
            -> Dynamics m Double                  
            -> Dynamics m Double                  
            -> Simulation m (Dynamics m Double)   
forecast x at hz =
  do y <- smooth x at
     return $ x * (1.0 + (x / y  1.0) / at * hz)
trend :: (MonadSD m, MonadFix m)
         => Dynamics m Double                  
         -> Dynamics m Double                  
         -> Dynamics m Double                  
         -> Simulation m (Dynamics m Double)   
trend x at i =
  do y <- smoothI x at (x / (1.0 + i * at))
     return $ (x / y  1.0) / at
diffsum :: (MonadSD m, MonadFix m,
            MU.MonadMemo m a, Num a)
           => Dynamics m a                  
           -> Dynamics m a                  
           -> Simulation m (Dynamics m a)   
diffsum (Dynamics diff) (Dynamics i) =
  mdo y <-
        MU.memo0Dynamics $
        Dynamics $ \p ->
        case pointIteration p of
          0 -> i p
          n -> do 
            let Dynamics m = y
                sc = pointSpecs p
                ty = basicTime sc (n  1) 0
                py = p { pointTime = ty, 
                         pointIteration = n  1, 
                         pointPhase = 0 }
            a <- m py
            b <- diff py
            let !v = a + b
            return v
      return y
diffsumEither :: (MonadSD m, MonadFix m,
                  MU.MonadMemo m a, Num a)
                 => Dynamics m (Either a a)
                 
                 -> Dynamics m a
                 
                 -> Simulation m (Dynamics m a)
                 
diffsumEither (Dynamics diff) (Dynamics i) =
  mdo y <-
        MU.memo0Dynamics $
        Dynamics $ \p ->
        case pointIteration p of
          0 -> i p
          n -> do 
            let Dynamics m = y
                sc = pointSpecs p
                ty = basicTime sc (n  1) 0
                py = p { pointTime = ty, 
                         pointIteration = n  1, 
                         pointPhase = 0 }
            b <- diff py
            case b of
              Left v ->
                return v
              Right b -> do
                a <- m py
                let !v = a + b
                return v
      return y
lookupDynamics :: Monad m => Dynamics m Double -> Array Int (Double, Double) -> Dynamics m Double
lookupDynamics (Dynamics m) tbl =
  Dynamics $ \p ->
  do a <- m p
     return $ tableLookup a tbl
lookupStepwiseDynamics :: Monad m => Dynamics m Double -> Array Int (Double, Double) -> Dynamics m Double
lookupStepwiseDynamics (Dynamics m) tbl =
  Dynamics $ \p ->
  do a <- m p
     return $ tableLookupStepwise a tbl
delay :: Monad m
         => Dynamics m a          
         -> Dynamics m Double     
         -> Dynamics m a          
delay (Dynamics x) (Dynamics d) = discreteDynamics $ Dynamics r 
  where
    r p = do 
      let t  = pointTime p
          sc = pointSpecs p
          n  = pointIteration p
      a <- d p
      let t' = t  a
          n' = fromIntegral $ floor $ (t'  spcStartTime sc) / spcDT sc
          y | n' < 0    = x $ p { pointTime = spcStartTime sc,
                                  pointIteration = 0, 
                                  pointPhase = 0 }
            | n' < n    = x $ p { pointTime = t',
                                  pointIteration = n',
                                  pointPhase = 1 }
            | n' > n    = error $
                          "Cannot return the future data: delay. " ++
                          "The lag time cannot be negative."
            | otherwise = error $
                          "Cannot return the current data: delay. " ++
                          "The lag time is too small."
      y
delayI :: MonadSD m
          => Dynamics m a                    
          -> Dynamics m Double               
          -> Dynamics m a                    
          -> Simulation m (Dynamics m a)     
delayI (Dynamics x) (Dynamics d) (Dynamics i) = M.memo0Dynamics $ Dynamics r 
  where
    r p = do 
      let t  = pointTime p
          sc = pointSpecs p
          n  = pointIteration p
      a <- d p
      let t' = t  a
          n' = fromIntegral $ floor $ (t'  spcStartTime sc) / spcDT sc
          y | n' < 0    = i $ p { pointTime = spcStartTime sc,
                                  pointIteration = 0, 
                                  pointPhase = 0 }
            | n' < n    = x $ p { pointTime = t',
                                  pointIteration = n',
                                  pointPhase = 1 }
            | n' > n    = error $
                          "Cannot return the future data: delay. " ++
                          "The lag time cannot be negative."
            | otherwise = error $
                          "Cannot return the current data: delay. " ++
                          "The lag time is too small."
      y
delayByDT :: Monad m
             => Dynamics m a
             
             -> Dynamics m Int
             
             
             -> Dynamics m a
             
delayByDT (Dynamics x) (Dynamics d) = discreteDynamics $ Dynamics r 
  where
    r p = do 
      let sc = pointSpecs p
          n  = pointIteration p
      a <- d p
      let p' = delayPoint p a
          n' = pointIteration p'
          y | n' < 0    = x $ p { pointTime = spcStartTime sc,
                                  pointIteration = 0, 
                                  pointPhase = 0 }
            | n' < n    = x p'
            | n' > n    = error $
                          "Cannot return the future data: delayByDT. " ++
                          "The lag time cannot be negative."
            | otherwise = error $
                          "Cannot return the current data: delayByDT. " ++
                          "The lag time is too small."
      y
      
delayIByDT :: MonadSD m
              => Dynamics m a
              
              -> Dynamics m Int
              
              
              -> Dynamics m a
              
              -> Simulation m (Dynamics m a)
              
delayIByDT (Dynamics x) (Dynamics d) (Dynamics i) = M.memoDynamics $ Dynamics r 
  where
    r p = do 
      let sc = pointSpecs p
          n  = pointIteration p
      a <- d p
      let p' = delayPoint p a
          n' = pointIteration p'
          y | n' < 0    = i $ p { pointTime = spcStartTime sc,
                                  pointIteration = 0, 
                                  pointPhase = 0 }
            | n' < n    = x p'
            | n' > n    = error $
                          "Cannot return the future data: delayIByDT. " ++
                          "The lag time cannot be negative."
            | otherwise = error $
                          "Cannot return the current data: delayIByDT. " ++
                          "The lag time is too small."
      y
npv :: (MonadSD m, MonadFix m)
       => Dynamics m Double                  
       -> Dynamics m Double                  
       -> Dynamics m Double                  
       -> Dynamics m Double                  
       -> Simulation m (Dynamics m Double)   
npv stream rate init factor =
  mdo let dt' = liftParameter dt
      df <- integ ( df * rate) 1
      accum <- integ (stream * df) init
      return $ (accum + dt' * stream * df) * factor
npve :: (MonadSD m, MonadFix m)
        => Dynamics m Double                  
        -> Dynamics m Double                  
        -> Dynamics m Double                  
        -> Dynamics m Double                  
        -> Simulation m (Dynamics m Double)   
npve stream rate init factor =
  mdo let dt' = liftParameter dt
      df <- integ ( df * rate / (1 + rate * dt')) (1 / (1 + rate * dt'))
      accum <- integ (stream * df) init
      return $ (accum + dt' * stream * df) * factor
step :: Monad m
        => Dynamics m Double
        
        -> Dynamics m Double
        
        -> Dynamics m Double
step h st =
  discreteDynamics $
  Dynamics $ \p ->
  do let sc = pointSpecs p
         t  = pointTime p
     st' <- invokeDynamics p st
     let t' = t + spcDT sc / 2
     if st' < t'
       then invokeDynamics p h
       else return 0
pulse :: Monad m
         => Dynamics m Double
         
         -> Dynamics m Double
         
         -> Dynamics m Double
pulse st w =
  discreteDynamics $
  Dynamics $ \p ->
  do let sc = pointSpecs p
         t  = pointTime p
     st' <- invokeDynamics p st
     let t' = t + spcDT sc / 2
     if st' < t'
       then do w' <- invokeDynamics p w
               return $ if t' < st' + w' then 1 else 0
       else return 0
pulseP :: Monad m
          => Dynamics m Double
          
          -> Dynamics m Double
          
          -> Dynamics m Double
          
          -> Dynamics m Double
pulseP st w period =
  discreteDynamics $
  Dynamics $ \p ->
  do let sc = pointSpecs p
         t  = pointTime p
     p'  <- invokeDynamics p period
     st' <- invokeDynamics p st
     let y' = if (p' > 0) && (t > st')
              then fromIntegral (floor $ (t  st') / p') * p'
              else 0
     let st' = st' + y'
     let t' = t + spcDT sc / 2
     if st' < t'
       then do w' <- invokeDynamics p w
               return $ if t' < st' + w' then 1 else 0
       else return 0
ramp :: Monad m
        => Dynamics m Double
        
        -> Dynamics m Double
        
        -> Dynamics m Double
        
        -> Dynamics m Double
ramp slope st e =
  discreteDynamics $
  Dynamics $ \p ->
  do let sc = pointSpecs p
         t  = pointTime p
     st' <- invokeDynamics p st
     if st' < t
       then do slope' <- invokeDynamics p slope
               e' <- invokeDynamics p e
               if t < e'
                 then return $ slope' * (t  st')
                 else return $ slope' * (e'  st')
       else return 0