{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RebindableSyntax #-}

-- | Numerical solution of the initial value problem,
--
--     y'(x)  = f(x, y(x))
--     y'(x0) = y0
--
--   for x in [x0, xN].
--

module ODE.IVP
where

import Misc (partition)
import NumericPrelude hiding (abs)
import Algebra.Absolute (abs)
import qualified Algebra.Field as Field
import qualified Algebra.ToInteger as ToInteger
import qualified Algebra.ToRational as ToRational
import qualified Algebra.RealField as RealField

-- | A single iteration of Euler's method over the interval
--   [$x0$, $x0$+$h$].
--
--   Examples:
--
--   >>> let x0 = 0.0
--   >>> let y0 = 1.0
--   >>> let h  = 1.0
--   >>> let f x y = y
--   >>> eulers_method1 x0 y0 f h
--   2.0
--
eulers_method1 :: (Field.C a, ToRational.C a, Field.C b)
               => a -- ^ x0, the initial point
               -> b -- ^ y0, the initial value at x0
               -> (a -> b -> b) -- ^ The function f(x,y)
               -> a -- ^ The step size h.
               -> b
eulers_method1 x0 y0 f h =
  y0 +  h'*y'
  where
    h' = fromRational'$ toRational h
    y' = (f x0 y0)


-- | Perform $n$ iterations of Euler's method over the interval [$x0$,
--   $xN$]. The step size `h` will be calculated automatically. A list
--   of y-values will be returned.
--
--   The explicit 'forall' in the type signature allows us to refer
--   back to the type variables 'a' and 'b' in the 'where' clause.
--
--   Examples:
--
--   >>> let x0 = 0.0
--   >>> let xN = 1.0
--   >>> let y0 = 1.0
--   >>> let f x y = y
--   >>> let ys = eulers_method x0 xN y0 f 10000
--   >>> let yN = head $ reverse ys
--   >>> abs ((exp 1) - yN) < 1/10^3
--   True
--   >>> head ys == y0
--   True
--
eulers_method :: forall a b c. (Field.C a,
                                ToRational.C a,
                                Field.C b,
                                ToInteger.C c,
                                Enum c)
               => a -- ^ x0, the initial point
               -> b -- ^ y0, the initial value at x0
               -> a -- ^ xN, the terminal point
               -> (a -> b -> b) -- ^ The function f(x,y)
               -> c -- ^ n, the number of intervals to use.
               -> [b]
eulers_method x0 y0 xN f n =
  y0 : go xs y0 f
  where
    xs = partition n x0 xN

    -- The 'go' function actually does all the work. It takes a list
    -- of intervals [(v0,v1), (v1, v2)...] and peels off the first
    -- one. It then runs the single-step Euler's method on that
    -- interval, and afterwards recurses down the rest of the list.
    --
    go :: [(a,a)] -> b -> (a -> b -> b) -> [b]
    go [] _ _ = []
    go ((v0,v1):rest) w0 g = w1 : (go rest w1 g)
      where
        w1 = eulers_method1 v0 w0 g (v1 - v0)


-- | Perform as many iterations of Euler's method over the interval
--   [$x0$, $xN$] as is necessary for the given step size $h$. The
--   number of subintervals `n` will be calculated automatically. A
--   list of y-values will be returned.
--
--   The implementation simply computes `n` from the length of the
--   interval and the given $h$, and then calls 'eulers_method'.
--
--   Examples:
--
--   >>> let x0 = 0.0
--   >>> let xN = 1.0
--   >>> let y0 = 1.0
--   >>> let f x y = y
--   >>> let ys = eulers_method x0 xN y0 f 10
--   >>> let ys' = eulers_methodH x0 xN y0 f 0.1
--   >>> ys == ys'
--   True
--
eulers_methodH :: (RealField.C a, ToRational.C a, Field.C b)
               => a -- ^ x0, the initial point
               -> b -- ^ y0, the initial value at x0
               -> a -- ^ xN, the terminal point
               -> (a -> b -> b) -- ^ The function f(x,y)
               -> a -- ^ h, the step size.
               -> [b]
eulers_methodH x0 y0 xN f h =
  eulers_method x0 y0 xN f n
  where
    n :: Integer
    n = floor $ (xN - x0) / h