{-# LANGUAGE TypeFamilies #-}

-- |
-- Module     : Simulation.Aivika.Trans.Internal.Specs
-- 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
--
-- It defines the simulation specs and related stuff.
module Simulation.Aivika.Trans.Internal.Specs
       (Specs(..),
        Method(..),
        Run(..),
        Point(..),
        basicTime,
        integIterationBnds,
        integIterationHiBnd,
        integIterationLoBnd,
        integPhaseBnds,
        integPhaseHiBnd,
        integPhaseLoBnd,
        integTimes,
        integPoints,
        integPointsStartingFrom,
        integStartPoint,
        integStopPoint,
        simulationStopPoint,
        timeGrid,
        pointAt,
        delayPoint) where

import Simulation.Aivika.Trans.Internal.Types

-- | Returns the integration iterations starting from zero.
integIterations :: Specs m -> [Int]
integIterations sc = [i1 .. i2] where
  i1 = integIterationLoBnd sc
  i2 = integIterationHiBnd sc

-- | Returns the first and last integration iterations.
integIterationBnds :: Specs m -> (Int, Int)
integIterationBnds sc = (i1, i2) where
  i1 = integIterationLoBnd sc
  i2 = integIterationHiBnd sc

-- | Returns the first integration iteration, i.e. zero.
integIterationLoBnd :: Specs m -> Int
integIterationLoBnd sc = 0

-- | Returns the last integration iteration.
integIterationHiBnd :: Specs m -> Int
integIterationHiBnd sc =
  let n = round ((spcStopTime sc -
                  spcStartTime sc) / spcDT sc)
  in if n < 0
     then
       error $
       "Either the simulation specs are incorrect, " ++
       "or a step time is too small, because of which " ++
       "a floating point overflow occurred on 32-bit Haskell implementation."
     else n

-- | Returns the phases for the specified simulation specs starting from zero.
integPhases :: Specs m -> [Int]
integPhases sc =
  case spcMethod sc of
    Euler -> [0]
    RungeKutta2 -> [0, 1]
    RungeKutta4 -> [0, 1, 2, 3]
    RungeKutta4b -> [0, 1, 2, 3]

-- | Returns the first and last integration phases.
integPhaseBnds :: Specs m -> (Int, Int)
integPhaseBnds sc =
  case spcMethod sc of
    Euler -> (0, 0)
    RungeKutta2 -> (0, 1)
    RungeKutta4 -> (0, 3)
    RungeKutta4b -> (0, 3)

-- | Returns the first integration phase, i.e. zero.
integPhaseLoBnd :: Specs m -> Int
integPhaseLoBnd sc = 0

-- | Returns the last integration phase, 0 for Euler's method, 1 for RK2 and 3 for RK4.
integPhaseHiBnd :: Specs m -> Int
integPhaseHiBnd sc =
  case spcMethod sc of
    Euler -> 0
    RungeKutta2 -> 1
    RungeKutta4 -> 3
    RungeKutta4b -> 3

-- | Returns a simulation time for the integration point specified by 
-- the specs, iteration and phase.
basicTime :: Specs m -> Int -> Int -> Double
{-# INLINE basicTime #-}
basicTime sc n ph =
  if ph < 0 then
    error "Incorrect phase: basicTime"
  else
    spcStartTime sc + n' * spcDT sc + delta (spcMethod sc) ph
      where n' = fromIntegral n
            delta Euler       0 = 0
            delta RungeKutta2 0 = 0
            delta RungeKutta2 1 = spcDT sc
            delta RungeKutta4 0 = 0
            delta RungeKutta4 1 = spcDT sc / 2
            delta RungeKutta4 2 = spcDT sc / 2
            delta RungeKutta4 3 = spcDT sc
            delta RungeKutta4b 0 = 0
            delta RungeKutta4b 1 = spcDT sc / 3
            delta RungeKutta4b 2 = 2 * spcDT sc / 3
            delta RungeKutta4b 3 = spcDT sc

-- | Return the integration time values.
integTimes :: Specs m -> [Double]
integTimes sc = map t [nl .. nu]
  where (nl, nu) = integIterationBnds sc
        t n = basicTime sc n 0

-- | Return the integration time points.
integPoints :: Run m -> [Point m]
integPoints r = points
  where sc = runSpecs r
        (nl, nu) = integIterationBnds sc
        points   = map point [nl .. nu]
        point n  = Point { pointSpecs = sc,
                           pointRun = r,
                           pointTime = basicTime sc n 0,
                           pointIteration = n,
                           pointPhase = 0 }

-- | Return the start time point.
integStartPoint :: Run m -> Point m
integStartPoint r = point nl
  where sc = runSpecs r
        (nl, nu) = integIterationBnds sc
        point n  = Point { pointSpecs = sc,
                           pointRun = r,
                           pointTime = basicTime sc n 0,
                           pointIteration = n,
                           pointPhase = 0 }

-- | Return the integration stop time point.
integStopPoint :: Run m -> Point m
integStopPoint r = point nu
  where sc = runSpecs r
        (nl, nu) = integIterationBnds sc
        point n  = Point { pointSpecs = sc,
                           pointRun = r,
                           pointTime = basicTime sc n 0,
                           pointIteration = n,
                           pointPhase = 0 }

-- | Return the simulation stop time point.
simulationStopPoint :: Run m -> Point m
simulationStopPoint r = pointAt r (spcStopTime $ runSpecs r)

-- | Return the point at the specified time.
pointAt :: Run m -> Double -> Point m
{-# INLINABLE pointAt #-}
pointAt r t = p
  where sc = runSpecs r
        t0 = spcStartTime sc
        dt = spcDT sc
        n  = fromIntegral $ floor ((t - t0) / dt)
        p = Point { pointSpecs = sc,
                    pointRun = r,
                    pointTime = t,
                    pointIteration = n,
                    pointPhase = -1 }

-- | Return the integration time points startin from the specified iteration.
integPointsStartingFrom :: Point m -> [Point m]
integPointsStartingFrom p = points
  where r  = pointRun p
        sc = runSpecs r
        (nl, nu) = integIterationBnds sc
        n0       = if pointPhase p == 0
                   then pointIteration p
                   else pointIteration p + 1
        points   = map point [n0 .. nu]
        point n  = Point { pointSpecs = sc,
                           pointRun = r,
                           pointTime = basicTime sc n 0,
                           pointIteration = n,
                           pointPhase = 0 }

-- | Return the indexed time values in the grid by specified size.
timeGrid :: Specs m -> Int -> [(Int, Double)]
timeGrid sc n =
  let t0 = spcStartTime sc
      t2 = spcStopTime sc
      n' = max (n - 1) 1
      dt = (t2 - t0) / (fromIntegral n')
      f i | i == 0    = (i, t0)
          | i == n'   = (i, t2)
          | otherwise = (i, t0 + (fromIntegral i) * dt)
  in map f [0 .. n']

-- | Delay the point by the specified positive number of iterations.
delayPoint :: Point m -> Int -> Point m
delayPoint p dn
  | dn <= 0   = error "Expected the positive number of iterations: delayPoint"
  | otherwise =
    let sc = pointSpecs p
        n  = pointIteration p
        ph = pointPhase p
    in if ph < 0
       then let t' = pointTime p - fromIntegral dn * spcDT sc
                n' = fromIntegral $ floor $ (t' - spcStartTime sc) / spcDT sc
            in p { pointTime = t',
                   pointIteration = n',
                   pointPhase = -1 }
       else let n' = n - dn
                t' = basicTime sc n' ph
            in p { pointTime = t',
                   pointIteration = n',
                   pointPhase = ph }