{-# 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 (
  eulers_method,
  eulers_method1,
  eulers_methodH )
where

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

-- | 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:
--
--   >>> import Algebra.Absolute (abs)
--   >>> 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