{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances, GADTs, TypeOperators, FlexibleContexts #-}

module Physics.Quantum.Algebra where

import Numeric.Field.Class          (Field)
import Numeric.Algebra.Class        (LeftModule((.*)), Multiplicative((*)), Monoidal(zero))
import Numeric.Additive.Class       (Additive((+)))
import Numeric.Algebra.Unital       (Unital(one))

import Prelude hiding ((+), (*), Complex)

import qualified Data.Map as M
import qualified Data.List as L
import Data.Maybe (catMaybes)

import Data.Complex                 (Complex, conjugate)

import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.Eval as R
import qualified Data.Array.Repa.Repr.Unboxed as R
import qualified Data.Array.Repa.Repr.Vector as R

import Data.Functor.Identity (runIdentity)

import qualified Data.Eigen.SparseMatrix as ES
import qualified Data.Eigen.Matrix as E

-- a vector space m over field k
class (Field k, Additive m) => VectorSpace k m | m -> k where
    (*^) :: k -> m -> m

{-
-- lhs type ‘m’ does not determine rhs type ‘k’
instance (Field k, LeftModule k m) => VectorSpace k m where
    (*^) = (.*)
-}

class VectorSpace k m => LieAlgebra k m | m -> k where
    lieBr :: m -> m -> m

class (VectorSpace k a, VectorSpace k b) => DualVectors k a b | a b -> k where
    dual :: a -> b

class VectorSpace k m => InnerProductSpace k m | m -> k where
    (<.>) :: m -> m -> k

-- Endo-morphism of vector space
-- Data structure for operator may be different for different use cases
-- A Map be suitable for sparse operator, matrix for discreet vector indices, and functions for continuous vector indices.
class VectorSpace k v => SimpleOperator k v o | o -> k, o -> v where
    op :: o -> v -> v

-- Once we have simple operator, we can extract the map information from that
-- and get addition and multiplication operations
data OperatorMap v where
    OperatorMap :: VectorSpace k v => (v -> v) -> OperatorMap v

unOperatorMap :: OperatorMap v -> (v -> v)
unOperatorMap (OperatorMap m) = m

extractOperatorMap :: SimpleOperator k v o => o -> OperatorMap v
extractOperatorMap = OperatorMap . op

instance Additive v => Additive (OperatorMap v) where
    (+) (OperatorMap operator1) (OperatorMap operator2) = OperatorMap $ \vec -> (operator1 vec) + (operator2 vec)

instance  Multiplicative (OperatorMap o) where
    (*) (OperatorMap operator1) (OperatorMap operator2) = OperatorMap $ \vec -> operator1 (operator2 vec)

class (DualVectors k cv v, SimpleOperator k v o) => Operator k v cv o | o -> k where
    (|><|) :: v -> cv -> o

(<|>) :: (InnerProductSpace k v, DualVectors k cv v) => cv -> v -> k
(<|>) covec vec = (dual covec) <.> vec

(|>) :: SimpleOperator k v o => o -> v -> v
(|>) = op

(<|) :: (InnerProductSpace k v, DualVectors k cv v, Operator k v cv o) =>  cv -> o -> (v -> k)
(<|) covector operator  = \vector -> covector <|> (operator |> vector)

class IsComplex a where
    complexConjugate :: a -> a

instance Num a => IsComplex (Complex a) where
    complexConjugate = conjugate

----------------------------------
-- Constructions of vector spaces

-- m is a set of vector indices indices
-- May or may not represent orthogonal vectors, write instance of InnerProductSpace separately
-- kets are represented as maps from indices to coefficients
-- sparse kets have finite number of non zero entries
data SparseKet k m where
    SparseKet :: (Field k, Ord m) => M.Map m k -> SparseKet k m

unSparseKet :: SparseKet k m -> M.Map m k
unSparseKet (SparseKet ketMap) = ketMap

data SparseBra k m where
     SparseBra :: (Field k, Ord m) => M.Map m k -> SparseBra k m

unSparseBra :: SparseBra k m -> M.Map m k
unSparseBra (SparseBra braMap) = braMap

instance (Show k, Show m) => Show (SparseKet k m) where
    show ket@(SparseKet map)  = show map

instance (Show k, Show m) => Show (SparseBra k m) where
    show ket@(SparseBra map)  = show map

instance (Ord m, Field k) => Additive (SparseKet k m) where
    (+) (SparseKet a) (SparseKet b) = SparseKet $ M.unionWith (+) a b

instance (Ord m, Field k) => Additive (SparseBra k m) where
    (+) (SparseBra a) (SparseBra b) = SparseBra $ M.unionWith (+) a b

instance (Ord m, Field k) => VectorSpace k (SparseKet k m) where
    (*^) s (SparseKet v) = SparseKet $ M.map (s *) v

instance (Field k, Ord m) => VectorSpace k (SparseBra k m) where
    (*^) s (SparseBra v) = SparseBra $ M.map (s *) v

instance (Field k, IsComplex k, Ord m) => DualVectors k (SparseKet k m) (SparseBra k m) where
    dual (SparseKet a) = SparseBra $ M.map complexConjugate a

instance (Field k, IsComplex k, Ord m) => DualVectors k (SparseBra k m) (SparseKet k m) where
    dual (SparseBra a) = SparseKet $ M.map complexConjugate a

-- very generic mappings
data InfKet k m = InfKet { unContKet :: (m -> k) }
data InfBra k m = InfBra { unContBra :: (m -> k) }

instance Field k => Additive (InfKet k m) where
    (+) (InfKet a) (InfKet b) = InfKet $ \x -> a x + b x

instance Field k => Additive (InfBra k m) where
    (+) (InfBra a) (InfBra b) = InfBra $ \x -> a x + b x

instance Field k => VectorSpace k (InfKet k m) where
    (*^) s (InfKet v) = InfKet $ \x -> s * (v x)

instance Field k => VectorSpace k (InfBra k m) where
    (*^) s (InfBra v) = InfBra $ \x -> s * (v x)

instance (Field k, IsComplex k) => DualVectors k (InfKet k m) (InfBra k m) where
    dual (InfKet a) = InfBra $ \x -> complexConjugate $ a x

instance (Field k, IsComplex k) => DualVectors k (InfBra k m) (InfKet k m) where
    dual (InfBra a) = InfKet $ \x -> complexConjugate $ a x

-- a vector of coefficients at lattice points
data LatticeKet k m = LatticeKet { unLatticeKet :: R.Array R.D R.DIM1 k }
data LatticeBra k m = LatticeBra { unLatticeBra :: R.Array R.D R.DIM1 k }

instance (Show k, Show m) => Show (LatticeKet k m) where
    -- Delayed representation doesn't have show instance
    show (LatticeKet vec) = show $! compute vec
        where compute :: Show k => R.Array R.D R.DIM1 k -> R.Array R.V R.DIM1 k
              compute = R.computeS

instance (Show k, Show m) => Show (LatticeBra k m) where
    -- Delayed representation doesn't have show instance
    show (LatticeBra vec) = show $! compute vec
        where compute :: Show k => R.Array R.D R.DIM1 k -> R.Array R.V R.DIM1 k
              compute = R.computeS

instance Field k => Additive (LatticeKet k m) where
   (+) (LatticeKet a) (LatticeKet b) = LatticeKet $! R.zipWith (+) a b

instance Field k => Additive (LatticeBra k m) where
   (+) (LatticeBra a) (LatticeBra b) = LatticeBra $! R.zipWith (+) a b

instance Field k => VectorSpace k (LatticeKet k m) where
    (*^) s (LatticeKet v) = LatticeKet $! R.map (s *) v

instance Field k => VectorSpace k (LatticeBra k m) where
    (*^) s (LatticeBra v) = LatticeBra $! R.map (s *) v

instance (Field k, IsComplex k) => DualVectors k (LatticeKet k m) (LatticeBra k m) where
    dual (LatticeKet a) = LatticeBra $! R.map complexConjugate a

instance (Field k, IsComplex k) => DualVectors k (LatticeBra k m) (LatticeKet k m) where
    dual (LatticeBra a) = LatticeKet $! R.map complexConjugate a

-- n x 1 matrix
data SparseMatKet k ck m where
    SparseMatKet :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatKet k ck m

unSparseMatKet :: SparseMatKet k ck m -> ES.SparseMatrix k ck
unSparseMatKet (SparseMatKet mat) = mat

data SparseMatBra k ck m where
    SparseMatBra :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatBra k ck m

unSparseMatBra :: SparseMatBra k ck m -> ES.SparseMatrix k ck
unSparseMatBra (SparseMatBra mat) = mat

instance (Show k) => Show (SparseMatKet k ck m) where
    show (SparseMatKet mat) = show mat

instance (Show k) => Show (SparseMatBra k ck m) where
    show (SparseMatBra mat) = show mat

instance Additive (SparseMatKet k ck m) where
    (+) (SparseMatKet a) (SparseMatKet b) = SparseMatKet $! ES.add a b

instance Additive (SparseMatBra k ck m) where
    (+) (SparseMatBra a) (SparseMatBra b) = SparseMatBra $! ES.add a b

instance Field k => VectorSpace k (SparseMatKet k ck m) where
    (*^) s (SparseMatKet v) = SparseMatKet $! ES.scale s v

instance Field k => VectorSpace k (SparseMatBra k ck m) where
    (*^) s (SparseMatBra v) = SparseMatBra $! ES.scale s v

-------------------------------------
-- Constructions of operators

-- Operator formed from outerproduct of two vectors
data OuterOperator k v cv where
    OuterOperator :: (InnerProductSpace k v, DualVectors k cv v) => [(v, cv)] -> OuterOperator k v cv

unOuterOperator :: OuterOperator k v cv -> [(v, cv)]
unOuterOperator (OuterOperator outerProducts) = outerProducts

instance ( InnerProductSpace k v
         , DualVectors k cv v
         )
         => SimpleOperator k v (OuterOperator k v cv) where
    op (OuterOperator outerProducts) vec = L.foldr1 (+) $ map (\(ovec,ocovec) -> (ocovec <|> vec) *^ ovec) $ outerProducts

instance ( InnerProductSpace k v
         , DualVectors k cv v
         )
         => Operator k v cv (OuterOperator k v cv) where
    (|><|) vec covec = OuterOperator [(vec, covec)]

instance Additive (OuterOperator k v cv) where
    (+) (OuterOperator o1) (OuterOperator o2) = OuterOperator $ o1 L.++ o2

-- a sparse operator represented as a map from (i,j) to coefficient of |i><j|
-- Note that (|i><j|) |k> = <j|k> |i> and therefore the operation requires inner product to be defined
data SparseOperator k m = SparseOperator { unSparseOperator :: M.Map (m, m) k }

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => SimpleOperator k (SparseKet k m) (SparseOperator k m) where
    op (SparseOperator opMap) ket = SparseKet $! M.fromListWith (+) $!
        map (\(veci, coveci) -> ( veci
                                , (SparseBra $! M.fromListWith (+) $! catMaybes $! [fmap ((,) coveci) $! M.lookup (veci, coveci) opMap]) <|> ket
                                )
            ) $!
            (M.keys opMap)

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Operator k (SparseKet k m) (SparseBra k m) (SparseOperator k m) where
    (|><|) (SparseKet vecMap) (SparseBra covecMap) = SparseOperator $! M.fromListWith (+) $!
        catMaybes [(,) (i,j) <$> (fmap (*) (M.lookup i vecMap) <*> (M.lookup j covecMap)) | i <- M.keys vecMap, j <- M.keys covecMap ]

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Additive (SparseOperator k m) where
    (+) (SparseOperator a) (SparseOperator b) = SparseOperator $! M.unionWith (+) a b

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => Multiplicative (SparseOperator k m) where
     (*) (SparseOperator a) (SparseOperator b) = SparseOperator $! M.fromListWith (+) $! multiplyMaps a b []
         where
             -- chaoose one index i and sum (over j and k) |j><i| |i><k|
             multiplyMaps :: ( Field k, Ord m,
                               InnerProductSpace k (SparseKet k m),
                               DualVectors k (SparseBra k m) (SparseKet k m)
                             )
                             => M.Map (m,m) k -> M.Map (m,m) k -> [((m,m), k)] -> [((m,m), k)]
             multiplyMaps mapA mapB oldResult = case null mapA || null mapB of
                 True -> oldResult
                 False ->
                   -- toList uses list fusion so this should not go through entire map
                   let oneElement = L.take 1 $ M.toList $ mapA
                   in case oneElement of
                          [] -> oldResult
                          elementToUse : _ -> let indexToUse = (snd $! fst elementToUse)

                                                  -- <i|i>
                                                  indexKet = SparseKet $! M.singleton indexToUse one
                                                  indexSquaredNorm =  indexKet <.> indexKet

                                                  -- take all keys from mapA which have this index in snd of key
                                                  (mapAToUse, mapAToPass) = M.partitionWithKey (\key _ -> snd key == indexToUse) mapA
                                                  -- take all keys from mapA which have this index in snd of key
                                                  (mapBToUse, mapBToPass) = M.partitionWithKey (\key _ -> fst key == indexToUse) mapB

                                                  -- we can discard second index from mapAToUse and first index from mapBToUse
                                                  -- we also know that the preserved indices will occur uniquely in index
                                                  -- so we can just take the cross product
                                                  listAToUse = M.toList $! mapAToUse
                                                  listBToUse = M.toList $! mapBToUse
                                                  newResult = [((a,b), va * vb * indexSquaredNorm) | ((a,_), va) <- listAToUse,
                                                                                ((_,b), vb) <- listBToUse
                                                                                ]

                                              in multiplyMaps mapAToPass mapBToPass $! newResult ++ oldResult

instance ( Field k, Ord m,
           InnerProductSpace k (SparseKet k m),
           DualVectors k (SparseBra k m) (SparseKet k m)
         )
         => LeftModule k (SparseOperator k m) where
    (.*) scalar (SparseOperator b) = SparseOperator $! M.map (scalar *) b

-- a more general operator represented as a map from j to Map i_n s_n where operator takes |i_n> to s_n |j>
-- this construct may not suffice when there are uncountable such i_n
data InfOperator k m = InfOperator { unInfOperator :: m -> M.Map m k }

instance (Field k) => SimpleOperator k (InfKet k m) (InfOperator k m) where
    op operator = \(InfKet coeffMap) -> InfKet $ \index -> foldr1 (+) $
        map ( \(invIndex, factor) -> factor * coeffMap invIndex ) $
            M.toList $ unInfOperator operator $ index

instance (Field k, Ord m) => Additive (InfOperator k m) where
    (+) (InfOperator op1) (InfOperator op2) = InfOperator $ \index -> M.unionWith (+) (op1 index) (op2 index)

instance (Field k, Ord m) => Multiplicative (InfOperator k m) where
    (*) (InfOperator op1) (InfOperator op2) = InfOperator $ \index ->
            foldr1 (M.unionWith (+)) $
                map (\(i,s) -> M.map (s *) $ op2 i) $
                    M.toList $ op1 index

-- A matrix to transfom LatticeKet vector
data LatticeOperator k m where
    LatticeOperator :: (Field k, R.Elt k, R.Unbox k, Num k) => R.Array R.D R.DIM2 k -> LatticeOperator k m

unLatticeOperator :: LatticeOperator k m -> R.Array R.D R.DIM2 k
unLatticeOperator (LatticeOperator mat) = mat

-- Num k is required because of sumP in method definition
instance Field k => SimpleOperator k (LatticeKet k m) (LatticeOperator k m) where
    op (LatticeOperator matrix) (LatticeKet vector) = LatticeKet $! R.delay $! runIdentity $! mvMultP matrix vector

instance Additive (LatticeOperator k m) where
    (+) (LatticeOperator op1) (LatticeOperator op2) = LatticeOperator $! R.zipWith (+) op1 op2

-- Num k is required because of sumP in method definition
instance Multiplicative (LatticeOperator k m) where
    (*) (LatticeOperator op1) (LatticeOperator op2) = LatticeOperator $! R.delay $! runIdentity $! mmMultP op1 op2

instance Field k => LeftModule k (LatticeOperator k m) where
    (.*) scalar (LatticeOperator matrix) = LatticeOperator $! R.map (scalar *) matrix

-- Taken from Repa wiki, extended it to Additive and Multiplicative structures
mmMultP :: (Monad m, R.Source r k, Additive k, Multiplicative k, R.Elt k, R.Unbox k, Num k)
        => R.Array r R.DIM2 k
        -> R.Array r R.DIM2 k
        -> m (R.Array R.U R.DIM2 k)
mmMultP a b = R.sumP (R.zipWith (*) aRepl bRepl)
    where
      -- Note: transpose transposes the lowest two dimensions
      t = R.transpose b
      aRepl = R.extend (R.Z R.:.R.All R.:.colsB R.:.R.All) a
      bRepl = R.extend (R.Z R.:.rowsA R.:.R.All R.:.R.All) t

      (R.Z R.:.colsA R.:.rowsA) = R.extent a
      (R.Z R.:.colsB R.:.rowsB) = R.extent b

-- Multiply a martix with a vector
-- based on the above matrix matrix multiplication
mvMultP :: (Monad m, R.Source r k, Additive k, Multiplicative k, R.Elt k, R.Unbox k, Num k)
        => R.Array r R.DIM2 k
        -> R.Array r R.DIM1 k
        -> m (R.Array R.U R.DIM1 k)
mvMultP a b = R.sumP (R.zipWith (*) aRepl bRepl)
    where
      aRepl = a
      bRepl = R.extend (R.Z R.:.rowsA R.:.R.All) b
      (R.Z R.:.rowsA R.:.colsA) = R.extent a

-- A matrix to transfom LatticeKet vector
data SparseMatOperator k ck m where
    SparseMatOperator :: E.Elem k ck => ES.SparseMatrix k ck -> SparseMatOperator k ck m

unSparseMatOperator :: SparseMatOperator k ck m -> ES.SparseMatrix k ck
unSparseMatOperator (SparseMatOperator mat) = mat

-- Num k is required because of sumP in method definition
instance Field k => SimpleOperator k (SparseMatKet k ck m) (SparseMatOperator k ck m) where
    op (SparseMatOperator matrix) (SparseMatKet vector) = SparseMatKet $! ES.mul matrix vector

instance Field k => Additive (SparseMatOperator k ck m) where
    (+) (SparseMatOperator op1) (SparseMatOperator op2) = SparseMatOperator $! ES.add op1 op2

instance Field k => Multiplicative (SparseMatOperator k ck m) where
    (*) (SparseMatOperator op1) (SparseMatOperator op2) = SparseMatOperator $! ES.mul op1 op2

instance Field k => LeftModule k (SparseMatOperator k ck m) where
    (.*) scalar (SparseMatOperator matrix) = SparseMatOperator $! ES.scale scalar matrix


import Test.HUnit
import FFT.FFTData
import FFT.Orig
import FFT.Samples
import System.Environment
import Data.Complex

testfftbFlySParZip = TestCase (assertEqual "FFT-blfys-parzip" expected actual)
    where
        expected = (sum $ fftbFlySParZip samples512)
        actual = (sum $ fft samples512)
        samples512 = samples 1 512

testfftbFlySParZipMap = TestCase (assertEqual "FFT-blfys-parzipmap" expected actual)
    where
        expected = (sum $ fftbFlySParZipMap samples512)
        actual = (sum $ fft samples512)
        samples512 = samples 1 512

testfftMapReduce = TestCase (assertEqual "FFT-map-reduce" expected actual)
    where
        expected = (sum $ fftMapReduce samples512)
        actual = (sum $ fft samples512)
        samples512 = samples 1 512

tests = TestList [TestLabel "testfftFlySParZip" testfftbFlySParZip,
                TestLabel "testfftbFlySParZipMap" testfftbFlySParZipMap,
		TestLabel "testfftMapReduce" testfftMapReduce]

main = runTestTT tests


module Typed(
    typed,

    testPropertyI,
    testPropertyD,
    testPropertyZ,
    testPropertyDZ,

    immutableVector,
    immutableMatrix,
    ) where

import Data.Complex( Complex )
import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.QuickCheck

import Numeric.LinearAlgebra.Vector( Vector )
import Numeric.LinearAlgebra.Matrix( Matrix )

typed :: e -> a e -> a e
typed _ = id

immutableVector :: Vector e -> Vector e
immutableVector = id

immutableMatrix :: Matrix e -> Matrix e
immutableMatrix = id

testPropertyI :: (Testable a)
              => TestName
              -> (Int -> a)
              -> Test
testPropertyI str prop =
    testProperty str $ prop undefined

testPropertyDZ :: (Testable a, Testable b)
               => TestName
               -> (Double -> a)
               -> (Complex Double -> b)
               -> Test
testPropertyDZ str propd propz =
    testGroup str
        [ testProperty "Double" $ propd undefined
        , testProperty "Complex Double" $ propz undefined
        ]

testPropertyD :: (Testable a)
              => TestName
              -> (Double -> a)
              -> Test
testPropertyD str prop =
    testProperty str $ prop undefined

testPropertyZ :: (Testable a)
              => TestName
              -> (Complex Double -> a)
              -> Test
testPropertyZ str prop =
    testProperty str $ prop undefined


module Numeric.NLOPTSpec (spec) where

import Numeric.LinearAlgebra
import Numeric.NLOPT
import Test.Hspec

--
-- Next we solve the alkyl problem (see)
--
-- @
--     var x2 := 1.745, >= 0, <= 2;
--     var x3 := 1.2, >= 0, <= 1.6;
--     var x4 := 1.1, >= 0, <= 1.2;
--     var x5 := 3.048, >= 0, <= 5;
--     var x6 := 1.974, >= 0, <= 2;
--     var x7 := 0.893, >= 0.85, <= 0.93;
--     var x8 := 0.928, >= 0.9, <= 0.95;
--     var x9 := 8, >= 3, <= 12;
--     var x10 := 3.6, >= 1.2, <= 4;
--     var x11 := 1.45, >= 1.45, <= 1.62;
--     var x12 := 1, >= 0.99, <= 1.01010101010101;
--     var x13 := 1, >= 0.99, <= 1.01010101010101;
--     var x14 := 1, >= 0.9, <= 1.11111111111111;
--     var x15 := 1, >= 0.99, <= 1.01010101010101;

--     minimize obj:  - 6.3*x5*x8 + 5.04*x2 + 0.35*x3 + x4 + 3.36*x6;

--     subject to

--     e2:  - 0.819672131147541*x2 + x5 - 0.819672131147541*x6 = 0;
--     e3: 0.98*x4 - x7*(0.01*x5*x10 + x4) = 0;
--     e4:  - x2*x9 + 10*x3 + x6 = 0;
--     e5: x5*x12 - x2*(1.12 + 0.13167*x9 - 0.0067*x9*x9) = 0;
--     e6: x8*x13 - 0.01*(1.098*x9 - 0.038*x9*x9) - 0.325*x7 = 0.57425;
--     e7: x10*x14 + 22.2*x11 = 35.82;
--     e8: x11*x15 - 3*x8 = -1.33;
-- @
--
-- Expected solution:
--
-- @
--    x(1) = 3.03582
--    x(2) =  0.95
--    x(3) =  1.7037
--    x(4) =  0.543084
--    x(5) =  0.901319
--    x(6) =  10.4755
--    x(7) =  1.56164
--    x(8) =  1.53535
--    x(9) =  0.99
--    x(10) =  0.99
--    x(11) =  1.11111
--    x(12) =  0.99
--    x(13) =  1.58471
--    x(14) =  2
--    obj =  -1.765
-- @
spec :: Spec
spec = do
{-   describe "Sequential Least-Squares Quadratic Programming (SQSLP)" $ do
    it "solves the Alkyl problem" $ do
      let lb = fromList [0, 0, 0, 0, 0, 0.85, 0.9, 3, 1.2, 1.45, 0.99, 0.99, 0.9, 0.99]
          ub = fromList [2, 1.6, 1.2, 5, 2, 0.93, 0.95, 12, 4, 1.62, 1.01010101010101, 1.01010101010101, 1.11111111111111, 1.01010101010101]
          objfD x =
            ( -6.3 * x ! 3 * x ! 6 + 5.04 * x ! 0 + 0.35 * x ! 1 + x ! 2 + 3.36 * x ! 4,
              fromList [5.04, 0.35, 1, -6.3 * x ! 6, 3.36, 0, -6.3 * x ! 3, 0, 0, 0, 0, 0, 0, 0]
            )
          stop = ObjectiveRelativeTolerance 1e-8 :| []
          eqCons =
            [ EqualityConstraint
                { eqConstraintFunctions =
                    Vector
                      7
                      ( \x _ ->
                          (fromList
                            [ -0.819672131147541 * x ! 0 + x ! 3 - 0.819672131147541 * x ! 4,
                              0.98 * x ! 2 - x ! 5 * (0.01 * x ! 3 * x ! 8 + x ! 2),
                              - x ! 0 * x ! 7 + 10 * x ! 1 + x ! 4,
                              x ! 3 * x ! 10 - x ! 0 * (1.12 + 0.13167 * x ! 7 - 0.0067 * x ! 7 * x ! 7),
                              x ! 6 * x ! 11 - 0.01 * (1.098 * x ! 7 - 0.038 * x ! 7 * x ! 7) - 0.325 * x ! 5 - 0.57425,
                              x ! 8 * x ! 12 + 22.2 * x ! 9 - 35.82,
                              x ! 9 * x ! 13 - 3 * x ! 6 + 1.33
                            ],
                          toDense [
                            ((0,0), -0.819672131147541),
                            ((0,3), 1),
                            ((0,4), -0.819672131147541),
                            ((1,2), 0.98),
                            ((1,5), -(0.01 * x ! 3 * x ! 8 + x ! 2)),
                            ((1,3, -0.01*x ! 5 * x ! 8),
                            ((1,8, -0.01*x ! 5 * x ! 3),
                            ((1,2, -0.01*x ! 5),
                            ((2,0), -x ! 7),
                            ((2,0), -x ! 7),
                            ((2,7), -x ! 0),
                            ((2,1), 10),
                            ((2,4), 1),
                            ((3,3), x ! 10),
                            ((3,10), x ! 3),
                            ((3,0), -(1.12 + 0.13167 * x ! 7 - 0.0067 * x ! 7 * x ! 7)),
                            ((3,7), -x!0*(0.13167 -2*0.0067 * x ! 7)),
                            ((4,))

                                ]
                      ),
                  eqConstraintTolerance = 1e-6
                }
            ]
          subproblem = LocalProblem (fromIntegral $ size lb) stop algorithm
          algorithm = SLSQP objfD [LowerBounds lb, UpperBounds ub] [{- InequalityConstraintsD -}] eqCons
          alkylX0 = [1.75, 1.2, 1.1, 3.048, 1.974, 0.893, 0.982, 8, 3.6, 1.45, 1, 1, 1, 1]
          res = minimizeLocal problem (fromList alkylX0)
      res
        `shouldBe` ( Right
                       ( Solution
                           { solutionCost = -1.765,
                             solutionParams =
                               fromList
                                 [ 3.03582,
                                   0.95,
                                   1.7037,
                                   0.543084,
                                   0.901319,
                                   10.4755,
                                   1.56164,
                                   1.53535,
                                   0.99,
                                   0.99,
                                   1.11111,
                                   0.99,
                                   1.58471,
                                   2
                                 ],
                             solutionResult = FTOL_REACHED
                           }
                       )
                   )
 -}
  describe "Augmented Lagrangian based on SBPLX (Subplex)" $ do
    it "solves the Alkyl problem" $ do
      let lb = fromList [0, 0, 0, 0, 0, 0.85, 0.9, 3, 1.2, 1.45, 0.99, 0.99, 0.9, 0.99]
          ub = fromList [2, 1.6, 1.2, 5, 2, 0.93, 0.95, 12, 4, 1.62, 1.01010101010101, 1.01010101010101, 1.11111111111111, 1.01010101010101]
          objf x = -6.3 * x ! 3 * x ! 6 + 5.04 * x ! 0 + 0.35 * x ! 1 + x ! 2 + 3.36 * x ! 4
          stop = ObjectiveRelativeTolerance 1e-8 :| []
          eqCons =
            [ EqualityConstraint
                { eqConstraintFunctions =
                    Vector
                      7
                      ( \x _ ->
                          fromList
                            [ -0.819672131147541 * x ! 0 + x ! 3 - 0.819672131147541 * x ! 4,
                              0.98 * x ! 2 - x ! 5 * (0.01 * x ! 3 * x ! 8 + x ! 2),
                              - x ! 0 * x ! 7 + 10 * x ! 1 + x ! 4,
                              x ! 3 * x ! 10 - x ! 0 * (1.12 + 0.13167 * x ! 7 - 0.0067 * x ! 7 * x ! 7),
                              x ! 6 * x ! 11 - 0.01 * (1.098 * x ! 7 - 0.038 * x ! 7 * x ! 7) - 0.325 * x ! 5 - 0.57425,
                              x ! 8 * x ! 12 + 22.2 * x ! 9 - 35.82,
                              x ! 9 * x ! 13 - 3 * x ! 6 + 1.33
                            ]
                      ),
                  eqConstraintTolerance = 1e-6
                }
            ]
          problem = AugLagProblem eqCons [] (AUGLAG_EQ_LOCAL subproblem)
          subproblem = LocalProblem (fromIntegral $ size lb) stop algorithm
          algorithm = SBPLX objf [LowerBounds lb, UpperBounds ub] Nothing
          alkylX0 = [1.75, 1.2, 1.1, 3.048, 1.974, 0.893, 0.982, 8, 3.6, 1.45, 1, 1, 1, 1]
          res = minimizeAugLag problem (fromList alkylX0)
      res
        `shouldBe` ( Right
                       ( Solution
                           { solutionCost = -1.765,
                             solutionParams =
                               fromList
                                 [ 3.03582,
                                   0.95,
                                   1.7037,
                                   0.543084,
                                   0.901319,
                                   10.4755,
                                   1.56164,
                                   1.53535,
                                   0.99,
                                   0.99,
                                   1.11111,
                                   0.99,
                                   1.58471,
                                   2
                                 ],
                             solutionResult = FTOL_REACHED
                           }
                       )
                   )


-- Filename: fft.hs
-- Created by: Daniel Winograd-Cort
-- Created on: unknown
-- Last Modified by: Daniel Winograd-Cort
-- Last Modified on: 12/12/2013

-- -- DESCRIPTION --
-- This code was inspired by Euterpea.  It uses UISF to present a GUI that
-- shows the sum of two waves (whose frequencies are specified by the user)
-- as well as the Fast Fourier Transform of that sum.
--
-- This module requires the array and pure-fft packages.

{-# LANGUAGE Arrows #-}
module FRP.UISF.Examples.FFT where
import FRP.UISF
import Numeric.FFT (fft)
import Data.Complex
import Data.Map (Map)
import Data.Maybe (listToMaybe, catMaybes)
import qualified Data.Map as Map

import Data.Array.Unboxed

-- | Alternative for working with Math.FFT instead of Numeric.FFT
--import qualified Math.FFT as FFT
--import Data.Array.IArray
--import Data.Array.CArray
--myFFT n lst = elems $ (FFT.dft) (listArray (0, n-1) lst)

--------------------------------------
-- Sine wave oscillators
--------------------------------------

-- Table definition
type Table = UArray Int Double

-- Sine table generator. Takes an integer representing the number of samples to generate
-- and a list of relative intensities for the overtones of the sine wave.

tableSinesN :: Int -> [Double] -> Table
tableSinesN size amps =
    let wave x   = sum (zipWith (*) [sin (2*pi*x*n) | n <- [1..]] amps)
        delta    = 1 / fromIntegral size
        waveform = take size $ map wave [0,delta..]
        divisor  = (maximum . map abs) waveform
     in listArray (0,size) (map (/divisor) waveform)

-- Two example sine tables.

tab1, tab2 :: Table
tab1 = tableSinesN 4096 [1] -- Basic sine wave
tab2 = tableSinesN 4096 [1.0,0.5,0.33]

-- Table-driven oscillator

osc :: ArrowCircuit a => Table -> Double -> a Double Double
osc table sr = proc freq -> do
    rec
      let delta = 1 / sr * freq
          phase = if next > 1 then frac next else next
      next <- delay 0 -< frac (phase + delta)
    returnA -< ((table!).(`mod` size).round.(*rate)) phase
  where (_,size)    = bounds table
        rate        = fromIntegral size
        frac x = if x > 1 then x - fromIntegral (truncate x) else x

--------------------------------------
-- Fast Fourier Transform
--------------------------------------

-- | Returns n samples of type b from the input stream at a time,
--   updating after k samples.  This function is good for chunking
--   data and is a critical component to fftA
quantize :: ArrowCircuit a => Int -> Int -> a b (SEvent [b])
quantize n k = proc d -> do
    rec (ds,c) <- delay ([],0) -< (take n (d:ds), c+1)
    returnA -< if c >= n && c `mod` k == 0 then Just ds else Nothing

-- | Converts the vector result of a dft into a map from frequency to magnitude.
--   One common use is:
--      fftA >>> arr (fmap $ presentFFT clockRate)
presentFFT :: Double -> [Double] -> Map Double Double
presentFFT clockRate a = Map.fromList $ zipWith (curry mkAssoc) [0..] a where
    mkAssoc (i,c) = (freq, c) where
        samplesPerPeriod = fromIntegral (length a)
        freq = fromIntegral i * (clockRate / samplesPerPeriod)

-- | Given a quantization frequency (the number of samples between each
--   successive FFT calculation) and a fundamental period, this will decompose
--   the input signal into its constituent frequencies.
--   NOTE: The fundamental period must be a power of two!
fftA :: ArrowCircuit a => Int -> Int -> a Double (SEvent [Double])
fftA qf fp = proc d -> do
    carray <- quantize fp qf -< d :+ 0
    returnA -< fmap (map magnitude . take (fp `div` 2) . fft) carray

--------------------------------------
-- UISF Example
--------------------------------------

-- This example shows off the histogram and realtimeGraph widgets by
-- summing two sin waves and displaying them.  Additionally, it makes
-- use of two horizontal sliders.
-- This example also shows off asyncUISFV and how to take a SigFun,
-- of the type used to create sound, and convert it to a UISF.
fftEx :: UISF () ()
fftEx = proc _ -> do
    f1 <- hSlider (1, 2000) 440 -< ()
    _ <- leftRight (label "Freq 1: " >>> display) -< f1
    f2 <- hSlider (1, 2000) 440 -< ()
    _ <- leftRight (label "Freq 2: " >>> display) -< f2
    d <- asyncVT sr 0.1 myAutomaton -< (f1, f2)
    let fft = listToMaybe $ catMaybes $ map (snd . fst) d
        s = map (\((s, _), t) -> (s,t)) d
    _ <- histogram (makeLayout (Stretchy 10) (Fixed 150)) -< fft
    _ <- realtimeGraph (makeLayout (Stretchy 10) (Fixed 150)) 2 Black -< s
    returnA -< ()
  where
    sr = 1000 -- signal rate
    myAutomaton = proc (f1, f2) -> do
        s1 <- osc tab1 sr -< f1
        s2 <- osc tab2 sr -< f2
        let s = (s1 + s2)/2
        fftData <- fftA 100 256 -< s
        returnA -< (s, fftData)

-- This test is run separately from the others.
main :: IO ()
main = runUI (defaultUIParams {uiSize=(500, 600), uiTitle="FFT Example"}) fftEx


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns     #-}
-- Tests for Statistics.Test.NonParametric
module Tests.NonParametric (tests) where

import Statistics.Distribution.Normal (standard)
import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import Statistics.Test.KruskalWallis
import Statistics.Test.WilcoxonT
import Statistics.Types (PValue,pValue,mkPValue)

import Test.Tasty                (testGroup)
import Test.Tasty.HUnit
import Tests.ApproxEq            (eq)
import Tests.Helpers             (testAssertion, testEquality)
import Tests.NonParametric.Table (tableKSD, tableKS2D)
import qualified Test.Tasty          as Tst
import qualified Data.Vector.Unboxed as U

tests :: Tst.TestTree
tests = testGroup "Nonparametric tests"
        $ concat [ mannWhitneyTests
                 , wilcoxonSumTests
                 , wilcoxonPairTests
                 , kruskalWallisRankTests
                 , kruskalWallisTests
                 , kolmogorovSmirnovDTest
                 ]

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

mannWhitneyTests :: [Tst.TestTree]
mannWhitneyTests = zipWith test [(0::Int)..] testData ++
  [ testEquality "Mann-Whitney U Critical Values, m=1"
      (replicate (20*3) Nothing)
      [mannWhitneyUCriticalValue (1,x) (mkPValue p) | x <- [1..20], p <- [0.005,0.01,0.025]]
  , testEquality "Mann-Whitney U Critical Values, m=2, p=0.025"
      (replicate 7 Nothing ++ map Just [0,0,0,0,1,1,1,1,1,2,2,2,2])
      [mannWhitneyUCriticalValue (2,x) (mkPValue 0.025) | x <- [1..20]]
  , testEquality "Mann-Whitney U Critical Values, m=6, p=0.05"
      (replicate 1 Nothing ++ map Just [0, 2,3,5,7,8,10,12,14,16,17,19,21,23,25,26,28,30,32])
      [mannWhitneyUCriticalValue (6,x) (mkPValue 0.05) | x <- [1..20]]
  , testEquality "Mann-Whitney U Critical Values, m=20, p=0.025"
      (replicate 1 Nothing ++ map Just [2,8,14,20,27,34,41,48,55,62,69,76,83,90,98,105,112,119,127])
      [mannWhitneyUCriticalValue (20,x) (mkPValue 0.025) | x <- [1..20]]
  ]
  where
    test n (a, b, c, d)
      = testCase "Mann-Whitney" $ do
          assertEqual ("Mann-Whitney U "     ++ show n) c us
          assertEqual ("Mann-Whitney U Sig " ++ show n) d ss
      where
        us = mannWhitneyU (U.fromList a) (U.fromList b)
        ss = mannWhitneyUSignificant SamplesDiffer (length a, length b) p005 us
    -- List of (Sample A, Sample B, (Positive Rank, Negative Rank))
    testData :: [([Double], [Double], (Double, Double), Maybe TestResult)]
    testData = [ ( [3,4,2,6,2,5]
                 , [9,7,5,10,6,8]
                 , (2, 34)
                 , Just Significant
                 )
               , ( [540,480,600,590,605]
                 , [760,890,1105,595,940]
                 , (2, 23)
                 , Just Significant
                 )
               , ( [19,22,16,29,24]
                 , [20,11,17,12]
                 , (17, 3)
                 , Just NotSignificant
                 )
               , ( [126,148,85,61, 179,93, 45,189,85,93]
                 , [194,128,69,135,171,149,89,248,79,137]
                 , (35,65)
                 , Just NotSignificant
                 )
               , ( [1..30]
                 , [1..30]
                 , (450,450)
                 , Just NotSignificant
                 )
               , ( [1 .. 30]
                 , [11.5 .. 40 ]
                 , (190.0,710.0)
                 , Just Significant
                 )
               ]

wilcoxonSumTests :: [Tst.TestTree]
wilcoxonSumTests = zipWith test [(0::Int)..] testData
  where
    test n (a, b, c) = testCase "Wilcoxon Sum"
                     $ assertEqual ("Wilcoxon Sum " ++ show n) c (wilcoxonRankSums (U.fromList a) (U.fromList b))
    -- List of (Sample A, Sample B, (Positive Rank, Negative Rank))
    testData :: [([Double], [Double], (Double, Double))]
    testData = [ ( [8.50,9.48,8.65,8.16,8.83,7.76,8.63]
                 , [8.27,8.20,8.25,8.14,9.00,8.10,7.20,8.32,7.70]
                 , (75, 61)
                 )
               , ( [0.45,0.50,0.61,0.63,0.75,0.85,0.93]
                 , [0.44,0.45,0.52,0.53,0.56,0.58,0.58,0.65,0.79]
                 , (71.5, 64.5)
                 )
               ]

wilcoxonPairTests :: [Tst.TestTree]
wilcoxonPairTests = zipWith test [(0::Int)..] testData ++
  -- Taken from the Mitic paper:
  [ testAssertion "Sig 16, 35" (to4dp 0.0467 $ wilcoxonMatchedPairSignificance 16 35)
  , testAssertion "Sig 16, 36" (to4dp 0.0523 $ wilcoxonMatchedPairSignificance 16 36)
  , testEquality   "Wilcoxon critical values, p=0.05"
      (replicate 4 Nothing ++ map Just [0,2,3,5,8,10,13,17,21,25,30,35,41,47,53,60,67,75,83,91,100,110,119])
      [wilcoxonMatchedPairCriticalValue x (mkPValue 0.05) | x <- [1..27]]
  , testEquality "Wilcoxon critical values, p=0.025"
      (replicate 5 Nothing ++ map Just [0,2,3,5,8,10,13,17,21,25,29,34,40,46,52,58,65,73,81,89,98,107])
      [wilcoxonMatchedPairCriticalValue x (mkPValue 0.025) | x <- [1..27]]
  , testEquality "Wilcoxon critical values, p=0.01"
      (replicate 6 Nothing ++ map Just [0,1,3,5,7,9,12,15,19,23,27,32,37,43,49,55,62,69,76,84,92])
      [wilcoxonMatchedPairCriticalValue x (mkPValue 0.01) | x <- [1..27]]
  , testEquality "Wilcoxon critical values, p=0.005"
      (replicate 7 Nothing ++ map Just [0,1,3,5,7,9,12,15,19,23,27,32,37,42,48,54,61,68,75,83])
      [wilcoxonMatchedPairCriticalValue x (mkPValue 0.005) | x <- [1..27]]
  ]
  where
    test n (a, b, c) = testEquality ("Wilcoxon Paired " ++ show n) c res
      where res = wilcoxonMatchedPairSignedRank (U.zip (U.fromList a) (U.fromList b))

    -- List of (Sample A, Sample B, (Positive Rank, Negative Rank))
    testData :: [([Double], [Double], (Int,Double, Double))]
    testData = [ ([1..10], [1..10], (0, 0, 0     ))
               , ([1..5],  [6..10], (5, 0, 5*(-3)))
               -- Worked example from the Internet:
               , ( [125,115,130,140,140,115,140,125,140,135]
                 , [110,122,125,120,140,124,123,137,135,145]
                 , ( 9
                   , sum $ filter (> 0) [7,-3,1.5,9,0,-4,8,-6,1.5,-5]
                   , sum $ filter (< 0) [7,-3,1.5,9,0,-4,8,-6,1.5,-5]
                   )
                 )
               -- Worked examples from books/papers:
               , ( [2.4,1.9,2.3,1.9,2.4,2.5]
                 , [2.0,2.1,2.0,2.0,1.8,2.0]
                 , (6, 18, -3)
                 )
               , ( [130,170,125,170,130,130,145,160]
                 , [120,163,120,135,143,136,144,120]
                 , (8, 27, -9)
                 )
               , ( [540,580,600,680,430,740,600,690,605,520]
                 , [760,710,1105,880,500,990,1050,640,595,520]
                 , (9, 3, -42)
                 )
               ]
    to4dp tgt (pValue -> x) = x >= tgt - 0.00005 && x < tgt + 0.00005

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

kruskalWallisRankTests :: [Tst.TestTree]
kruskalWallisRankTests = zipWith test [(0::Int)..] testData
  where
    test n (a, b) = testCase "Kruskal-Wallis Ranking"
                  $ assertEqual ("Kruskal-Wallis " ++ show n) (map U.fromList b) (kruskalWallisRank $ map U.fromList a)
    testData :: [([[Int]],[[Double]])]
    testData = [ ( [ [68,93,123,83,108,122]
                   , [119,116,101,103,113,84]
                   , [70,68,54,73,81,68]
                   , [61,54,59,67,59,70]
                   ]
                 , [ [8.0,14.0,16.0,19.0,23.0,24.0]
                   , [15.0,17.0,18.0,20.0,21.0,22.0]
                   , [1.5,8.0,8.0,10.5,12.0,13.0]
                   , [1.5,3.5,3.5,5.0,6.0,10.5]
                   ]
                 )
               ]

kruskalWallisTests :: [Tst.TestTree]
kruskalWallisTests = zipWith test [(0::Int)..] testData
  where
    test n (a, b, c) = testCase "Kruskal-Wallis" $ do
        assertEqual ("Kruskal-Wallis " ++ show n) (round100 b) (round100 kw)
        assertEqual ("Kruskal-Wallis Sig " ++ show n) c kwt
      where
        kw  = kruskalWallis $ map U.fromList a
        kwt = isSignificant p005 `fmap` kruskalWallisTest (map U.fromList a)
        round100 :: Double -> Integer
        round100 = round . (*100)

    testData :: [([[Double]], Double, Maybe TestResult)]
    testData = [ ( [ [68,93,123,83,108,122]
                   , [119,116,101,103,113,84]
                   , [70,68,54,73,81,68]
                   , [61,54,59,67,59,70]
                   ]
                 , 16.03
                 , Just Significant
                 )
               , ( [ [5,5,3,5,5,5,5]
                   , [5,5,5,5,7,5,5]
                   , [5,5,6,5,5,5,5]
                   , [4,5,5,5,6,5,5]
                   ]
               , 2.24
               , Just NotSignificant
               )
               , ( [ [36,48,5,67,53]
                   , [49,33,60,2,55]
                   , [71,31,140,59,42]
                   ]
                 , 1.22
                 , Just NotSignificant
                 )
               , ( [ [6,38,3,17,11,30,15,16,25,5]
                   , [34,28,42,13,40,31,9,32,39,27]
                   , [13,35,19,4,29,0,7,33,18,24]
                   ]
                 , 6.10
                 , Just Significant
                 )
               ]

----------------------------------------------------------------
-- K-S test
----------------------------------------------------------------

kolmogorovSmirnovDTest :: [Tst.TestTree]
kolmogorovSmirnovDTest =
  [ testAssertion "K-S D statistics" $
    and [ eq 1e-6 (kolmogorovSmirnovD standard (toU sample)) reference
        | (reference,sample) <- tableKSD
        ]
  , testAssertion "K-S 2-sample statistics" $
    and [ eq 1e-6 (kolmogorovSmirnov2D (toU xs) (toU ys)) reference
        | (reference,xs,ys) <- tableKS2D
        ]
  , testAssertion "K-S probability" $
    and [ eq 1e-14 (kolmogorovSmirnovProbability n d) p
        | (d,n,p) <- testData
        ]
  ]
  where
    toU = U.fromList
    -- Test data for the calculation of cumulative probability
    -- P(D[n] < d).
    --
    -- Test data is:
    --    (D[n], n, p)
    -- Table is generated using sample program from paper
    testData :: [(Double,Int,Double)]
    testData =
      [ (0.09           ,    3, 0                   )
      , (0.2            ,    3, 0.00177777777777778 )
      , (0.301          ,    3, 0.116357025777778   )
      , (0.392          ,    3, 0.383127210666667   )
      , (0.5003         ,    3, 0.667366306558667   )
      , (0.604          ,    3, 0.861569877333333   )
      , (0.699          ,    3, 0.945458198         )
      , (0.802          ,    3, 0.984475216         )
      , (0.9            ,    3, 0.998               )
      , (0.09           ,    5, 0                   )
      , (0.2            ,    5, 0.0384              )
      , (0.301          ,    5, 0.33993786080016    )
      , (0.392          ,    5, 0.66931908083712    )
      , (0.5003         ,    5, 0.888397260183794   )
      , (0.604          ,    5, 0.971609957879808   )
      , (0.699          ,    5, 0.994331075994008   )
      , (0.802          ,    5, 0.999391366368064   )
      , (0.9            ,    5, 0.99998             )
      , (0.09           ,    8, 3.37615237575e-06   )
      , (0.2            ,    8, 0.151622071801758   )
      , (0.301          ,    8, 0.613891042670582   )
      , (0.392          ,    8, 0.871491561427005   )
      , (0.5003         ,    8, 0.977534089199071   )
      , (0.604          ,    8, 0.997473116268255   )
      , (0.699          ,    8, 0.999806082005123   )
      , (0.802          ,    8, 0.999995133786947   )
      , (0.9            ,    8, 0.99999998          )
      , (0.09           ,   10, 3.89639433093119e-05)
      , (0.2            ,   10, 0.25128096          )
      , (0.301          ,   10, 0.732913126355935   )
      , (0.392          ,   10, 0.932185254518767   )
      , (0.5003         ,   10, 0.992276179340446   )
      , (0.604          ,   10, 0.999495533516769   )
      , (0.699          ,   10, 0.999979691783985   )
      , (0.802          ,   10, 0.999999801409237   )
      , (0.09           ,   20, 0.00794502217168886 )
      , (0.2            ,   20, 0.647279826376584   )
      , (0.301          ,   20, 0.958017466965765   )
      , (0.392          ,   20, 0.997206424843499   )
      , (0.5003         ,   20, 0.999962641414228   )
      , (0.09           ,   30, 0.0498147538075168  )
      , (0.2            ,   30, 0.842030838984526   )
      , (0.301          ,   30, 0.993403560017612   )
      , (0.392          ,   30, 0.99988478803318    )
      , (0.09           ,  100, 0.629367974413669   )
      ]

p005 :: PValue Double
p005 = mkPValue 0.05


module Matrix.StateTest
  (
    testsMatrix
  ) where

import Test.HUnit

import Numeric.LinearAlgebra.HMatrix hiding (corr)
import Matrix.State

mat :: Matrix Double
mat = reshape 3 $ fromList $ take 9 [1..]

vec :: Vector Double
vec = fromList $ take 3 [10..]

test1RE :: Test
test1RE = TestCase (assertEqual "replaceElement - first position"
                    (fromList [1,11,12])
                    (replaceElement 0 1 vec))

test2RE :: Test
test2RE = TestCase (assertEqual "replaceElement - last position"
                    (fromList [10,11,1])
                    (replaceElement 2 1 vec))
test3RE :: Test
test3RE = TestCase (assertEqual "replaceElement - in vector"
                    (fromList [10,2,12])
                    (replaceElement 1 2 vec))

test4RE :: Test
test4RE = TestCase (assertEqual "replaceElement - out of bounds low"
                    (fromList [10,11,12])
                    (replaceElement (-1) 2 vec))

test5RE :: Test
test5RE = TestCase (assertEqual "replaceElement - out of bounds high"
                    (fromList [10,11,12])
                    (replaceElement 4 2 vec))

testsReplaceElement :: Test
testsReplaceElement = TestList [  TestLabel "test1" test1RE
                                , TestLabel "test2" test2RE
                                , TestLabel "test3" test3RE
                                , TestLabel "test4" test4RE
                                , TestLabel "test5" test5RE
                                ]

test1RV :: Test
test1RV = TestCase (assertEqual "replaceVector - first row"
                    (reshape 3 $ fromList [10,11,12,4,5,6,7,8,9])
                    (replaceVector 0 vec mat))

test2RV :: Test
test2RV = TestCase (assertEqual "replaceVector - inner row"
                    (reshape 3 $ fromList [1,2,3,10,11,12,7,8,9])
                    (replaceVector 1 vec mat))

test3RV :: Test
test3RV = TestCase (assertEqual "replaceVector - last row"
                    (reshape 3 $ fromList [1,2,3,4,5,6,10,11,12])
                    (replaceVector 2 vec mat))

test4RV :: Test
test4RV = TestCase (assertEqual "replaceVector - index out of bounds - low"
                    mat
                    (replaceVector (-1) vec mat))

test5RV :: Test
test5RV = TestCase (assertEqual "replaceVector - index out of bounds - high"
                    mat
                    (replaceVector 10 vec mat))

testsReplaceVector :: Test
testsReplaceVector = TestList [ TestLabel "test1" test1RV
                              , TestLabel "test2" test2RV
                              , TestLabel "test3" test3RV
                              , TestLabel "test4" test4RV
                              , TestLabel "test5" test5RV
                              ]

test1RME :: Test
test1RME = TestCase (assertEqual "replaceMatrixElement - first row"
                     (reshape 3 $ fromList [0,2,3,4,5,6,7,8,9])
                     (replaceMatrixElement (0,0) 0 mat))

test2RME :: Test
test2RME = TestCase (assertEqual "replaceMatrixElement - middle row"
                     (reshape 3 $ fromList [1,2,3,0,5,6,7,8,9])
                     (replaceMatrixElement (1,0) 0 mat))

test3RME :: Test
test3RME = TestCase (assertEqual "replaceMatrixElement - end first row"
                     (reshape 3 $ fromList [1,2,0,4,5,6,7,8,9])
                     (replaceMatrixElement (0,2) 0 mat))

test4RME :: Test
test4RME = TestCase (assertEqual "replaceMatrixElement - last row"
                     (reshape 3 $ fromList [1,2,3,4,5,6,7,8,0])
                     (replaceMatrixElement (2,2) 0 mat))

test5RME :: Test
test5RME = TestCase (assertEqual "replaceMatrixElement - out of bounds"
                     mat
                     (replaceMatrixElement (0,4) 0 mat))

test6RME :: Test
test6RME = TestCase (assertEqual "replaceMatrixElement - out of bounds"
                     mat
                     (replaceMatrixElement ((-1),0) 0 mat))

testsReplaceMatrixElement :: Test
testsReplaceMatrixElement = TestList [ TestLabel "test1" test1RME
                                     , TestLabel "test2" test2RME
                                     , TestLabel "test3" test3RME
                                     , TestLabel "test4" test4RME
                                     , TestLabel "test5" test5RME
                                     , TestLabel "test6" test6RME]

testsMatrix :: Test
testsMatrix = TestList [ TestLabel "replaceElementTests" testsReplaceElement
                       , TestLabel "replaceVectorTests" testsReplaceVector
                       , TestLabel "replaceMatrixElementTests" testsReplaceMatrixElement]


{-# LANGUAGE NoMonomorphismRestriction, ImpredicativeTypes ,RankNTypes, TypeOperators #-}
import Kalman
import qualified Data.Vector as V
import Vectorization.Instances
import Local
import Linear.V1
import Data.Distributive
import qualified Data.Foldable as F
import Linear.V2
import Data.Complex
import CartesianProduct
import System.Random.MWC.Distributions
import Vectorization
import Foreign.Storable
import Foreign.Cholesky
import Space.Class
import Control.Applicative
import Control.Monad
import System.Random.MWC
import Graphics.Gnuplot.Simple
import Statistics.Autocorrelation
import Statistics.Sample
import Data.Functor.Compose
import Foreign.FStorable
import Foreign.FStorable

type Model = V1 :|: V1 :|: V1 :|: V1

predSin :: V1 Double -> Model Double -> Model Double
predSin  dt s@(theta :|: omega :|: amp :|: accel)   = s |+|  ((omega*dt )  :|: 0 :|: (dt*accel ) :|: 0 )

projSin :: Model Double -> V1 Double
projSin (theta :|: omega :|: amp :|: accel) = amp * (fmap sin theta )

nx :: Model (Model Double)
nx = diag $ fmap (*1e-3)((1/3) :|: 1 :|: 0.1 :|: 0.1)

getMean :: Functor f => (forall b . g b -> V1 b) -> f (g b , g (g b)) -> f b
getMean s = fmap ( unV1 . s . fst )

getCov :: (Num b ,Functor f )=> (forall b .  g b -> V1 b) -> f (g b , g (g b)) -> f b
getCov s = fmap ( abs .unV1 . s .  unV1 . s . snd )

sigmaBound :: Num b => (forall b . g b -> V1 b) -> [(g b , g (g b))] -> [[b]]
sigmaBound s f =  [zipWith (+) (getMean s f) (getCov s f), getMean s  f,zipWith (-) (getMean s f) (getCov s f)]

sigmas :: Num b => (forall b . g b -> V1 b) -> [[(g b , g (g b))]] -> [[b]]
sigmas s = concat . map (sigmaBound s)

allSigmas :: Num b => [forall b . g b -> V1 b] -> [[(g b , g (g b))]] -> [[[b]]]
allSigmas s d =  map (flip sigmas $ d) s

unV1 (V1 x) = x
execute = do
   seed <- create
   let n = 500
   v <- replicateM n (normal 0 0.25 seed)
   let noisedData = testeSet v
       cleanData = testeSet (replicate n 0)
   let ny :: (V1 (V1 Double))
       ny = 1
       xo :: Model Double
       xo = 0.1 :|: 0.1 :|: 1 :|: 1e-3
       alpha = 1e-2
       iteration  r = full (predSin 1) nx projSin ny r
       backwardIteration = smoothing (predSin 1 ) nx
       forwards =   scanl (flip iteration)  (xo,nx) (fmap V1 $ testeSet v)
       backwards =  reverse $ scanl1 (flip backwardIteration)  $ reverse forwards
       project = fmap (unV1 . projSin . fst )
       getAmp = fmap ( unV1 . amp . fst )
       getAmpCov = fmap ( unV1 .  amp . unV1 . amp . snd)
       get x = fmap ( unV1 . x . fst )
   -- print projections
   plotLists [] [testeSet (replicate n 0) ,testeSet v,project backwards  ]
   -- print all parameters
   mapM_ (plotLists []) (allSigmas [amp ,accel,omega] [backwards])
   let ampCov = V.fromList noiseSignalBackward
       noiseSignalBackward = zipWith (-)  cleanData  (project backwards)
       noiseSignalForward = zipWith (-)  cleanData (project forwards)
   -- Print residuals
   plotLists [] [noiseSignalForward , noiseSignalBackward ,v]
   print $ mean ampCov
   print $ stdDev ampCov
   print $ mean $ V.fromList $  zipWith (-) (getAmp forwards) (getAmp backwards)
   print $ stdDev $ V.fromList $  zipWith (-) (getAmp forwards) (getAmp backwards)

amp :: Model a -> V1 a
amp ( x :|: _ :|: amp :|: _ ) = amp
accel ( x :|: _ :|: amp :|: accel ) = accel
omega ( x :|: omega  :|: amp :|: accel ) = omega
theta ( theta :|: _ :|: amp :|: accel ) = theta

matI :: Num a => Int -> [[a]]
matI n = [ [fromIntegral $ fromEnum $ i == j | i <- [1..n]] | j <- [1..n]]

diag = liftIdent
liftIdent ::(R f,Num a,Storable a,Vectorize f,Functor f,Applicative f)=> f a -> f ( f a)
liftIdent x = fmap (liftA2 (*) x) (fromLists $ matI n )
    where n = dimM x

testeSet v =  zipWith (+)  (map (function.fromIntegral) [1..length v]) v where
  function x =  (1+x/500) * sin(16 * x/500 * pi)


{-# LANGUAGE BangPatterns #-}

module BBrotRender(showCells, render) where

import Codec.Picture
import Control.Monad
import Data.Aeson(decode)
import Data.Array.IO
import Data.Array.Unboxed
import qualified Data.ByteString.Lazy as BS
import Data.Complex
import Data.Maybe
import Data.List
import Data.Word(Word32, Word8)
import System.Console.CmdArgs(whenNormal)
import System.Exit
import System.IO
import Text.Read(readMaybe)

import BBrotCompute
import BBrotConf
import BBrotSelection

xmin, xmax, ymin, ymax :: Double
xmin = realPart loCorner
xmax = realPart hiCorner
ymin = imagPart loCorner
ymax = imagPart hiCorner

xrange, yrange :: Double
xrange = xmax - xmin
yrange = ymax - ymin

orbs :: Complex Double -> [Complex Double]
orbs !(x0 :+ y0) = unfoldr f (x0, y0)
  where f :: (Double, Double) -> Maybe (Complex Double, (Double, Double))
        f (!x, !y) =
          let !x2 = x * x
              !y2 = y * y
              !newx = (x2 - y2 + x0)
              !newy = (2 * x * y + y0)
          in
           if x2 + y2 > 4
           then Nothing
           else Just (newx :+ newy, (newx, newy))

inWindow :: Complex Double -> Bool
inWindow z = x >= xmin && x < xmax && y >= ymin && y < ymax
  where x = realPart z
        y = imagPart z

norm :: (Double -> Double) -> Word32 -> Word32 -> Word32 -> Word8
norm curveFunc minV maxV cnt = v
  where t = fromIntegral (cnt - minV) / fromIntegral (maxV - minV)
        u = curveFunc $ min 1.0 (2.0 * t)
        v = floor $ u * 255 :: Word8

gray :: Word8 -> PixelRGB8
gray x = PixelRGB8 x x x

reddish :: Word8 -> PixelRGB8
reddish x = PixelRGB8 r g 0
  where xr = min x 127
        xg = min (x - xr) 127
        r = 2 * xr
        g = 2 * xg

flames :: Word8 -> PixelRGB8
flames x = PixelRGB8 r g b
  where xr = min 120 x
        xg = min 120 (x - xr)
        xb = min 15 (x - xr - xg)
        r = xr * 2 + xr `div` 10
        g = xg * 2 + xg `div` 10
        b = xb * 17

toPixel :: (Double -> Double) -> Word32 -> Word32 -> (Word8 -> PixelRGB8) -> Word32 -> PixelRGB8
toPixel curveFunc smallest biggest scheme = scheme . norm curveFunc smallest biggest

plotPix :: IOUArray (Int, Int) Word32 -> (Int, Int) -> IO ()
plotPix !img (!x, !y) = do
  v <- readArray img (x, y)
  writeArray img (x, y) (v + 1)

rel :: Double -> Double -> Double -> Double
rel !lo !realrange !v = (v - lo) / realrange

toPix :: Int -> Double -> Double -> Double -> Int
toPix !pixrange !realmin !realrange !a =
  floor $ (fromIntegral pixrange) * (rel realmin realrange a)

toImgCoords :: Int -> Int -> Complex Double -> (Int, Int)
toImgCoords !xres !yres !z =
    (toPix xres xmin xrange $ realPart z,
     toPix yres ymin yrange $ imagPart z)

toPlaneCoords :: Int -> Int -> Int -> Int -> Complex Double
toPlaneCoords !xres !yres !i !j = x :+ y
    where x = xmin + xrange * fromIntegral i / fromIntegral xres
          y = ymin + yrange * fromIntegral j / fromIntegral yres

emptyPS :: PointSelection
emptyPS = PointSelection { pointList   = []
                         , commandLine = Nothing
                         , randGen     = Nothing
                         , timeStamp   = Nothing
                         }

loadPointsJson :: String -> IO (Maybe PointSelection)
loadPointsJson filepath = do
  contents <- BS.readFile filepath
  return (decode contents :: Maybe PointSelection)

loadPointsComplex :: String -> IO (Maybe PointSelection)
loadPointsComplex filepath = do
  contents <- readFile filepath
  let zs = concatMap g $ lines contents
        where g = maybeToList . readMaybe :: String -> [Complex Double]
      ps = map f zs
        where f (x :+ y) = BBPoint x y 1
  return $ if null ps
           then Nothing
           else Just (emptyPS { pointList = ps })

render :: BBrotConf -> IO ()
render conf = do
  -- Load points of interest, render their orbits into a PNG image.
  whenNormal $ putStrLn $ "Loading cache " ++ icachepath conf ++ " ..."
  maybePS <- if isComplex conf
             then loadPointsComplex $ icachepath conf
             else loadPointsJson $ icachepath conf
  when (isNothing maybePS) $ do
    hPutStrLn stderr $ "failed to parse " ++ icachepath conf
    unless (isComplex conf) $
      hPutStrLn stderr "try with flag -z"
    exitFailure

  let psel = pointList $ fromMaybe emptyPS maybePS
      selected = map (\(BBPoint x y _) -> x :+ y) psel
      orbits = concatMap orbs selected
      result = filter inWindow $ if dontRender conf
                                 then selected
                                 else orbits
  whenNormal $ putStrLn $ "selected points: " ++ show (length selected)

  let xres = xpixels conf
      yres = ypixels conf
      coords = map (toImgCoords xres yres) result
  img <- newArray ((0, 0), (xres - 1, yres - 1)) 0 :: IO (IOUArray (Int, Int) Word32)
  mapM_ (plotPix img) coords
  whenNormal $ putStrLn "done plotting"

  img2 <- freeze img :: IO (UArray (Int, Int) Word32)
  let values = elems img2
      !v0 = head values
      (!total, !smallest, !biggest) = foldl' f (0, v0, v0) values
        where f (!x, !y, !z) !a = (a + x, min a y, max a z)

  whenNormal $ do
    putStrLn $ "img points: " ++ show total
    putStrLn $ "value range: " ++ show smallest ++ "-" ++ show biggest

  let outfile = fromMaybe defPath (imagepath conf)
        where defPath = icachepath conf ++ ".png"
  whenNormal $ putStrLn $ "Writing " ++ outfile ++ " ..."
  let colorScheme = case palette conf of
        Flames -> flames
        Gray -> gray
        Reddish -> reddish
      curveFunc = case curve conf of
        Line -> id
        Root -> sqrt
        Square -> (**2)
      pixFunc = toPixel curveFunc smallest biggest colorScheme
      renderer i j = pixFunc $ img2!(i,j)
  writePng outfile $ generateImage (flip $ renderer) yres xres

showCells :: BBrotConf -> IO ()
showCells conf = do
  -- Compute selection of cells (squares in the complex plane) that
  -- are close to the edge of the Mandelbrot set, and render them as
  -- animated GIF.
  let step = gridStep conf
      bailout = maxIters conf
      xres = 1000
      yres = 1000
      !cells = selectCells step bailout
  whenNormal $ putStrLn $ "cell count: " ++ show (length cells)

  let red   = PixelRGB8 255 0 0
      black = PixelRGB8 0 0 0
      grey  = PixelRGB8 64 64 64

  whenNormal $ putStrLn "rendering cells"
  cellMap <- newArray ((0, 0), (xres - 1, yres - 1)) False :: IO (IOUArray (Int, Int) Bool)
  forM_ cells $ \(x, y) -> do
    let cornerA = (x - step / 2) :+ (y - step / 2)
        cornerB = (x + step / 2) :+ (y + step / 2)
        (imin, jmin) = toImgCoords xres yres cornerA
        (imax, jmax) = toImgCoords xres yres cornerB
        coords = [ (i, j) | i <- [imin..imax], j <- [jmin..jmax],
                                 i >= 0 && i < xres,
                                 j >= 0 && j < yres ]
    forM_ coords $ \(i, j) ->
      writeArray cellMap (i, j) True

  whenNormal $ putStrLn "rendering mandel"
  let imgMandel = generateImage (flip $ mandelRenderer) yres xres
      inMandelbrotSet z = inSet 0 bailout z
      mandelRenderer i j = if inMandelbrotSet $ toPlaneCoords xres yres i j
                           then grey else black

  cellMapPure <- freeze cellMap :: IO (UArray (Int, Int) Bool)
  let imgCells = generateImage (flip $ cellRenderer) yres xres
      cellRenderer i j = if cellMapPure!(i,j)
                         then red else mandelRenderer i j

  whenNormal $ putStrLn $ "writing " ++ animpath conf ++ " ..."
  case writeGifAnimation (animpath conf) 100 LoopingForever [imgMandel, imgCells] of
    Left err -> putStrLn $ "error generating gif animation: " ++ err
    Right action -> action


module Parser (SchemeData, readExpr) where

import Text.Parsec
import Text.Parsec.String

import Numeric
import Data.Ratio
import Data.Complex

-- Scheme definition

-- Scheme Data

data SchemeData = Symbol String
          | Bool Bool
          | Character Char
          | String String
          | Number Integer
          | Float Double
          | Ratio Rational
          | Complex (Complex Double)
          | List [SchemeData]
          | DottedList [SchemeData] SchemeData

-- SchemeValue to String for print
showValue :: SchemeData -> String
showValue (Symbol s) = s
showValue (Bool b) = case b of
    True -> "#t"
    False -> "#f"
showValue (Character c) = ['#', '\\', c]
showValue (String s) = "\"" ++ s ++ "\""
showValue (Number n) = (show n)
showValue (Float f) = (show f)
showValue (Ratio r) = (show r)
showValue (Complex (x :+ y)) = (show x) ++ "+" ++ (show y) ++ "i"
showValue (List l) = "(" ++ unwords (map showValue l) ++ ")"
showValue (DottedList l a) = (showValue (List l)) ++ " . " ++ (showValue a)

instance Show SchemeData where show = showValue

-- Symbol Character
-- not scheme symbol
symbol :: Parser Char
symbol = oneOf "#!$%&|*+-/:<=>?@^_~"

-- Scheme Escaped Character
--- \n, \r, \t, \\, \"
escapedChar :: Parser Char
escapedChar = do
    char '\\'
    x <- oneOf "\"\\nrt"
    return $ case x of
        'n' -> '\n'
        'r' -> '\r'
        't' -> '\t'
        _ -> x

-- Scheme String
-- " String "
parseString :: Parser SchemeData
parseString = do
    char '"'
    x <- many1 (escapedChar <|> noneOf "\"")
    char '"'
    return $ String x

-- Scheme Symbol
-- a-z , A-Z, 0-9, _
parseSymbol :: Parser SchemeData
parseSymbol = do
    head <- letter <|> char '_' <|> symbol
    tail <- many (letter <|> digit <|> char '_')
    return $ Symbol (head:tail)

-- Scheme Bool
-- #t, #f
parseBool :: Parser SchemeData
parseBool = do
    char '#'
    x <- oneOf "tf"
    return $ case x of
        't' -> Bool True
        'f' -> Bool False

-- Scheme Number
-- Bin, "#b-"
parseBin :: Parser Integer
parseBin = do
    try $ string "#b"
    x <- many1 $ oneOf "01"
    return $ bin2dec x

-- Oct, "#o-"
parseOct :: Parser Integer
parseOct = do
    try $ string "#o"
    x <- many1 digit
    return $ fst $ (readOct x) !! 0

-- Hex, "#x-"
parseHex :: Parser Integer
parseHex = do
    try $ string "#x"
    x <- many1 (digit <|> letter)
    return $ fst $ (readHex x) !! 0

-- Dec, "#d-", "0-9"
parseDec :: Parser Integer
parseDec = do
    many $ string "#d"
    x <- many1 digit
    return $ read x

-- Number Bin, Hex, Oct, Dec
parseNumber :: Parser SchemeData
parseNumber = do
    num <- parseBin <|> parseHex <|> parseOct <|> parseDec
    return $ Number num

-- Scheme Character
-- "#\a-z"
parseCharacter :: Parser SchemeData
parseCharacter = do
    string "#\\"
    v <- try (string "newline" <|> string "space") <|> do
        x <- anyChar
        notFollowedBy alphaNum
        return [x]
    return $ Character $ case v of
        "newline" -> '\n'
        "space" -> ' '
        _ -> v !! 0

-- Scheme Float
-- "-.-"
parseFloat :: Parser SchemeData
parseFloat = do
    x <- many1 digit
    char '.'
    y <- many1 digit
    return $ Float $ fst $ (readFloat $ x ++ "." ++ y) !! 0

-- Scheme Rational
-- "-/-"
parseRatio :: Parser SchemeData
parseRatio = do
    x <- many1 digit
    char '/'
    y <- many1 digit
    return $ Ratio $ (read x) % (read y)

-- Scheme Complex
parseComplex :: Parser SchemeData
parseComplex = do
    x <- try parseFloat <|> parseNumber
    char '+'
    y <- try parseFloat <|> parseNumber
    char 'i'
    return $ Complex $ (toDouble x) :+ (toDouble y)

-- Parser Combinator

parsecNum :: Parser SchemeData
parsecNum = (try parseComplex) <|> (try parseFloat) <|> (try parseRatio) <|> parseNumber

parsecSym :: Parser SchemeData
parsecSym = (try parseBool) <|> (try parseCharacter) <|> parseSymbol

-- Scheme Expr

parseExpr :: Parser SchemeData
parseExpr = parseString
        <|> try parseQuotes
        <|> try parsecNum
        <|> try parsecSym
        <|> try parseQuoted
        <|> try parseDottedList
        <|> try parseList

readExpr :: String -> SchemeData
readExpr input = case (parse parseExpr "" input) of
     Left x -> String $ "Error message is " ++ (show x)
     Right y -> y

parseList :: Parser SchemeData
parseList = do
    char '('
    x <- sepBy parseExpr spaces
    char ')'
    return $ List x

parseDottedList :: Parser SchemeData
parseDottedList = do
    char '('
    head <- endBy parseExpr spaces
    tail <- do
        char '.'
        spaces
        parseExpr
    char ')'
    return $ DottedList head tail

parseQuoted :: Parser SchemeData
parseQuoted = do
    char '\''
    x <- parseExpr
    return $ List [(Symbol "quote"), x]

parseUnQuoted :: Parser SchemeData
parseUnQuoted = do
    char ','
    x <- parseExpr
    return $ List [(Symbol "unquote"), x]

parseBackQuoted :: Parser SchemeData
parseBackQuoted = do
    char '`'
    x <- parseExpr
    return $ List [(Symbol "quasiquoted"), x]

parseQuotes :: Parser SchemeData
parseQuotes = parseQuoted <|> parseUnQuoted <|> parseBackQuoted

-- Util function

bin2dec :: String -> Integer
bin2dec [] = 0
bin2dec t@(x:xs) = (read [x]) * 2^(len -1) + (bin2dec xs)
    where len = length t

toDouble :: SchemeData -> Double
toDouble (Number x) = fromIntegral x
toDouble (Float x) = x


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ViewPatterns #-}
module Algorithm.Training where

import Prelude hiding (foldl, foldr)
-- import Numeric.AD
import Numeric.LinearAlgebra as LA hiding (fromList, (<>))
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V hiding (foldr, modify)
import Data.Functor
import Data.Bifunctor
import Data.Monoid as M
import Data.Foldable as F
import Data.Ord
import Control.Monad
import Pipes
import Pipes.Prelude as P
import Control.DeepSeq
-- import Control.Monad.RWS
-- import Control.Monad.RWS as RWS
import Data.Network.Activation
import Data.Network
import Data.Network.Utils
import Algorithm.GradientDescent

-- type TrainingSet = [Sample]

-- type TestSet = [Sample]

-- type ResultSet = [(Output, Output)]

-- Loss function parameterized by sample type (a), output loss type(l) and parameter space (p)
data LossFunction a l = Evaluator { evaluate :: (a, a) -> l, evaluate' :: (a, a) -> l }

instance Differentiable (a,a) l (LossFunction a l) where
  derivate = evaluate'

-- The type of a loss/error evaluation
type Loss = Vector R
type LossBatch = Matrix R

-- y is target, y' is actual
mse :: (Floating a) => LossFunction a a
mse = let f (y,y') = let γ = y'-y in (/ 2) $ γ**2
          f' (y,y') = (y'-y)
      in  Evaluator f f'

mseBatch :: LossFunction Batch Batch
mseBatch = let fb (b, b')  = (/ 2) $ (**2) (b' - b)
               fb' (b, b') = (b' - b)
           in Evaluator fb fb'

type NetworkEvaluator a = LossFunction a a

type LearningRate = Double
-- Training parameters = (epochs, learning rate)
type TrainingParameters = (Int, Double)

class (Network i n, Differentiable (i, i) i lf, Monad m) => Backpropagation i lf n m | lf -> i, n -> i where
  -- given a neural network and an input, compute each layer's activation and derivative
  -- forward :: n -> i -> m [(i,i)]
  -- given a loss function, a network, an output and target pair and a list of activation derivatives, compute the deltas
  backward :: lf -> n -> (i, i) -> [i] -> m [i]

data TrainingState m  = TS {
  tsmodel :: m,      -- the model(network) state
  tseta   :: Double, -- the learning rate η
  tsiteration :: Int,-- the iteration number(one per gradient move)
  tsloss  :: Double,  -- the loss of the model over the last training batch/sample
  tssampleid :: Int
  } deriving Show

-- type DescentStopper i n = StopCondition (LossFunction i i) (n, (i,i))

class AsDouble a where
  toDouble :: a -> Double

instance AsDouble (Vector R) where
  toDouble = norm_2

instance AsDouble (Matrix R) where
  toDouble m = (F.sum . (fmap norm_2) . toRows) m / (fromIntegral $ rows m)

-- A pipe that takes a sample and makes one step in the parameter space
stepP :: (Monad m, AsDouble i, Network i n, GradientDescent (LossFunction i i) (n, (i,i)) m) =>
  n -> LossFunction i i -> Int -> (Int -> Double) -> (Int, (i,i)) -> Producer (TrainingState n) m r
stepP m lf n lrf (sid, (i, t)) = loop m n where
  loop m n = do
    let !η = lrf n -- compute learning rate for iteration
    (m', _) <- lift $ descend lf η (m, (i,t)) -- step
    let o = fst $ runNetwork m' i -- test
        loss = evaluate lf (t,o) -- evaluate
    m' `deepseq` yield $ TS m' η n (toDouble loss) sid -- yield new state
    loop m' (n+1)

-- A pipe that produce a reporting action on a training state
reportP :: (Monad m, AsDouble i, Network i n) => Int -> (TrainingState n -> m ()) -> Pipe (TrainingState n) (TrainingState n) m r
reportP interval action = P.mapM $ \ts -> do
  when (tsiteration ts `mod` interval == 0) (action ts)
  return ts

type Stopper n x m = (TrainingState n, TrainingState n) -> m (Maybe x)

-- Does gradient descent given a stop condition and a stream of "steps"(training states containing an updated model)
descendP :: (Monad m, AsDouble i, Network i n) =>
    Stopper n x m-> Consumer (TrainingState n) m x
descendP stopper = loop
  where loop = do
          ts <- await
          ts' <- await
          mx <- lift (stopper (ts, ts'))
          case mx of
            Just x -> return x
            Nothing -> loop

-- A stop condition: stop if the loss doesn't change significantly enough,
-- returning the state with the minimum loss
closeEnough :: (Network i n, Monad m) => Double -> Stopper n (TrainingState n) m
closeEnough tolerance (ts, ts') = do
  if abs (tsloss ts' - tsloss ts) < tolerance then
    return (Just $ minimumBy (comparing tsloss) [ts, ts']) else return Nothing

-- stops after a given number of iterations
longEnough :: (Network i n, Monad m) => Int -> Stopper n (TrainingState n) m
longEnough iteration (ts, ts') = if tsiteration ts' >= iteration then
  return (Just $ minimumBy (comparing tsloss) [ts, ts']) else return Nothing

-- stops when there's not enough improvement(i.e. loss reduction) between the two states
improvement :: (Network i n, Monad m) => Double -> Stopper n (TrainingState n) m
improvement tolerance (ts, ts') = if ((tsloss ts' - tsloss ts) > tolerance) then
  return (Just $ minimumBy (comparing tsloss) [ts, ts']) else return Nothing


module Statistics.Classification.ConfusionMatrix.Multi (
    MultiConfusionMatrix(..)
  , multiClassCount
  , multiTotalCount
  , truePositiveCount, falsePositiveCount, falseNegativeCount, trueNegativeCount
  , confusionMatrix
  , PosNeg
  , posVal, negVal, posVals, negVals
  , trueVal, falseVal
  , isPos, isNeg
  ) where

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Statistics.Classification.Types (ClassificationResult(..))

data MultiConfusionMatrix a = MultiConfusionMatrix {
    confCounts :: Map.Map (ClassificationResult a) Int
  , confClasses :: Set.Set a
  } deriving (Eq, Show, Read)

instance Ord a => Semigroup (MultiConfusionMatrix a) where
  (MultiConfusionMatrix mcm0 mcc0) <> (MultiConfusionMatrix mcm1 mcc1) =
    MultiConfusionMatrix (Map.unionWith (+) mcm0 mcm1) (Set.union mcc0 mcc1)

instance Ord a => Monoid (MultiConfusionMatrix a) where
  mempty = MultiConfusionMatrix Map.empty Set.empty
  mappend = (<>)

multiClassCount :: Ord a => MultiConfusionMatrix a -> ClassificationResult a -> Int
multiClassCount (MultiConfusionMatrix conf _) r = Map.findWithDefault 0 r conf

truePositiveCount,falsePositiveCount,falseNegativeCount,trueNegativeCount :: Ord a => MultiConfusionMatrix a -> PosNeg a -> Int
truePositiveCount (MultiConfusionMatrix conf _) posneg = sum $
  Map.filterWithKey (\(ClassificationResult a p) _-> isPos posneg a && isPos posneg p) conf
falsePositiveCount (MultiConfusionMatrix conf _) posneg = sum $
  Map.filterWithKey (\(ClassificationResult a p) _-> isNeg posneg a && isPos posneg p) conf
falseNegativeCount (MultiConfusionMatrix conf _) posneg = sum $
  Map.filterWithKey (\(ClassificationResult a p) _-> isPos posneg a && isNeg posneg p) conf
trueNegativeCount (MultiConfusionMatrix conf _) posneg = sum $
  Map.filterWithKey (\(ClassificationResult a p) _-> isNeg posneg a && isNeg posneg p) conf

multiTotalCount :: MultiConfusionMatrix a -> Int
multiTotalCount = sum . confCounts

confusionMatrix :: Ord a => [ClassificationResult a] -> MultiConfusionMatrix a
confusionMatrix = foldMap go
  where
    go k = MultiConfusionMatrix {
        confCounts = Map.singleton k 1
      , confClasses = Set.fromList [classifiedActual k, classifiedPredicted k]
      }

data PosNeg a = PosVal (Set.Set a) | NegVal (Set.Set a)
  deriving (Eq, Show, Read)

posVal,negVal :: a -> PosNeg a
posVal = PosVal . Set.singleton
negVal = NegVal . Set.singleton

posVals,negVals :: Ord a => [a] -> PosNeg a
posVals = PosVal . Set.fromList
negVals = NegVal . Set.fromList

trueVal,falseVal :: PosNeg Bool
trueVal = posVal True
falseVal = posVal False

isPos,isNeg :: Ord a => PosNeg a -> a -> Bool
isPos (PosVal pos) v = Set.member v pos
isPos (NegVal neg) v = Set.notMember v neg
isNeg (PosVal pos) v = Set.notMember v pos
isNeg (NegVal neg) v = Set.member v neg


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module JuliaC where

import Control.Monad (replicateM_)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.State.Strict

import Data.Complex
import Data.Foldable (maximumBy)
import Data.Function (on)

import Firestarter

import Linear

import Numeric.AD

data Julia = Julia (V2 Double) Double

instance Metric Complex

instance Firestarter Julia V2 Double where
  isInside (Julia _ e) v = juliaDistance v <= e
  distanceAndDirection (Julia _ _) v =
    let
      (d,dir) = grad' juliaDistance v
      n = fmap signum $ maximumBy (compare `on` norm) $ map (negate dir*) basis
    in
      (d,n)
  startingLocation _ = V2 2 2

juliaFunction
  :: (Floating a, RealFloat a, Ord a)
  => Complex a
  -> StateT (Complex a,Complex a) (Either (Complex a, Complex a)) ()
juliaFunction c = do
  (q,dq) <- get
  if quadrance q > 200
    then lift $ Left (q,dq)
    else put (q*q+c, 2*^ (q*dq))

juliaDistanceEstimate :: Floating a => (Complex a, Complex a) -> a
juliaDistanceEstimate (q,dq) = 0.5 * r * log r / dr
  where r = norm q
        dr = norm dq

juliaDistance :: (Floating a, Ord a, RealFloat a) => V2 a -> a
juliaDistance v = juliaDistanceEstimate . either id id $ execStateT s qs
  where
    s = replicateM_ 100 $ juliaFunction c
    c = initC
    qs = fromVec v

initC :: Floating a => Complex a
initC = (-0.125) :+ (-0.256)
--initComplex = (-0.4) :+ 0.6

fromVec :: Floating a => V2 a -> (Complex a, Complex a)
fromVec (V2 x y) = (x :+ y, 1 :+ 0)


{-# LANGUAGE BangPatterns #-}
module Util where
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util
import Data.List
import Data.Ord (comparing)
import Control.Applicative
import System.Random
import System.IO.Unsafe
import Data.IORef
import qualified Data.Map as Map

partitionList :: [Double] -> [a] -> [[a]]
partitionList ps xs = iterate (proportional ps) xs
  where
    sumP = sum ps
    lenX = length xs
    proportional = map (/sumP)
    iterate [] leftover = [leftover]
    iterate (p:rest) li = (take proportion li):iterate rest (drop proportion li)
      where
        proportion = floor $ p * (fromIntegral lenX)

{- shuffle the input and partition into training, test and validation sets in the
   given proportions
   -}
shuffleAndPartitionInput :: RandomGen g => g -> (Matrix Double,Matrix Double) ->
                                           (Double,Double,Double) ->
                                           ((Matrix Double,Matrix Double),
                                            (Matrix Double,Matrix Double),
                                            (Matrix Double,Matrix Double))
shuffleAndPartitionInput gen (inputX,inputY) (trainProp,valProp,testProp) = (train,val,test)
  where
   m = rows inputY
   (trainIndices:valIndices:testIndices:rest) = partitionList [trainProp,valProp,testProp] (shuffle gen m [0..])
   train   = get trainIndices
   val     = get valIndices
   test    = get testIndices
   get ind = (extractRows ind inputX,extractRows ind inputY)

-- |shuffle a list
shuffle :: RandomGen g => g -> Int -> [a] -> [a]
shuffle gen 0   _  = []
shuffle gen len xs =
        let
                n = fst $ randomR (0, len - 1) gen
                (y, ys) =  choose n xs
                ys' = shuffle gen (len - 1) ys
        in y:ys'

choose _ [] = error "choose: index out of range"
choose 0 (x:xs) = (x, xs)
choose i (x:xs) = let (y, ys) = choose (i - 1) xs in (y, x:ys)

-- | choose a random sample from x and y
randomSample :: RandomGen g => g -> (Matrix Double,Matrix Double) -> Int -> (Matrix Double, Matrix Double)
randomSample gen (x,y) size = (newX,newY)
  where
    m = rows y
    indices = sort $ uniqueRandoms size m gen
    newX = extractRows indices x
    newY = extractRows indices y

uniqueRandoms :: RandomGen g => Int -> Int -> g -> [Int]
uniqueRandoms size max gen | max < size = error "size greater than max"
                           | otherwise = take size $ shuffle gen max [0..(max-1)]

-- fatureNormalize(X)
-- Normalizes the features in X
-- FEATURENORMALIZE(X) returns a normalized version of X where
-- the mean value of each feature is 0 and the standard deviation
-- is 1. This is often a good preprocessing step to do when
-- working with learning algorithms.
featureNormalize :: Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
featureNormalize x = ((doNormalize x mu sigma), mu, sigma)
                      where
                        mu = fromRows $ [fromList $ map mean $ map toList $ toColumns x]
                        sigma = fromRows $ [fromList $ map stddev $ map toList $ toColumns x]

-- normalize an input matrix given sigma and mu
doNormalize :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
doNormalize x mu sigma = (x - mu) / sigma

-- split a list into words by a function
wordsBy :: (a -> Bool) -> [a] -> [[a]]
wordsBy p s = case dropWhile p s of
    []      -> []
    s':rest -> (s':w) : wordsBy p (drop 1 s'')
        where (w, s'') = break p rest

-- read a csv file of numbers into a matrix
readMatrixCSV :: FilePath -> IO (Matrix Double)
readMatrixCSV path = do
    s <- map (wordsBy (\ x -> x == ',' )) <$> lines <$> readFile path
    let d = map (map read) $ s :: [[Double]]
    let cols = length $ head d
    return $ reshape cols $ fromList (concat d)

-- sum up the rows in a matrix, returning a column vector of sums
sumRows m = fromColumns $ [ sum $ toColumns m ]

-- sum up the columns in a matrix, returning a row vector of sums
sumColumns m = fromRows $ [ sum $ toRows m ]

-- pad a matrix on the left with a value
padLeft :: Matrix Double -> Double -> Matrix Double
padLeft !m num = fromBlocks [[scalar num, m]]

-- convenience function similar to above
addBias !m = fromBlocks [[1,m]]

-- |Numerically stable mean
mean :: Floating a => [a] -> a
mean x = fst $ foldl' (\(!m, !n) x -> (m+(x-m)/(n+1),n+1)) (0,0) x

-- |Same as 'mean'
average :: Floating a => [a] -> a
average = mean

-- |Median
median :: (Floating a, Ord a) => [a] -> a
median x | odd n  = head  $ drop (n `div` 2) x'
         | even n = mean $ take 2 $ drop i x'
                  where i = (length x' `div` 2) - 1
                        x' = sort x
                        n  = length x
-- |Range
range :: (Num a, Ord a) => [a] -> a
range xs = maximum xs - minimum xs

-- |Standard deviation of sample
stddev :: (Floating a) => [a] -> a
stddev xs = sqrt $ var xs

-- |Sample variance
var xs = (var' 0 0 0 xs) / (fromIntegral $ length xs - 1)
    where
      var' _ _ s [] = s
      var' m n s (x:xs) = var' nm (n + 1) (s + delta * (x - nm)) xs
         where
           delta = x - m
           nm = m + delta/(fromIntegral $ n + 1)

-- |Replace items in a list
replace :: Eq a => a -> a -> [a] -> [a]
replace _ _ [] = []
replace target replacement (x:xs) | target == x = replacement:(replace target replacement xs)
                                  | otherwise   = x:(replace target replacement xs)

-- The Sigmoid function
sigmoid x = 1.0 / (1 + exp (-x))

sigmoidGradient z = (sigmoid z) * (1 - (sigmoid z))

-- |Find the maximum value of each row and return a column vector of their indices
rowMaxIndex :: Matrix Double -> Matrix Double
rowMaxIndex m = asColumn $ fromList $ map maxIndex $ toLists m
  where
    maxIndex xs = snd $ maximum $ zip xs [0,1..]

-- |Turn a list of matrices into a vector
flattenConcat ms = fromList $ concat $ map (toList . flatten) ms

-- |Memoizes a function with unsafePerformIO
{-# NOINLINE unsafeMemoize #-}
unsafeMemoize :: Ord a => (a -> b) -> (a -> b)
unsafeMemoize f = unsafePerformIO $ do
    r <- newIORef Map.empty
    return $ \ x -> unsafePerformIO $ do
        m <- readIORef r
        case Map.lookup x m of
            Just y  -> return y
            Nothing -> do
                    let y = f x
                    writeIORef r (Map.insert x y m)
                    return y

{- Generates the train and cross validation set errors needed
  to plot a learning curve
   [error_train, error_val] = ...
       LEARNINGCURVE(X, y, Xval, yval, lambda) returns the train and
       cross validation set errors for a learning curve. In particular,
       it returns two vectors of the same length - error_train and
       error_val. Then, error_train(i) contains the training error for
       i examples (and similarly for error_val(i)).

   In this function, you will compute the train and test errors for
   dataset sizes from 1 up to m. In practice, when working with larger
   datasets, you might want to do this in larger intervals.
-}

learningCurves :: (Matrix Double -> Matrix Double -> Matrix Double) ->
                  (Matrix Double -> Matrix Double -> Matrix Double -> Double) ->
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   Int ->
                   [(Double,Double,Int)]
learningCurves train cost x y xval yval step_size = iterate 1
  where
    m = rows y
    iterate i | i >= m    = []
              | otherwise = (train_cost,cval_cost,i):iterate (i+step_size)
      where
        x' = takeRows i x
        y' = takeRows i y
        theta = train x' y'
        train_cost = cost x' y' theta
        cval_cost = cost xval yval theta

{- POLYFEATURES Maps X (1D vector) into the p-th power
   [X_poly] = POLYFEATURES(X, p) takes a data matrix X (size m x 1) and
   maps each example into its polynomial features where
   X_poly(i, :) = [X(i) X(i).^2 X(i).^3 ...  X(i).^p];
-}
oneDPolyFeatures :: Matrix Double -> Int -> Matrix Double
oneDPolyFeatures x p = fromColumns $ iterate 1
  where
    xVect = (toColumns x) !! 0
    iterate i | i > p     = []
              | otherwise = (xVect ^ i):iterate (i+1)

{- VALIDATIONCURVE Generate the train and validation errors needed to
   plot a validation curve that we can use to select lambda
   [lambda_vec, error_train, error_val] = ...
   VALIDATIONCURVE(X, y, Xval, yval) returns the train
   and validation errors (in error_train, error_val)
   for different values of lambda. You are given the training set (X,
   y) and validation set (Xval, yval).
-}
validationCurve :: (Double -> Matrix Double) -> (Matrix Double -> Double) -> (Matrix Double -> Double) -> [Double] -> [(Double,Double)]
validationCurve trainF trainErrorF valErrorF candidates = iterate candidates
  where
    iterate [] = []
    iterate (c:cs) = (trainEr,valEr):(iterate cs)
      where
        theta = trainF c
        trainEr = trainErrorF theta
        valEr = valErrorF theta

{- mapOutputsVectorToTargetMatrix
   Takes an m-dimension column vector of outputs in space {0:(k-1)} and maps it to an m><k target
   matrix in space {0,1} where each row represents an output and a 1 is used to indicate which
   class (column) an output is in.  Eg if there are 3 output classes and the input vector looks
   like this:

   [ 0,
     0,
     1,
     2 ]

  The output matrix will look like this:

  [ 1, 0, 0;
    1, 0, 0;
    0, 1, 0;
    0, 0, 1 ]
-}
mapOutputsVectorToTargetMatrix y = (targetYsMatrix,targets)
  where
    targets = nub $ toList $ flatten y
    targetYsList = map (\t -> cond y t 0 1 0) (map scalar targets)
    targetYsMatrix = fromBlocks [targetYsList]


{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE BangPatterns #-}

{-|
Module      : Statistics.FastBayes.Linear
Description : Bayesian linear regression via maximum marginal likelihood.
Copyright   : (c) Melinae, 2014
                  Chad Scherrer, 2014
License     : MIT
Maintainer  : chad.scherrer@gmail.com
Stability   : experimental
Portability : POSIX

This module gives an implementation of Bayesian linear regression, with the scale of the prior chosen by marginal likelihood.

The inputs for a Bayesian linear model are identical to those of a classical linear model, except that in addition to a design matrix and response, we must also specify a prior distribution on the weights and the noise. This leaves us with an open question of how these should be specified.

In his book /Pattern Recognition and Machine Learning/, Christopher Bishop provides details for an approach that simplifies the situation significantly, and allows for much faster inference. The structure of the linear model allows us to integrate the posterior over the weights, resulting in the /marginal likelihood/, expressed as a function of the prior precision and noise precision. This, in turn, can be easily optimized.
-}

module Statistics.FastBayes.Linear
  ( test
  , Fit
  , fit0
  , fit1
  , priorPrecision
  , noisePrecision
  , effectiveNumParameters
  , logEvidence
  , mapWeights
  , intercept
  , hessian
  , features
  , Predict
  , predict
  , studentize
  , studentizeM
  , inputs
  , fitted
  , fittedSD
  )
  where

import qualified Data.Vector.Storable as V
import Data.Packed.Matrix
import Numeric.LinearAlgebra

data Fit a = Fit
  { priorPrecision         :: Double             -- ^The precision (inverse variance) of the prior distribution, determined by maximizing the marginal likelihood
  , noisePrecision         :: Double             -- ^The precision (inverse variance) of the noise
  , effectiveNumParameters :: Double             -- ^The effective number of parameters in the model
  , logEvidence            :: Double             -- ^The log of the evidence, which is useful for model comparison (different features, same response)
  , mapWeights             :: Vector Double      -- ^The MAP (maximum a posteriori) values for the parameter weights
  , intercept              :: Double             -- ^The y-intercept
  , hessian                :: Matrix Double      -- ^The Hessian (matrix of second derivatives) for the posterior distribution
  , features               :: a → Vector Double  -- ^Extract features from a data point
  }

-- instance Show (Fit a) where

-- |@fit0 lim x y@ fits a Bayesian linear model to a design matrix @x@ and response vector @y@. This is an iterative algorithm, resulting in a sequence (list) of (α,β) values. Here α is the prior precision, and β is the noise precision. The @lim@ function passed in is used to specify how the limit of this sequence should be computed.
fit0 ::
  ([(Double, Double)] → (Double, Double)) -- ^How to take the limit of the (α,β) sequence. A simple approach is, /e.g./, @fit (!! 1000) x y@
  → Bool                                  -- ^Whether to studentize (translate and re-scale) the features
  → (a → Vector Double)                   -- ^The features
  → [a]                                   -- ^The input data
  → Vector Double                         -- ^The response vector
  → Fit a
fit0 lim student f xs y = Fit α β γ logEv m 0 h f'
  where
  (x, f') | not student = (fromRows $ map f xs,      f)
          | student     = (                 x0, fs . f)
            where
            (x0, fs) = studentizeM . fromRows . map f $ xs
  n = rows x
  p = cols x
  α0 = 1.0
  β0 = 1.0

  -- A vector of the eigenvalues of xtx
  (_,sqrtEigs,_) = compactSVD x
  eigs = V.map square sqrtEigs

  xtx = trans x <> x
  xty = trans x <> y
  getHessian a b = diag (V.replicate p a) + scale b xtx

  m = scale β $ h <\> xty

  go :: Double → Double → [(Double, Double)]
  go a0 b0 = (a0, b0) : go a b
    where
    h0 = getHessian a0 b0
    m0 = scale b0 $ h0 <\> xty
    c = V.sum $ V.map (\x' → x' / (a0 + x')) eigs
    a = c / (m0 <.> m0)
    b = recip $ (normSq $ y - x <> m0) / (fromIntegral n - c)

  γ = V.sum $ V.map (\x' → x' / (α + x')) eigs

  h = getHessian α β
  logEv = 0.5 *
    ( fromIntegral p * log α
    + fromIntegral n * log β
    - (β * normSq (y - x <> m) + α * (m <.> m))
    - logDetH
    - fromIntegral n * log (2*pi)
    )
    where
    (_,(logDetH, _)) = invlndet h

  (α, β) = lim $ go α0 β0

fit1 ::
  ([(Double, Double)] → (Double, Double)) -- ^How to take the limit of the (α,β) sequence. A simple approach is, /e.g./, @fit (!! 1000) x y@
  → Bool                                  -- ^Whether to studentize (translate and re-scale) the features
  → (a → Vector Double)                   -- ^The features
  → [a]                                   -- ^The input data
  → Vector Double                         -- ^The response vector
  → Fit a
fit1 lim student f xs y = f0 {intercept = μ}
  where
  μ = mean y
  n = V.length y
  f0 = fit0 lim student f xs (y - konst μ n)

square :: Double → Double
square x = x * x

normSq :: Vector Double → Double
normSq x = x <.> x

mean :: Vector Double -> Double
mean v = V.sum v / fromIntegral (V.length v)

studentizeM :: Matrix Double -> (Matrix Double, Vector Double → Vector Double)
studentizeM m = (x, f)
  where
  (xs, fs) = unzip . map studentize . toColumns $ m
  x = fromColumns xs
  f = fromList . zipWith ($) fs . toList

studentize :: Vector Double -> (Vector Double, Double → Double)
studentize v = (scale (1/σ) v0, f)
  where
  n   = V.length v
  μ   = mean v
  v0  = v - konst μ n
  σ   = sqrt $ normSq v0 / fromIntegral n
  f x = (x - μ) / σ

test = myFit 11
  where
  x = [1,1.1..pi]
  y = fromList $ map sin x
  f n x = fromList . map (x **) $ [1..n]
  myFit n = fit1 (!!100) True (f n) x y

data Predict a = Predict
  { inputs          :: [a]
  , fitted          :: [Double]
  , fittedSD        :: [Double]
  }

predict :: Fit a -> [a] -> Predict a
predict myFit xs = Predict xs ys (toList . V.map sqrt $ konst (1/β) n + takeDiag cov)
  where
  n = rows x
  α  = priorPrecision myFit
  β  = noisePrecision myFit
  w  = mapWeights     myFit
  y0 = intercept      myFit
  h  = hessian        myFit
  f  = features       myFit
  x  = fromRows $ map f xs
  cov = x <> (h <\> trans x)
  y = x <> w + konst y0 n
  ys = toList y