{-# LANGUAGE CPP                    #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE RecordWildCards        #-}
-- |
--
-- @regression-simple@ provides (hopefully) simple regression functions.
--
-- The @'linear' :: Foldable f => (a -> (Double, Double)) -> f a -> 'V2'@
-- is the simplest one.
--
-- There are variants with weights, y-errors, and x and y-errors.
-- In addition, package includes Levenberg–Marquardt algorithm implementation
-- to fit arbitrary functions (with one, two or three parameters),
-- as long as you can give their partial derivatives as well (@ad@ package is handy for that).
--
-- For multiple independent variable ordinary least squares
-- or Levenberg-Marquard with functions with \> 3 parameter you should look elsewhere.
--
-- Package has been tested to return similar results as @fit@ functionality in @gnuplot@
-- (L-M doesn't always converge to exactly the same points in parameter space).
--
module Math.Regression.Simple (
    -- * Linear regression
    linear,
    linearFit,
    linearWithWeights,
    linearWithYerrors,
    linearWithXYerrors,
    -- ** Step-by-step interface
    linearFit',
    LinRegAcc (..),
    zeroLinRegAcc,
    addLinReg,
    addLinRegW,
    -- * Quadratic regression
    quadratic,
    quadraticFit,
    quadraticWithWeights,
    quadraticWithYerrors,
    quadraticWithXYerrors,
    -- ** Step-by-step interface
    quadraticFit',
    QuadRegAcc (..),
    zeroQuadRegAcc,
    addQuadReg,
    addQuadRegW,
    quadRegAccToLin,
    -- * Levenberg–Marquardt algorithm
    -- ** One parameter
    levenbergMarquardt1,
    levenbergMarquardt1WithWeights,
    levenbergMarquardt1WithYerrors,
    levenbergMarquardt1WithXYerrors,
    -- ** Two parameters
    levenbergMarquardt2,
    levenbergMarquardt2WithWeights,
    levenbergMarquardt2WithYerrors,
    levenbergMarquardt2WithXYerrors,
    -- ** Three parameters
    levenbergMarquardt3,
    levenbergMarquardt3WithWeights,
    levenbergMarquardt3WithYerrors,
    levenbergMarquardt3WithXYerrors,
    -- * Auxiliary types
    Fit (..),
    V2 (..),
    V3 (..),
) where

import Control.DeepSeq (NFData (..))

import qualified Data.Foldable      as F
import qualified Data.List.NonEmpty as NE

import Math.Regression.Simple.LinAlg
import Numeric.KBN

-- $setup
-- >>> :set -XDeriveFunctor -XDeriveFoldable -XDeriveTraversable
--
-- >>> import Numeric (showFFloat)
-- >>> import Data.List.NonEmpty (NonEmpty (..))
-- >>> import qualified Data.List.NonEmpty as NE
--
-- Don't show too much decimal digits
--
-- >>> let showDouble x = showFFloat (Just (min 5 (5 - ceiling (logBase 10 x)))) (x :: Double)
-- >>> showDouble 123.456 ""
-- "123.46"
--
-- >>> showDouble 1234567 ""
-- "1234567"
--
-- >>> showDouble 123.4567890123456789 ""
-- "123.46"
--
-- >>> showDouble 0.0000000000000012345 ""
-- "0.00000"
--
-- >>> newtype PP a = PP a
-- >>> class Show' a where showsPrec' :: Int -> a -> ShowS
-- >>> instance Show' a => Show (PP a) where showsPrec d (PP x) = showsPrec' d x
-- >>> instance Show' Double where showsPrec' d x = if x < 0 then showParen (d > 6) (showChar '-' . showDouble (negate x)) else showDouble x
-- >>> instance Show' Int where showsPrec' = showsPrec
-- >>> instance (Show' a, Show' b) => Show' (a, b) where showsPrec' d (x, y) = showParen True $ showsPrec' 0 x . showString ", " . showsPrec' 0 y
-- >>> instance Show' v => Show' (Fit v) where showsPrec' d (Fit p e ndf wssr) = showParen (d > 10) $ showString "Fit " . showsPrec' 11 p . showChar ' ' . showsPrec' 11 e . showChar ' ' . showsPrec' 11 ndf . showChar ' ' . showsPrec' 11 wssr
-- >>> instance Show' V2 where showsPrec' d (V2 x y)   = showParen (d > 10) $ showString "V2 " . showsPrec' 11 x . showChar ' ' . showsPrec' 11 y
-- >>> instance Show' V3 where showsPrec' d (V3 x y z) = showParen (d > 10) $ showString "V3 " . showsPrec' 11 x . showChar ' ' . showsPrec' 11 y . showChar ' ' . showsPrec' 11 z
-- >>> instance Show' a => Show' [a] where showsPrec' _ [] = id; showsPrec' _ [x] = showsPrec' 0 x; showsPrec' _ (x:xs) = showsPrec' 0 x . showChar '\n' . showsPrec' 0 xs
-- >>> instance Show' a => Show' (NonEmpty a) where showsPrec' _ (x :| []) = showsPrec' 0 x; showsPrec' _ (x:|xs) = showsPrec' 0 x . showChar '\n' . showsPrec' 0 xs
-- >>>
--
-- Inputs:
-- >>> let input1 = [(0, 1), (1, 3), (2, 5)]
-- >>> let input2 = [(0.1, 1.2), (1.3, 3.1), (1.9, 4.9), (3.0, 7.1), (4.1, 9.0)]
-- >>> let input3 = [(0, 2), (1, 3), (2, 6), (3, 11)]
--
-- >>> let sq z = z * z
--

-------------------------------------------------------------------------------
-- Linear
-------------------------------------------------------------------------------

-- | Linear regression.
--
-- >>> let input1 = [(0, 1), (1, 3), (2, 5)]
-- >>> PP $ linear id input1
-- V2 2.0000 1.00000
--
-- >>> let input2 = [(0.1, 1.2), (1.3, 3.1), (1.9, 4.9), (3.0, 7.1), (4.1, 9.0)]
-- >>> PP $ linear id input2
-- V2 2.0063 0.88685
--
linear :: F.Foldable f => (a -> (Double, Double)) -> f a -> V2
linear f = fitParams . linearFit f

-- | Like 'linear' but returns complete 'Fit'.
--
-- To get confidence intervals you should multiply the errors
-- by @quantile (studentT (n - 2)) ci'@ from @statistics@ package
-- or similar.
-- For big @n@ using value 1 gives 68% interval and using value 2 gives 95% confidence interval.
-- See https://en.wikipedia.org/wiki/Student%27s_t-distribution#Table_of_selected_values
-- (@quantile@ calculates one-sided values, you need two-sided, thus adjust @ci@ value).
--
-- The first input is perfect fit:
--
-- >>> let fit = linearFit id input1
-- >>> PP fit
-- Fit (V2 2.0000 1.00000) (V2 0.00000 0.00000) 1 0.00000
--
-- The second input is quite good:
--
-- >>> PP $ linearFit id input2
-- Fit (V2 2.0063 0.88685) (V2 0.09550 0.23826) 3 0.25962
--
-- But the third input isn't so much,
-- standard error of a slope parameter is 20%.
--
-- >>> let input3 = [(0, 2), (1, 3), (2, 6), (3, 11)]
-- >>> PP $ linearFit id input3
-- Fit (V2 3.0000 1.00000) (V2 0.63246 1.1832) 2 4.0000
--
linearFit :: F.Foldable f => (a -> (Double, Double)) -> f a -> Fit V2
linearFit f = linearFit' . linRegAcc f

-- | Weighted linear regression.
--
-- >>> let input2 = [(0.1, 1.2), (1.3, 3.1), (1.9, 4.9), (3.0, 7.1), (4.1, 9.0)]
-- >>> PP $ linearFit id input2
-- Fit (V2 2.0063 0.88685) (V2 0.09550 0.23826) 3 0.25962
--
-- >>> let input2w = [(0.1, 1.2, 1), (1.3, 3.1, 1), (1.9, 4.9, 1), (3.0, 7.1, 1/4), (4.1, 9.0, 1/4)]
-- >>> PP $ linearWithWeights id input2w
-- Fit (V2 2.0060 0.86993) (V2 0.12926 0.23696) 3 0.22074
--
linearWithWeights :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> Fit V2
linearWithWeights f = linearFit' . linRegAccW f

-- | Linear regression with y-errors.
--
-- >>> let input2y = [(0.1, 1.2, 0.12), (1.3, 3.1, 0.31), (1.9, 4.9, 0.49), (3.0, 7.1, 0.71), (4.1, 9.0, 1.9)]
-- >>> let fit = linearWithYerrors id input2y
-- >>> PP fit
-- Fit (V2 1.9104 0.98302) (V2 0.13006 0.10462) 3 2.0930
--
-- When we know actual y-errors, we can calculate the Q-value using @statistics@ package:
--
-- >>> import qualified Statistics.Distribution            as S
-- >>> import qualified Statistics.Distribution.ChiSquared as S
-- >>> S.cumulative (S.chiSquared (fitNDF fit)) (fitWSSR fit)
-- 0.446669639443138
--
-- or using @math-functions@
--
-- >>> import Numeric.SpecFunctions (incompleteGamma)
-- >>> incompleteGamma (fromIntegral (fitNDF fit) / 2) (fitWSSR fit / 2)
-- 0.446669639443138
--
-- It is not uncommon to deem acceptable on equal terms any models with, say, Q > 0.001.
-- If Q is too large, too near to 1 is most likely caused by overestimating
-- the y-errors.
--
--
linearWithYerrors :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> Fit V2
linearWithYerrors f = linearWithWeights f' where
    f' a = case f a of
        (x, y, dy) -> (x, y, recip (dy * dy))

-- | Iterative linear regression with x and y errors.
--
-- /Orear, J. (1982). Least squares when both variables have uncertainties. American Journal of Physics, 50(10), 912–916. doi:10.1119\/1.12972/
--
-- >>> let input2xy = [(0.1, 1.2, 0.01, 0.12), (1.3, 3.1, 0.13, 0.31), (1.9, 4.9, 0.19, 0.49), (3.0, 7.1, 0.3, 0.71), (4.1, 9.0, 0.41, 1.9)]
-- >>> let fit :| fits = linearWithXYerrors id input2xy
--
-- First fit is done using 'linearWithYerrors':
--
-- >>> PP fit
-- Fit (V2 1.9104 0.98302) (V2 0.13006 0.10462) 3 2.0930
--
-- After that the effective variance is used to refine the fit,
-- just a few iterations is often enough:
--
-- >>> PP $ take 3 fits
-- Fit (V2 1.9092 0.99251) (V2 0.12417 0.08412) 3 1.2992
-- Fit (V2 1.9092 0.99250) (V2 0.12418 0.08414) 3 1.2998
-- Fit (V2 1.9092 0.99250) (V2 0.12418 0.08414) 3 1.2998
--
linearWithXYerrors
    :: F.Foldable f
    => (a -> (Double, Double, Double, Double))  -- ^ \(x_i, y_i, \delta x_i, \delta y_i\)
    -> f a                                      -- ^ data
    -> NE.NonEmpty (Fit V2)
linearWithXYerrors f xs = iterate1 go fit0 where
    fit0   = linearWithYerrors (\a -> case f a of (x,y,_,dy) -> (x,y,dy)) xs
    go fit = linearWithWeights (\a -> case f a of (x,y,dx,dy) -> (x,y,recip $ sq (param1 * dx) + sq dy)) xs where
        V2 param1 _ = fitParams fit

-- >>> import qualified Numeric.AD.Mode.Reverse.Double as AD
-- >>> data H3 a = H3 a a a deriving (Functor, Foldable, Traversable)
-- >>> let linearF (H3 a b x) = a * x + b
-- >>> let lin' (V2 a b) (x, y, dx, dy) = case AD.grad' linearF (H3 a b x) of (f, H3 da db f') -> (y, f, V2 da db, recip $ sq (f' * dx) + sq dy)
--
-- >>> PP $ NE.last $ levenbergMarquardt2WithWeights lin' (V2 1 1) input2xy
-- Fit (V2 1.9092 0.99250) (V2 0.12418 0.08414) 3 1.2998

-- | Calculate linear fit from 'LinRegAcc'.
linearFit' :: LinRegAcc -> Fit V2
linearFit' LinRegAcc {..} = Fit params errors ndf wssr where
    matrix@(SM22 a11 _ a22) = inv (SM22 x2 x w)
    params@(V2 a b)         = mult matrix (V2 xy y)

    errors = V2 sa sb

    -- ensure that error is always non-negative.
    -- Due rounding errors, in perfect fit situations it can be slightly negative.
    wssr = max 0 (y2 - a * xy - b * y)
    ndf  = lra_n - 2
    ndf' = fromIntegral ndf :: Double

    sa   = sqrt (a11 * wssr / ndf')
    sb   = sqrt (a22 * wssr / ndf')

    w  = getKBN lra_w
    x  = getKBN lra_x
    x2 = getKBN lra_x2
    y  = getKBN lra_y
    xy = getKBN lra_xy
    y2 = getKBN lra_y2

    -- is it useful?
    -- r2 = (n * xy - x*y) / (sqrt (n * x2 - x*x) * sqrt (n * y2 - y*y))

linRegAcc :: F.Foldable f => (a -> (Double, Double)) -> f a -> LinRegAcc
linRegAcc f = F.foldl' (\acc a -> case f a of (x,y) -> addLinReg acc x y) zeroLinRegAcc

linRegAccW :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> LinRegAcc
linRegAccW f = F.foldl' (\acc a -> case f a of (x,y,w) -> addLinRegW acc x y w) zeroLinRegAcc

-------------------------------------------------------------------------------
-- Quadractic
-------------------------------------------------------------------------------

-- | Quadratic regression.
--
-- >>> let input1 = [(0, 1), (1, 3), (2, 5)]
-- >>> quadratic id input1
-- V3 0.0 2.0 1.0
--
-- >>> let input2 = [(0.1, 1.2), (1.3, 3.1), (1.9, 4.9), (3.0, 7.1), (4.1, 9.0)]
-- >>> PP $ quadratic id input2
-- V3 (-0.00589) 2.0313 0.87155
--
-- >>> let input3 = [(0, 2), (1, 3), (2, 6), (3, 11)]
-- >>> PP $ quadratic id input3
-- V3 1.00000 0.00000 2.0000
--
quadratic :: F.Foldable f => (a -> (Double, Double)) -> f a -> V3
quadratic f = fitParams . quadraticFit f

-- | Like 'quadratic' but returns complete 'Fit'.
--
-- >>> PP $ quadraticFit id input2
-- Fit (V3 (-0.00589) 2.0313 0.87155) (V3 0.09281 0.41070 0.37841) 2 0.25910
--
-- >>> PP $ quadraticFit id input3
-- Fit (V3 1.00000 0.00000 2.0000) (V3 0.00000 0.00000 0.00000) 1 0.00000
--
quadraticFit :: F.Foldable f => (a -> (Double, Double)) -> f a -> Fit V3
quadraticFit f = quadraticFit' . quadRegAcc f

-- | Weighted quadratic regression.
--
-- >>> let input2w = [(0.1, 1.2, 1), (1.3, 3.1, 1), (1.9, 4.9, 1), (3.0, 7.1, 1/4), (4.1, 9.0, 1/4)]
-- >>> PP $ quadraticWithWeights id input2w
-- Fit (V3 0.02524 1.9144 0.91792) (V3 0.10775 0.42106 0.35207) 2 0.21484
--
quadraticWithWeights :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> Fit V3
quadraticWithWeights f = quadraticFit' . quadRegAccW f

-- | Quadratic regression with y-errors.
--
-- >>> let input2y = [(0.1, 1.2, 0.12), (1.3, 3.1, 0.31), (1.9, 4.9, 0.49), (3.0, 7.1, 0.71), (4.1, 9.0, 0.9)]
-- >>> PP $ quadraticWithYerrors id input2y
-- Fit (V3 0.08776 1.6667 1.0228) (V3 0.10131 0.31829 0.11917) 2 1.5398
--
quadraticWithYerrors :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> Fit V3
quadraticWithYerrors f = quadraticWithWeights f' where
    f' a = case f a of
        (x, y, dy) -> (x, y, recip (dy * dy))

-- | Iterative quadratic regression with x and y errors.
--
-- /Orear, J. (1982). Least squares when both variables have uncertainties. American Journal of Physics, 50(10), 912–916. doi:10.1119\/1.12972/
--
quadraticWithXYerrors
    :: F.Foldable f
    => (a -> (Double, Double, Double, Double))  -- ^ \(x_i, y_i, \delta x_i, \delta y_i\)
    -> f a                                      -- ^ data
    -> NE.NonEmpty (Fit V3)
quadraticWithXYerrors f xs = iterate1 go fit0 where
    fit0   = quadraticWithYerrors (\a -> case f a of (x,y,_,dy) -> (x,y,dy)) xs
    go fit = quadraticWithWeights (\a -> case f a of (x,y,dx,dy) -> (x,y,recip $ sq ((2 * p1 * x + p2) * dx) + sq dy)) xs where
        V3 p1 p2 _ = fitParams fit

-- | Calculate quadratic fit from 'QuadRegAcc'.
quadraticFit' :: QuadRegAcc -> Fit V3
quadraticFit' QuadRegAcc {..} = Fit params errors ndf wssr where
    matrix@(SM33 a11
                 _   a22
                 _   _   a33) = inv (SM33 x4
                                          x3 x2
                                          x2  x w)

    params@(V3 a b c) = mult matrix (V3 x2y xy y)

    errors = V3 sa sb sc

    wssr = max 0 (y2 - a * x2y - b * xy - c * y)
    ndf  = qra_n - 3
    ndf' = fromIntegral ndf :: Double

    sa  = sqrt (a11 * wssr / ndf')
    sb  = sqrt (a22 * wssr / ndf')
    sc  = sqrt (a33 * wssr / ndf')

    w   = getKBN qra_w
    x   = getKBN qra_x
    x2  = getKBN qra_x2
    x3  = getKBN qra_x3
    x4  = getKBN qra_x4
    y   = getKBN qra_y
    xy  = getKBN qra_xy
    x2y = getKBN qra_x2y
    y2  = getKBN qra_y2

    -- is it useful?
    -- r2 = (n * xy - x*y) / (sqrt (n * x2 - x*x) * sqrt (n * y2 - y*y))

quadRegAcc :: F.Foldable f => (a -> (Double, Double)) -> f a -> QuadRegAcc
quadRegAcc f = F.foldl' (\acc a -> case f a of (x,y) -> addQuadReg acc x y) zeroQuadRegAcc

quadRegAccW :: F.Foldable f => (a -> (Double, Double, Double)) -> f a -> QuadRegAcc
quadRegAccW f = F.foldl' (\acc a -> case f a of (x,y,w) -> addQuadRegW acc x y w) zeroQuadRegAcc

-------------------------------------------------------------------------------
-- Levenberg–Marquardt 1
-------------------------------------------------------------------------------

-- | Levenberg–Marquardt for functions with one parameter.
--
-- See 'levenbergMarquardt2' for examples, this is very similar.
--
-- For example we can fit \(f = x \mapsto \beta x + 1\), its derivative is \(\partial_\beta f = x \mapsto x\).
--
-- >>> let scale a (x, y) = (y, a * x + 1, x)
-- >>> PP $ NE.last $ levenbergMarquardt1 scale 1 input2
-- Fit 1.9685 0.04735 4 0.27914
--
-- Not bad, but worse then linear fit which fits the intercept point too.
--
levenbergMarquardt1
    :: F.Foldable f
    => (Double -> a -> (Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \partial_\beta f(\beta, x)\)
    -> Double                                     -- ^ initial parameter, \(\beta_0\)
    -> f a                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit Double)                   -- ^ non-empty list of iteration results
levenbergMarquardt1 f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    lambda0 = c11
      where
        c11 = getKBN $ lm1_c11 acc0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d) -> addLM1Acc acc y g d) zeroLM1Acc xs

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (lambda1 c11)
        delta     = mult matrix z1

        beta' = add beta delta
        acc'  = calcAcc beta'

        a11    = inv c11
        errors = sa

        wssr  = max 0 $ getKBN $ lm1_wssr acc
        wssr' =         getKBN $ lm1_wssr acc'

        ndf   = lm1_n acc - 1
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')

        c11 = getKBN $ lm1_c11 acc
        z1  = getKBN $ lm1_z1 acc

-- | 'levenbergMarquardt1' with weights.
levenbergMarquardt1WithWeights
    :: F.Foldable f
    => (Double -> a -> (Double, Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \partial_\beta f(\beta, x), w_i\)
    -> Double                                             -- ^ initial parameter, \(\beta_0\)
    -> f a                                                -- ^ data, \(d\)
    -> NE.NonEmpty (Fit Double)                           -- ^ non-empty list of iteration results
levenbergMarquardt1WithWeights f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    lambda0 = c11
      where
        c11 = getKBN $ lm1_c11 acc0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d, w) -> addLM1AccW acc y g d w) zeroLM1Acc xs

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (lambda1 c11)
        delta     = mult matrix z1

        beta' = add beta delta
        acc'  = calcAcc beta'

        a11    = inv c11
        errors = sa

        wssr  = max 0 $ getKBN $ lm1_wssr acc
        wssr' =         getKBN $ lm1_wssr acc'

        ndf   = lm1_n acc - 1
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')

        c11 = getKBN $ lm1_c11 acc
        z1  = getKBN $ lm1_z1 acc

-- | 'levenbergMarquardt1' with Y-errors.
levenbergMarquardt1WithYerrors
    :: F.Foldable f
    => (Double -> a -> (Double, Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \partial_\beta f(\beta, x), \delta y_i\)
    -> Double                                             -- ^ initial parameter, \(\beta_0\)
    -> f a                                                -- ^ data, \(d\)
    -> NE.NonEmpty (Fit Double)                           -- ^ non-empty list of iteration results
levenbergMarquardt1WithYerrors f = levenbergMarquardt1WithWeights f' where
    f' beta x = case f beta x of (y, fbetax, grad, dy) -> (y, fbetax, grad, recip $ sq dy)

-- | 'levenbergMarquardt1' with XY-errors.
levenbergMarquardt1WithXYerrors
    :: F.Foldable f
    => (Double -> a -> (Double, Double, Double, Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \partial_\beta f(\beta, x), \partial_x f(\beta, x), \delta x_i, \delta y_i\)
    -> Double                                                             -- ^ initial parameter, \(\beta_0\)
    -> f a                                                                -- ^ data, \(d\)
    -> NE.NonEmpty (Fit Double)                                           -- ^ non-empty list of iteration results
levenbergMarquardt1WithXYerrors g = levenbergMarquardt1WithWeights g' where
    g' beta x = case g beta x of (y, fbetax, grad, f', dx, dy) -> (y, fbetax, grad, recip $ sq (f' * dx) + sq dy)

data LM1Acc = LM1Acc
    { lm1_n    :: !Int
    , lm1_c11  :: !KBN
    , lm1_z1   :: !KBN
    , lm1_wssr :: !KBN
    }
  deriving Show

zeroLM1Acc :: LM1Acc
zeroLM1Acc = LM1Acc 0 zeroKBN zeroKBN zeroKBN

addLM1Acc :: LM1Acc -> Double -> Double -> Double -> LM1Acc
addLM1Acc LM1Acc {..} y f d1 = LM1Acc
    { lm1_n    = lm1_n + 1
    , lm1_c11  = addKBN lm1_c11  (d1 * d1)
    , lm1_z1   = addKBN lm1_z1   (d1 * res)
    , lm1_wssr = addKBN lm1_wssr (res * res)
    }
  where
    res = y - f

addLM1AccW :: LM1Acc -> Double -> Double -> Double -> Double -> LM1Acc
addLM1AccW LM1Acc {..} y f d1 w = LM1Acc
    { lm1_n    = lm1_n + 1
    , lm1_c11  = addKBN lm1_c11  (w * d1 * d1)
    , lm1_z1   = addKBN lm1_z1   (w * d1 * res)
    , lm1_wssr = addKBN lm1_wssr (w * res * res)
    }
  where
    res = y - f

-------------------------------------------------------------------------------
-- Levenberg–Marquardt 2
-------------------------------------------------------------------------------

-- | Levenberg–Marquardt for functions with two parameters.
--
-- You can use this sledgehammer to do a a linear fit:
--
-- >>> let lin (V2 a b) (x, y) = (y, a * x + b, V2 x 1)
--
-- We can then use 'levenbergMarquardt2' to find a fit:
--
-- >>> PP $ levenbergMarquardt2 lin (V2 1 1) input2
-- Fit (V2 1.00000 1.00000) (V2 1.0175 2.5385) 3 29.470
-- Fit (V2 1.0181 1.0368) (V2 0.98615 2.4602) 3 27.681
-- Fit (V2 1.1557 1.2988) (V2 0.75758 1.8900) 3 16.336
-- Fit (V2 1.5463 1.6577) (V2 0.29278 0.73043) 3 2.4400
-- Fit (V2 1.9129 1.1096) (V2 0.11033 0.27524) 3 0.34645
-- Fit (V2 2.0036 0.89372) (V2 0.09552 0.23830) 3 0.25970
-- Fit (V2 2.0063 0.88687) (V2 0.09550 0.23826) 3 0.25962
-- Fit (V2 2.0063 0.88685) (V2 0.09550 0.23826) 3 0.25962
--
-- This is the same result what 'linearFit' returns:
--
-- >>> PP $ linearFit id input2
-- Fit (V2 2.0063 0.88685) (V2 0.09550 0.23826) 3 0.25962
--
-- == Using AD
--
-- You can use @ad@ to calculate derivatives for you.
--
-- >>> import qualified Numeric.AD.Mode.Reverse.Double as AD
--
-- We need a ('Traversable') homogenic triple to represent the two parameters and @x@:
--
-- >>> data H3 a = H3 a a a deriving (Functor, Foldable, Traversable)
--
-- Then we define a function @ad@ can operate with:
--
-- >>> let linearF (H3 a b x) = a * x + b
--
-- which we can use to fit the curve in generic way:
--
-- >>> let lin' (V2 a b) (x, y) = case AD.grad' linearF (H3 a b x) of (f, H3 da db _f') -> (y, f, V2 da db)
-- >>> PP $ levenbergMarquardt2 lin' (V2 1 1) input2
-- Fit (V2 1.00000 1.00000) (V2 1.0175 2.5385) 3 29.470
-- Fit (V2 1.0181 1.0368) (V2 0.98615 2.4602) 3 27.681
-- Fit (V2 1.1557 1.2988) (V2 0.75758 1.8900) 3 16.336
-- Fit (V2 1.5463 1.6577) (V2 0.29278 0.73043) 3 2.4400
-- Fit (V2 1.9129 1.1096) (V2 0.11033 0.27524) 3 0.34645
-- Fit (V2 2.0036 0.89372) (V2 0.09552 0.23830) 3 0.25970
-- Fit (V2 2.0063 0.88687) (V2 0.09550 0.23826) 3 0.25962
-- Fit (V2 2.0063 0.88685) (V2 0.09550 0.23826) 3 0.25962
--
-- == Non-polynomial example
--
-- We can fit other curves too, for example an example from Wikipedia
-- https://en.wikipedia.org/wiki/Gauss%E2%80%93Newton_algorithm#Example
--
-- >>> let rateF (H3 vmax km s) = (vmax * s) / (km + s)
-- >>> let rateF' (V2 vmax km) (x, y) = case AD.grad' rateF (H3 vmax km x) of (f, H3 vmax' km' _) -> (y, f, V2 vmax' km')
-- >>> let input = zip [0.038,0.194,0.425,0.626,1.253,2.500,3.740] [0.050,0.127,0.094,0.2122,0.2729,0.2665,0.3317]
-- >>> PP $ levenbergMarquardt2 rateF' (V2 0.9 0.2) input
-- Fit (V2 0.90000 0.20000) (V2 0.43304 0.43936) 5 1.4455
-- Fit (V2 0.83306 0.25278) (V2 0.39164 0.49729) 5 1.0055
-- Fit (V2 0.59437 0.43508) (V2 0.21158 0.53403) 5 0.18832
-- Fit (V2 0.39687 0.56324) (V2 0.05723 0.25666) 5 0.01062
-- Fit (V2 0.36289 0.56104) (V2 0.04908 0.24007) 5 0.00784
-- Fit (V2 0.36190 0.55662) (V2 0.04887 0.23843) 5 0.00784
-- Fit (V2 0.36184 0.55629) (V2 0.04885 0.23830) 5 0.00784
-- Fit (V2 0.36184 0.55627) (V2 0.04885 0.23829) 5 0.00784
--
-- We get the same result as in the article: 0.362 and 0.556
--
-- The algorithm terminates when a scaling parameter \(\lambda\) becomes larger than 1e20 or smaller than 1e-20, or relative WSSR change is smaller than 1e-10, or sum-of-squared-residuals candidate becomes @NaN@ (i.e. when it would start to produce garbage).
-- You may want to terminate sooner, Numerical Recipes suggest to stop when WSSR decreases by a neglible amount absolutely or fractionally.
--
levenbergMarquardt2
    :: F.Foldable f
    => (V2 -> a -> (Double, Double, V2))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x)\)
    -> V2                                 -- ^ initial parameters, \(\beta_0\)
    -> f a                                -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V2)               -- ^ non-empty list of iteration results
levenbergMarquardt2 f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d) -> addLM2Acc acc y g d) zeroLM2Acc xs

    lambda0 = max l1 l2
      where
        V2 l1 l2 = eigenSM22 c11 c12 c22
        c11 = getKBN $ lm2_c11 acc0
        c12 = getKBN $ lm2_c11 acc0
        c22 = getKBN $ lm2_c22 acc0

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (SM22 (lambda1 c11) c12 (lambda1 c22))
        delta     = mult matrix (V2 z1 z2)

        beta' = add beta delta
        acc'  = calcAcc beta'

        SM22 a11 _ a22 = inv (SM22 c11 c12 c22)
        errors = V2 sa sb

        wssr  = max 0 $ getKBN $ lm2_wssr acc
        wssr' =         getKBN $ lm2_wssr acc'

        ndf   = lm2_n acc - 2
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')
        sb   = sqrt (a22 * wssr / ndf')

        c11 = getKBN $ lm2_c11 acc
        c12 = getKBN $ lm2_c12 acc
        c22 = getKBN $ lm2_c22 acc
        z1  = getKBN $ lm2_z1 acc
        z2  = getKBN $ lm2_z2 acc

-- | 'levenbergMarquardt2' with weights.
--
-- Because 'levenbergMarquardt2' is an iterative algorithm,
-- not only we can use it to fit curves with known y-errors ('levenbergMarquardt2WithYerrors'),
-- but also with both x and y-errors ('levenbergMarquardt2WithXYerrors').
--
levenbergMarquardt2WithWeights
    :: F.Foldable f
    => (V2 -> a -> (Double, Double, V2, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), w_i\)
    -> V2                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V2)                       -- ^ non-empty list of iteration results
levenbergMarquardt2WithWeights f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    lambda0 = max l1 l2
      where
        V2 l1 l2 = eigenSM22 c11 c12 c22
        c11 = getKBN $ lm2_c11 acc0
        c12 = getKBN $ lm2_c11 acc0
        c22 = getKBN $ lm2_c22 acc0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d, w) -> addLM2AccW acc y g d w) zeroLM2Acc xs

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (SM22 (lambda1 c11) c12 (lambda1 c22))
        delta     = mult matrix (V2 z1 z2)

        beta' = add beta delta
        acc'  = calcAcc beta'

        SM22 a11 _ a22 = inv (SM22 c11 c12 c22)
        errors = V2 sa sb

        wssr  = max 0 $ getKBN $ lm2_wssr acc
        wssr' =         getKBN $ lm2_wssr acc'

        ndf   = lm2_n acc - 2
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')
        sb   = sqrt (a22 * wssr / ndf')

        c11 = getKBN $ lm2_c11 acc
        c12 = getKBN $ lm2_c12 acc
        c22 = getKBN $ lm2_c22 acc
        z1  = getKBN $ lm2_z1 acc
        z2  = getKBN $ lm2_z2 acc

-- | 'levenbergMarquardt2' with Y-errors.
levenbergMarquardt2WithYerrors
    :: F.Foldable f
    => (V2 -> a -> (Double, Double, V2, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), \delta y_i\)
    -> V2                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V2)                       -- ^ non-empty list of iteration results
