{-|

 Module      :  HF.Gauss
 Copyright   :  Copyright (c) David Schlegel
 License     :  BSD
 Maintainer  :  David Schlegel
 Stability   :  experimental
 Portability :  Haskell

Gaussian Integral evaluation plays an important role in quantum chemistry. Here functions will be provided to compute the most important integrals involving gaussian-type orbitals (GTOs).

An unnormalized primitive cartesian Gaussian Orbital has the form

<<centered_gaussian.svg centered_gaussian>>

where /A/ denotes the center of the orbital.

An unnormalized contracted Gaussian Orbital is a linear combination of primitive Gaussian:

<<centered_contracted_gaussian.svg centered_contracted_gaussian>>

"Gauss" uses the datatype contstructors 'PG' (Primitive Gaussian) and 'Ctr' (Contraction), defined in "HF.Data".

Notice: Formulas of Gaussian Integrals are taken from:
	Fundamentals of Molecular Integral Evaluation
	by Justin T. Fermann and Edward F. Valeev

-}

module HF.Gauss (
-- * Basic functions
distance, frac, center, itersum, erf, f, factorial2, general_erf,
-- * Normalization
normCtr, normPG,
-- * Integral Evaluation
s_12, t_12,
-- * Contraction operations
zipContractionWith, constr_matrix

) where

import HF.Data
import Numeric.Container hiding (linspace)
import Numeric.LinearAlgebra.Data hiding (linspace)
import Data.Maybe
import Control.DeepSeq
import Debug.Trace

---------------------
---------------------
--Helper Functions---
---------------------
---------------------

itersum function range
	| range == [] = 0
	| otherwise = (function $ head range) + itersum function (tail range)

-- |Calculates the factorial f(x) = x!
factorial n
    | n < 0     = 0
    | otherwise = product [1..n]

-- |Calculates the Doublefactorial f(x) = x!!
factorial2 :: (Show a, Eq a, Num a) => a -> a
factorial2 0 = 1
factorial2 1 = 1
factorial2 (-1) = 1
factorial2 n = n * factorial2 (n-2)

-- |Calculates the binomialcoefficient n over k
--binom :: Integer -> Integer -> Double
binom n 0 = 1
binom 0 k = 0
binom n k = binom (n-1) (k-1) * n `div` k

-- |Computes the distance between to vectors:
--
-- <<norm.svg norm>>
distance :: Vector Double -> Vector Double -> Double
distance rA rB = norm2(rA - rB)

-- | Calculates effective coefficient:
--
-- >>> frac alpha beta = alpha*beta/(alpha + beta)
frac ::  Double -> Double -> Double
frac alpha beta = alpha*beta/(alpha + beta)

-- | Calculates weighted radius of two gaussians
center :: Double -- ^ alpha
			-> Vector Double -- ^ rA
			-> Double -- ^ beta
			-> Vector Double -- ^ rB
			-> Vector Double
center a rA b rB = ((scalar b * rB) * (scalar a * rA))/scalar (a+b)

-- |error function calculation
erf :: Double -> Double
erf x = 2/sqrt pi * sum [x/(2*n+1) * product [-x^2/k |  k <- [1..n] ] | n <- [0..100]]

--error function calculation
f_0 :: Double -> Double
f_0 = erf

-- | General error function. See Eq. 4.118 Thijssen.
general_erf nu u
	| nu == 0 = sqrt pi * erf (sqrt u) * 1/(2*sqrt u)
	| otherwise = 1/(2*u) * (2*nu - 1) * (general_erf (nu-1) u ) - exp (-u)

--------------------------
---Important Functions----
--------------------------

-- |Calculates normalization factor of contracted Gaussian with arbitrary angular momentum
normCtr :: Ctr -> Double
normCtr contr = (prefactor * summation)**(-1.0/2.0)
	--See also Eq. 2.11
	where
		(l, m, n) = lmncontract contr
		--(l,m,n) = (fromIntegral l_, fromIntegral m_, fromIntegral n_) --This looks quite dirty
		n_sum = lengthcontract contr -1
		a = toList $ coefflist contr
		alp = [alpha prim | prim <- (gaussians contr)]
		mom = fromIntegral $ momentumctr contr
		prefactor = 1.0/(2.0**mom)*pi**(3.0/2.0)* factorial2 (2*l-1) * factorial2 (2*m-1)* factorial2 (2*n-1)
		summation = sum $ concat $ [[(a !! i)*(a !! j)/((alp !! i +alp !! j)**(mom + 3.0/2.0)) | i <-[0..n_sum]]| j <- [0..n_sum]]

-- |Calculates normalization factor of primitive Gaussian with arbitrary angular momentum
normPG :: PG -> Double
normPG (PG lmn alpha pos) = (2*alpha/pi)**(3/4) * (a/b)**(1/2)
	where
		(l, m, n) = (fromIntegral $ lmn !! 0, fromIntegral $ lmn !! 1, fromIntegral $ lmn !! 2)
		a = (8*alpha)**(l+m+n) * factorial l * factorial m * factorial n
		b = factorial (2*l) * factorial (2*m) * factorial (2*n)

-- |Calculates f_k (l1 l2 pa_x pb_x) used in Gaussian Product Theorem
f :: Int -> Int -> Int -> Double -> Double -> Double
--trace ("f" ++ " " ++  show k ++ " " ++ show l1_ ++ " " ++ show l2_ ++ " " ++ show pa_x ++ " " ++ show pb_x) $
f k l1_ l2_ pa_x pb_x = sum $ concat $ [[pa_x**(fromIntegral (l1_-i))  *  pb_x**(fromIntegral(l2_-j)) * fromIntegral ((binom l1_ i) * (binom l2_ j))| i <- [0..l1_], i+j <= k]| j <-[0..l2_]]
--See also Eq. 2.45

-- |Evaluates a given function for two contractions. The operation is symmetric in contraction arguments.
zipContractionWith :: (PG -> PG -> Double) -- ^ Function of two primitive Gaussians
							-> Ctr -- ^ Contraction
							-> Ctr -- ^ Contraction
							-> Double
zipContractionWith zipfunction (Ctr pglist1 coeffs1) (Ctr pglist2 coeffs2) = pr1 * pr2 * value
	where
		coefflist1 = toList coeffs1
		coefflist2 = toList coeffs2
		n1 = length coefflist1 -1
		n2 = length coefflist2 -1
		pr1 = normCtr (Ctr pglist1 coeffs1) --I am not so sure about this
		pr2 = normCtr (Ctr pglist2 coeffs2) --I am not so sure about this
		value = sum $ [(coefflist1 !! i)* (coefflist2 !! j)* (zipfunction (pglist1 !! i) (pglist2 !! j)) | i <- [0..n1], j <- [0..n2]]

-- |Construct a matrix out of a list of Contractions and a function for two primitive gaussians
constr_matrix :: [Ctr] -- ^ List of Contraction
						-> (PG -> PG -> Double) -- ^ Function of two primitive Gaussians
						-> Matrix Double -- ^ Matrix of dimensions (length [Ctr]) x (length [Ctr])
constr_matrix contractionlist function = buildMatrix (length contractionlist) (length contractionlist) (\(i,j) -> zipContractionWith function (contractionlist !! i) (contractionlist !! j) )

-- |Calculates overlap integral of two primitive gaussians

