```----------------------------------------------------
-- |
-- Module     :  Math.LevenbergMarquardt
--
-- Maintainer :  Kiet Lam <ktklam9@gmail.com>
--
--
-- The Levenberg-Marquardt algorithm is a minimization
-- algorithm for functions expressed as a sum of squared
-- errors
--
-- This can be used for curve-fitting, multidimensional
-- function minimization, and neural networks training
--
--
----------------------------------------------------

module Math.LevenbergMarquardt (
Function,
Jacobian,
lmMinimize
) where

import Data.Packed.Vector
import Data.Packed.Matrix
import Numeric.LinearAlgebra.Algorithms
import Numeric.Container

-- | Type that represents the function that can calculate the residues
type Function = Vector Double -> Vector Double

-- | Type that represents the function that can calculate the jacobian matrix
--   of the residue with respect to each parameter
type Jacobian = Vector Double -> Matrix Double

-- | Name says it all
sumSquared :: Vector Double -> Double
sumSquared = foldVector (+) 0 . mapVector (**2)

-- | Evolves the parameter x for f(x) = sum-square(e(x)) so that f(x)
--   will be minimized, where:
--
--   f     = real-valued error function,
--   e(x)  = {e1(x),e2(x),..,eN(x)}, where
--   e1(x) = the residue at the vector x
--
--   NOTE: eN(x) is usually represented as (sample - hypothesis(x))
--
--         e.g.: In training neural networks, hypothesis(x) would be the
--               network's output for a training set, and sample would be the
--               expected output for that training set
--
--   NOTE: The dampening constant(lambda) should be set to 0.01 and the dampening
--         update value (beta) should be set to be 10
lmMinimize :: Function         -- ^ Multi-dimensional function that will
--   return a vector of residues
-> Jacobian      -- ^ The function that calculate the Jacobian
--   matrix of each residue with respect to
--   each parameter
-> Vector Double -- ^ The initial guess for the parameter
-> Double        -- ^ Dampening constant (usually lambda in most literature)
-> Double        -- ^ Dampening update value (usually beta in most literature)
-> Double        -- ^ The precision desired
-> Int           -- ^ The maximum iteration
-> (Vector Double, Matrix Double) -- ^ Returns the optimal parameter
--   and the matrix path
lmMinimize func jacob param lambda beta prec iter =
let errValue = sumSquared \$ func param
(v, m) = lmMinimizeInternal func jacob param prec errValue lambda beta iter [param]
in
(v, (fromLists . reverse . map toList) m)

lmMinimizeInternal :: Function           -- ^ Multi-dimensional function that will
--   return a vector of residues
-> Jacobian        -- ^ The Jacobian matrix of each residue
--   with respect to each parameter
-> Vector Double   -- ^ The current guess for the parameter
-> Double          -- ^ The precision desired
-> Double          -- ^ The current value for the error function
--   (which is the sum of the residues squared)
-> Double          -- ^ The dampening constant (lambda)
-> Double          -- ^ The beta constant to update lambda
-> Int             -- ^ The maximum iteration
-> [Vector Double] -- ^ The current matrix path
-> (Vector Double, [Vector Double]) -- ^ Returns the optimal parameter
--   and the matrix path
lmMinimizeInternal _ _ v _ _ _ _ 0 m = (v,m)
lmMinimizeInternal func jacob param prec errValue lambda beta iter mat =
let jacobian = jacob param
residues = func param

gradients = (trans jacobian) `mXv` residues
hessian = (trans jacobian) `multiply` jacobian

diagonalHessian = diagRect 0 (takeDiag hessian) (cols hessian) (cols hessian)
diagonalLambda = mapMatrix (*lambda) (ident (cols hessian))