levenbergMarquardt2WithYerrors f = levenbergMarquardt2WithWeights f' where
    f' beta x = case f beta x of (y, fbetax, grad, dy) -> (y, fbetax, grad, recip $ sq dy)

-- | 'levenbergMarquardt2' with XY-errors.
levenbergMarquardt2WithXYerrors
    :: F.Foldable f
    => (V2 -> a -> (Double, Double, V2, Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), \partial_x f(\beta, x), \delta x_i, \delta y_i\)
    -> V2                                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V2)                                       -- ^ non-empty list of iteration results
levenbergMarquardt2WithXYerrors g = levenbergMarquardt2WithWeights g' where
    g' beta x = case g beta x of (y, fbetax, grad, f', dx, dy) -> (y, fbetax, grad, recip $ sq (f' * dx) + sq dy)

data LM2Acc = LM2Acc
    { lm2_n    :: !Int
    , lm2_c11  :: !KBN
    , lm2_c12  :: !KBN
    , lm2_c22  :: !KBN
    , lm2_z1   :: !KBN
    , lm2_z2   :: !KBN
    , lm2_wssr :: !KBN
    }
  deriving Show

zeroLM2Acc :: LM2Acc
zeroLM2Acc = LM2Acc 0 zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN

addLM2Acc :: LM2Acc -> Double -> Double -> V2 -> LM2Acc
addLM2Acc LM2Acc {..} y f (V2 d1 d2) = LM2Acc
    { lm2_n    = lm2_n + 1
    , lm2_c11  = addKBN lm2_c11  (d1 * d1)
    , lm2_c12  = addKBN lm2_c12  (d1 * d2)
    , lm2_c22  = addKBN lm2_c22  (d2 * d2)
    , lm2_z1   = addKBN lm2_z1   (d1 * res)
    , lm2_z2   = addKBN lm2_z2   (d2 * res)
    , lm2_wssr = addKBN lm2_wssr (res * res)
    }
  where
    res = y - f

addLM2AccW :: LM2Acc -> Double -> Double -> V2 -> Double -> LM2Acc
addLM2AccW LM2Acc {..} y f (V2 d1 d2) w = LM2Acc
    { lm2_n    = lm2_n + 1
    , lm2_c11  = addKBN lm2_c11  (w * d1 * d1)
    , lm2_c12  = addKBN lm2_c12  (w * d1 * d2)
    , lm2_c22  = addKBN lm2_c22  (w * d2 * d2)
    , lm2_z1   = addKBN lm2_z1   (w * d1 * res)
    , lm2_z2   = addKBN lm2_z2   (w * d2 * res)
    , lm2_wssr = addKBN lm2_wssr (w * res * res)
    }
  where
    res = y - f

-------------------------------------------------------------------------------
-- Levenberg–Marquardt 3
-------------------------------------------------------------------------------

-- | Levenberg–Marquardt for functions with three parameters.
--
-- See 'levenbergMarquardt2' for examples, this is very similar.
--
-- >>> let quad (V3 a b c) (x, y) = (y, a * x * x + b * x + c, V3 (x * x) x 1)
-- >>> PP $ NE.last $ levenbergMarquardt3 quad (V3 2 2 2) input3
-- Fit (V3 1.00000 0.00000 2.0000) (V3 0.00000 0.00000 0.00000) 1 0.00000
--
-- Same as quadratic fit, just less direct:
--
-- >>> PP $ quadraticFit id input3
-- Fit (V3 1.00000 0.00000 2.0000) (V3 0.00000 0.00000 0.00000) 1 0.00000
--
levenbergMarquardt3
    :: F.Foldable f
    => (V3 -> a -> (Double, Double, V3))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x)\)
    -> V3                                 -- ^ initial parameters, \(\beta_0\)
    -> f a                                -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V3)               -- ^ non-empty list of iteration results
