-- | Defines a structure to contain scientifically-relevant information about an ODE.
module Language.Drasil.Data.ODEInfo (
  ODEInfo(..), odeInfo, odeInfo', ODEOptions(..), odeOptions, ODEMethod(..)
) where

import Language.Drasil.Chunk.Code (CodeVarChunk)
import Language.Drasil.CodeExpr.Development
import Language.Drasil(makeAODESolverFormat, formEquations,
  DifferentialModel(..), ODESolverFormat(..), InitialValueProblem(..))
import Language.Drasil.Chunk.CodeBase (quantvar)

-- This may be temporary, but need a structure to hold ODE info for now.
-- Goal will be for this info to be populated by the instance model for the ODE and the Choices structure.
-- Probably doesn't belong here, but where?
-- | Structure to hold ODE information.
data ODEInfo = ODEInfo {
  -- | Independent variable.
  ODEInfo -> CodeVarChunk
indepVar :: CodeVarChunk,
  -- | Dependent variable.
  ODEInfo -> CodeVarChunk
depVar :: CodeVarChunk,
  -- | Other variables in the ODE.
  ODEInfo -> [CodeVarChunk]
otherVars :: [CodeVarChunk],
  ODEInfo -> CodeExpr
tInit :: CodeExpr,
  ODEInfo -> CodeExpr
tFinal :: CodeExpr,
  -- | Initial value of an ODE.
  ODEInfo -> [CodeExpr]
initVal :: [CodeExpr],
  -- | ODE equations.
  ODEInfo -> [CodeExpr]
odeSyst :: [CodeExpr],
  -- | Various options related to the ODE, including solution method, step size, initial value of a second order ODE, etc.
  ODEInfo -> ODEOptions
odeOpts :: ODEOptions
}

-- | Basic 'ODEInfo' constructor.
odeInfo :: CodeVarChunk -> CodeVarChunk -> [CodeVarChunk] -> CodeExpr -> CodeExpr ->
  [CodeExpr] -> [CodeExpr] -> ODEOptions -> ODEInfo
odeInfo :: CodeVarChunk
-> CodeVarChunk
-> [CodeVarChunk]
-> CodeExpr
-> CodeExpr
-> [CodeExpr]
-> [CodeExpr]
-> ODEOptions
-> ODEInfo
odeInfo = CodeVarChunk
-> CodeVarChunk
-> [CodeVarChunk]
-> CodeExpr
-> CodeExpr
-> [CodeExpr]
-> [CodeExpr]
-> ODEOptions
-> ODEInfo
ODEInfo

-- | Create ODEInfo with Other variables, ODEOptions, DifferentialModel, and InitialValueProblem
odeInfo' :: [CodeVarChunk] -> ODEOptions -> DifferentialModel -> InitialValueProblem -> ODEInfo
odeInfo' :: [CodeVarChunk]
-> ODEOptions
-> DifferentialModel
-> InitialValueProblem
-> ODEInfo
odeInfo' [CodeVarChunk]
ovs ODEOptions
opt DifferentialModel
dm InitialValueProblem
ivp = CodeVarChunk
-> CodeVarChunk
-> [CodeVarChunk]
-> CodeExpr
-> CodeExpr
-> [CodeExpr]
-> [CodeExpr]
-> ODEOptions
-> ODEInfo
ODEInfo
  (UnitalChunk -> CodeVarChunk
forall c. (Quantity c, MayHaveUnit c) => c -> CodeVarChunk
quantvar (UnitalChunk -> CodeVarChunk) -> UnitalChunk -> CodeVarChunk
forall a b. (a -> b) -> a -> b
$ DifferentialModel -> UnitalChunk
_indepVar DifferentialModel
dm)
  (ConstrConcept -> CodeVarChunk
forall c. (Quantity c, MayHaveUnit c) => c -> CodeVarChunk
quantvar (ConstrConcept -> CodeVarChunk) -> ConstrConcept -> CodeVarChunk
forall a b. (a -> b) -> a -> b
$ DifferentialModel -> ConstrConcept
_depVar DifferentialModel
dm)
  [CodeVarChunk]
ovs
  (Expr -> CodeExpr
expr (Expr -> CodeExpr) -> Expr -> CodeExpr
forall a b. (a -> b) -> a -> b
$ InitialValueProblem -> Expr
initTime InitialValueProblem
ivp)
  (Expr -> CodeExpr
expr (Expr -> CodeExpr) -> Expr -> CodeExpr
forall a b. (a -> b) -> a -> b
$ InitialValueProblem -> Expr
finalTime InitialValueProblem
ivp)
  ((Expr -> CodeExpr) -> [Expr] -> [CodeExpr]
forall a b. (a -> b) -> [a] -> [b]
map Expr -> CodeExpr
expr ([Expr] -> [CodeExpr]) -> [Expr] -> [CodeExpr]
forall a b. (a -> b) -> a -> b
$ InitialValueProblem -> [Expr]
initValues InitialValueProblem
ivp)
  (DifferentialModel -> [CodeExpr]
createFinalExpr DifferentialModel
dm)
  ODEOptions
opt

-- | Other parameters for solving the ODE numerically
data ODEOptions = ODEOpts {
  -- | Solution method.
  ODEOptions -> ODEMethod
solveMethod :: ODEMethod,
  -- | Absolute tolerance.
  ODEOptions -> CodeExpr
absTol :: CodeExpr,
  -- | Relative tolerance.
  ODEOptions -> CodeExpr
relTol :: CodeExpr,
  -- | Step size.
  ODEOptions -> CodeExpr
stepSize :: CodeExpr
}

-- | Basic 'ODEOptions' constructor
odeOptions :: ODEMethod -> CodeExpr -> CodeExpr -> CodeExpr -> ODEOptions
odeOptions :: ODEMethod -> CodeExpr -> CodeExpr -> CodeExpr -> ODEOptions
odeOptions = ODEMethod -> CodeExpr -> CodeExpr -> CodeExpr -> ODEOptions
ODEOpts

-- | Methods for solving ODEs. Includes Runge-Kutta 4-5, Backwards Differentiation Formula, or Adams' method.
data ODEMethod = RK45 | BDF | Adams

-- | Create well-formatted ODE equations which the ODE solvers can solve.
createFinalExpr :: DifferentialModel -> [CodeExpr]
createFinalExpr :: DifferentialModel -> [CodeExpr]
createFinalExpr DifferentialModel
dm = (Expr -> CodeExpr) -> [Expr] -> [CodeExpr]
forall a b. (a -> b) -> [a] -> [b]
map Expr -> CodeExpr
expr ([Expr] -> [CodeExpr]) -> [Expr] -> [CodeExpr]
forall a b. (a -> b) -> a -> b
$ [[Expr]] -> [Unknown] -> [Expr] -> ConstrConcept -> [Expr]
formEquations (ODESolverFormat -> [[Expr]]
coeffVects ODESolverFormat
ode) (ODESolverFormat -> [Unknown]
unknownVect ODESolverFormat
ode) (ODESolverFormat -> [Expr]
constantVect ODESolverFormat
ode) (DifferentialModel -> ConstrConcept
_depVar DifferentialModel
dm)
  where ode :: ODESolverFormat
ode = DifferentialModel -> ODESolverFormat
makeAODESolverFormat DifferentialModel
dm