{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- | Classical iterative methods to solve the system Ax = b.

module Linear.Iteration
where

import Data.List (find)
import Data.Maybe (fromJust)
import Data.Vector.Fixed (Arity, N1, S)
import NumericPrelude hiding ((*))
import qualified Algebra.Algebraic as Algebraic
import qualified Algebra.Field as Field
import qualified Algebra.RealField as RealField
import qualified Algebra.ToRational as ToRational
import qualified Prelude as P

import Linear.Matrix
import Linear.System
import Normed

-- | Perform one Jacobi iteration,
--
--     x1 = M^(-1) * (N*x0 + b)
--
--   Examples:
--
--   >>> let m  = fromList [[4,2],[2,2]] :: Mat2 Double
--   >>> let x0 = vec2d (0, 0::Double)
--   >>> let b  = vec2d (1, 1::Double)
--   >>> jacobi_iteration m b x0
--   ((0.25),(0.5))
--   >>> let x1 = jacobi_iteration m b x0
--   >>> jacobi_iteration m b x1
--   ((0.0),(0.25))
--
jacobi_iteration :: (Field.C a, Arity m)
                 => Mat m m  a
                 -> Mat m N1 a
                 -> Mat m N1 a
                 -> Mat m N1 a
jacobi_iteration matrix b x_current =
  x_next
  where
    big_m = diagonal matrix
    big_n = big_m - matrix
    rhs = big_n*x_current + b
    x_next = forward_substitute big_m rhs


-- | Compute an infinite list of Jacobi iterations starting with the
--   vector x0.
jacobi_iterations :: (Field.C a, Arity m)
                  => Mat m m  a
                  -> Mat m N1 a
                  -> Mat m N1 a
                  -> [Mat m N1 a]
jacobi_iterations matrix b =
  iterate (jacobi_iteration matrix b)


-- | Solve the system Ax = b using the Jacobi method. This will run
--   forever if the iterations do not converge.
--
--   Examples:
--
--   >>> let m  = fromList [[4,2],[2,2]] :: Mat2 Double
--   >>> let x0 = vec2d (0, 0::Double)
--   >>> let b  = vec2d (1, 1::Double)
--   >>> let epsilon = 10**(-6)
--   >>> jacobi_method m b x0 epsilon
--   ((0.0),(0.4999995231628418))
--
jacobi_method :: forall m n a b.
                 (RealField.C a,
                  Algebraic.C a, -- Normed instance
                  ToRational.C a, -- Normed instance
                  Algebraic.C b,
                  RealField.C b,
                  Arity m,
                  Arity n, -- Normed instance
                  m ~ S n)
              => Mat m m  a
              -> Mat m N1 a
              -> Mat m N1 a
              -> b
              -> Mat m N1 a
jacobi_method matrix b x0 epsilon =
  -- fromJust is "safe," because the list is infinite. If the
  -- algorithm doesn't converge, 'find' will search forever and never
  -- return Nothing.
  fst' $ fromJust $ find error_small_enough diff_pairs
  where
    x_n = jacobi_iterations matrix b x0

    pairs :: [(Mat m N1 a, Mat m N1 a)]
    pairs = zip (tail x_n) x_n

    append_diff :: (Mat m N1 a, Mat m N1 a)
                -> (Mat m N1 a, Mat m N1 a, b)
    append_diff (cur,prev) =
      (cur,prev,diff)
      where
        diff = norm (cur - prev)

    diff_pairs :: [(Mat m N1 a, Mat m N1 a, b)]
    diff_pairs = map append_diff pairs

    fst' :: (c, d, e) -> c
    fst' (x,_,_) = x

    error_small_enough :: (Mat m N1 a, Mat m N1 a, b)-> Bool
    error_small_enough (_,_,err) = err < epsilon