{-|
Module      : Integrators
Description : Implementations of Runge-Kutta 8 and Gauss-Jackson integrators

Implements fourth, fifth, and eighth order Runge-Kutta integration methods for
time independent second order differential equations using the classic fourth
order, Dormand-Prince fifth order, and Cooper-Verner eigth order coefficients
respectively.
The Gauss-Jackson coefficients are generated to arbitrary order, but using RKCV8
for the startup conditions necessarily restricts the order to eigth order at
most, though correction steps can more or less minimise that loss. For some
inputs RK4 instead of RKCV8 or RKDP5 may make more sense for startup.

-}

module Integrators
       (
         rk4
       , rkdp5
       , rkcv8
       , timeIndepRK
       , Integrator(..)
       ) where

import PhysicsVectors


-- | Time independent Butcher Tableu with a_ij in backwards row order
data BTabTIndep n = BTabTIndep [[n]] [n]
  deriving (Show)

rkcv8t :: Floating t => BTabTIndep t
rkcv8t = BTabTIndep
  [
    [1/2]
  , [1/4, 1/4]
  , [(21 + 5*s21)/49, (-7 - 3*s21)/98, 1/7]
  , [(21 - s21)/252, (18 + 4*s21)/63, 0, (11 + s21)/84]
  , [(63 - 7*s21)/80, (-231 + 14*s21)/360, (9 + s21)/36, 0, (5 + s21)/48]
  , [(63 - 13*s21)/35, (-504 + 115*s21)/70, (633 - 145*s21)/90
    , (-432 + 92*s21)/315, 0, (10 - s21)/42]
  , [1/9, (13 - 3*s21)/63, (14 - 3*s21)/126, 0, 0, 0, 1/14]
  , [(63 + 13*s21)/128, (-385 - 75*s21)/1152, 11/72, (91 - 21*s21)/576, 0, 0, 0, 1/32]
  , [(132 + 28*s21)/245, (-51 - 11*s21)/56, (515 + 111*s21)/504
    , (-733 - 147*s21)/2205, 1/9, 0, 0, 0, 1/14]
  , [(49 - 7*s21)/18, (28 - 28*s21)/45, (301 + 53*s21)/72, (-273 - 53*s21)/72
    , (-18 + 28*s21)/45, (-42 + 7*s21)/18, 0, 0, 0, 0]
  ]
  [ 1/20, 49/180, 16/45, 49/180, 0, 0, 0, 0, 0, 0, 1/20]
  where
    s21 = sqrt 21

rk4t :: Floating t => BTabTIndep t
rk4t = BTabTIndep [[1/2],[1/2,0],[1,0,0]] [1/6,1/3,1/3,1/6]

rkdp5t :: Floating t => BTabTIndep t
rkdp5t = BTabTIndep
  [
    [1/5]
  , [9/40, 3/40]
  , [32/9, -56/15, 44/45]
  , [-212/729, 64448/6561, -25360/2187, 19372/6561]
  , [-5103/18656, 49/176, 46732/5247, -355/33, 9017/3168]
  , [11/84, -2187/6784, 125/192, 500/1113, 0, 35/384]
  ]
  [0, 11/84, -2187/6784, 125/192, 500/1113, 0, 35/384]

{-|
  'timeIndepRK' implements the general Runge-Kutta form for a time independent
  differential equation \(f\).

  'timeIndepRK' @btab f yn dt@ steps one step forward in the described
  time-independent Runge-Kutta method where @btab@ describes the implementation
  Butcher-Tableu, @f@ the differetial function, @rn@ the current position and
  velocity, and @dt@ the integration time step.

  Note that \(f\) operates on both position and velocity, i.e. it returns both
  \(\dot{\vec{r}}\) and \(\ddot{\vec{r}}\) and that the Butcher-Tableu's used
  must have the \(a_{ij}\) and \(b_i\) coefficients in /reverse/ row order. That
  is, each row is in reverse order.
-}
timeIndepRK
    :: (Floating t, Eq t)
    => BTabTIndep t
    -> (Vec2 t -> Vec2 t)
    -> Vec2 t
    -> t
    -> Vec2 t
timeIndepRK _ _ rn 0 = rn
timeIndepRK (BTabTIndep a b) f rn dt = step ks b
  where
    coeffSum [] _ = 0
    coeffSum kl zetas = sum [ zi `scalarMult` ki | (ki, zi) <- zip kl zetas ]
    ks = foldl (\kl aj-> f (step kl aj) : kl) [f rn] a
    step kl zi = rn + (dt `scalarMult` coeffSum kl zi)

-- | Data class encapsulating the concept of an integrator function
data Integrator t =
    SingStepIntegrator ((Vec2 t -> Vec2 t) -> Vec2 t -> t -> Vec2 t)
  | MultStepIntegrator ((Vec2 t -> Vec2 t) -> [(Vec2 t, Vec3 t)] -> t -> [(Vec2 t, Vec3 t)])


-- | The classic Runge-Kutta fourth order method
rk4 :: (Floating t, Eq t) => Integrator t
rk4 = SingStepIntegrator $ timeIndepRK rk4t

{-|
  Dormand-Prince fifth order method. A less general RK method could avoid the
  final b calculation and the first f calculation of the next step with this one
  due to the coefficient choice.
-}
rkdp5 :: (Floating t, Eq t) => Integrator t
rkdp5 = SingStepIntegrator $ timeIndepRK rkdp5t

-- | Cooper-Verner eigth order method
rkcv8 :: (Floating t, Eq t) => Integrator t
rkcv8 = SingStepIntegrator $ timeIndepRK rkcv8t