{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Control.Monad  (sequence_)
import LevMar.Chart   (PlotValue, levmarChart)
import LevMar.Fitting
import NFunction      (($*))
import System.Random  (Random, mkStdGen, randoms)

import qualified SizedList as SL

-------------------------------------------------------------------------------
-- Some handy type level naturals

type N0 = Z
type N1 = S N0
type N2 = S N1
type N3 = S N2
type N4 = S N3
type N5 = S N4

-------------------------------------------------------------------------------
-- Model functions

constant  :: Num r => SimpleModel N1 r
linear    :: Num r => SimpleModel N2 r
quadratic :: Num r => SimpleModel N3 r
cubic     :: Num r => SimpleModel N4 r

constant        a _ = a
linear        b a x = b * x     + constant      a x
quadratic   c b a x = c * x*x   + linear      b a x
cubic     d c b a x = d * x*x*x + quadratic c b a x

trig  :: Floating r => SimpleModel N2 r
trig a b x = a * cos x + b * sin x

-------------------------------------------------------------------------------
-- Jacobians

constantJacob  :: Num r => SimpleJacobian N1 r
linearJacob    :: Num r => SimpleJacobian N2 r
quadraticJacob :: Num r => SimpleJacobian N3 r
cubicJacob     :: Num r => SimpleJacobian N4 r

constantJacob        _ _ = 1     ::: Nil
linearJacob        _ a x = x     ::: constantJacob      a x
quadraticJacob   _ b a x = x*x   ::: linearJacob      b a x
cubicJacob     _ c b a x = x*x*x ::: quadraticJacob c b a x

trigJacob :: Floating r => SimpleJacobian N2 r
trigJacob _ _ x = cos x ::: sin x ::: Nil

-------------------------------------------------------------------------------
-- Test utility function

testChart :: forall n r a.
             ( Nat n
             , PlotValue r, Fractional r, LevMarable r, Random r
             , PlotValue a, Fractional a
             )
          => (Model n r a)
          -> Maybe (Jacobian n r a)
          -> SizedList n r
          -> [a]
          -> r
          -> IO ()
testChart f j ps xs noise = levmarChart f
                                        j
                                        initPs
                                        samples'
                                        1000
                                        opts
                                        Nothing
                                        Nothing
                                        noLinearConstraints
                                        Nothing
                           >> return ()
    where
      ns         = take (length xs) $ randoms $ mkStdGen rndGenSeed
      initPs     = SL.replicate 0 :: SizedList n r
      samples    = zip xs $ map (f $* ps) xs
      samples'   = zipWith (\(x, y) n -> (x, y + (n - 0.5) * 2 * noise)) samples ns
      rndGenSeed = 123456
      opts       = defaultOpts { optStopNormInfJacTe = 1e-15
                               , optStopNorm2Dp      = 1e-15
                               , optStopNorm2E       = 1e-20
                               }

-------------------------------------------------------------------------------
-- Examples

testConstant :: IO ()
testConstant = testChart constant (Just constantJacob)
                         (42 ::: Nil)
                         [-100 .. 100]
                         (10 :: Double)

testLinear :: IO ()
testLinear = testChart linear (Just linearJacob)
                       (3.14 ::: -8 ::: Nil)
                       [-100 .. 100]
                       (300 :: Double)

testQuadratic :: IO ()
testQuadratic = testChart quadratic (Just quadraticJacob)
                          (0.2 ::: 8 ::: 23 ::: Nil)
                          [-100 .. 100]
                          (1000 :: Double)

testCubic :: IO ()
testCubic = testChart cubic (Just cubicJacob)
                      (-0.05 ::: 0.5 ::: -12 ::: 10 ::: Nil)
                      [-100 .. 100]
                      (8000 :: Double)

testTrig :: IO ()
testTrig = testChart trig (Just trigJacob)
                     (100 ::: 102 ::: Nil)
                     (map (/4) [0..400])
                     (50 :: Double)

-------------------------------------------------------------------------------
-- Main

main :: IO ()
main = sequence_ [ testConstant
                 , testLinear
                 , testQuadratic
                 , testCubic
                 , testTrig
                 ]