-- | <<s_12.svg s_12>>
s_12 :: PG -> PG -> Double
s_12 (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = trace ("s_12:   range " ++ show range  ++ "   f 0 =" ++ show (factor 0) ++ "   f 1 = " ++ show (factor 1)++ "   f 2 =" ++ show (factor 2)) $ prefactor * (factor 0) * (factor 1) * (factor 2) --pref * I_x * I_y * I_z
	where
		(l1, m1, n1) = (lmn1 !! 0, lmn1 !! 1, lmn1 !! 2)
		(l2, m2, n2) = (lmn2 !! 0, lmn2 !! 1, lmn2 !! 2)
		g = alpha1 + alpha2
		prefactor = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		p = center alpha1 pos1 alpha2 pos2
		pa = toList $ p  - pos1
		pb = toList $ p  - pos2
		range = [0..(fromIntegral (l1+l2) /2)] :: (Enum a, Fractional a) => [a]
		-- See also Eq. 3.15
		function k i = force $ f (round (2*i)) l1 l2 (pa !! k) (pb !! k) *  (factorial2 (2.0* i -1.0)) * ((2*g)** (- i)) * (pi/g)**(1/2)
		factor k = itersum (function k) range
		--i k = sum $ [f (round (2*i)) l1 l2 (pa !! k) (pb !! k) * factorial2 (2.0* i -1.0) * ((2*g)** (- i)) * (pi/g)**(1/2) | i <- range ]

-- |Calculates kinetic energy integral of two primitive gaussians

-- | <<t_12.svg t_12>>
t_12 :: PG -> PG -> Double
t_12 (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) =  (i 0) + (i 1) + (i 2) -- I_x + I_y + I_z
	where
		--lifts or lowers l, m or n, indicated by index k for a lmn-list, respectively
		pp1 list k = [if i == k then list !! i + 1 else list !! i | i<-[0..2]]	--lifts
		mm1 list k = [if i == k then list !! i - 1 else list !! i | i<-[0..2]]	--lowers
		--overlaps where parts of l, m or n ar lifted +1 or lowered -1
		m1m1 k = s_12 (PG (mm1 lmn1 k) alpha1 pos1) (PG (mm1 lmn2 k) alpha2 pos2) -- <-1|-1>_k
		p1p1 k = s_12 (PG (pp1 lmn1 k) alpha1 pos1) (PG (pp1 lmn2 k) alpha2 pos2) -- <+1|+1>_k
		p1m1 k = s_12 (PG (pp1 lmn1 k) alpha1 pos1) (PG (mm1 lmn2 k) alpha2 pos2) -- <+1|-1>_k
		m1p1 k = s_12 (PG (mm1 lmn1 k) alpha1 pos1) (PG (pp1 lmn2 k) alpha2 pos2) -- <-1|+1>_k
		--See also Eq.4.13
		i k = 1/2 * (fromIntegral (lmn1 !! k)) * (fromIntegral (lmn2 !! k)) * (m1m1 k)
			+ 2 * alpha1 * alpha2 * (p1p1 k)
			- alpha1 * (fromIntegral (lmn2 !! k)) * (p1m1 k)
			- alpha2 * (fromIntegral (lmn1 !! k)) * (m1p1 k)

{-
-- |Calculates nuclear attraction integral of two primitive gaussians
v_12 :: Double -> Vecotr Double -> PG -> PG -> Double
v_12 z r_c (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = z*pi*pre/g * 2* summation
	where
		ax, ay, az = (fromList pos1) !! 0, (fromList pos1) !! 1, (fromList pos1) !! 2
		bx, by, bz = (fromList pos2) !! 0, (fromList pos2) !! 1, (fromList pos2) !! 2
		r_cx, r_cy, r_cz = (fromList r_c) !! 0, (fromList r_c) !! 1, (fromList r_c) !! 2
		pref = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		g = alpha1 + alpha2
		p = 1/g * (alpha1*pos1 + alpha2*pos2)
		px, py, pz = (fromList p) !! 0, (fromList p) !! 1, (fromList p) !! 2
		eta = (alpha1 + alpha2)/g
		mu_x = l1 + l2 - 2*(ijk1 + ijk2) - (opq1 + opq2)
		mu_y = l1 + l2 - 2*(j1 + j2) - (p1 + p2)
		mu_z = l1 + l2 - 2*(k1 + k2) - (q1 + q2)
		nu = mu_x + mu_y + mu_z - (u + v + w)
		summation = general_erf nu (g * distance p r_c) * a_x * a_y * a_z

		--ranges:
		ijk1_range = [0..floor $ (fromIntegral l1 /2)] :: (Integral a, Enum a) => [a]
		ijk2_range = [0..floor $ (fromIntegral l2 /2)] :: (Integral a, Enum a) => [a]
		opq1_range ijk1 = [0..floor $ (fromIntegral l1 - 2*ijk1)] :: (Integral a, Enum a) => [a]
		opq2_range ijk2 = [0..floor $ (fromIntegral l2 - 2*ijk2)] :: (Integral a, Enum a) => [a]
		rst_range opq1 opq2 = [0..floor $ fromIntegral (opq1+l2) /2] :: (Integral a, Enum a) => [a]
		u_range mu_x = [0..floor $ (fromIntegral mu_x /2)] :: (Integral a, Enum a) => [a]

		a_x = (-1)**(l1+l2) * (factorial l1) * (factorial l2) *
		(-1)**(opq2 + r) * factorial (opq1 + opq2) * 1/(4**(ijk1+ijk2+r) * factorial ijk1 * factorial ijk2 * factorial opq1 * factorial opq2 * factorial r) *
		alpha1**(opq2-ijk1-r) * alpha2**(opq1-ijk2-r) * (ax-bx)**(opq1+opq2-2*r) * 1/(factorial (l1-2*ijk1-opq1) * factorial (l2-2*ijk2-opq2) * factorial (opq1-2*opq2-2*r)) *
		* [ (-1)**u * factorial mu_x * (px -r_cx)**(mu_x - 2*u) * 1/ ((4)**u * factorial u * factorial (mu_x-2*u) * g**(opq1 + opq2 - r + u)) | u <- [0..1/2 * mu_x]]
-}

--v_12_test ::  Double -> Vector Double -> PG -> PG -> [Double]
{- v_12_test z r_c (PG lmn1 alpha1 pos1) (PG lmn2 alpha2 pos2) = summation 0
	where
		--(ax, ay, az) = ((toList pos1) !! 0, (toList pos1) !! 1, (toList pos1) !! 2)
		--(bx, by, bz) = ((toList pos2) !! 0, (toList pos2) !! 1, (toList pos2) !! 2)
		--(r_cx, r_cy, r_cz) = ((toList r_c) !! 0, (toList r_c) !! 1, (toList r_c) !! 2)
		--lmn1_ = fromList lmn1
		--lmn2_ = fromList lmn2

		pref = exp (-alpha1*alpha2 * (distance pos1 pos2) /g)
		g = alpha1 + alpha2
		p = scale (1/g) (add (scale alpha1 pos1) (scale alpha2 pos2))
		--(px, py, pz) = ((toList p) !! 0, (toList p) !! 1, (toList p) !! 2)

		--f_nu ijk1 ijk2 opq1 opq2 = general_erf (nu ijk1 ijk2 opq1 opq2) (g * distance p r_c)

		eta = (alpha1 + alpha2)/g
		mu ijk1 ijk2 opq1 opq2 s = fromIntegral (atIndex lmn1 s) + fromIntegral (atIndex lmn2 s) - 2*(ijk1 + ijk2) - (opq1 + opq2)
		--nu ijk1 ijk2 opq1 opq2 = (mu ijk1 ijk2 opq1 opq2 s) + (mu ijk1 ijk2 opq1 opq2 s) + (mu ijk1 ijk2 opq1 opq2 s) - (u v q)

		--ranges alle ranges (X,Y,Z) mit einer klappe schlagen, dafür index s:
		ijk1_range s = [0..fromIntegral $ floor $ fromIntegral (lmn1 !! s) /2] :: [Double]
		ijk2_range s = [0..fromIntegral $ floor $ fromIntegral (lmn2 !! s) /2] :: [Double]
		opq1_range ijk1 s = [0..fromIntegral $ floor $ (fromIntegral (lmn1 !! s) - 2*ijk1)] :: [Double]
		opq2_range ijk2 s = [0..fromIntegral $ floor $ (fromIntegral (lmn2 !! s) - 2*ijk2)] :: [Double]
		rst_range opq1 opq2 s = [0..fromIntegral $ floor $ (opq1+opq2) /2] :: [Double]
		uvw_range ijk1 ijk2 opq1 opq2 s = [0..fromIntegral $ floor $ ((mu ijk1 ijk2 opq1 opq2 s) /2)] :: [Double]

		--a_x :: Num a => a -> a -> a -> a -> a -> Double

		--(-1)**(l1+l2) * (factorial fromIntegral (lmn1 !! s)) * (factorial fromIntegral (lmn2 !! s)) *
		function ijk1 ijk2 opq1 opq2 rst s = (-1)**(opq2 + rst) * factorial (opq1 + opq2) * 1/(4**(ijk1+ijk2+rst) * factorial ijk1 * factorial ijk2 * factorial opq1 * factorial opq2 * factorial rst) * alpha1**(opq2-ijk1-rst) * alpha2**(opq1-ijk2-rst) * (atIndex pos1 s - atIndex pos2 s)**(opq1+opq2-2*rst) * 1/(factorial ((lmn1 !! s)-2*ijk1-opq1) * factorial ((lmn2 !! s)-2*ijk2-opq2) * factorial (opq1-2*opq2-2*rst)) * sum $ [ (-1)**uvw * factorial (mu ijk1 ijk2 opq1 opq2 s) * ((atIndex p s) - (atIndex r_c s))**((mu ijk1 ijk2 opq1 opq2 s) - 2*uvw) * 1/ ((4)**uvw * factorial uvw * factorial ((mu ijk1 ijk2 opq1 opq2 s)-2*uvw) * g**(opq1 + opq2 - rst + uvw)) | uvw <- (uvw_range ijk1 ijk2 opq1 opq2 s)]

		summation s = (-1)**(lmn1 !! s + lmn2 !! s) * (factorial fromIntegral (lmn1 !! s)) * (factorial fromIntegral (lmn2 !! s)) * sum $ concat $ concat $ [[[function ijk1 ijk2 opq1 opq2 rst s | rst <- (rst_range opq1 opq2 s)] | opq1 <- (opq1_range ijk1 s), opq2 <- (opq2_range ijk2 s)]| ijk1 <- ijk1_range s, ijk2 <- ijk2_range s]
-}

-----------------
--Deprecated-----
-----------------

-- Here all importand integrals involving gaussian functions will be evaluated.

--Overlap integral for s orbitals
--calculates <1s,alpha, A | 1s, beta, B>
overlaps :: Double -> Double -> Vector Double -> Vector Double -> Double
overlaps alpha beta rA rB = prefactor * exp exponent
	where 	prefactor = (pi/(alpha + beta))**(3/2)
		exponent = - (frac alpha beta * distance rA rB)

--Kinetic integral
--calculates <1s,alpha, A | - Laplace | 1s, beta, B>
kinetic :: Double -> Double -> Vector Double -> Vector Double -> Double
kinetic alpha beta rA rB = prefactor * exp(exponent)
	where 	prefactor = (frac alpha beta)
			 *(6 -4 *(frac alpha beta) * distance rA rB)
			 *(pi/(alpha + beta))**(3/2)
		exponent = - ((frac alpha beta) * distance rA rB)

--Nuclear interaction integral
--calculates <1s,alpha, A | - Z/r_C | 1s, beta, B>
nuclear :: Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Double -> Double
nuclear alpha beta rA rB rC z = pref1 * pref2 * exp(exponent)
	where 	pref1 = -2*pi*z/(alpha + beta)
		pref2 = f_0 arg
		arg   = (distance (center alpha rA beta rB) rC )   * (alpha*beta)
		exponent = - ((frac alpha beta) * (distance rA rB))

--two-electron integral
--important!
--calculates <1s,alpha, A ; 1s, beta, B | 1s,gamma, C ; 1s, delta, D >
twoelectron :: Double -> Double -> Double -> Double -> Vector Double -> Vector Double -> Vector Double -> Vector Double -> Double
twoelectron a b g d rA rB rC rD = pref * exp(exponent) * (f_0 arg)
	where	pref = 			(2*pi**(5/2))/
				( (a+g)*(b+d)*(a+b+g+d)**(1/2)   )
		exponent = 	-a*g/(a+g) * (distance rA rC)
				-b*d/(b+d) * (distance rB rD)
		arg = (a+g)*(b+d)/(a+b+g+d) * ( distance (center a rA g rC) (center b rB d rD) )

--Overlap Matrix S_pq = <p|q>
--overlap p q  = t
--where
--	t = buildMatrix (n*n) (n*n) (\(i,j) -> overlaps alpha beta rA rB )


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Random.Distribution.Static.MultivariateNormal
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- Sample from the multivariate normal distribution with a given
-- vector-valued \(\mu\) and covariance matrix \(\Sigma\). For
-- example, the chart below shows samples from the bivariate normal
-- distribution. The dimension of the mean \(n\) is statically checked
-- to be compatible with the dimension of the covariance matrix \(n \times n\).
--
-- <<diagrams/src_Data_Random_Distribution_Static_MultivariateNormal_diagMS.svg#diagram=diagMS&height=600&width=500>>
--
-- Example code to generate the chart:
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- >
-- > import Data.Random.Distribution.Static.MultivariateNormal
-- >
-- > import qualified Data.Random as R
-- > import Data.Random.Source.PureMT
-- > import Control.Monad.State
-- > import Numeric.LinearAlgebra.Static
-- >
-- > nSamples :: Int
-- > nSamples = 10000
-- >
-- > sigma1, sigma2, rho :: Double
-- > sigma1 = 3.0
-- > sigma2 = 1.0
-- > rho = 0.5
-- >
-- > singleSample :: R.RVarT (State PureMT) (R 2)
-- > singleSample = R.sample $ Normal (vector [0.0, 0.0])
-- >                (sym $ matrix [ sigma1, rho * sigma1 * sigma2
-- >                              , rho * sigma1 * sigma2, sigma2])
-- >
-- > multiSamples :: [R 2]
-- > multiSamples = evalState (replicateM nSamples $ R.sample singleSample) (pureMT 3)
-- >
-- > pts = map f multiSamples
-- >   where
-- >     f z = (x, y)
-- >       where
-- >         (x, t) = headTail z
-- >         (y, _) = headTail t
-- >
-- > chartPoint pointVals n = C.toRenderable layout
-- >   where
-- >
-- >     fitted = C.plot_points_values .~ pointVals
-- >               $ C.plot_points_style  . C.point_color .~ opaque red
-- >               $ C.plot_points_title .~ "Sample"
-- >               $ def
-- >
-- >     layout = C.layout_title .~ "Sampling Bivariate Normal (" ++ (show n) ++ " samples)"
-- >            $ C.layout_y_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >            $ C.layout_x_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >
-- >            $ C.layout_plots .~ [C.toPlot fitted]
-- >            $ def
-- >
-- > diagMS = do
-- >   denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >   return $ fst $ runBackend denv (C.render (chartPoint pts nSamples) (500, 500))
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE MultiParamTypeClasses        #-}
{-# LANGUAGE TypeFamilies                 #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE DataKinds                    #-}

module Data.Random.Distribution.Static.MultivariateNormal
    ( Normal(..)
    ) where

import           Data.Random hiding ( StdNormal, Normal )
import qualified Data.Random as R
import           Control.Monad.State ( replicateM )
import qualified Numeric.LinearAlgebra.HMatrix as H
import           Numeric.LinearAlgebra.Static as S
                 ( R, vector, extract, Sq, Sym, col,
                   tr, linSolve, uncol, chol, (<.>),
                   ℝ, (<>), diag, (#>), eigensystem
                 )
import          GHC.TypeLits ( KnownNat, natVal )
import          Data.Maybe ( fromJust )

normalMultivariate :: KnownNat n =>
                      R n -> Sym n -> RVarT m (R n)
normalMultivariate mu bigSigma = do
  z <- replicateM (fromIntegral $ natVal mu) (rvarT R.StdNormal)
  return $ mu + bigA #> (vector z)
  where
    (vals, bigU) = eigensystem bigSigma
    lSqrt = diag $ mapVector sqrt vals
    bigA = bigU S.<> lSqrt

mapVector :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapVector f = vector . H.toList . H.cmap f . extract

sumVector :: KnownNat n => R n -> ℝ
sumVector x = x <.> 1

data family Normal k :: *

data instance Normal (R n) = Normal (R n) (Sym n)

instance KnownNat n => Distribution Normal (R n) where
  rvar (Normal m s) = normalMultivariate m s

normalLogPdf :: KnownNat n =>
                R n -> Sym n -> R n -> Double
normalLogPdf mu bigSigma x = - sumVector (mapVector log (diagonals dec))
                             - 0.5 * (fromIntegral $ natVal mu) * log (2 * pi)
                             - 0.5 * s
  where
    dec = chol bigSigma
    t = uncol $ fromJust $ linSolve (tr dec) (col $ x - mu)
    u = mapVector (\x -> x * x) t
    s = sumVector u

normalPdf :: KnownNat n =>
             R n -> Sym n -> R n -> Double
normalPdf mu sigma x = exp $ normalLogPdf mu sigma x

diagonals :: KnownNat n => Sq n -> R n
diagonals = vector . H.toList . H.takeDiag . extract

instance KnownNat n => PDF Normal (R n) where
  pdf (Normal m s) = normalPdf m s
  logPdf (Normal m s) = normalLogPdf m s


{- |
This module provides normalized versions of the transforms in @fftw@.

All of the transforms are normalized so that

 - Each transform is unitary, i.e., preserves the inner product and the sum-of-squares norm of its input.

 - Each backwards transform is the inverse of the corresponding forwards transform.

(Both conditions only hold approximately, due to floating point precision.)

For more information on the underlying transforms, see
<http://www.fftw.org/fftw3_doc/What-FFTW-Really-Computes.html>.
--
-- @since 0.2
-}

module Numeric.FFT.Vector.Unitary.Multi
  (
        -- * Creating and executing 'Plan's
        run,
        plan,
        execute,
        -- * Complex-to-complex transforms
        dft,
        idft,
        -- * Real-to-complex transforms
        dftR2C,
        dftC2R,
  ) where

import Control.Exception (assert)
import Control.Monad (forM_)
import Numeric.FFT.Vector.Base
import qualified Numeric.FFT.Vector.Unnormalized.Multi as U
import Data.Complex
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MS
import Control.Monad.Primitive(RealWorld)

-- | A discrete Fourier transform. The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(-2pi i j k\/n)@
dft :: TransformND (Complex Double) (Complex Double)
dft = U.dft {normalizationND = \ns -> constMultOutput $ 1 / sqrt (toEnum (VS.product ns))}

-- | An inverse discrete Fourier transform.  The output and input sizes are the same (@n@).
--
-- @y_k = (1\/sqrt n) sum_(j=0)^(n-1) x_j e^(2pi i j k\/n)@
idft :: TransformND (Complex Double) (Complex Double)
idft = U.idft {normalizationND = \ns -> constMultOutput $ 1 / sqrt (toEnum (VS.product ns))}

-- | A forward discrete Fourier transform with real data.  If the input size is @n@,
-- the output size will be @n \`div\` 2 + 1@.
dftR2C :: TransformND Double (Complex Double)
dftR2C = U.dftR2C {normalizationND = \ns -> modifyOutput $
                    complexR2CScaling (sqrt 2) ns (outputSizeND U.dftR2C $ VS.last ns)
        }

-- | A normalized backward discrete Fourier transform which is the left inverse of
-- 'U.dftR2C'.  (Specifically, @run dftC2R . run dftR2C == id@.)
--
-- This 'Transform' behaves differently than the others:
--
--  - Calling @plan dftC2R n@ creates a 'Plan' whose /output/ size is @n@, and whose
--    /input/ size is @n \`div\` 2 + 1@.
--
--  - If @length v == n@, then @length (run dftC2R v) == 2*(n-1)@.
--
dftC2R :: TransformND (Complex Double) Double
dftC2R = U.dftC2R {normalizationND = \ns -> modifyInput $
                    complexR2CScaling (sqrt 0.5) ns (inputSizeND U.dftC2R $ VS.last ns)
        }

complexR2CScaling :: Double -> VS.Vector Int -> Int -> MS.MVector RealWorld (Complex Double) -> IO ()
complexR2CScaling !t !ns !len !a = assert (MS.length a == VS.product (VS.init ns) * len) $ do
    let !s1 = sqrt (1/toEnum (VS.product ns))
    let !s2 = t * s1
    -- Justification for the use of unsafeModify:
    -- The output size is 2n+1; so if n>0 then the output size is >=1;
    -- and if n even then the output size is >=3.
    forM_ [0.. VS.product (VS.init ns) - 1] $ \idx -> do
      unsafeModify a (idx * len) $ scaleByD s1
      if odd (VS.last ns)
        then multC s2 (MS.unsafeSlice (idx * len + 1) (len-1) a)
        else do
            unsafeModify a (idx * len + len - 1) $ scaleByD s1
            multC s2 (MS.unsafeSlice (idx * len + 1) (len-2) a)


{-# LANGUAGE RankNTypes #-}

-- |
--  Author: Benedikt Spies
--
--  Utility functions to generate simple Images
--  Image rows are generated in parallel
module Utils.ImageGenerator (generateRainbowImage, generateMandelbrotImage) where

import Codec.Picture
  ( Image,
    Pixel8,
    PixelRGB8 (PixelRGB8),
  )
import Data.Complex (Complex ((:+)), magnitude)
import Utils.ParallelImageUtils (parGenerateImage)

generateRainbowImage :: Image PixelRGB8
generateRainbowImage = parGenerateImage f 256 256
  where
    f :: Int -> Int -> PixelRGB8
    f x y = PixelRGB8 r g b
      where
        r = fromIntegral $ 255 - x
        g = fromIntegral $ 255 - y
        b = fromIntegral x

generateMandelbrotImage :: Int -> Int -> Image PixelRGB8
generateMandelbrotImage width height = parGenerateImage (generateMandelbrotPixel width height) width height

generateMandelbrotPixel :: Int -> Int -> Int -> Int -> PixelRGB8
generateMandelbrotPixel imageWidth imageHeight x y = color $ mandelbrot (x' / w' * 3 - 2) (y' / h' * 3 - 1.5)
  where
    x' = fromIntegral x
    y' = fromIntegral y
    w' = fromIntegral imageWidth
    h' = fromIntegral imageHeight
    maxIter = 300
    color :: Int -> PixelRGB8
    color iter = PixelRGB8 colorVal colorVal colorVal
      where
        colorVal = floor $ 255 - amt * 255 :: Pixel8
        amt = fromIntegral iter / fromIntegral maxIter :: Float
    mandelbrot :: Float -> Float -> Int
    mandelbrot r i =
      length . takeWhile (\z -> magnitude z <= 2)
        . take maxIter
        $ iterate (\z -> z ^ (2 :: Int) + (r :+ i)) 0


-- | Dynamic Programming module.
module DP where

import Data.Map (Map)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Numeric.LinearAlgebra.Data
       (Matrix, Vector, (!), (¿), (><), maxIndex)
import Numeric.LinearAlgebra.HMatrix as H

import MDP
import Planning

import Debug.Trace

data MDPM s a = MDPM
    { statesm :: Map s Int
    , actionsm :: Map a Int
    , transitionm :: Map a (Matrix Double)
    , rewardm :: Matrix Double
    , gammam :: Double             -- discount
    , sizeStates :: Int           -- S
    , sizeActions :: Int          -- A
    , actionsr :: Map Int a
    }

prepareMDP :: (Ord s, Ord a) => MDP s a -> MDPM s a
prepareMDP mdp =
    MDPM
    { statesm = statesM
    , actionsm = actionsM
    , transitionm = transitionM
    , rewardm = rewardM
    , gammam = gamma mdp
    , sizeStates = S.size (states mdp)
    , sizeActions = S.size (actions mdp)
    , actionsr = M.fromList (zip [0..] (S.toList (actions mdp)))
    }
  where
    statesM = M.fromList (zip (S.toList (states mdp)) [0 ..])
    actionsM = M.fromList (zip (S.toList (actions mdp)) [0 ..])
    transitionM = transitionMap
    transitionMap =
        M.fromList
            (map
                 (\a ->
                       (a, createMatrix a))
                 (S.toList (actions mdp)))
    createMatrix a =
        H.fromLists
            (map
                 (\s ->
                       map (transition mdp a s) (S.toList (states mdp)))
                 (S.toList (states mdp)))
    rewardM = createRewardMatrix
    createRewardMatrix =
        H.fromLists
            (map
                 (\s ->
                       map
                           (\a ->
                                 reward mdp a s)
                           (S.toList (actions mdp)))
                 (S.toList (states mdp)))

valueIterationM :: (Ord a) => MDPM s a -> Int -> Matrix Double -> Matrix Double
valueIterationM mdpm k v
  | k == 0 = v
  | otherwise = valueIterationM mdpm (k - 1) (calc v)
  where
    calc v' =
        rewardm mdpm +
        gammam mdpm `H.scale`
        fromColumns
            (map
                 (\a ->
                       flatten
                           (transitionm mdpm M.! a H.<>
                            (v' ¿ [actionsm mdpm M.! a])))
                 (M.keys (actionsm mdpm)))

valueIteration :: (Ord a, Ord s) => MDP s a -> Int -> s -> a
valueIteration mdp k s = actionsr mdpm M.! maxIndex (vk ! (statesm mdpm M.! s))
  where
    mdpm = prepareMDP mdp
    v0 = (sizeStates mdpm >< sizeActions mdpm) [0,0 ..]
    vk = valueIterationM mdpm k v0


{- |
Description :  Convenience functions to simulate Gamma rate heterogeneity
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

Sites evolve at different speeds. One way to model this phenomenon is to use
Gamma rate heterogeneity. This module provides the necessary definitions.

* Changelog

-}

module GammaRate where

import           Numeric.Integration.TanhSinh
import           Statistics.Distribution
import           Statistics.Distribution.Gamma

-- | For a given number of rate categories 'n' and a shape parameter 'alpha'
-- (the rate or scale is set such that the mean is 1.0), return a list of rates
-- that represent the respective categories. Use the mean rate for each
-- category.
getMeans :: Int -> Double -> [Double]
getMeans n alpha = means ++ lastMean
  where gamma = gammaDistr alpha (1.0/alpha)
        quantiles = [ quantile gamma (fromIntegral i / fromIntegral n) | i <- [0..n] ]
        -- Calculate the mean rate. Multiplication with the number of rate
        -- categories 'n' is necessary because in each n-quantile the
        -- probability mass is 1/n.
        meanFunc x = fromIntegral n * x * density gamma x
        -- Only calculate the first (n-1) categories with normal integration.
        means = [ intAToB meanFunc (quantiles !! i) (quantiles !! (i+1)) | i <- [0..n-2] ]
        -- The last category has to be calculated with an improper integration.
        lastMean = [intAToInf meanFunc (quantiles !! (n-1))]

-- | The error of integration.
eps :: Double
eps = 1e-6

-- | The integration method to use
method :: (Double -> Double ) -> Double -> Double -> [Result]
method = parSimpson

-- | Helper function for a normal integral from 'a' to 'b'.
intAToB :: (Double -> Double) -> Double -> Double -> Double
intAToB f a b = result . absolute eps $ method f a b

-- | Helper function for an improper integral from 'a' to infinity.
intAToInf :: (Double -> Double) -> Double -> Double
intAToInf f a = (result . absolute eps $ nonNegative method f) - intAToB f 0 a


{-# LANGUAGE BangPatterns #-}
-- file: Astro/Misc.hs
-- Miscellaneous utilities

module Astro.Misc where

import Data.List
import Debug.Trace
import Numeric.LinearAlgebra
import Numeric (showEFloat)
import Text.Printf

{-
 - General math
 -}

type Polynomial = [Double]

-- Evaluate real polynomial given as list of coefficients [a0,a1,...]
evalPoly :: Polynomial->Double->Double
evalPoly p x = sum $ zipWith (*) p $ map (x^) [0..]

-- Derivative of a (real) polynomial
diffPoly :: Polynomial->Polynomial
diffPoly [] = []
diffPoly (a0:[]) = []
diffPoly (a0:as) = zipWith (*) as [1..]

-- Degree of a polynomial
degPoly :: Polynomial->Int
degPoly [] = 0
degPoly (a:[]) = if a /= 0 then 1 else 0
degPoly p = if l == [] then 0 else last l
	where
	l = Data.List.findIndices (/= 0) p

-- Laguerre polynomial root finder
-- Returns Nothing if fails to converge to a real root within maxIt
rootLaguerre (tol,maxIt) p x0  = result
	where
	result = helper x0 0
	p' = diffPoly p
	p'' = diffPoly p'
	deg = fromIntegral $ degPoly p
	helper x n
		| isNaN x = Nothing
		| abs (evalPoly p x) <= tol = Just x
		| n >= maxIt = Nothing -- `debug` ("maxit x,nx:"++show(x,nx))
		| otherwise = helper nx (n+1)
                {-
                        `debug` (printf ("g:n %g h: %g sqa: %g sq: %g x:"
                                ++"%g y: %g") g h sqarg sq nx y)
                                -}
		where
		g = evalPoly p' x / evalPoly p x
		h = g^2 - evalPoly p'' x / evalPoly p x
		sqarg = (deg-1)*(deg*h-g^2)
		sq = sqrt sqarg
		denom = if g>0 then g+sq else g-sq
		nx = x - deg/denom -- `debug` ("n,denom:"++show(n,denom))
                y = evalPoly p nx

{-
 - Vector Math
 -}

-- Standard vector norms
norm :: Vector Double -> Double
--norm v = pnorm PNorm2 v
-- XXX: For 3-vectors this is quite a bit faster
norm !v = sqrt $ v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

norm3D :: Vector Double -> Double
norm3D = sqrt . sqrNorm3D

sqrNorm3D :: Vector Double -> Double
sqrNorm3D v = v@>0*v@>0 + v@>1*v@>1 + v@>2*v@>2

-- Normalize vector
unitV v = scale (recip $ norm v) v

-- 3-vector cross product
cross :: Vector Double->Vector Double->Vector Double
cross !a !b =  3|>[a@>1 * b@>2 - a@>2 * b@>1
		,-a@>0 * b@>2 + a@>2 * b@>0
		,a@>0 * b@>1 - a@>1 * b@>0]

-- Rodrigues' rotation formula
rodr :: Vector Double->Double->Vector Double->Vector Double
rodr z th vec = scale (cos th) vec + scale (sin th) (cross z vec)
		+ scale ((z<.>vec)*(1-cos th)) z

{-
 - List and tuple manipulations
 -}

thrd (_,_,a) = a
frth (_,_,_,a) = a

-- Give all pairings, excluding (p,p) and when (p1,p2) excluding (p2,p1)
-- Courtesy of Jussi Knuuttila
chooseTwo :: [a] -> [(a,a)]
chooseTwo as = concatMap pairs (tails as)
    where pairs [] = []
          pairs (a:as) = zip (repeat a) as

-- Take every nth of list
takenth 0 _ = []
takenth _ [] = []
takenth n list =
	iter list []
	where
	iter [] acc = acc
	iter list acc =
		iter (tail newList) (acc ++ [head newList])
		where newList = drop (n-1) list

-- Debugging and printing functions
debug = flip trace
showFloat f = showEFloat (Just 4) f ""


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics.Information
-- Copyright   :  (c) A. V. H. McPhail 2010, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Shannon entropy
--
-----------------------------------------------------------------------------

module Numeric.Statistics.Information (
                                   entropy
                                   , mutual_information
                                  ) where

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

import Numeric.Statistics.PDF

import qualified Numeric.LinearAlgebra as LA
--import Numeric.LinearAlgebra.Data hiding(Vector)

--import qualified Data.Vector as DV
import qualified Data.Vector.Storable as V

import Prelude hiding(map,zip)

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

zeroToOne x
    | x == 0.0  = 1.0
    | otherwise = x

logE = V.map (log . zeroToOne)

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

-- | the entropy \sum p_i l\ln{p_i} of a sequence
entropy :: PDF a Double
        => a                       -- ^ the underlying distribution
        -> LA.Vector Double           -- ^ the sequence
        -> Double                  -- ^ the entropy
entropy p x = let ps = probability p x
              in negate $ (LA.dot ps (logE ps))

-- | the mutual information \sum_x \sum_y p(x,y) \ln{\frac{p(x,y)}{p(x)p(y)}}
mutual_information :: (PDF a Double, PDF b (Double,Double))
                   => b                                          -- ^ the underlying distribution
                   -> a                                          -- ^ the first dimension distribution
                   -> a                                          -- ^ the second dimension distribution
                   -> (LA.Vector Double, LA.Vector Double)             -- ^ the sequence
                   -> Double         -- ^ the mutual information
mutual_information p px py (x,y) = let ps = probability p $ V.zipWith (,) x y
                                       xs = probability px x
                                       ys = probability py y
                                   in (LA.dot ps (logE ps - logE (xs*ys)))

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


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Vector.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Vector.Statistics (
    -- * Immutable interface
    sum,
    mean,
    weightedSum,
    weightedMean,

    -- * Mutable interface
    addSumTo,
    meanTo,
    addWeightedSumTo,
    weightedMeanTo,

    ) where

import Prelude hiding ( sum )

import Control.Monad( forM_ )
import Control.Monad.ST( ST )

import Numeric.LinearAlgebra.Types

import Numeric.LinearAlgebra.Vector.Base( Vector )
import Numeric.LinearAlgebra.Vector.STBase( RVector, STVector )
import qualified Numeric.LinearAlgebra.Vector.STBase as V

-- | Returns the sum of the vectors.  The first argument gives the dimension
-- of the vectors.
sum :: (BLAS1 e) => Int -> [Vector e] -> Vector e
sum p xs = V.create $ do
    s <- V.new_ p
    V.clear s
    addSumTo s xs
    return s

-- | Returns the mean of the vectors.  The first argument gives the dimension
-- of the vectors.
mean :: (BLAS1 e) => Int -> [Vector e] -> Vector e
mean p xs = V.create $ do
      m <- V.new_ p
      meanTo m xs
      return m

-- | Returns the weighted sum of the vectors.  The first argument gives the
-- dimension of the vectors.
weightedSum :: (BLAS1 e) => Int -> [(e, Vector e)] -> Vector e
weightedSum p wxs = V.create $ do
    s <- V.new_ p
    V.clear s
    addWeightedSumTo s wxs
    return s

-- | Returns the weighted mean of the vectors.  The first argument gives the
-- dimension of the vectors.
weightedMean :: (BLAS1 e)
             => Int -> [(Double, Vector e)] -> Vector e
weightedMean p wxs = V.create $ do
    m <- V.new_ p
    weightedMeanTo m wxs
    return m

-- | Adds the sum of the vectors to the target vector.
addSumTo :: (RVector v, BLAS1 e) => STVector s e -> [v e] -> ST s ()
addSumTo dst = addWeightedSumTo dst . zip (repeat 1)

-- | Sets the target vector to the mean of the vectors.
meanTo :: (RVector v, BLAS1 e)
       => STVector s e -> [v e] -> ST s()
meanTo dst = weightedMeanTo dst . zip (repeat 1)

-- | Adds the weigthed sum of the vectors to the target vector.
addWeightedSumTo :: (RVector v, BLAS1 e)
                 => STVector s e ->  [(e, v e)] -> ST s ()
addWeightedSumTo s wxs = do
    n <- V.getDim s
    err <- V.new n 0
    old_s <- V.new_ n
    diff <- V.new_ n
    val <- V.new_ n

    forM_ wxs $ \(w,x) -> do
        V.unsafeCopyTo old_s s -- old_s := s

        V.unsafeCopyTo val x   -- val := w * x
        V.scaleM_ w val

        V.addTo err err val    -- err := err + val
        V.addTo s s err        -- s := s + err

        V.subTo diff old_s s   -- diff := old_s - s
        V.addTo err diff val   -- err := diff + val

-- | Sets the target vector to the weighted mean of the vectors.
weightedMeanTo :: (RVector v, BLAS1 e)
               =>  STVector s e -> [(Double, v e)] -> ST s ()
weightedMeanTo m wxs = let
    go _ _ [] = return ()
    go diff w_sum ((w,x):wxs') | w == 0    = go diff w_sum wxs'
                               | otherwise = let w_sum' = w_sum + w
                                             in do
                                    V.subTo diff x m
                                    V.addWithScaleM_
                                        (realToFrac $ w/w_sum') diff m
                                    go diff w_sum' wxs'
    in do
        n <- V.getDim m
        diff <- V.new_ n
        V.clear m
        go diff 0 wxs


{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeSynonymInstances #-}

module PLA
  ( Weight(..)
  , Point'(..)
  , viewX
  , trainPLA
  , trainPLAWithInitial
  , isMisclassified
  ) where

import           Control.Lens
import           Control.Lens.Combinators
import           Control.Lens.Operators
import           Control.Monad.Reader
import           Data.List
import           Data.Traversable           (for)
import           Numeric.LinearAlgebra.Data
import           System.Random

--------------------------------
--  Solutions to homework 1, problems 7-10 of
--  the 'Learning from Data' course
-- http://work.caltech.edu/homeworks.html
-------------------------------------
-------------------------------
-- Types and Typeclass Instances
--------------------------------
type Weight = (R, R, R)

type Point = (R, R)

-- | To be used at a later stage - currently just a test
data Point' =
  Point'
    { _x :: R
    , _y :: R
    }

makeLenses ''Point'

instance Num Weight where
  (+) (x0, x1, x2) (w0, w1, w2) = (w0 + x0, w1 + x1, w2 + x2)
  (-) (x0, x1, x2) (w0, w1, w2) = (x0 - w0, x1 - w1, x2 - w2)
  (*) (x0, x1, x2) (w0, w1, w2) =
    (w2 * x1 - w1 * x2, w2 * x0 - w0 * x2, w1 * x0 - w0 * x1)
  negate (w0, w1, w2) = (negate w0, negate w1, negate w2)
  abs (w0, w1, w2) = (abs w0, abs w1, abs w2)
  fromInteger x = (fromInteger x, 0.0, 0.0)
  signum (w0, w1, w2) = signum (w0, w1, w2)

--------------------------------------------
-- Creating the line for the target function
--------------------------------------------
m :: Point -> Point -> R
m (x1, x2) (z1, z2) = (x2 - z2) / (x1 - z1)

b :: Point -> Point -> R
b (x1, x2) (z1, z2) = x2 - x1 * (x2 - z2) / (x1 - z1)

-- | Takes 2 Points and returns the function of a line connecting them
line :: Point -> Point -> (R -> R)
line point1 point2 = \x -> m point1 point2 * x + b point1 point2

-----------------------------------
-- functions for the Weight type
-----------------------------------
-- | The dot product
dotProd :: Weight -> Weight -> R
dotProd (x0, x1, x2) (w0, w1, w2) = w0 * x0 + w1 * x1 + w2 * x2

-- | Scalar multiplication of the weight vecotr
prod :: R -> Weight -> Weight
prod lambda (x0, x1, x2) = (x0 * lambda, x1 * lambda, x2 * lambda)

-- | Converting a point to a weight by adding the first coordinate (1)
point2weight :: Point -> Weight
point2weight (x1, x2) = (1, x1, x2)

-- | looks up the label of the point by applying the target function
yn :: (R -> R) -> Point -> R
yn yLine (x1, x2) =
  if yLine x1 < x2
    then (-1)
    else 1

-- | Checking if a point is misclassified
isMisclassified yLine weight point = classified /= label
  where
    classified = signum $ dotProd (point2weight point) weight
    label = yn yLine point

updateWeight' yLine w xn = w + prod (yn yLine xn) (point2weight xn)

-- | Performs 1 training epoch (one run through all training points)
-- | performs the perceptron learning algorithm
pla yLine w niter gen xs
  | misclassifiedXs == [] = (w, niter)
  | otherwise = pla yLine wUpdated (niter + 1) nextRandom xs
  where
    wUpdated = updateWeight' yLine w $ misclassifiedXs !! number
    misclassifiedXs = filter (isMisclassified yLine w) xs
    randomRange = (0, length misclassifiedXs - 1)
    (number, nextRandom) = randomR randomRange gen

-- | stores the target function and performs training, returning the weight and the number of epochs performed before convergence
weightReader' :: [Point] -> StdGen -> Reader (R -> R) (Weight, Int)
weightReader' points gen = do
  targetFunc <- ask
  let weightsEpochs = pla targetFunc (0, 0, 0) 0 gen points
  return weightsEpochs

-- | stores the target function and performs training, returning the weight and the number of epochs performed before convergence
weightReader :: [Point] -> StdGen -> Weight -> Reader (R -> R) (Weight, Int)
weightReader points gen initialWeights = do
  targetFunc <- ask
  let weightsEpochs = pla targetFunc initialWeights 0 gen points
  return weightsEpochs

-- |Creates a random Point between (-1,-1) and (1,1)
createRandomPoint = do
  rand1 :: R <- randomRIO (-1, 1)
  rand2 :: R <- randomRIO (-1, 1)
  return (rand1, rand2)

-- |Creates a list of n random Points between (-1,-1) and (1,1)
createRandomPoints n = do
  randList1 :: [R] <- for [1 .. n] $ \_ -> randomRIO (-1, 1)
  randList2 :: [R] <- for [1 .. n] $ \_ -> randomRIO (-1, 1)
  return $ zip randList1 randList2-- | creates an infinite list of permutations of the list provided as an argument

pointPermutations = cycle . permutations

-- | Just a test for the point lens
viewX point = view x point

---------------------------------------------------------------
-- The final loop that creates the target function and training and testing data,
-- performs training and testing
-------------------------------------------------------------------
-- | performs the PLA
trainPLA :: Int -> IO ()
trainPLA nPoints = do
  point1 <- createRandomPoint
  point2 <- createRandomPoint
  stdgen <- getStdGen
  let target = line point1 point2
  trainpoints <- createRandomPoints nPoints
  let (weights, epochs) = runReader (weightReader' trainpoints stdgen) target
  putStr "Weights: "
  print weights
  putStr "Epochs: "
  print epochs
  testpoints <- createRandomPoints 1000
  let misclassifiedTestPoints =
        length $ filter (isMisclassified target weights) testpoints
  let prob = fromIntegral misclassifiedTestPoints / 1000
  putStr "Probability f /= g: "
  print prob

-- | performs the PLA with a set of initial weights, training points and a target function
trainPLAWithInitial :: (R, R, R) -> [Point] -> (R -> R) -> IO (Weight, Int)
trainPLAWithInitial initialWeights trainPoints target = do
  stdgen <- getStdGen
  let (weights, epochs) =
        runReader (weightReader trainPoints stdgen initialWeights) target
  return (weights, epochs)


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE ParallelListComp #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Newton
  (
  -- * Newton's Method (Forward AD)
    findZero
  , findZeroNoEq
  , inverse
  , inverseNoEq
  , fixedPoint
  , fixedPointNoEq
  , extremum
  , extremumNoEq
  -- * Gradient Ascent/Descent (Reverse AD)
  , gradientDescent, constrainedDescent, CC(..), eval
  , gradientAscent
  , conjugateGradientDescent
  , conjugateGradientAscent
  , stochasticGradientDescent
  ) where

import Data.Foldable (all, sum)
import Data.Reflection (Reifies)
import Data.Traversable
import Numeric.AD.Internal.Combinators
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Or
import Numeric.AD.Internal.Reverse (Reverse, Tape)
import Numeric.AD.Internal.Type (AD(..))
import Numeric.AD.Mode
import Numeric.AD.Mode.Reverse as Reverse (gradWith, gradWith', grad')
import Numeric.AD.Rank1.Kahn as Kahn (Kahn, grad)
import qualified Numeric.AD.Rank1.Newton as Rank1
import Prelude hiding (all, mapM, sum)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,2.5,2.05,2.000609756097561,2.0000000929222947,2.000000000000002,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
findZero f = Rank1.findZero (runAD.f.AD)
{-# INLINE findZero #-}

-- | The 'findZeroNoEq' function behaves the same as 'findZero' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
findZeroNoEq :: Fractional a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
findZeroNoEq f = Rank1.findZeroNoEq (runAD.f.AD)
{-# INLINE findZeroNoEq #-}

-- | The 'inverse' function inverts a scalar function using
-- Newton's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- Example:
--
-- >>> last $ take 10 $ inverse sqrt 1 (sqrt 10)
-- 10.0
inverse :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a -> [a]
inverse f = Rank1.inverse (runAD.f.AD)
{-# INLINE inverse  #-}

-- | The 'inverseNoEq' function behaves the same as 'inverse' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
inverseNoEq :: Fractional a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a -> [a]
inverseNoEq f = Rank1.inverseNoEq (runAD.f.AD)
{-# INLINE inverseNoEq #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Newton's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
fixedPoint f = Rank1.fixedPoint (runAD.f.AD)
{-# INLINE fixedPoint #-}

-- | The 'fixedPointNoEq' function behaves the same as 'fixedPoint' except that
-- it doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
fixedPointNoEq :: Fractional a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> [a]
fixedPointNoEq f = Rank1.fixedPointNoEq (runAD.f.AD)
{-# INLINE fixedPointNoEq #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Newton's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream
-- becomes constant ("it converges"), no further elements are returned.
--
-- >>> last $ take 10 $ extremum cos 1
-- 0.0
extremum :: (Fractional a, Eq a) => (forall s. AD s (On (Forward (Forward a))) -> AD s (On (Forward (Forward a)))) -> a -> [a]
extremum f = Rank1.extremum (runAD.f.AD)
{-# INLINE extremum #-}

-- | The 'extremumNoEq' function behaves the same as 'extremum' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
extremumNoEq :: Fractional a => (forall s. AD s (On (Forward (Forward a))) -> AD s (On (Forward (Forward a)))) -> a -> [a]
extremumNoEq f = Rank1.extremumNoEq (runAD.f.AD)
{-# INLINE extremumNoEq #-}

-- | The 'gradientDescent' function performs a multivariate
-- optimization, based on the naive-gradient-descent in the file
-- @stalingrad\/examples\/flow-tests\/pre-saddle-1a.vlad@ from the
-- VLAD compiler Stalingrad sources.  Its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- It uses reverse mode automatic differentiation to compute the gradient.
gradientDescent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a]
gradientDescent f x0 = go x0 fx0 xgx0 0.1 (0 :: Int)
  where
    (fx0, xgx0) = Reverse.gradWith' (,) f x0
    go x fx xgx !eta !i
      | eta == 0     = [] -- step size is 0
      | fx1 > fx     = go x fx xgx (eta/2) 0 -- we stepped too far
      | zeroGrad xgx = [] -- gradient is 0
      | otherwise    = x1 : if i == 10
                            then go x1 fx1 xgx1 (eta*2) 0
                            else go x1 fx1 xgx1 eta (i+1)
      where
        zeroGrad = all (\(_,g) -> g == 0)
        x1 = fmap (\(xi,gxi) -> xi - eta * gxi) xgx
        (fx1, xgx1) = Reverse.gradWith' (,) f x1
{-# INLINE gradientDescent #-}

data SEnv (f :: * -> *) a = SEnv { sValue :: a, origEnv :: f a }
  deriving (Functor, Foldable, Traversable)

-- | Convex constraint, CC, is a GADT wrapper that hides the existential
-- ('s') which is so prevalent in the rest of the API.  This is an
-- engineering convenience for managing the skolems.
data CC f a where
  CC :: forall f a. (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> CC f a

-- |@constrainedDescent obj fs env@ optimizes the convex function @obj@
-- subject to the convex constraints @f <= 0@ where @f `elem` fs@. This is
-- done using a log barrier to model constraints (i.e. Boyd, Chapter 11.3).
-- The returned optimal point for the objective function must satisfy @fs@,
-- but the initial environment, @env@, needn't be feasible.
constrainedDescent :: forall f a. (Traversable f, RealFloat a, Floating a, Ord a)
                   => (forall s. Reifies s Tape => f (Reverse s a)
                                                -> Reverse s a)
                   -> [CC f a]
                   -> f a
                   -> [(a,f a)]
constrainedDescent objF [] env =
  map (\x -> (eval objF x, x)) (gradientDescent objF env)
constrainedDescent objF cs env =
    let s0       = 1 + maximum [eval c env | CC c <- cs]
        -- ^ s0 = max ( f_i(0) )
        cs'      = [CC (\(SEnv sVal rest) -> c rest - sVal) | CC c <- cs]
        -- ^ f_i' = f_i - s0  and thus f_i' <= 0
        envS     = SEnv s0 env
        -- feasible point for f_i', use gd to find feasiblity for f_i
        cc       = constrainedConvex' (CC sValue) cs' envS ((<=0) . sValue)
    in case dropWhile ((0 <) . fst) (take (2^(20::Int)) cc) of
        []                  -> []
        (_,envFeasible) : _ ->
            constrainedConvex' (CC objF) cs (origEnv envFeasible) (const True)
{-# INLINE constrainedDescent #-}

eval :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> a
eval f e = fst (grad' f e)
{-# INLINE eval #-}

-- | Like 'constrainedDescent' except the initial point must be feasible.
constrainedConvex' :: forall f a. (Traversable f, RealFloat a, Floating a, Ord a)
                   => CC f a
                   -> [CC f a]
                   -> f a
                   -> (f a -> Bool)
                   -> [(a,f a)]
constrainedConvex' objF cs env term =
  -- 1. Transform cs using a log barrier with increasing t values.
  let os   = map (mkOpt objF cs) tValues
  -- 2. Iteratively run gradientDescent on each os.
      envs =  [(undefined,env)] :
              [gD (snd $ last e) o
                          | o  <- os
                          | e  <- limEnvs
                          ]
      -- Obtain a finite number of elements from the initial len tValues - 1 lists.
      limEnvs = zipWith id nrSteps envs
  in dropWhile (not . term . snd) (concat $ drop 1 limEnvs)
 where
  tValues = map realToFrac $ take 64 $ iterate (*2) (2 :: a)
  nrSteps = [take 20 | _ <- [1..length tValues]] ++ [id]
  -- | `gD f e` is gradient descent with the evaulated result
  gD e (CC f)  = (eval f e, e) :
                 map (\x -> (eval f x, x)) (gradientDescent f e)
{-# INLINE constrainedConvex' #-}

-- @mkOpt u fs t@ converts an inequality convex problem (@u,fs@) into an
-- unconstrained convex problem using log barrier @u + -(1/t)log(-f_i)@.
-- As @t@ increases the approximation is more accurate but the gradient
-- decreases, making the gradient descent more expensive.
mkOpt :: forall f a. (Traversable f, RealFloat a, Floating a, Ord a)
      => CC f a -> [CC f a]
      -> a -> CC f a
mkOpt (CC o) xs t = CC (\e -> o e + sum (map (\(CC c) -> iHat t c e) xs))
{-# INLINE mkOpt #-}

iHat :: forall a f. (Traversable f, RealFloat a, Floating a, Ord a)
     => a
     -> (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a)
     -> (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a)
iHat t c e =
   let r = c e
   in if r >= 0 || isNaN r
        then 1  / 0
        else (-1 / auto t) * log( - (c  e))
{-# INLINE iHat #-}

-- | The 'stochasticGradientDescent' function approximates
-- the true gradient of the constFunction by a gradient at
-- a single example. As the algorithm sweeps through the training
-- set, it performs the update for each training example.
--
-- It uses reverse mode automatic differentiation to compute the gradient
-- The learning rate is constant through out, and is set to 0.001
stochasticGradientDescent :: (Traversable f, Fractional a, Ord a)
  => (forall s. Reifies s Tape => e -> f (Reverse s a) -> Reverse s a)
  -> [e]
  -> f a
  -> [f a]
stochasticGradientDescent errorSingle d0 x0 = go xgx0 0.001 dLeft
  where
    dLeft = tail $ cycle d0
    xgx0 = Reverse.gradWith (,) (errorSingle (head d0)) x0
    go xgx !eta d
      | eta ==0       = []
      | otherwise     = x1 : go xgx1 eta (tail d)
      where
        x1 = fmap (\(xi, gxi) -> xi - eta * gxi) xgx
        (_, xgx1) = Reverse.gradWith' (,) (errorSingle (head d)) x1
{-# INLINE stochasticGradientDescent #-}

-- | Perform a gradient descent using reverse mode automatic differentiation to compute the gradient.
gradientAscent :: (Traversable f, Fractional a, Ord a) => (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> f a -> [f a]
gradientAscent f = gradientDescent (negate . f)
{-# INLINE gradientAscent #-}

-- | Perform a conjugate gradient descent using reverse mode automatic differentiation to compute the gradient, and using forward-on-forward mode for computing extrema.
--
-- >>> let sq x = x * x
-- >>> let rosenbrock [x,y] = sq (1 - x) + 100 * sq (y - sq x)
-- >>> rosenbrock [0,0]
-- 1
-- >>> rosenbrock (conjugateGradientDescent rosenbrock [0, 0] !! 5) < 0.1
-- True
conjugateGradientDescent
  :: (Traversable f, Ord a, Fractional a)
  => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a))
  -> f a -> [f a]
conjugateGradientDescent f = conjugateGradientAscent (negate . f)
{-# INLINE conjugateGradientDescent #-}

lfu :: Functor f => (f (Or F a b) -> Or F a b) -> f a -> a
lfu f = runL . f . fmap L

rfu :: Functor f => (f (Or T a b) -> Or T a b) -> f b -> b
rfu f = runR . f . fmap R

-- | Perform a conjugate gradient ascent using reverse mode automatic differentiation to compute the gradient.
conjugateGradientAscent
  :: (Traversable f, Ord a, Fractional a)
  => (forall s. Chosen s => f (Or s (On (Forward (Forward a))) (Kahn a)) -> Or s (On (Forward (Forward a))) (Kahn a))
  -> f a -> [f a]
conjugateGradientAscent f x0 = takeWhile (all (\a -> a == a)) (go x0 d0 d0 delta0)
  where
    dot x y = sum $ zipWithT (*) x y
    d0 = Kahn.grad (rfu f) x0
    delta0 = dot d0 d0
    go xi _ri di deltai = xi : go xi1 ri1 di1 deltai1
      where
        ai = last $ take 20 $ Rank1.extremum (\a -> lfu f $ zipWithT (\x d -> auto x + a * auto d) xi di) 0
        xi1 = zipWithT (\x d -> x + ai*d) xi di
        ri1 = Kahn.grad (rfu f) xi1
        deltai1 = dot ri1 ri1
        bi1 = deltai1 / deltai
        di1 = zipWithT (\r d -> r + bi1 * d) ri1 di
{-# INLINE conjugateGradientAscent #-}


module Numeric.Layer.Sigmoid where
import Numeric.Layer
import Dependent.Size
import Numeric.Vector.Sized
import qualified Numeric.Vector.Sized as Sized
import GHC.TypeLits(KnownNat,Nat)

newtype Sigmoid (n :: Nat) = Sigmoid ()

instance Show (Sigmoid n) where
  showsPrec d (Sigmoid _) = showString "Sigmoid"

sigmoid :: Floating x => x -> x
sigmoid x = 1 / (1 + exp (-x))

sigmoid' :: Num a => a -> a
sigmoid' y = (1 - y) * y

instance KnownNat n => Layer (Sigmoid n) where
  type Inputs (Sigmoid n) = ZZ ::. n
  type Outputs (Sigmoid n) = ZZ ::. n
  type Tape (Sigmoid n) = SizedArray (ZZ ::. n)
  type Gradient (Sigmoid n) = ()
  forward _ x = (y, tape)
    where
      y = Sized.map sigmoid x
      tape = y
  backward _ tape dy = (gradient, dx)
    where
      gradient = ()
      dx = Sized.zipWith (*) dy (Sized.map sigmoid' tape)
  applyGradient layer _ = layer


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
module Numeric.Layer.Affine where

import Dependent.Size
import Numeric.Layer
import Numeric.Vector.Sized
import qualified Numeric.Vector.Sized as Sized
import GHC.TypeLits
import Data.Array.Accelerate.LLVM.Native

data Affine (n :: Nat) (m :: Nat) = Affine (SizedMatrix' n m) (SizedVector' m)
  deriving (Show)

instance (KnownNat n, KnownNat m) => Layer (Affine n m) where
  type Inputs (Affine n m) = ZZ '::. n
  type Outputs (Affine n m) = ZZ '::. m
  type Tape (Affine n m) = SizedArray (ZZ ::. n)
  type Gradient (Affine n m) = Affine n m
  forward (Affine ws b) x = (affine (use ws) (use b) x, x)
  backward (Affine ws b) x dy = (Affine dw db, dx)
    where
      dx = use ws #> dy
      dw = runSized run (dy `outer` x)
      db = runSized run (dy)
  applyGradient (Affine ws b) (Affine dw db)
    = Affine
      (runSized run $ Sized.zipWith (-) (use ws) (use dw))
      (runSized run $ Sized.zipWith (-) (use b) (use db))

affine
  :: (KnownNat n, KnownNat m)
  => SizedMatrix n m
  -> SizedVector m
  -> SizedVector n
  -> SizedVector m
affine ws b x = (x <# ws) .+. b

-- affine_dw dedy = dedw
-- affine_db dedy = dedb
-- affine_dx dedy = dedx


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Random.Distribution.MultivariateNormal
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- Sample from the multivariate normal distribution with a given
-- vector-valued \(\mu\) and covariance matrix \(\Sigma\). For example,
-- the chart below shows samples from the bivariate normal
-- distribution.
--
-- <<diagrams/src_Data_Random_Distribution_MultivariateNormal_diagM.svg#diagram=diagM&height=600&width=500>>
--
-- Example code to generate the chart:
--
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- >
-- > import Data.Random.Distribution.MultivariateNormal
-- >
-- > import qualified Data.Random as R
-- > import Data.Random.Source.PureMT
-- > import Control.Monad.State
-- > import qualified Numeric.LinearAlgebra.HMatrix as LA
-- >
-- > nSamples :: Int
-- > nSamples = 10000
-- >
-- > sigma1, sigma2, rho :: Double
-- > sigma1 = 3.0
-- > sigma2 = 1.0
-- > rho = 0.5
-- >
-- > singleSample :: R.RVarT (State PureMT) (LA.Vector Double)
-- > singleSample = R.sample $ Normal (LA.fromList [0.0, 0.0])
-- >                (LA.sym $ (2 LA.>< 2) [ sigma1, rho * sigma1 * sigma2
-- >                                      , rho * sigma1 * sigma2, sigma2])
-- >
-- > multiSamples :: [LA.Vector Double]
-- > multiSamples = evalState (replicateM nSamples $ R.sample singleSample) (pureMT 3)
-- > pts = map (f . LA.toList) multiSamples
-- >   where
-- >     f [x, y] = (x, y)
-- >     f _      = error "Only pairs for this chart"
-- >
-- >
-- > chartPoint pointVals n = C.toRenderable layout
-- >   where
-- >
-- >     fitted = C.plot_points_values .~ pointVals
-- >               $ C.plot_points_style  . C.point_color .~ opaque red
-- >               $ C.plot_points_title .~ "Sample"
-- >               $ def
-- >
-- >     layout = C.layout_title .~ "Sampling Bivariate Normal (" ++ (show n) ++ " samples)"
-- >            $ C.layout_y_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >            $ C.layout_x_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >
-- >            $ C.layout_plots .~ [C.toPlot fitted]
-- >            $ def
-- >
-- > diagM = do
-- >   denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >   return $ fst $ runBackend denv (C.render (chartPoint pts nSamples) (500, 500))
--
-----------------------------------------------------------------------------

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

module Data.Random.Distribution.MultivariateNormal
    ( Normal(..)
    ) where

import           Data.Random.Distribution
import qualified Numeric.LinearAlgebra.HMatrix as H
import           Control.Monad
import qualified Data.Random as R
import           Foreign.Storable ( Storable )
import           Data.Maybe ( fromJust )

normalMultivariate :: H.Vector Double -> H.Herm Double -> R.RVarT m (H.Vector Double)
normalMultivariate mu bigSigma = do
  z <- replicateM (H.size mu) (rvarT R.StdNormal)
  return $ mu + bigA H.#> (H.fromList z)
  where
    (vals, bigU) = H.eigSH bigSigma
    lSqrt = H.diag $ H.cmap sqrt vals
    bigA = bigU H.<> lSqrt

data family Normal k :: *

data instance Normal (H.Vector Double) = Normal (H.Vector Double) (H.Herm Double)

instance Distribution Normal (H.Vector Double) where
  rvar (Normal m s) = normalMultivariate m s

normalPdf :: (H.Numeric a, H.Field a, H.Indexable (H.Vector a) a, Num (H.Vector a)) =>
             H.Vector a -> H.Herm a -> H.Vector a -> a
normalPdf mu sigma x = exp $ normalLogPdf mu sigma x

normalLogPdf :: (H.Numeric a, H.Field a, H.Indexable (H.Vector a) a, Num (H.Vector a)) =>
                 H.Vector a -> H.Herm a -> H.Vector a -> a
normalLogPdf mu bigSigma x = - H.sumElements (H.cmap log (diagonals dec))
                              - 0.5 * (fromIntegral (H.size mu)) * log (2 * pi)
                              - 0.5 * s
  where
    dec = fromJust $ H.mbChol bigSigma
    t = fromJust $ H.linearSolve (H.tr dec) (H.asColumn $ x - mu)
    u = H.cmap (\v -> v * v) t
    s = H.sumElements u

diagonals :: (Storable a, H.Element t, H.Indexable (H.Vector t) a) =>
             H.Matrix t -> H.Vector a
diagonals m = H.fromList (map (\i -> m H.! i H.! i) [0..n-1])
  where
    n = max (H.rows m) (H.cols m)

instance PDF Normal (H.Vector Double) where
  pdf (Normal m s) = normalPdf m s
  logPdf (Normal m s) = normalLogPdf m s


-- ch03 example 3.7

import Statistics.Distribution
import Statistics.Distribution.Binomial

b :: BinomialDistribution
b = binomial 5 0.05

main = do
  putStrLn $ "If 5 samples is tested, the probabiliy of observing at least one defective: "
          ++ show (1 - probability b 0)


module Fractals
( order
, floatOrder
, mandelbrot
, bmp
, tangentPoint
, mandelbrotCardioid
, mag2
, toBitmap
) where

import Data.Complex (Complex(..), mkPolar)
import Graphics.Gloss
import qualified Data.ByteString as B
import Data.Word
import Control.Parallel.Strategies

import Palette

mag2:: Complex Double -> Double
mag2 (a:+b) = a^(2::Int) + b^(2::Int)

-- | Compute the order of a given complex function.
order :: (Complex Double -> Complex Double) -- ^ a function on the complex plane
      -> Double                             -- ^ the escape radius squared
      -> Int                                -- ^ the maximum order
      -> Complex Double                     -- ^ a starting point
      -> (Int, Complex Double)              -- ^ the order
order f radius2 nMax z0 = let f' = \ (i, z) -> (i+1, f z)
                              fs = dropWhile (\ (i, z) -> i<nMax && mag2 z < radius2) $ iterate f' (0, z0)
                           in head fs

-- | Compute the floating order of a given complex function.
floatOrder :: (Complex Double -> Complex Double) -- ^ a function on the complex plane
           -> Double                             -- ^ the escape radius squared
           -> Int                                -- ^ the maximum order
           -> Complex Double                     -- ^ a starting point
           -> Double                             -- ^ the order
floatOrder f radius2 nMax z0 = let (nu, z) = order f radius2 nMax z0
                                in fromIntegral nu - log (0.5 * (log $ mag2 z)) / log 2.0

-- | the Mandelbrot function: z^2 + c.
mandelbrot :: Complex Double -- ^ the c parameter
           -> Complex Double -- ^ the z variable
           -> Complex Double -- ^ the resulting value
mandelbrot (cr:+ci) (zr:+zi) = (cr + zr^(2::Int) - zi^(2::Int)) :+ (ci + 2*zr*zi)
--mandelbrot c z = z**2 + c

-- | Yield the point tangent to the main cardioid and to the bulb with period q
--   and combinatorial rotation number r=p/q.
tangentPoint :: Rational -> Complex Double
tangentPoint r = mandelbrotCardioid $ mkPolar 1.0 (2.0 * pi * fromRational r)

-- | The equation of the main cardioid.
mandelbrotCardioid :: Complex Double -> Complex Double
mandelbrotCardioid mu = 0.5 * mu * (1.0 - 0.5 * mu)

bitmapFormat :: BitmapFormat
bitmapFormat = BitmapFormat BottomToTop PxRGBA

escapeRadius2 :: Double
escapeRadius2 = 100000.0

maxIter :: Int
maxIter = 255

bmpValues :: Int -> Int -> Double -> Double -> Double -> [Double]
bmpValues w h xc yc picScale =
    let wf = fromIntegral w * picScale
        hf = fromIntegral h * picScale
        floatOrderFromIJ :: (Int, Int) -> Double
        floatOrderFromIJ (i,j) = let x0 = xc - 0.5*wf + (fromIntegral j)*wf/(fromIntegral (w-1))
                                     y0 = yc - 0.5*hf + (fromIntegral i)*hf/(fromIntegral (h-1))
                                     z0 = x0 :+ y0
                                  in floatOrder (mandelbrot z0) escapeRadius2 maxIter z0
        indices = (,) <$> [0..h-1] <*> [0..w-1]
        vals = map floatOrderFromIJ indices
     in vals `using` parListChunk h rseq

toBitmap :: Palette -> [Double] -> [Word8]
toBitmap _ [] = []
toBitmap palette l = let maxL = maximum l
                         toBitmap' _ [] = []
                         toBitmap' maxX (x:xs) = let (r,g,b) = palette (x/maxX)
                                                     in r:g:b:255:(toBitmap' maxX xs)
                      in toBitmap' maxL l

bmpByteString :: Int -> Int -> Double -> Double -> Double -> B.ByteString
bmpByteString w h xc yc picScale = B.pack vals
    where vals = toBitmap hsv $ bmpValues w h xc yc picScale

bmp :: Int -> Int -> Complex Double -> Double -> Picture
bmp w h (xc :+ yc) picScale = bitmapOfByteString w h bitmapFormat (bmpByteString w h xc yc picScale) False


{-# LANGUAGE NoMonomorphismRestriction, ScopedTypeVariables, FlexibleContexts, TypeFamilies #-}

module Math.Probably.PCA where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix ((#>))
import Math.Probably.FoldingStats
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as VS
import Foreign.Storable.Tuple ()
import Control.Spoon
import qualified Data.Map.Strict as M
import           Data.Maybe (fromJust)
import           Data.List
import           Data.Ord
import           Data.Number.Erf
import Debug.Trace

empiricalCovariance :: [Vector Double] -> Matrix Double
empiricalCovariance xs
   = let k = length xs
         --barxk = L.scale (recip $ realToFrac $ k+1) $ sum xs
         --m1 = sum $ map (\xv-> L.outer xv xv) xs
         --m2 = L.scale (realToFrac $ k+1) $ L.outer barxk barxk
         xmn = empiricalMean xs
         f xi = outerSame $ xi - xmn
--     in L.scale (recip $ realToFrac k) $ m1 - m2
     in scale (recip $ realToFrac k - 1 ) $ sum $ map f xs

empiricalCovariance' :: Vector Double -> [Vector Double] -> Matrix Double
empiricalCovariance' xmn xs
   = let k = length xs
         --barxk = L.scale (recip $ realToFrac $ k+1) $ sum xs
         --m1 = sum $ map (\xv-> L.outer xv xv) xs
         --m2 = L.scale (realToFrac $ k+1) $ L.outer barxk barxk
         f xi = outerSame $ xi - xmn
--     in L.scale (recip $ realToFrac k) $ m1 - m2
     in scale (recip $ realToFrac k - 1 ) $ sumWith f xs

outerSame v = outer v v

--sumWithVs :: (Vector Double -> Vector Double) -> [Vector Double] -> Vector Double
--sumWithVs f = foldl1' (\acc v -> acc + f v)

sumWith :: Num a => (Vector Double -> a) -> [Vector Double] -> a
sumWith f (v:vs) = foldl' (\acc v -> acc + f v) (f v) vs

empiricalMean :: [Vector Double] -> Vector Double
empiricalMean vecs = scale (recip $ realToFrac $ length vecs) $ sumWith id vecs

-- copied from https://github.com/albertoruiz/hmatrix/blob/master/examples/pca2.hs

type Vec = Vector Double
type Mat = Matrix Double

-- Vector with the mean value of the columns of a matrix
mean :: Mat -> Vec
mean a = constant (recip . fromIntegral . rows $ a) (rows a) <> a

-- covariance matrix of a list of observations stored as rows
cov :: Mat -> Mat
cov x = (trans xc <> xc) / fromIntegral (rows x - 1)
    where xc = x - asRow (mean x)

covN :: Mat -> Mat
covN x = (trans x <> x) / fromIntegral (rows x)

cov0 :: Mat -> Mat
cov0 x = (trans x <> x) / fromIntegral (rows x - 1)
    --where xc = x - asRow (mean x)

type Stat = (Vec, [Double], Mat)
-- 1st and 2nd order statistics of a dataset (mean, eigenvalues and eigenvectors of cov)
stat :: Mat -> Stat
stat x = (m, toList s, trans v) where
    m = mean x
    (s,v) = eigSH' (cov x)

stat0 :: Mat -> Stat
stat0 x = (m, toList s, trans v) where
    m = VS.map (const 0) s
    c = cov0 x
    (s,v) = eigSH' $ trace ("cov00="++(show (c!0!0))) $ c

statSVD :: Mat -> Stat
statSVD x = (m, [], evecs) where
    m = mean x
    c = cov x
    (_, evals, evecCols) = svd c
    evecs = fromRows $ map evecSigns $ toColumns evecCols
    evecSigns ev = let maxelidx = maxIndex $ cmap abs ev
                       sign = signum (ev ! maxelidx)
                   in cmap (sign *) ev

statThinSVD :: Mat -> Stat
statThinSVD x = (m, [], evecs) where
    m = mean x
    c = cov x
    (_, evals, evecCols) = thinSVD c
    evecs = fromRows $ map evecSigns $ toColumns evecCols
    evecSigns ev = let maxelidx = maxIndex $ cmap abs ev
                       sign = signum (ev ! maxelidx)
                   in cmap (sign *) ev

pcaNSVD :: Int -> Stat -> (Vec -> Vec , Vec -> Vec)
pcaNSVD  n (m,_,evecs) = (encode,decode)
  where
    encode x = VS.take n $ evecs #> (x - m)
    decode x = error "not sure right now"

statVs :: [Vector Double] -> Stat
statVs vs = (m, toList s, trans v) where
    m = empiricalMean vs
    (s,v) = eigSH' (empiricalCovariance' m vs)

mbStat :: Mat -> Maybe Stat
mbStat = spoon . stat

pca :: Double -> Stat -> (Vec -> Vec , Vec -> Vec)
pca prec (m,s,v) = (encode,decode)
  where
    encode x = vp <> (x - m)
    decode x = x <> vp + m
    vp = takeRows n v
    n = 1 + (length $ fst $ span (< (prec'*sum s)) $ cumSum s)
    cumSum = tail . scanl (+) 0.0
    prec' = if prec <=0.0 || prec >= 1.0
                then error "the precision in pca must be 0<prec<1"
                else prec

pcaN :: Int -> Stat -> (Vec -> Vec , Vec -> Vec)
pcaN n (m,s,v) = (encode,decode)
  where
    encode x = vp <> (x - m)
    decode x = x <> vp + m
    vp = takeRows n v

type Lens a b = ( (a->b), (a-> b -> a) )

get :: Lens a b -> a -> b
get (f,_) x = f x

set :: Lens a b -> a -> b -> a
set (_,f) x y = f x y

idlens :: Lens a a
idlens = (id,  (\_ x -> x))

--, pcaFeatures
normaliseBy :: (Eq b, Ord b) => Lens a Vec -> Lens a b -> Int -> [a] -> [a]
normaliseBy veclens idlens ncomponents = equaliseVars veclens
                                         . rankTransform veclens idlens
                                         . pcaFeatures veclens ncomponents
                                         . equaliseVars veclens

normaliseNoRank :: Lens a Vec -> Int -> [a] -> [a]
normaliseNoRank veclens  ncomponents = equaliseVars veclens
                                         . pcaFeatures veclens ncomponents
                                         . equaliseVars veclens

equaliseVars ::Lens a Vec -> [a] -> [a]
equaliseVars veclens xs = equal_vars where
  decorr_vars :: VS.Vector (Double,Double)
  decorr_vars = uncurry (VS.zipWith (,)) $ runStat meanSDF $ map (get veclens) xs

  equal_vars =  map (\x -> set veclens x (g $ get veclens x)) xs

  g :: Vec -> Vec
  g fl = VS.zipWith (\x (mn,sd) -> (x-mn)/sd) fl decorr_vars

vecMeanSds :: [Vector Double] -> Vector (Double, Double)
vecMeanSds xs = uncurry (VS.zipWith (,)) $ runStat meanSDF  xs

vecMeans :: [Vector Double] -> Vector Double
vecMeans xs = scale (recip $ realToFrac $ length xs) $ sumWith id xs

rankTransform :: (Eq b, Ord b) => Lens a Vec -> Lens a b -> [a] -> [a]
rankTransform veclens idlens nmfeatures = rankTransformed where
  nfeatures = VS.length $ get veclens $ head $ nmfeatures
  nposters = realToFrac $ length nmfeatures
  sortations = flip map [0..(nfeatures-1)] $ \ix -> toRankMap $ map (get idlens) $ sortBy (comparing ((VS.! ix) . get veclens)) nmfeatures -- decorrelated
  rankTransformed = flip map nmfeatures $ \x ->
                      let nm = get idlens x
                          f sortation = invnormcdf $ realToFrac ((getRank sortation nm)+1) / (nposters+1)
                          newv =  VS.fromList $ map f sortations
                      in set veclens x newv

pcaFeatures :: Lens a Vec -> Int -> [a] -> [a]
pcaFeatures veclens ncomp nmfeatures = pca'd where
  input_data = map (get veclens ) nmfeatures
  pcaStats@(_,eigVals,_) = stat $ fromRows input_data
  (enc,_) = trace ("eigenVals ="++show eigVals ) $ pcaN ncomp pcaStats
--  pca_red = PCA.pcaReduceN input_data 5
  setFeatures x =
      let oldv = get veclens x
          newNormV = enc oldv
      in set veclens x newNormV
  pca'd = map setFeatures nmfeatures

type RankMap a = M.Map a Int

toRankMap :: (Ord a, Eq a) => [a] -> RankMap a
toRankMap = M.fromList . (`zip` [0..])

getRank :: (Ord a, Eq a) => RankMap a -> a -> Int
getRank m x = fromJust $ M.lookup x m

newtype NumUV a = NumUV { unNumUV :: U.Vector a } deriving (Eq, Ord)

numUmap f = NumUV . U.map f . unNumUV
numUlength = U.length . unNumUV

instance (U.Unbox a, Num a) => Num (NumUV a) where
    (+) = binvop (+)
    (-) = binvop (-)
    (*) = binvop (*)
    abs = numUmap abs
    signum = numUmap signum
    fromInteger x = NumUV $ U.fromList [fromInteger x]

instance (U.Unbox a, Fractional a) => Fractional (NumUV a) where
    (/) = binvop (/)
    fromRational x = NumUV $ U.fromList [fromRational x]
instance (U.Unbox a, Floating a) => Floating (NumUV a) where
    pi = NumUV $ U.fromList [pi]
    sqrt = numUmap sqrt
    log = numUmap log
    exp = numUmap exp
    sin = numUmap sin
    cos = numUmap cos
    tan = numUmap tan
    asin = numUmap asin
    acos = numUmap acos
    atan = numUmap atan
    sinh = numUmap sinh
    cosh = numUmap cosh
    tanh = numUmap tanh
    asinh = numUmap asinh
    acosh = numUmap acosh
    atanh = numUmap atanh

binvop f xs ys
  | numUlength ys == 1 = let y = (unNumUV ys) U.! 0 in numUmap (\x-> f x y) xs
  | numUlength xs == 1 = let x = (unNumUV xs) U.! 0 in numUmap (\y-> f x y) ys
  | otherwise = NumUV $ U.zipWith f (unNumUV xs) (unNumUV ys)


module Rotation where

import Control.Monad
import Numeric.LinearAlgebra ((<>), (#>), Z, Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import qualified Data.List as List
import qualified Data.Maybe as Maybe

import Cube (Cube, cubeToVector, vectorToCube)
import qualified Cube as Cube

type CubeRotation = Cube -> Cube

toCubeRotation :: VectorRotation -> CubeRotation
toCubeRotation rotation = vectorToCube . rotation . cubeToVector

type VectorRotation = Vector Z -> Vector Z

toVectorRotation :: MatrixRotation -> VectorRotation
toVectorRotation = (#>)

rotate :: MatrixRotation -> Cube -> Cube
rotate rotation = toCubeRotation $ toVectorRotation rotation

type MatrixRotation = Matrix Z

showMatrixRotation :: MatrixRotation -> String
showMatrixRotation = show . fmap (go 0 . LA.toList) . LA.toColumns
  where
    go i [] = i
    go i (v:vs)
      | v == 1 = i
      | otherwise = go (i + 1) vs

identity :: MatrixRotation
identity = LA.ident $ 9 * 6

topToFront :: MatrixRotation
topToFront =
  let vs = LA.toColumns identity
      mapping = [  9, 10, 11, 12, 13, 14, 15, 16, 17
                , 45, 46, 47, 48, 49, 50, 51, 52, 53
                , 24, 21, 18, 25, 22, 19, 26, 23, 20
                ,  8,  7,  6,  5,  4,  3,  2,  1,  0
                , 38, 41, 44, 37, 40, 43, 36, 39, 42
                , 35, 34, 33, 32, 31, 30, 29, 28, 27 ]
  in LA.fromColumns $ fmap (vs !!) mapping

topToBack :: MatrixRotation
topToBack = topToFront <> topToFront <> topToFront

topToRight :: MatrixRotation
topToRight =
  let vs = LA.toColumns identity
      mapping = [ 20, 23, 26, 19, 22, 25, 18, 21, 24
                , 11, 14, 17, 10, 13, 16,  9, 12, 15
                , 47, 50, 53, 46, 49, 52, 45, 48, 51
                , 33, 30, 27, 34, 31, 28, 35, 32, 29
                ,  2,  5,  8,  1,  4,  7,  0,  3,  6
                , 38, 41, 44, 37, 40, 43, 36, 39, 42 ]
  in LA.fromColumns $ fmap (vs !!) mapping

topToLeft :: MatrixRotation
topToLeft = topToRight <> topToRight <> topToRight

topTwistRight :: MatrixRotation
topTwistRight = topToBack <> topToRight <> topToFront

topTwistLeft :: MatrixRotation
topTwistLeft = topTwistRight <> topTwistRight <> topTwistRight

rightL :: MatrixRotation
rightL =
  let vs = LA.toColumns identity
      mapping = [  0,  1, 11,  3,  4, 14,  6,  7, 17
                ,  9, 10, 47, 12, 13, 50, 15, 16, 53
                , 24, 21, 18, 25, 22, 19, 26, 23, 20
                ,  8, 28, 29,  5, 31, 32,  2, 34, 35
                , 36, 37, 38, 39, 40, 41, 42, 43, 44
                , 45, 46, 33, 48, 49, 30, 51, 52, 27 ]
  in LA.fromColumns $ fmap (vs !!) mapping

rightR :: MatrixRotation
rightR = rightL <> rightL <> rightL

leftL :: MatrixRotation
leftL =
  let twist = topTwistRight <> topTwistRight
  in twist <> rightL <> twist

leftR :: MatrixRotation
leftR = leftL <> leftL <> leftL

topL :: MatrixRotation
topL = topToLeft <> rightL <> topToRight

topR :: MatrixRotation
topR = topL <> topL <> topL

bottomL :: MatrixRotation
bottomL =
  let twist = topToRight <> topToRight
  in twist <> topL <> twist

bottomR :: MatrixRotation
bottomR = bottomL <> bottomL <> bottomL

frontL :: MatrixRotation
frontL = topTwistRight <> rightL <> topTwistLeft

frontR :: MatrixRotation
frontR = frontL <> frontL <> frontL

backL :: MatrixRotation
backL =
  let twist = topTwistLeft <> topTwistLeft
  in twist <> frontL <> twist

backR :: MatrixRotation
backR = backL <> backL <> backL

rotations :: [(MatrixRotation, String)]
rotations = [
    (identity, "identity")
  , (topToFront, "topToFront")
  , (topToBack, "topToBack")
  , (topToRight, "topToRight")
  , (topToLeft, "topToLeft")
  , (topTwistRight, "topTwistRight")
  , (topTwistLeft, "topTwistLeft")
  , (rightL, "rightL")
  , (rightR, "rightR")
  , (leftL, "leftL")
  , (leftR, "leftR")
  , (topL, "topL")
  , (topR, "topR")
  , (bottomL, "bottomL")
  , (bottomR, "bottomR")
  , (frontL, "frontL")
  , (frontR, "frontR")
  , (backL, "backL")
  , (backR, "backR")
  ]

nameMatrixRotation :: MatrixRotation -> String
nameMatrixRotation = Maybe.fromMaybe "unknown" . flip List.lookup rotations

inverse :: MatrixRotation -> Maybe MatrixRotation
inverse r = Maybe.listToMaybe $ do
  r' <- map fst rotations
  guard $ r <> r' == identity
  return r'


module Main where

import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning.Types as T
import qualified MachineLearning as ML
import qualified MachineLearning.Classification.Binary as CB
import qualified MachineLearning.Classification.OneVsAll as OVA
import qualified MachineLearning.TerminalProgress as TP

processFeatures :: T.Matrix -> T.Matrix
processFeatures = ML.addBiasDimension . (ML.mapFeatures 2)

calcAccuracy :: T.Matrix -> T.Vector -> [T.Vector] -> Double
calcAccuracy x y thetas = OVA.calcAccuracy y yPredicted
  where yPredicted = OVA.predict x thetas

main = do
  putStrLn "\n=== Optical Recognition of Handwritten Digits Data Set ===\n"
  -- Step 1. Data loading.
  -- Step 1.1 Training Data loading.
  (x, y) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tra"
  -- Step 1.1 Testing Data loading.
  (xTest, yTest) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tes"

  -- Step 2. Outputs and features preprocessing.
  let numLabels = 10
      x1 = processFeatures x
      initialTheta = LA.konst 0 (LA.cols x1)
      initialThetas = replicate numLabels initialTheta
  -- Step 3. Learning.
  (thetas, optPath) <- TP.learnOneVsAllWithProgressBar (CB.learn (OVA.BFGS2 0.01 0.1) 0.001 30 (CB.L2 30) x1) y initialThetas 1

  -- Step 4. Prediction and checking accuracy
  let accuracyTrain = calcAccuracy x1 y thetas
      accuracyTest = calcAccuracy (processFeatures xTest) yTest thetas

  -- Step 5. Printing results.
  putStrLn $ "\nNumber of iterations to learn: " ++ show (map LA.rows optPath)

  putStrLn $ "Accuracy on train set (%): " ++ show (accuracyTrain*100)
  putStrLn $ "Accuracy on test set (%): " ++ show (accuracyTest*100)


module Trade.Statistics.Algorithm where

import Data.Time.Clock (NominalDiffTime)

import qualified Data.Vector as Vec
import Data.Vector (Vector)

import qualified Statistics.Sample as Stat

import Trade.Type.Equity (Equity(..))
import Trade.Type.Percent (Percent(..))
import Trade.Type.Price (Price(..))

class Statistics a where
  mean :: Vector a -> a
  stdDev :: Vector a -> a
  skewness :: Vector a -> Double
  kurtosis :: Vector a -> Double
  volatility :: Vector a -> Percent

instance Statistics Double where
  mean = Stat.mean
  stdDev = Stat.stdDev
  skewness = Stat.skewness
  kurtosis = Stat.kurtosis

  volatility xs =
    let as = Vec.zipWith (\a b -> log (a/b)) (Vec.tail xs) xs
        s = sqrt (fromIntegral (Vec.length as))
    in Percent (Stat.stdDev as * s)

instance Statistics Price where
  mean = Price . Stat.mean . Vec.map unPrice
  stdDev = Price . Stat.stdDev . Vec.map unPrice
  skewness = Stat.skewness . Vec.map unPrice
  kurtosis = Stat.kurtosis . Vec.map unPrice
  volatility = volatility . Vec.map unPrice

instance Statistics Equity where
  mean = Equity . Stat.mean . Vec.map unEquity
  stdDev = Equity . Stat.stdDev . Vec.map unEquity
  skewness = Stat.skewness . Vec.map unEquity
  kurtosis = Stat.kurtosis . Vec.map unEquity
  volatility = volatility . Vec.map unEquity

instance Statistics NominalDiffTime where
  mean = realToFrac . Stat.mean . Vec.map realToFrac
  stdDev = realToFrac . Stat.stdDev . Vec.map realToFrac
  skewness = Stat.skewness . Vec.map realToFrac
  kurtosis = Stat.kurtosis . Vec.map realToFrac
  volatility = volatility . Vec.map (realToFrac :: NominalDiffTime -> Double)


import Data.Complex
import Data.Maybe
import Data.Word
import qualified Data.Array.Repa as R
import Data.Array.Repa hiding (map,zip,foldr)

f :: Complex Double -> Complex Double -> Complex Double
f c z = z*z + c

isMandelbrot :: Int -> Complex Double -> Int
isMandelbrot n c = case (n' < n) of
                      True ->   n'
                      False -> (-1)
  where
    m = iterate (f c) 0
    n' = length. take n . takeWhile (\z -> magnitude z <2) $ m

array :: Double -> Complex Double -> Int -> Array D DIM2 (Complex Double)
array size centre n = R.fromFunction (Z:.n:.n) posToVal
  where
    dx = size / (fromIntegral n)
    posToVal (Z:.i:.j) = (((fromIntegral i)*dx - size/2) :+ ((fromIntegral j)*dx - size/2)) + centre

genMandelbrot :: Int -> Int -> Double -> Complex Double -> Array D DIM2 (Word8,Word8,Word8)
genMandelbrot res n size centre = cs
  where
    zs = array size centre res
    cs = R.map (toColor . isMandelbrot n) zs

toColor :: Int -> (Word8,Word8,Word8)
toColor n = case (n<0) of
               True = (0,0,0)
               False = (255 - (fromIntegral n),fromIntegral n,50)

main :: IO ()
main = do
  [i',j',size',res',n',filename] <- getArgs
  let res    = read res'     :: Int
      n      = read n'       :: Int
      size   = read size'    :: Double
      i      = read i'       :: Double
      j      = read j'       :: Double
      centre = i :+ j
      cs' = genMandelbrot res n size centre
  cs <- computeP cs' :: IO (Array U DIM2 (Word8,Word8,Word8))
  writeImageToBMP filename cs


-- | Planning by Dynamic Programming (Matrix Version)

module PlanningM where

import           Data.Map                      (Map)
import qualified Data.Map                      as M
import qualified Data.Set                      as S
import           Numeric.LinearAlgebra.Data    (Matrix, Vector, (!))
import           Numeric.LinearAlgebra.HMatrix as H

import           MDP
import           Planning

data MDPM s a = MDPM
    { statesm     :: Map s Int
    , actionsm    :: Map a Int
    , transitionm :: Matrix Double -- matrix S x S
    , rewardm     :: Vector Double     -- vector 1 x S
    , gammam      :: Double             -- discount
    , sizeStates  :: Int           -- S
    , sizeActions :: Int          -- A
    }

type PolicyM = Matrix Double

policyEvaluationM :: (Ord s, Ord a) => MDP s a -> ValueFunction s a
policyEvaluationM mdp policy k s =
    matrixPolicyEvaluation mdpm k ! (M.!) (statesm mdpm) s
  where
    mdpm = prepareMDP mdp (preparePolicy mdp policy)

preparePolicy :: MDP s a -> Policy s a -> PolicyM
preparePolicy mdp policy = H.fromLists (
  map (
    \s -> map (policy s) (S.toList (actions mdp)))
  (S.toList (states mdp)))

prepareMDP :: (Ord s, Ord a) => MDP s a -> PolicyM -> MDPM s a
prepareMDP mdp policy =
    MDPM
    { statesm = statesM
    , actionsm = actionsM
    , transitionm = transitionM
    , rewardm = rewardM
    , gammam = gamma mdp
    , sizeStates = S.size (states mdp)
    , sizeActions = S.size (actions mdp)
    }
  where
    statesM = M.fromList (zip (S.toList (states mdp)) [0 ..])
    actionsM = M.fromList (zip (S.toList (actions mdp)) [0 ..])
    transitionM =
        H.fromRows
            (map
                 (\s ->
                       takeDiag (policy H.<> (transitionMap M.! s)))
                 (S.toList (states mdp)))
    transitionMap =
        M.fromList
            (map
                 (\s ->
                       (s, createMatrix s))
                 (S.toList (states mdp)))
    createMatrix s' =
        H.fromLists
            (map
                 (\a ->
                       map
                           (\s ->
                                 transition mdp a s s')
                           (S.toList (states mdp)))
                 (S.toList (actions mdp)))
    rewardM = takeDiag (policy H.<> createRewardMatrix)
    createRewardMatrix =
        H.fromLists
            (map
                 (\a ->
                       map
                           (
                                 reward mdp a)
                           (S.toList (states mdp)))
                 (S.toList (actions mdp)))

matrixPolicyEvaluation :: MDPM s a -> Int -> Vector Double
matrixPolicyEvaluation mdpm = evalVF (konst 0.0 (sizeStates mdpm))
  where
    evalVF v 0 = v
    evalVF v k =
        evalVF
            (rewardm mdpm + gammam mdpm `scale` (transitionm mdpm H.#> v))
            (k - 1)

matrixPolicyImprove :: MDPM s a -> Vector Double -> Matrix Double
matrixPolicyImprove mdpm vf = undefined


module Measurable.Measures where

import Measurable.Core
import Numeric.SpecFunctions (choose)
import Statistics.Distribution
import qualified Statistics.Distribution.Beta as Statistics
import qualified Statistics.Distribution.Binomial as Statistics
import qualified Statistics.Distribution.ChiSquared as Statistics
import qualified Statistics.Distribution.Gamma as Statistics
import qualified Statistics.Distribution.Exponential as Statistics
import qualified Statistics.Distribution.Normal as Statistics

standard :: Measure Double
standard = fromDensityFunction pdf where
  pdf = density Statistics.standard

normal :: Double -> Double -> Measure Double
normal m s = fromDensityFunction pdf where
  pdf = density $ Statistics.normalDistr m s

logNormal :: Double -> Double -> Measure Double
logNormal m s = fmap exp (normal m s)

exponential :: Double -> Measure Double
exponential r = fromDensityFunction pdf where
  pdf = density $ Statistics.exponential r

gamma :: Double -> Double -> Measure Double
gamma a b = fromDensityFunction pdf where
  pdf = density $ Statistics.gammaDistr a b

inverseGamma :: Double -> Double -> Measure Double
inverseGamma a b = fmap recip (gamma a b)

chiSquare :: Int -> Measure Double
chiSquare k = fromDensityFunction pdf where
  pdf = density $ Statistics.chiSquared k

beta :: Double -> Double -> Measure Double
beta a b = fromDensityFunction pdf where
  pdf = density $ Statistics.betaDistr a b

binomial :: Int -> Double -> Measure Int
binomial n p = fromMassFunction pmf [0..n] where
  pmf = binomialMass n p

bernoulli :: Double -> Measure Int
bernoulli = binomial 1

binomialMass :: Int -> Double -> Int -> Double
binomialMass n p k = bc * p ^ k * (1 - p) ^ (n - k) where
  bc = n `choose` k


-- perceptron simulation code
-- created by Tomohiro KIDA 2018/05/08/Tue

module Main
    ( main
    , linearFunc
    , linearFunc'
    , lossFunc
    , gradientDescent
    , plotDataPng
    , simPerceptron
    ) where

import Graphics.Gnuplot.Simple
import System.Random
import Numeric.LinearAlgebra

xxs :: [[Double]]
xxs = [[1,4,6],[1,1,4],[1,3,2],[1,7,3]]
ts :: [Int]
ts = [-1,-1,1,1]
a :: (Fractional frac) => frac
a  = 0.1

-- f (x) = w (tr x)
linearFunc :: Matrix R -> Matrix R -> Double
linearFunc w x
    | fst (size w) == 1 && fst (size x) == 1 = head . head . toLists $ w <> (tr x)
    | otherwise                              = -1 -- error

-- linear function for plot
-- the left hand side is x2
linearFunc' :: Matrix R -> Double -> Double
linearFunc' ws x =
    let w = (toLists ws)!!0
    in (-x)*(w!!1)/(w!!2) - ((w!!0)/(w!!2))

lossFunc :: Double -> Double -> Double
lossFunc t f = max 0.0 (-t*f)

-- return one itaration of gradientDescent
gradientDescent :: Int -> Matrix R -> Matrix R
gradientDescent rn ws =
    let ti      = realToFrac $ ts !! rn
        xi      = fromLists [xxs!!rn] :: Matrix R
        li      = lossFunc ti $ linearFunc ws xi
        wsn     = if li > 0 then ws - ( a*(realToFrac (-ti))*xi ) else ws
    in wsn

-- plot data and png
plotDataPng :: Matrix R -> IO ()
plotDataPng ws = do
    plotPathsStyle atribute plotstyle
    where
        lines     = map (\x->(x,(linearFunc' ws x))) [0..10]
        points_t  = map ((\[x,y]->(x,y)).tail) $ (\xs->[(xs!!0),(xs!!1)]) xxs
        points_f  = map ((\[x,y]->(x,y)).tail) $ (\xs->[(xs!!2),(xs!!3)]) xxs
        label     = [(XLabel "x"), (YLabel "y")]
        range     = [(XRange (0,8)), (YRange (0,8))]
        size      = [Aspect (Ratio 0.7)]
        title     = [(Title "simPerceptron")]
        png       = [(PNG   "simPerceptron.png")]
        atribute  = (label++range++title++png)
        plotstyle = [(defaultStyle {plotType = Points }, points_t),
                     (defaultStyle {plotType = Points }, points_f),
                     (defaultStyle {plotType = Lines  }, lines )]

-- simulation perceptron by grandient descent
simPerceptron :: StdGen -> Matrix R -> Int -> IO ()
simPerceptron gen ws count = do
    let (rn,ng) = randomR (0,3) gen :: (Int, StdGen) -- get one rand data of 4
    let wsn = gradientDescent rn ws -- update ws
    putStrLn $ show count ++ show wsn
    if (count > 0)
        then simPerceptron ng wsn (count-1) -- return recursion
        else plotDataPng wsn -- plot

main :: IO ()
main = do
    gen <- getStdGen
    let ws = (1><3) [-1,2,3] :: Matrix R
    simPerceptron gen ws 100 -- gradientDescent simulation


module Section_3_4_3 where

import Numeric.LinearAlgebra

type Network = [(Matrix Double, Vector Double, Double -> Double)]

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

initNetwork :: Network
initNetwork =
  [ ( (2><3) [0.1, 0.3, 0.5, 0.2, 0.4, 0.6]
    , vector [0.1, 0.2, 0.3]
    , sigmoid
    )
  , ( (3><2) [0.1, 0.4, 0.2, 0.5, 0.3, 0.6]
    , vector [0.1, 0.2]
    , sigmoid
    )
  , ( (2><2) [0.1, 0.3, 0.2, 0.4]
    , vector [0.1, 0.2]
    , id
    )
  ]

forward :: Network -> Vector Double -> Vector Double
forward nw x = foldl (\x (w, b, h) -> cmap h (x <# w + b)) x nw


-- root finding examples
import Numeric.GSL
import Numeric.LinearAlgebra
import Text.Printf(printf)

rosenbrock a b [x,y] = [ a*(1-x), b*(y-x^2) ]

test method = do
    print method
    let (s,p) = root method 1E-7 30 (rosenbrock 1 10) [-10,-5]
    print s -- solution
    disp p -- evolution of the algorithm

jacobian a b [x,y] = [ [-a    , 0]
                     , [-2*b*x, b] ]

testJ method = do
    print method
    let (s,p) = rootJ method 1E-7 30 (rosenbrock 1 10) (jacobian 1 10) [-10,-5]
    print s
    disp p

disp = putStrLn . format "  " (printf "%.3f")

main = do
    test Hybrids
    test Hybrid
    test DNewton
    test Broyden

    mapM_ testJ [HybridsJ .. GNewton]


-- |
-- Module      :  Spiral.RootOfUnity
-- Copyright   :  (c) 2016-2017 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.RootOfUnity (
    RootOfUnity(..)
  ) where

import Data.Complex
import Data.Complex.Cyclotomic

-- | Types that have primitive roots of unity.
class Fractional a => RootOfUnity a where
    -- | @rootOfUnity n k@ computes the primitive @n@th root of unity raised to
    -- the @k@th power.
    rootOfUnity :: Int -> Int -> a
    rootOfUnity n k = omega n ^^ k

    -- | Computes the primitive @n@th root of unity.
    omega :: Int -> a
    omega n = rootOfUnity n 1

instance RealFloat a => RootOfUnity (Complex a) where
    -- | $e^{-2 \pi i \frac{k}{n}$
    rootOfUnity n k = exp (-2*pi*i/fromIntegral n)^^k
      where
        i = 0:+1

instance RootOfUnity Cyclotomic where
    -- 'Data.Complex.Cyclotomic.e' computes $e^{2 \pi i/n}$, so we need to take
    -- the reciprocal.
    rootOfUnity n k = (1/e (fromIntegral n))^^k