levenbergMarquardt3 f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d) -> addLM3Acc acc y g d) zeroLM3Acc xs

    -- frobenius norm is larger than largest eigen value.
    -- calculating the eigen values for 3x3 (symmetric) matrix is becoming complicated.
    lambda0 = sqrt $ sumKBN [ sq c11
                            , 2 * sq c12, sq c22
                            , 2 * sq c13, 2 * sq c23, sq c33]
      where
        c11 = getKBN $ lm3_c11 acc0
        c12 = getKBN $ lm3_c12 acc0
        c13 = getKBN $ lm3_c13 acc0
        c22 = getKBN $ lm3_c22 acc0
        c23 = getKBN $ lm3_c23 acc0
        c33 = getKBN $ lm3_c33 acc0

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (SM33 (lambda1 c11)
                               c12          (lambda1 c22)
                               c13          c23           (lambda1 c33))
        delta     = mult matrix (V3 z1 z2 z3)

        beta' = add beta delta
        acc'  = calcAcc beta'

        SM33 a11
             _  a22
             _  _   a33 = inv (SM33 c11
                                    c12 c22
                                    c13 c23 c33)
        errors = V3 sa sb sc

        wssr  = max 0 $ getKBN $ lm3_wssr acc
        wssr' =         getKBN $ lm3_wssr acc'

        ndf   = lm3_n acc - 3
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')
        sb   = sqrt (a22 * wssr / ndf')
        sc   = sqrt (a33 * wssr / ndf')

        c11 = getKBN $ lm3_c11 acc
        c12 = getKBN $ lm3_c12 acc
        c13 = getKBN $ lm3_c13 acc
        c22 = getKBN $ lm3_c22 acc
        c23 = getKBN $ lm3_c23 acc
        c33 = getKBN $ lm3_c33 acc
        z1  = getKBN $ lm3_z1 acc
        z2  = getKBN $ lm3_z2 acc
        z3  = getKBN $ lm3_z3 acc

-- | 'levenbergMarquardt3' with weights.
levenbergMarquardt3WithWeights
    :: F.Foldable f
    => (V3 -> a -> (Double, Double, V3, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), w_i\)
    -> V3                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V3)                       -- ^ non-empty list of iteration results
levenbergMarquardt3WithWeights f b0 xs = loop lambda0 b0 acc0 where
    acc0 = calcAcc b0

    lambda0 = sqrt $ sumKBN [ sq c11
                            , 2 * sq c12, sq c22
                            , 2 * sq c13, 2 * sq c23, sq c33]
      where
        c11 = getKBN $ lm3_c11 acc0
        c12 = getKBN $ lm3_c12 acc0
        c13 = getKBN $ lm3_c13 acc0
        c22 = getKBN $ lm3_c22 acc0
        c23 = getKBN $ lm3_c23 acc0
        c33 = getKBN $ lm3_c33 acc0

    calcAcc beta = F.foldl' (\acc p -> case f beta p of (y, g, d, w) -> addLM3AccW acc y g d w) zeroLM3Acc xs

    loop lambda beta acc
        | lmStop lambda wssr wssr'
        = Fit beta errors ndf wssr NE.:| []

        | wssr' >= wssr
        = loop (lambda * 10) beta acc

        | otherwise
        = Fit beta errors ndf wssr `NE.cons` loop (lambda / 10) beta' acc'

      where
        lambda1 z = (1 + lambda) * z
        matrix    = inv (SM33 (lambda1 c11)
                               c12          (lambda1 c22)
                               c13          c23           (lambda1 c33))
        delta     = mult matrix (V3 z1 z2 z3)

        beta' = add beta delta
        acc'  = calcAcc beta'

        SM33 a11 _ _ a22 _ a33 = inv (SM33 c11 c12 c13 c22 c23 c33)
        errors = V3 sa sb sc

        wssr  = max 0 $ getKBN $ lm3_wssr acc
        wssr' =         getKBN $ lm3_wssr acc'

        ndf   = lm3_n acc - 3
        ndf'  = fromIntegral ndf :: Double

        sa   = sqrt (a11 * wssr / ndf')
        sb   = sqrt (a22 * wssr / ndf')
        sc   = sqrt (a33 * wssr / ndf')

        c11 = getKBN $ lm3_c11 acc
        c12 = getKBN $ lm3_c12 acc
        c13 = getKBN $ lm3_c13 acc
        c22 = getKBN $ lm3_c22 acc
        c23 = getKBN $ lm3_c23 acc
        c33 = getKBN $ lm3_c33 acc
        z1  = getKBN $ lm3_z1 acc
        z2  = getKBN $ lm3_z2 acc
        z3  = getKBN $ lm3_z3 acc

-- | 'levenbergMarquardt3' with Y-errors.
levenbergMarquardt3WithYerrors
    :: F.Foldable f
    => (V3 -> a -> (Double, Double, V3, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), \delta y_i\)
    -> V3                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V3)                       -- ^ non-empty list of iteration results
levenbergMarquardt3WithYerrors f = levenbergMarquardt3WithWeights f' where
    f' beta x = case f beta x of (y, fbetax, grad, dy) -> (y, fbetax, grad, recip $ sq dy)

-- | 'levenbergMarquardt3' with XY-errors.
levenbergMarquardt3WithXYerrors
    :: F.Foldable f
    => (V3 -> a -> (Double, Double, V3, Double, Double, Double))  -- ^ \(\beta, d_i \mapsto y_i, f(\beta, x_i), \nabla_\beta f(\beta, x), \partial_x f(\beta, x), \delta x_i, \delta y_i\)
    -> V3                                                         -- ^ initial parameters, \(\beta_0\)
    -> f a                                                        -- ^ data, \(d\)
    -> NE.NonEmpty (Fit V3)                                       -- ^ non-empty list of iteration results
levenbergMarquardt3WithXYerrors g = levenbergMarquardt3WithWeights g' where
    g' beta x = case g beta x of (y, fbetax, grad, f', dx, dy) -> (y, fbetax, grad, recip $ sq (f' * dx) + sq dy)

data LM3Acc = LM3Acc
    { lm3_n    :: !Int
    , lm3_c11  :: !KBN
    , lm3_c12  :: !KBN
    , lm3_c13  :: !KBN
    , lm3_c22  :: !KBN
    , lm3_c23  :: !KBN
    , lm3_c33  :: !KBN
    , lm3_z1   :: !KBN
    , lm3_z2   :: !KBN
    , lm3_z3   :: !KBN
    , lm3_wssr :: !KBN
    }
  deriving Show

zeroLM3Acc :: LM3Acc
zeroLM3Acc = LM3Acc 0 zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN

addLM3Acc :: LM3Acc -> Double -> Double -> V3 -> LM3Acc
addLM3Acc LM3Acc {..} y f (V3 d1 d2 d3) = LM3Acc
    { lm3_n    = lm3_n + 1
    , lm3_c11  = addKBN lm3_c11  (d1 * d1)
    , lm3_c12  = addKBN lm3_c12  (d1 * d2)
    , lm3_c13  = addKBN lm3_c12  (d1 * d3)
    , lm3_c22  = addKBN lm3_c22  (d2 * d2)
    , lm3_c23  = addKBN lm3_c22  (d2 * d3)
    , lm3_c33  = addKBN lm3_c22  (d3 * d3)
    , lm3_z1   = addKBN lm3_z1   (d1 * res)
    , lm3_z2   = addKBN lm3_z2   (d2 * res)
    , lm3_z3   = addKBN lm3_z3   (d3 * res)
    , lm3_wssr = addKBN lm3_wssr (res * res)
    }
  where
    res = y - f

addLM3AccW :: LM3Acc -> Double -> Double -> V3 -> Double -> LM3Acc
addLM3AccW LM3Acc {..} y f (V3 d1 d2 d3) w = LM3Acc
    { lm3_n    = lm3_n + 1
    , lm3_c11  = addKBN lm3_c11  (w * d1 * d1)
    , lm3_c12  = addKBN lm3_c12  (w * d1 * d2)
    , lm3_c13  = addKBN lm3_c12  (w * d1 * d3)
    , lm3_c22  = addKBN lm3_c22  (w * d2 * d2)
    , lm3_c23  = addKBN lm3_c22  (w * d2 * d3)
    , lm3_c33  = addKBN lm3_c22  (w * d3 * d3)
    , lm3_z1   = addKBN lm3_z1   (w * d1 * res)
    , lm3_z2   = addKBN lm3_z2   (w * d2 * res)
    , lm3_z3   = addKBN lm3_z3   (w * d3 * res)
    , lm3_wssr = addKBN lm3_wssr (w * res * res)
    }
  where
    res = y - f

-------------------------------------------------------------------------------
-- Output
-------------------------------------------------------------------------------

-- | Result of a curve fit.
data Fit v = Fit
    { fitParams :: !v       -- ^ fit parameters
    , fitErrors :: !v       -- ^ asympotic standard errors, /assuming a good fit/
    , fitNDF    :: !Int     -- ^ number of degrees of freedom
    , fitWSSR   :: !Double  -- ^ sum of squares of residuals
    }
  deriving Show

instance (NFData v) => NFData (Fit v) where
    rnf (Fit p e _ _) = rnf p `seq` rnf e

-------------------------------------------------------------------------------
-- LinRegAcc
-------------------------------------------------------------------------------

-- | Linear regression accumulator.
data LinRegAcc = LinRegAcc
    { lra_n  :: {-# UNPACK #-} !Int  -- ^ \(n\)
    , lra_w  :: {-# UNPACK #-} !KBN  -- ^ \(\sum w_i\)
    , lra_x  :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i \)
    , lra_x2 :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i^2 \)
    , lra_y  :: {-# UNPACK #-} !KBN  -- ^ \(\sum y_i \)
    , lra_xy :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i y_i \)
    , lra_y2 :: {-# UNPACK #-} !KBN  -- ^ \(\sum y_i^2 \)
    }
  deriving Show

instance NFData LinRegAcc where
    rnf LinRegAcc {} = ()

-- | All-zeroes 'LinRegAcc'.
zeroLinRegAcc :: LinRegAcc
zeroLinRegAcc = LinRegAcc 0 zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN

-- | Add a point to linreg accumulator.
addLinReg
    :: LinRegAcc
    -> Double   -- ^ x
    -> Double   -- ^ y
    -> LinRegAcc
addLinReg LinRegAcc {..} x y = LinRegAcc
    { lra_n  = lra_n + 1
    , lra_w  = addKBN lra_w  1
    , lra_x  = addKBN lra_x  x
    , lra_x2 = addKBN lra_x2 (x * x)
    , lra_y  = addKBN lra_y  y
    , lra_xy = addKBN lra_xy (x * y)
    , lra_y2 = addKBN lra_y2 (y * y)
    }

-- | Add a weighted point to linreg accumulator.
addLinRegW
    :: LinRegAcc
    -> Double   -- ^ x
    -> Double   -- ^ y
    -> Double   -- ^ w
    -> LinRegAcc
addLinRegW LinRegAcc {..} x y w = LinRegAcc
    { lra_n  = lra_n + 1
    , lra_w  = addKBN lra_w  w
    , lra_x  = addKBN lra_x  (w * x)
    , lra_x2 = addKBN lra_x2 (w * x * x)
    , lra_y  = addKBN lra_y  (w * y)
    , lra_xy = addKBN lra_xy (w * x * y)
    , lra_y2 = addKBN lra_y2 (w * y * y)
    }

-------------------------------------------------------------------------------
-- QuadRegAcc
-------------------------------------------------------------------------------

-- | Quadratic regression accumulator.
data QuadRegAcc = QuadRegAcc
    { qra_n   :: {-# UNPACK #-} !Int  -- ^ \(n\)
    , qra_w   :: {-# UNPACK #-} !KBN  -- ^ \(\sum w_i\)
    , qra_x   :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i \)
    , qra_x2  :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i^2 \)
    , qra_x3  :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i^3 \)
    , qra_x4  :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i^4 \)
    , qra_y   :: {-# UNPACK #-} !KBN  -- ^ \(\sum y_i \)
    , qra_xy  :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i y_i \)
    , qra_x2y :: {-# UNPACK #-} !KBN  -- ^ \(\sum x_i^2 y_i \)
    , qra_y2  :: {-# UNPACK #-} !KBN  -- ^ \(\sum y_i^2 \)
    }
  deriving Show

instance NFData QuadRegAcc where
    rnf QuadRegAcc {} = ()

-- | All-zeroes 'QuadRegAcc'.
zeroQuadRegAcc :: QuadRegAcc
zeroQuadRegAcc = QuadRegAcc 0 zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN zeroKBN

-- | Add a point to quadreg accumulator.
addQuadReg
    :: QuadRegAcc
    -> Double  -- ^ x
    -> Double  -- ^ y
    -> QuadRegAcc
addQuadReg QuadRegAcc {..} x y = QuadRegAcc
    { qra_n    = qra_n + 1
    , qra_w    = addKBN qra_w   1
    , qra_x    = addKBN qra_x   x
    , qra_x2   = addKBN qra_x2  x2
    , qra_x3   = addKBN qra_x3  (x * x2)
    , qra_x4   = addKBN qra_x4  (x2 * x2)
    , qra_y    = addKBN qra_y   y
    , qra_xy   = addKBN qra_xy  (x * y)
    , qra_x2y  = addKBN qra_x2y (x2 * y)
    , qra_y2   = addKBN qra_y2  (y * y)
    }
  where
    x2 = x * x

-- | Add a weighted point to quadreg accumulator.
addQuadRegW
    :: QuadRegAcc
    -> Double  -- ^ x
    -> Double  -- ^ y
    -> Double  -- ^ w
    -> QuadRegAcc
addQuadRegW QuadRegAcc {..} x y w = QuadRegAcc
    { qra_n    = qra_n + 1
    , qra_w    = addKBN qra_w   w
    , qra_x    = addKBN qra_x   (w * x)
    , qra_x2   = addKBN qra_x2  (w * x2)
    , qra_x3   = addKBN qra_x3  (w * x * x2)
    , qra_x4   = addKBN qra_x4  (w * x2 * x2)
    , qra_y    = addKBN qra_y   (w * y)
    , qra_xy   = addKBN qra_xy  (w * x * y)
    , qra_x2y  = addKBN qra_x2y (w * x2 * y)
    , qra_y2   = addKBN qra_y2  (w * y * y)
    }
  where
    x2 = x * x

-- | Convert 'QuadRegAcc' to 'LinRegAcc'.
--
-- Using this we can try quadratic and linear fits with a single data scan.
--
quadRegAccToLin :: QuadRegAcc -> LinRegAcc
quadRegAccToLin QuadRegAcc {..} = LinRegAcc
    { lra_n  = qra_n
    , lra_w  = qra_w
    , lra_x  = qra_x
    , lra_x2 = qra_x2
    , lra_y  = qra_y
    , lra_xy = qra_xy
    , lra_y2 = qra_y2
    }

-------------------------------------------------------------------------------
-- utils
-------------------------------------------------------------------------------

sq :: Num a => a -> a
sq x = x * x
{-# INLINE sq #-}

iterate1 :: (b -> b) -> b -> NE.NonEmpty b
iterate1 g x = NE.cons x (iterate1 g (g x))

eigenSM22 :: Double -> Double -> Double -> V2
eigenSM22 a b c = V2 ((ac + discr) / 2) ((ac - discr) / 2)
  where
    ac = a + c
    discr = sqrt (sq (a - c) + 4 * sq b)

-- | Levenberg-Marquard stop condition
lmStop :: Double -> Double -> Double -> Bool
lmStop lambda wssr wssr' =
    lambda < 1e-20 || lambda > 1e20 || isNaN wssr' || relDiff < 1e-10
  where
    relDiff = abs (wssr' - wssr) / wssr


import Data.Complex

type CD = Complex Double

quadraticRoots :: (CD, CD, CD) -> (CD, CD)
quadraticRoots (a, b, c) =
    if   realPart b > 0
    then ((2*c) / (-b - d), (-b - d) / (2*a))
    else ((-b + d) / (2*a), (2*c) / (-b + d))
  where d = sqrt $ b^2 - 4*a*c


--{-# LANGUAGE BangPatterns #-}
module NeuralNetworks where
import Util
import Data.List
import System.Random
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util
import Numeric.GSL.Minimization
import Control.Parallel (par,pseq)
import Debug.Trace
import System.IO
import System.Directory

readThetaList :: [Int] -> IO [Matrix Double]
readThetaList layers = do
                          let readThetaListInternal fname = do
                                                              contents <- readFile fname
                                                              let thetas = contents `seq` read contents :: [Matrix Double]
                                                              return thetas
                          let generateThetaList = do
                                                    rgen <- getStdGen
                                                    return $ randomInitializeAllWeights rgen layers
                          let fileName = show layers
                          fexist <- doesFileExist fileName
                          if fexist
                          then
                            readThetaListInternal fileName
                          else
                            generateThetaList

writeThetaList :: [Int] -> [Matrix Double] -> IO ()
writeThetaList layers thetas = do
                                let filename = show layers
                                let contents = show thetas
                                writeFile filename contents

{- RANDINITIALIZEWEIGHTS Randomly initialize the weights of a layer with L_in
   incoming connections and L_out outgoing connections
   W = RANDINITIALIZEWEIGHTS(L_in, L_out) randomly initializes the weights
   of a layer with L_in incoming connections and L_out outgoing
   connections.
   Note that W should be set to a matrix of size(L_out, 1 + L_in) as
   the first row of W handles the "bias" terms -}
randomInitializeWeights :: RandomGen g => g -> Int -> Int -> Matrix Double
randomInitializeWeights gen l_in l_out = randommatrix * 2 * epsilon_init - epsilon_init
     where
      randomlist = randoms gen :: [Double]
      epsilon_init = 0.12
      randommatrix = (l_out><(1+l_in)) randomlist

randomInitializeAllWeights :: RandomGen g => g -> [Int] -> [Matrix Double]
randomInitializeAllWeights gen (outputs:[]) = []
randomInitializeAllWeights gen (inputs:outputs:rest) = (randomInitializeWeights gen inputs outputs):(randomInitializeAllWeights gen (outputs:rest))

{- Implements the neural network cost function for a neural network which performs
   classification.
   [J, grad] = NNCOSTFUNCTON(nn_params, hidden_layer_size, num_labels, ...
   X, y, lambda) computes the cost and gradient of the neural network. The
   parameters for the neural network are "unrolled" into the vector
   nn_params and need to be converted back into the weight matrices.
   The returned parameter grad should be a "unrolled" vector of the
   partial derivatives of the neural network.

   NOTE: cost does not work on NNs with no hidden layers.  TODO fix this if
   I ever need it.

   -}
costFunction :: Vector Double -> [Int] -> Matrix Double -> Matrix Double -> Double -> (Double, Vector Double)
costFunction thetaVector layers x targetYsMatrix lambda = (j_reg, gradientVector)
    where
      m = fromIntegral $ rows targetYsMatrix
      thetaList = reshapeVector layers thetaVector
      (inputMatrices, activationMatrices, output_layer) = forwardPropagate x thetaList
      --j = trace (show $ -((1/m) * sumElements (targetYsMatrix * (log output_layer) + (1 - targetYsMatrix)  * (log (1 - output_layer))))) (-((1/m) * ( sumElements (targetYsMatrix * (log output_layer) + (1 - targetYsMatrix) * (log (1 - output_layer))))))
      --j = trace (show $ ((targetYsMatrix * (log output_layer) + (1 - targetYsMatrix) * (log (1 - output_layer))))) (-((1/m) * ( sumElements (targetYsMatrix * (log output_layer) + (1 - targetYsMatrix) * (log (1 - output_layer))))))
      j = -((1/m) * ( sumElements (targetYsMatrix * (log output_layer) + (1 - targetYsMatrix) * (log (1 - output_layer)))))
      j_reg = j + (lambda/(2*m)) * ( sum ( map (\ theta -> sumElements ((dropColumns 1 theta) ^ 2)) thetaList ))
      deltaList = backPropagate thetaList inputMatrices activationMatrices output_layer targetYsMatrix
      gradientVector = createGradientVector thetaList deltaList
        where
          -- create list of gradient matrices and then flatten them into a single vector
          createGradientVector :: [Matrix Double] -> [Matrix Double] -> Vector Double
          createGradientVector tlist dlist = flattenConcat $ iterate tlist dlist
            where
              iterate [] []         = []
              iterate (t:ts) (d:ds) = grad:(iterate ts ds)
                where
                  -- liminate j = 0 column from theta for purposes of regularization
                  tmpTheta = padLeft (dropColumns 1 t) 0
                  -- regularized gradients
                  grad =  (1/(scalar m)) * d + (scalar (lambda/m)) * tmpTheta

-- |Performs forward propagation on a neural network, returning the resulting output layer
forwardPropagate :: Matrix Double -> [Matrix Double] -> ([Matrix Double],[Matrix Double],Matrix Double)
--forwardPropagate inputMatrix _ | trace ("Forward propagate input matrix dimensions " ++ (show (rows inputMatrix)) ++ "x" ++ (show (cols inputMatrix)) ) False = undefined
forwardPropagate inputMatrix thetaList = (inputMatrices, activationMatrices, output_layer)
  where
    resultList' = doForwardProp inputMatrix thetaList
    --add x to the start of the input list, for both input and activation... testing
    resultList = (inputMatrix,inputMatrix):resultList'
    (inputMatrices, activationMatrices) = unzip resultList
    output_layer = last activationMatrices
    doForwardProp input []         = []
    doForwardProp input (t:thetas) = (initial_value,activation):(doForwardProp activation thetas)
      where
        input_with_bias = addBias input
        initial_value = input_with_bias <> (ctrans t)
        activation = sigmoid initial_value

-- |Performs backpropagation on a neural network, returning the resulting Delta matrices
backPropagate :: [Matrix Double] -> [Matrix Double] -> [Matrix Double] -> Matrix Double -> Matrix Double -> [Matrix Double]
backPropagate thetaList inputMatrices activationMatrices output_layer targetYs = output_deltas `par` doBackProp 0 deltaList
  where
    deltaList = makeZeroDeltaList thetaList
    m = rows targetYs
    -- reverse the lists we'll be iterating through, since we're going backwards from the output layer
    revInputs = reverse inputMatrices
    revActivations = reverse activationMatrices
    revThetaList = reverse thetaList
    output_deltas = (head revActivations) - targetYs
    --for each example in training set:
    doBackProp index ds | index >= m = ds
                        | otherwise  = doBackProp (index +1) updatedDs
      where
        -- reverse the lists we'll be iterating through, since we're going backwards from the output layer
        revDs = reverse ds
        --do output layer separately:
        output_delta = ctrans $ extractRows [index] output_deltas
        --do the rest of the layers:
        updatedDs = reverse $ backPropagateOneExample revThetaList output_delta (tail revInputs) (tail revActivations) revDs
        backPropagateOneExample _ _ _ _ [] = []
        backPropagateOneExample (thetaI:thetas) previous_delta (inputI:inputs) (activationI:activations) (deltaI:deltas) = current_delta `pseq` updatedDelta:(backPropagateOneExample thetas current_delta inputs activations deltas)
          where
            current_delta_with_bias = (ctrans thetaI) <> previous_delta * (sigmoidGradient (ctrans (addBias (extractRows [index] inputI))))
            current_delta = updatedDelta `par` dropRows 1 current_delta_with_bias
            -- capital Delta accumulation
            updatedDelta = deltaI + previous_delta <> addBias (extractRows [index] activationI)

-- |Makes a list of 0-filled Delta matrices from a list of theta matrices
makeZeroDeltaList :: [Matrix Double] -> [Matrix Double]
makeZeroDeltaList [] = []
makeZeroDeltaList (m:ms) = (zeros (rows m) (cols m)):(makeZeroDeltaList ms)

-- |Reshapes a vector into a list of weight matrices for a network of the given number and size of layers
reshapeVector :: [Int] -> Vector Double -> [Matrix Double]
reshapeVector layers vector = iterate layers (toList vector)
  where
    iterate _ [] = []
    iterate [] _ = []
    iterate (x:[]) _ = []
    iterate (l_in:l_out:xs) vlist = m:(iterate (l_out:xs) (drop msize vlist))
      where
        msize = l_out * (l_in + 1)
        m = reshape (l_in + 1) (fromList (take msize vlist))

-- |The prediction function
predict :: [Matrix Double] -> Matrix Double -> [Double] -> Matrix Double
predict [] input targetMapping = selectTarget
  where
    maxIndex xs = snd $ maximum $ zip xs targetMapping
    selectTarget = asColumn $ fromList $ map maxIndex $ toLists input
predict (theta:thetas) input targetMapping = predict thetas layer targetMapping
  where
    layer = sigmoid $ (addBias input) <> (ctrans theta)

{-
Minimization methods:
ConjugateFR
ConjugatePR
VectorBFGS
VectorBFGS2
SteepestDescent
-}

minimizeVDGeneric method x targetYsMatrix layers num_iters lambda precision initial_thetas = minimizeVD method precision num_iters first_step tol traceCostFunc gradientFunc (flattenConcat initial_thetas)
    where
      first_step = precision * 100
      tol = 0.1 -- see http://www.gnu.org/software/gsl/manual/html_node/Multimin-Algorithms-with-Derivatives.html
      costFunc thetaVect = fst $ costGradFunction thetaVect
      gradientFunc thetaVect = snd $ costGradFunction thetaVect
      costGradFunction t = (costFunction t layers x targetYsMatrix lambda)
      traceGradientFunc thetaVect = traceShow ((sumElements grad)) grad
        where grad = snd $ costGradFunction thetaVect
      traceCostFunc thetaVect = traceShow cost cost
        where cost = fst $ costGradFunction thetaVect
      --memoizedCostGradFunction = unsafeMemoize costGradFunction

-- Generate learning curves
learningCurvesNeuralNetwork::
  RandomGen g =>   MinimizeMethodD ->
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   [Int] ->
                   Double ->
                   Int ->
                   Double ->
                   g ->
                   Int ->
                   [(Double,Double,Int)]
learningCurvesNeuralNetwork method x y xval yval layers lambda iter precision gen step_size = learningCurves trainFunc costFunc x y xval yval step_size
  where
    initial_thetas = randomInitializeAllWeights gen layers
    trainFunc tx ty = asColumn $ fst $ minimizeVDGeneric method tx ty layers iter lambda precision initial_thetas
    costFunc tx ty tt = fst $ costFunction (flatten tt) layers tx ty  0


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.QuantumMat
Copyright   :  (c) Scott N. Walck 2016-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains state vectors and matrices
for quantum mechanics.
-}

-- Using only Complex Double here, no cyclotomic

module Physics.Learn.QuantumMat
    (
    -- * Complex numbers
      C
    -- * State Vectors
    , xp
    , xm
    , yp
    , ym
    , zp
    , zm
    , np
    , nm
    , dim
    , scaleV
    , inner
    , norm
    , normalize
    , probVector
    , gramSchmidt
    , conjV
    , fromList
    , toList
    -- * Matrices (operators)
    , sx
    , sy
    , sz
    , scaleM
    , (<>)
    , (#>)
    , (<#)
    , conjugateTranspose
    , fromLists
    , toLists
    , size
    , matrixFunction
    -- * Density matrices
    , couter
    , dm
    , trace
    , normalizeDM
    , oneQubitMixed
    -- * Quantum Dynamics
    , timeEvMat
    , timeEv
    , timeEvMatSpec
    -- * Composition
    , Kronecker(..)
    -- * Measurement
    , possibleOutcomes
    , outcomesProjectors
    , outcomesProbabilities
    -- * Vector and Matrix
    , Vector
    , Matrix
    )
    where

import Numeric.LinearAlgebra
    ( C
    , Vector
    , Matrix
    , Herm
    , iC        -- square root of negative one
    , (><)      -- matrix definition
    , ident
    , scale
    , norm_2
    , inv
    , (<\>)
    , sym
    , eigenvaluesSH
    , eigSH
    , cmap
    , takeDiag
    , conj
    , dot
    , tr
    )
--    , (<>)      -- matrix product (not * !!!!)
--    , (#>)      -- matrix-vector product
--    , fromList  -- vector definition

import qualified Numeric.LinearAlgebra as H
-- because H.outer does not conjugate
import Data.Complex
    ( Complex(..)
    , magnitude
    , realPart
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

-- | The state resulting from a measurement of
--   spin angular momentum in the x direction
--   on a spin-1/2 particle
--   when the result of the measurement is hbar/2.
xp :: Vector C
xp = normalize $ fromList [1, 1]

-- | The state resulting from a measurement of
--   spin angular momentum in the x direction
--   on a spin-1/2 particle
--   when the result of the measurement is -hbar/2.
xm :: Vector C
xm = normalize $ fromList [1, -1]

-- | The state resulting from a measurement of
--   spin angular momentum in the y direction
--   on a spin-1/2 particle
--   when the result of the measurement is hbar/2.
yp :: Vector C
yp = normalize $ fromList [1, iC]

-- | The state resulting from a measurement of
--   spin angular momentum in the y direction
--   on a spin-1/2 particle
--   when the result of the measurement is -hbar/2.
ym :: Vector C
ym = normalize $ fromList [1, -iC]

-- | The state resulting from a measurement of
--   spin angular momentum in the z direction
--   on a spin-1/2 particle
--   when the result of the measurement is hbar/2.
zp :: Vector C
zp = normalize $ fromList [1, 0]

-- | The state resulting from a measurement of
--   spin angular momentum in the z direction
--   on a spin-1/2 particle
--   when the result of the measurement is -hbar/2.
zm :: Vector C
zm = normalize $ fromList [0, 1]

-- | The state resulting from a measurement of
--   spin angular momentum in the direction
--   specified by spherical angles theta (polar angle)
--   and phi (azimuthal angle)
--   on a spin-1/2 particle
--   when the result of the measurement is hbar/2.
np :: Double -> Double -> Vector C
np theta phi = fromList [ cos (theta/2) :+ 0
                        , exp(0 :+ phi) * (sin (theta/2) :+ 0) ]

-- | The state resulting from a measurement of
--   spin angular momentum in the direction
--   specified by spherical angles theta (polar angle)
--   and phi (azimuthal angle)
--   on a spin-1/2 particle
--   when the result of the measurement is -hbar/2.
nm :: Double -> Double -> Vector C
nm theta phi = fromList [ sin (theta/2) :+ 0
                        , -exp(0 :+ phi) * (cos (theta/2) :+ 0) ]

-- | Dimension of a vector.
dim :: Vector C -> Int
dim = H.size

-- | Scale a complex vector by a complex number.
scaleV :: C -> Vector C -> Vector C
scaleV = scale

-- | Complex inner product.  First vector gets conjugated.
inner :: Vector C -> Vector C -> C
inner = dot

-- | Length of a complex vector.
norm :: Vector C -> Double
norm = norm_2

-- | Return a normalized version of a given state vector.
normalize :: Vector C -> Vector C
normalize v = scale (1 / norm_2 v :+ 0) v

-- | Return a vector of probabilities for a given state vector.
probVector :: Vector C       -- ^ state vector
           -> Vector Double  -- ^ vector of probabilities
probVector = cmap (\c -> magnitude c**2)

-- | Conjugate the entries of a vector.
conjV :: Vector C -> Vector C
conjV = conj

-- | Construct a vector from a list of complex numbers.
fromList :: [C] -> Vector C
fromList = H.fromList

-- | Produce a list of complex numbers from a vector.
toList :: Vector C -> [C]
toList = H.toList

--------------
-- Matrices --
--------------

-- | The Pauli X matrix.
sx :: Matrix C
sx = (2><2) [ 0, 1
            , 1, 0 ]

-- | The Pauli Y matrix.
sy :: Matrix C
sy = (2><2) [  0, -iC
            , iC,   0 ]

-- | The Pauli Z matrix.
sz :: Matrix C
sz = (2><2) [ 1,  0
            , 0, -1 ]

-- | Scale a complex matrix by a complex number.
scaleM :: C -> Matrix C -> Matrix C
scaleM = scale

-- | Matrix product.
(<>) :: Matrix C -> Matrix C -> Matrix C
(<>) = (H.<>)

-- | Matrix-vector product.
(#>) :: Matrix C -> Vector C -> Vector C
(#>) = (H.#>)

-- | Vector-matrix product
(<#) :: Vector C -> Matrix C -> Vector C
(<#) = (H.<#)

-- | Conjugate transpose of a matrix.
conjugateTranspose :: Matrix C -> Matrix C
conjugateTranspose = tr

-- | Construct a matrix from a list of lists of complex numbers.
fromLists :: [[C]] -> Matrix C
fromLists = H.fromLists

-- | Produce a list of lists of complex numbers from a matrix.
toLists :: Matrix C -> [[C]]
toLists = H.toLists

-- | Size of a matrix.
size :: Matrix C -> (Int,Int)
size = H.size

-- | Apply a function to a matrix.
--   Assumes the matrix is a normal matrix (a matrix
--   with an orthonormal basis of eigenvectors).
matrixFunction :: (C -> C) -> Matrix C -> Matrix C
matrixFunction f m
    = let (valv,vecm) = H.eig m
          fvalv = fromList [f val | val <- toList valv]
      in vecm <> H.diag fvalv <> tr vecm

----------------------
-- Density Matrices --
----------------------

-- | Complex outer product
couter :: Vector C -> Vector C -> Matrix C
couter v w = v `H.outer` conj w

-- | Build a pure-state density matrix from a state vector.
dm :: Vector C -> Matrix C
dm cvec = cvec `couter` cvec

-- | Trace of a matrix.
trace :: Matrix C -> C
trace = sum . toList . takeDiag

-- | Normalize a density matrix so that it has trace one.
normalizeDM :: Matrix C -> Matrix C
normalizeDM rho = scale (1 / trace rho) rho

-- | The one-qubit totally mixed state.
oneQubitMixed :: Matrix C
oneQubitMixed = normalizeDM $ ident 2

----------------------
-- Quantum Dynamics --
----------------------

-- | Given a time step and a Hamiltonian matrix,
--   produce a unitary time evolution matrix.
--   Unless you really need the time evolution matrix,
--   it is better to use 'timeEv', which gives the
--   same numerical results without doing an explicit
--   matrix inversion.  The function assumes hbar = 1.
timeEvMat :: Double -> Matrix C -> Matrix C
timeEvMat dt h
    = let ah = scale (0 :+ dt / 2) h
          (l,m) = size h
          n = if l == m then m else error "timeEv needs square Hamiltonian"
          identity = ident n
      in inv (identity + ah) <> (identity - ah)

-- | Given a time step and a Hamiltonian matrix,
--   advance the state vector using the Schrodinger equation.
--   This method should be faster than using 'timeEvMat'
--   since it solves a linear system rather than calculating
--   an inverse matrix.  The function assumes hbar = 1.
timeEv :: Double -> Matrix C -> Vector C -> Vector C
timeEv dt h v
    = let ah = scale (0 :+ dt / 2) h
          (l,m) = size h
          n = if l == m then m else error "timeEv needs square Hamiltonian"
          identity = ident n
      in (identity + ah) <\> ((identity - ah) #> v)

-- | Given a Hamiltonian matrix, return a function from time
--   to evolution matrix.  Uses spectral decomposition.
--   Assumes hbar = 1.
timeEvMatSpec :: Matrix C -> Double -> Matrix C
timeEvMatSpec m t = matrixFunction (\h -> exp(-iC * h * (t :+ 0))) m

-----------------
-- Composition --
-----------------

class Kronecker a where
    kron :: a -> a -> a

instance H.Product t => Kronecker (Vector t) where
    kron v1 v2 = H.fromList [c1 * c2 | c1 <- H.toList v1, c2 <- H.toList v2]

instance H.Product t => Kronecker (Matrix t) where
    kron = H.kronecker

-----------------
-- Measurement --
-----------------

-- | The possible outcomes of a measurement
--   of an observable.
--   These are the eigenvalues of the matrix
--   of the observable.
possibleOutcomes :: Matrix C -> [Double]
possibleOutcomes observable
    = H.toList $ eigenvaluesSH (sym observable)

-- From a Hermitian matrix, a list of pairs of eigenvalues and eigenvectors.
valsVecs :: Herm C -> [(Double,Vector C)]
valsVecs h = let (valv,m) = eigSH h
                 vals = H.toList valv
                 vecs = map (conjV . fromList) $ toLists (conjugateTranspose m)
             in zip vals vecs

-- From a Hermitian matrix, a list of pairs of eigenvalues and projectors.
valsPs :: Herm C -> [(Double,Matrix C)]
valsPs h = [(val,couter vec vec) | (val,vec) <- valsVecs h]

combineFst :: (Eq a, Num b) => [(a,b)] -> [(a,b)]
combineFst [] = []
combineFst [p] = [p]
combineFst ((x1,m1):(x2,m2):ps)
    = if x1 == x2
      then combineFst ((x1,m1+m2):ps)
      else (x1,m1):combineFst ((x2,m2):ps)

-- | Given an obervable, return a list of pairs
--   of possible outcomes and projectors
--   for each outcome.
outcomesProjectors :: Matrix C -> [(Double,Matrix C)]
outcomesProjectors m = combineFst (valsPs (sym m))

-- | Given an observable and a state vector, return a list of pairs
--   of possible outcomes and probabilites
--   for each outcome.
outcomesProbabilities :: Matrix C -> Vector C -> [(Double,Double)]
outcomesProbabilities m v
    = [(a,realPart (inner v (p #> v))) | (a,p) <- outcomesProjectors m]

------------------
-- Gram-Schmidt --
------------------

-- | Form an orthonormal list of complex vectors
--   from a linearly independent list of complex vectors.
gramSchmidt :: [Vector C] -> [Vector C]
gramSchmidt [] = []
gramSchmidt (v:vs) = let nvs = gramSchmidt vs
                         nv = normalize (v - sum [scale (inner w v) w | w <- nvs])
                     in nv:nvs

-- To Do
--   Generate higher spin operators and state vectors
--   eigenvectors
--   projection operators


{-# LANGUAGE RecordWildCards #-}

module Regression.Logistic
  ( computeCost
  , gradientDescent
  , minimizeBFGS2
  ) where

import Numeric.GSL.Minimization (MinimizeMethodD(VectorBFGS2), minimizeVD)
import Numeric.LinearAlgebra (( #> ), (<.>), add)
import Numeric.LinearAlgebra.Data
  ( Matrix
  , R
  , Vector
  , (?)
  , atIndex
  , cmap
  , cols
  , flatten
  , fromList
  , scalar
  , size
  , toList
  , tr'
  )
import Common (MinimizationOpts(..), sigmoid, sigmoidMatrix)

computeCost :: Matrix R -> Vector R -> Vector R -> R
computeCost x y theta =
  let s = sigmoidMatrix (x #> theta)
      m = size y
   in ((y <.> log (s)) + ((1 - y) <.> log (1 - s))) / fromIntegral (-m)

gradientDescent :: Matrix R -> Vector R -> Vector R -> R -> Int -> R -> Vector R
gradientDescent _ _ _ alpha numIters regFactor
  | regFactor < 0 = error "Regularization factor cannot be < 0"
  | alpha < 0 = error "Learning factor cannot be < 0"
  | numIters < 0 = error "Number of iterations cannot be < 0"
gradientDescent x y theta alpha depth regFactor = go 0 theta
  where
    alpha' = scalar alpha
    --
    m = fromIntegral $ size y
    --
    go k accum
      | k >= depth = accum
      | otherwise =
        let delta = computeDelta accum
         in go (k + 1) (accum - (alpha' * delta))
    --
    penalizedTheta = fromList $ 0 : (tail $ ((* regFactor) <$> toList theta))
    --
    computeDelta th =
      let delta = (tr' x) #> ((sigmoidMatrix (x #> th)) - y)
       in (delta + penalizedTheta) / scalar m

minimizeBFGS2 ::
     Matrix R
  -> Vector R
  -> Vector R
  -> Int
  -> MinimizationOpts
  -> R
  -> Vector R
minimizeBFGS2 _ _ _ numIters _ regFactor
  | regFactor < 0 = error "Regularization factor cannot be < 0"
  | numIters < 0 = error "Number of iterations cannot be < 0"
minimizeBFGS2 x y theta numIters MinimizationOpts {..} regFactor =
  let (params, _) =
        minimizeVD
          VectorBFGS2
          precision
          numIters
          sizeOfFirstTrialStep
          tolerance
          (computeCost x y)
          computeGradients
          theta
   in params
  where
    m = fromIntegral $ size y
    --
    xTr = tr' x
    --
    penalizedTheta = fromList $ 0 : (tail $ ((* regFactor) <$> toList theta))
    --
    computeGradients :: Vector R -> Vector R
    computeGradients th =
      let delta = (tr' x) #> ((sigmoidMatrix (x #> th)) - y)
       in (delta + penalizedTheta) / scalar m


{-# LANGUAGE TemplateHaskell #-}

module FractalsTest
( runTests
) where

import Test.QuickCheck
import Data.Complex

import Fractals
import Palette

prop_toGrayscaleLength :: [Double] -> Bool
prop_toGrayscaleLength ws = length (toBitmap grayscale ws) == 4 * (length ws)

inMandelbrotSet :: Int -> Complex Double -> Bool
inMandelbrotSet nMax z0 = nMax == (fst $ order (mandelbrot z0) 2.0 nMax z0)

prop_orderMandelbrotOutside :: Positive Int -> Complex Double -> Property
prop_orderMandelbrotOutside (Positive nMax) z0 =
    (mag2 z0 > 4.0) ==> not $ inMandelbrotSet nMax z0

toUnitCircle :: RealFloat a => Complex a -> Complex a -> Complex a
toUnitCircle w z = if magW>0.0 && magZ>0.0 then mu else (0.0 :+ 0.0)
    where magW = magnitude w
          magZ = magnitude z
          mu = if magW < magZ then w/(magZ:+0.0) else z/(magW:+0.0)

prop_orderMandelbrotCardioid :: Positive Int    -- ^ the max number of iterations
                             -> Complex Double  -- ^ required to generate a value in the unit disk
                             -> Complex Double  -- ^ required to generate a value in the unit disk
                             -> Bool
prop_orderMandelbrotCardioid (Positive nMax) w z = inMandelbrotSet nMax z0
    where z0 = mandelbrotCardioid $ toUnitCircle w z

prop_orderMandelbrotCircle :: Positive Int    -- ^ the max number of iterations
                           -> Complex Double  -- ^ required to generate a value in the unit disk
                           -> Complex Double  -- ^ required to generate a value in the unit disk
                           -> Bool
prop_orderMandelbrotCircle (Positive nMax) w z = inMandelbrotSet nMax z0
    where z0 = (-1.0) + mu * 0.25
          mu = toUnitCircle w z

return []
runTests :: IO Bool
runTests = $quickCheckAll


{-# LANGUAGE FlexibleContexts #-}

module Stastics where

import           Data.Heap
import           Data.Maybe                  (fromJust)

import qualified Data.Vector.Unboxed         as U
import           Graphics.EasyPlot
import           Statistics.LinearRegression

-- Moving/Running Mean

-- convert raw data to Double
clean raw = map (\s -> read s :: Double) (lines raw)

-- Calculate moving average -- hardcoded smoothening factor = 0.95
avg :: [Double] -> Double
avg xs = case xs of
    (x:xs) -> a * x + (1 - a) * (avg xs)
    []     -> 0
    where a = 0.95

movingMean = do
  rawInput <- readFile "numbers.txt"
  let input = avg . clean $ rawInput
  print input

-- Calculating a Moving Median
strToInt raw = map (\s -> read s :: Int) (lines raw)
-- put smaller nums in maxheap and big nums in minheap
median (x:xs) maxheap minheap = case viewHead maxheap of
    Just theMax -> if x < theMax
            then median xs (insert x maxheap) minheap
            else median xs maxheap (insert x minheap)
    Nothing -> median xs (insert x maxheap) minheap

median [] maxheap minheap
    | size maxheap + 1 < size minheap =
        median [] (insert minelem maxheap) $ (snd.fromJust.view) minheap
    | size minheap + 1 < size maxheap =
        median [] ((snd.fromJust.view) maxheap) $ insert maxelem minheap
    |  size maxheap == size minheap =
        (fromIntegral maxelem + fromIntegral minelem) / 2.0
    | size maxheap > size minheap = fromIntegral maxelem
    | otherwise = fromIntegral minelem
    where
        maxelem = fromJust (viewHead maxheap)
        minelem = fromJust (viewHead minheap)

movingMedian = do
    raw <- readFile "numbers.txt"
    let input = strToInt raw
    print $ median input (empty :: MaxHeap Int) (empty :: MinHeap Int)

-- APPROXIMATING A LINEAR REGRESSION
linearR =do
    let xs = U.fromList [1.0, 2.0, 3.0, 4.0, 5.0] :: U.Vector Double
    let ys = U.fromList [1.0, 2.0, 1.3, 3.75, 2.25] :: U.Vector Double
    let (b,m) = linearRegression xs ys
    print $ concat ["y = ", show m, " x + ", show b]

--
main = do
        let values = [4,5,16,15,14,13,13,17]
        let values2 = [2,5,7,10,21, 18, 15]
        -- plot X11 $
        --    Data2D [ Title "Line Graph"
        --           , Style Linespoints
        --           , Color Blue]
        --    [] (zip [1..] values)

        -- plot X11 $ [ Data2D [Color Red] [] (zip [1..] values)
        --             , Data2D [Color Blue] [] (zip [1..] values2) ]

        -- 3D
        -- plot' [Interactive] X11
        --                [ Data3D [Color Red] [] [(1,1,1), (1,2,1), (0,1,1), (1,1,0)]
        --                , Data3D [Color Blue] [] [(4,3,2), (3,3,2), (3,2,3), (4,4,3)] ]

        -- regression

        let winEstimate = map (\x -> x * 0.425 + 0.7850000000000001) [1.3, 2.4 .. 4.7]
        let regressionLine = zip [1.3, 2.4 .. 4.7] winEstimate
        plot X11 [Data2D [Title "Runs Per Game VS Win % in 2014" ] [] (zip [1.0, 2.0, 3.0, 4.0, 5.0] [1.0, 2.0, 1.3, 3.75, 2.25]), Data2D [Title "Regression Line", Style Lines, Color Blue] [] regressionLine]


module LearnHaskDA.Chapter5 where
--import Math.Combinatorics.Exact.Binomial
import Numeric.SpecFunctions (choose)
import LearnHaskDA.Chapter2
import LearnHaskDA.Chapter4
import Graphics.EasyPlot
import System.Random
import Data.Number.Erf

{- Using Math.Combinatorics.Exact.Binomial in the exact-combinatorics
package we could have a more general type here and an exact, yet
efficient, computation of choose. But exact-combinatorics is not on
stackage so I changed to the approximation (from Numeric.SpecFunctions
which is on stackage). We only use it for plotting anyway.

-- probabilityMassFunction :: (Integral i, Num n) => i -> i -> n -> n
-}
probabilityMassFunction :: Int -> Int -> Double -> Double
probabilityMassFunction k n p = (n `choose` k)
                              * p^k * (1-p)^(n-k)
plot81 :: IO Bool
plot81 = plot (PNG "coinflips.png") $
              Function2D [Title "Coin Flip Probabilities"]
                         [Range 0 1000]
                         (\k -> probabilityMassFunction (floor k) 1000 0.5)

-- almost equal would be enough
sanitycheck :: Bool
sanitycheck = 1 == sum (map (\k -> probabilityMassFunction k 1000 0.5) [0..1000])

-- sumMid :: (Fractional a, Integral i) => i -> a
sumMid :: Int -> Double
sumMid d = sum $ map (\k -> probabilityMassFunction k 1000 0.5) [500-d .. 500+d]

----------------

standardDeviation :: (Floating a, Real a) => [a] -> a
standardDeviation values = sqrt (sum [(x-mu)^2 | x <- values]) / sqrt_nm1
  where  mu = average values
         sqrt_nm1 = sqrt $ fromIntegral (length values - 1)

select place = "SELECT "++place++"team, SUM("++place++"score) "++
                 "FROM winloss "++
                 "GROUP BY "++ place ++"team " ++
                 "ORDER BY "++ place ++"team"

test91 = do
  runsAtHome <- queryDatabase "winloss.sql" (select "home")
  runsAway   <- queryDatabase "winloss.sql" (select "away")
  let runsHomeAway = zip (readDoubleColumn runsAtHome 1) (readDoubleColumn runsAway 1)
  plot (PNG "HomeScoreAwayScore.png") $ Data2D [Title "Runs at Home (x axis) and Runs Away (y axis)"] [] runsHomeAway
  let runsHomeAwayDiff = map (\(a,b) -> a-b) runsHomeAway
  plot (PNG "HomeScoreAwayScoreDiff.png") $ Data2D [Title "Difference in Runs at Home and Runs Away"] [] $ zip [1..] runsHomeAwayDiff
  let n = fromIntegral (length runsHomeAwayDiff)
      avg = average runsHomeAwayDiff
      stdDev = standardDeviation runsHomeAwayDiff
      stdErr = stdDev / sqrt n
  return (avg, stdDev, stdErr)

--

from = invnormcdf ((1-95/100)/2 :: Double)
to = - from

test95 = test91 >>= \ (avg, stdDev, stdErr) ->
         return (avg + from*stdErr, avg + to*stdErr)

-- The 95% confidence interval is (-14.9, 26.6) which contains zero so
-- the null hypothesis could not be rejected.


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module NeuralNetwork where

import Control.Lens
import Numeric.LinearAlgebra

stepFunction :: Double -> Double
stepFunction x | x > 0 = 1.0
               | otherwise = 0.0

sigmoidFunction :: Double -> Double
sigmoidFunction x = 1 / (1 + (exp (-x)))

reluFunction :: Double -> Double
reluFunction = max 0

-- |
-- >>> :set -XOverloadedLists
-- >>> softmax [1010, 1000, 990]
-- [0.999954600070331,4.539786860886666e-5,2.061060046209062e-9]
-- >>> softmax [0.3, 2.9, 4.0]
-- [1.821127329554753e-2,0.24519181293507392,0.7365969137693786]
-- >>> sumElements $ softmax [0.3, 2.9, 4.0]
-- 1.0
softmax :: Vector R -> Vector R
softmax xs = cmap (/denom) ys
  where ys = cmap calc xs
        calc x = exp (x - c)
        c = maxElement xs
        denom = sumElements ys

-- |
-- >>> :set -XOverloadedLists
-- >>> meanSquaredError [0, 0, 1, 0, 0, 0, 0, 0, 0, 0] [0.1, 0.05, 0.6, 0, 0.05, 0.1, 0, 0.1, 0, 0]
-- 9.750000000000003e-2
-- >>> meanSquaredError [0, 0, 1, 0, 0, 0, 0, 0, 0, 0] [0.1, 0.05, 0.1, 0, 0.05, 0.1, 0, 0.6, 0.5, 0]
-- 0.7224999999999999
meanSquaredError :: Vector R -> Vector R -> Double
meanSquaredError t y = (/2) $ sumElements $ cmap (^2) $ t - y

-- |
-- >>> :set -XOverloadedLists
-- >>> crossEntropyError [0, 0, 1, 0, 0, 0, 0, 0, 0, 0] [0.1, 0.05, 0.6, 0, 0.05, 0.1, 0, 0.1, 0, 0]
-- 0.510825457099338
-- >>> crossEntropyError [0, 0, 1, 0, 0, 0, 0, 0, 0, 0] [0.1, 0.05, 0.1, 0, 0.05, 0.1, 0, 0.6, 0.5, 0]
-- 2.302584092994546
crossEntropyError :: Vector R -> Vector R -> Double
crossEntropyError t y = negate $ dot t $ cmap (log . (+delta)) y
  where delta = 1e-7 :: Double

-- |
-- >>> function1 x = 0.01 * x * x + 0.1 * x
-- >>> numericalDiff function1 5
-- 0.1999999999990898
-- >>> numericalDiff function1 10
-- 0.2999999999997449
-- >>> function2 x y = x * x + y * y
-- >>> numericalDiff (flip function2 4) 3
-- 6.00000000000378
-- >>> numericalDiff (function2 3) 4
-- 7.999999999999119
numericalDiff :: (Double -> Double) -> Double -> Double
numericalDiff f x = (f (x + h) - f (x - h)) / (2 * h)
  where h = 1e-4 :: Double

-- |
-- >>> function2 v = dot v v
-- >>> numericalGradient function2 $ vector [3.0, 4.0]
-- [6.00000000000378,7.999999999999119]
-- >>> numericalGradient function2 $ vector [0.0, 2.0]
-- [0.0,4.000000000004]
-- >>> numericalGradient function2 $ vector [3.0, 0.0]
-- [6.000000000012662,0.0]
numericalGradient
  :: (Gradible c)
  => ((c R) -> R) -> c R -> c R
numericalGradient f v = assoc (size v) 0 xs
  where
    h = 1e-4 :: R
    size' = size v
    xs = do
      idx <- indicesOf v
      let v1 = v - assoc size' 0 [(idx, h)]
          v2 = v + assoc size' 0 [(idx, h)]
          x' = f v2 - f v1
      return (idx, x' / (2 * h))

-- |
-- >>> function2 v = dot v v
-- >>> gradientDescent function2 (vector [-3.0, 4.0]) 0.1 100
-- [-6.111107928998789e-10,8.148143905314271e-10]
gradientDescent :: ((Vector R) -> Double) -> Vector R -> Double -> Int -> Vector R
gradientDescent f v lr steps = gradientDescents f lr v !! steps

-- |
-- >>> function2 v = dot v v
-- >>> take 3 $ gradientDescents function2 0.1 (vector [-3.0, 4.0])
-- [[-3.0,4.0],[-2.399999999999622,3.200000000000088],[-1.9199999999982538,2.5599999999976717]]
gradientDescents :: ((Vector R) -> Double) -> Double -> Vector R -> [Vector R]
gradientDescents f lr v = iterate desc v
  where desc v = v - cmap (* lr) (numericalGradient f v)

class (Container c R, Num (c R)) => Gradible c where
  indicesOf :: c R -> [IndexOf c]

instance Gradible Vector where
  indicesOf v = take (size v) [0..]

instance Gradible Matrix where
  indicesOf w = (,) <$> (take (rows w) [0..]) <*> (take (cols w) [0..])


{-# language OverloadedStrings #-}
{-# language TypeApplications #-}

import Geometry
import Data.Complex

iterateN g f = zipWith g [1..] . iterate f

main = do
  let romashka =
        group [ aSegment # rotate (18 * fromIntegral i) #: label (show i) <> white
              | i <- [1..20] ]
      clock = aCircle <+> group (modularScale 24 (aCircle # rotate 90 ))
      hand h = mkTriangle @CN [0.03:+0, 0:+0.9, (-0.03):+0] #: fill "red" # rotate (-15 * h)
               <+> aCircle # scale 0.05 #: fill "red"

  writeSVG 450 "figures/1.5/26.svg" $ let
    seg = aSegment # scale (1/8) #: "1" <+> point (1/8,0)
    segs = aPoint <+> (group $ take 8 $ iterate (translate (1/8, 0)) seg)
    rays = aSegment <+>
           origin #: "O" <+>
           (group $ take 7 $ tail $ iterate (rotate 8) (aSegment #: white)) <+>
           aSegment # rotate 64 #: "1"
    in rays <+> segs # translate (0,-0.25)

  writeSVG 500 "figures/1.5/28.svg" $
    group [ aSegment # rotate (18 * fromIntegral i) #: label (show i)
          | i <- [1..20] ]

  writeSVG 500 "figures/1.5/30.svg" $
    (group [ aSegment # rotate (18 * fromIntegral i) #: label (show i) <> white
           | i <- [1..6] ]) <+>
    aSegment # rotate (18 * 7) #: "7" <+>
    aSegment <+>
    aCircle #: invisible

  writeSVG 500 "figures/1.5/31.svg" $
    (group [ aSegment # rotate (18 * fromIntegral i) #: label (show i) <> white
           | i <- [1..6] ]) <+>
    aSegment # rotate (18 * 7) #: "7" <+>
    (group [ aSegment # rotate (18 * fromIntegral i) #: label (show (i - 7)) <> white
           | i <- [8..19] ]) <+>
    aSegment # rotate (18 * 20) #: "13" <+>
    aCircle #: invisible

  writeSVG 500 "figures/1.5/33.svg" $
    romashka <+> aSegment # rotate (18 * 13)

  writeSVG 500 "figures/1.5/34.svg" $
    romashka <+> aSegment # rotate (18 * 3)

  writeSVG 300 "figures/1.5/35.svg" $
    clock <+> hand 11

  writeSVG 300 "figures/1.5/37.svg" $
    clock <+> hand 17

  writeSVG 200 "figures/1.5/38.svg" $
    let dow l i = aLabel
                  # on aCircle (-i * 1/7) # rotate 90
                  #: label l
    in group $ zipWith dow ["Пн","Вт","Ср","Чт","Пт","Сб","Вс"] [1..]

  writeSVG 300 "figures/1.5/51.svg" $
    let dec = regularPoly 10
        sc = group $ modularScale 10 dec
    in dec <+> sc

  writeSVG 450 "figures/1.5/53.svg" $
    group [ aSegment # scale 0.9 # rotate i <+>
            aLabel # at' i #: label (show i)
          | i <- asDeg <$> [10,20..360] ]

  writeSVG 450 "figures/1.5/54.svg" $
    anAngle 180 #: "#" <+> aPoint <+>
    aSegment <+> aSegment # rotate 180

  writeSVG 450 "figures/1.5/55.svg" $
    anAngle 90 #: "#" <+> aPoint <+>
    aSegment <+> aSegment # rotate 180 <+>
    aSegment # rotate 90


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Exponential
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The exponential distribution.  This is the continunous probability
-- distribution of the times between events in a poisson process, in
-- which events occur continuously and independently at a constant
-- average rate.

module Statistics.Distribution.Exponential
    (
      ExponentialDistribution
    -- * Constructors
    , exponential
    , exponentialFromSample
    -- * Accessors
    , edLambda
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_neg_inf)
import qualified Statistics.Distribution         as D
import qualified Statistics.Sample               as S
import qualified System.Random.MWC.Distributions as MWC
import Statistics.Types (Sample)
import Data.Binary (put, get)

newtype ExponentialDistribution = ED {
      edLambda :: Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON ExponentialDistribution
instance ToJSON ExponentialDistribution

instance Binary ExponentialDistribution where
    put = put . edLambda
    get = fmap ED get

instance D.Distribution ExponentialDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.ContDistr ExponentialDistribution where
    density (ED l) x
      | x < 0     = 0
      | otherwise = l * exp (-l * x)
    logDensity (ED l) x
      | x < 0     = m_neg_inf
      | otherwise = log l + (-l * x)
    quantile = quantile

instance D.Mean ExponentialDistribution where
    mean (ED l) = 1 / l

instance D.Variance ExponentialDistribution where
    variance (ED l) = 1 / (l * l)

instance D.MaybeMean ExponentialDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance ExponentialDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy ExponentialDistribution where
  entropy (ED l) = 1 - log l

instance D.MaybeEntropy ExponentialDistribution where
  maybeEntropy = Just . D.entropy

instance D.ContGen ExponentialDistribution where
  genContVar = MWC.exponential . edLambda

cumulative :: ExponentialDistribution -> Double -> Double
cumulative (ED l) x | x <= 0    = 0
                    | otherwise = 1 - exp (-l * x)

complCumulative :: ExponentialDistribution -> Double -> Double
complCumulative (ED l) x | x <= 0    = 1
                         | otherwise = exp (-l * x)

quantile :: ExponentialDistribution -> Double -> Double
quantile (ED l) p
  | p == 1          = 1 / 0
  | p >= 0 && p < 1 = -log (1 - p) / l
  | otherwise       =
    error $ "Statistics.Distribution.Exponential.quantile: p must be in [0,1] range. Got: "++show p

-- | Create an exponential distribution.
exponential :: Double            -- ^ λ (scale) parameter.
            -> ExponentialDistribution
exponential l
  | l <= 0 =
    error $ "Statistics.Distribution.Exponential.exponential: scale parameter must be positive. Got " ++ show l
  | otherwise = ED l

-- | Create exponential distribution from sample. No tests are made to
-- check whether it truly is exponential.
exponentialFromSample :: Sample -> ExponentialDistribution
exponentialFromSample = ED . S.mean