{-|
Module      : Qubism.Algebra
Description : Typeclasses and utilities for Vector and Hilbert spaces.
Copyright   : (c) Keith Pearson, 2018
License     : MIT
Maintainer  : keith@qubitrot.org
-}

module Qubism.Algebra where

import Data.Complex
import Data.Monoid

type C = Complex Double

-- | A vector space over the complex numbers
class Eq v => VectorSpace v where
  -- | The additive identity
  zero :: v
  -- | Scalar multiplication
  (.:) :: C -> v -> v
  -- | Vector addition
  (+:) :: v -> v -> v
  -- | Vector subtraction
  (-:) :: v -> v -> v
  -- | Inverse element
  neg  :: v -> v
  a -: b = a +: (neg  b)

-- | A Hilbert space over the complex numbers
class VectorSpace v => HilbertSpace v where
  -- | Sesquilinear inner product
  (<.>) :: v -> v -> C
  -- | Norm induced by the inner product
  norm  :: v -> Double
  norm a = realPart $ a <.> a

class VectorSpace v => Algebra v where
  -- | Bilinear product
  (*:) :: v -> v -> v

commutator :: Algebra v  => v -> v -> v
commutator a b = (a *: b) -: (b *: a)

anticommutator :: Algebra v => v -> v -> v
anticommutator a b = (a *: b) +: (b *: a)


module Math.Ray where

import Numeric.LinearAlgebra

import Codec.Picture

type Vec3 = Vector Double
vec3 :: Double -> Double -> Double -> Vec3
vec3 x y z= vector $! [x,y,z]

data Ray = Ray
    {-# UNPACK #-} !Vec3
    {-# UNPACK #-} !Vec3
    deriving Show

parametric :: Ray -> (Double -> Vec3)
parametric (Ray p v) k = p + scale k v
{-# INLINE parametric #-}


module FFT (fft) where
import Data.Complex (Complex, cis, conjugate)
import Data.List (head, zip, take)

fft :: [Complex Float] -> [Complex Float]
fft xs = fft' (length xs) xs

fft' :: Int -> [Complex Float] -> [Complex Float]
fft' 1 [x] = [x]
fft' n coeffs = eval [0..n-1] halfN nThPrimitiveRoots halfFFT
  where
        halfN                   = n `div` 2
        (evenCoeffs, oddCoeffs) = split coeffs
        evenFFT                 = fft' halfN evenCoeffs
        oddFFT                  = fft' halfN oddCoeffs
        nThPrimitiveRoots       = duplicate $ take halfN (rootsOfUnity n)
        halfFFT                 = duplicate $ zip evenFFT oddFFT

rootsOfUnity :: Int -> [Complex Float]
rootsOfUnity = rootsOfUnity' . fromIntegral

rootsOfUnity' :: Float -> [Complex Float]
rootsOfUnity' n = map conjugate $ [ cis (2 * pi * k / n) | k <- [0 .. n - 1] ]

split :: [Complex Float] -> ([Complex Float], [Complex Float])
split coeffs = (evenCoeffs, oddCoeffs)
  where indexedCoeffs = zip coeffs [0..]
        evenCoeffs    = [ coeff | (coeff, index) <- indexedCoeffs, even index ]
        oddCoeffs     = [ coeff | (coeff, index) <- indexedCoeffs, odd index ]

eval :: [Int] -> Int -> [Complex Float] -> [(Complex Float, Complex Float)] -> [Complex Float]
eval [] _ _ _ = []
eval (power:restPowers) halfPower (root:restRoots) (coeffs:restCoeffs)
  | power < halfPower = (evenCoeff + root * oddCoeff) : (eval restPowers halfPower restRoots restCoeffs)
  | otherwise         = (evenCoeff - root * oddCoeff) : (eval restPowers halfPower restRoots restCoeffs)
  where (evenCoeff, oddCoeff) = coeffs

duplicate :: [a] -> [a]
duplicate xs = xs ++ xs


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

The forwards transforms in this module are identical to those in "Numeric.FFT.Vector.Unnormalized".
The backwards transforms are normalized to be their inverse operations (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>.
-}
module Numeric.FFT.Vector.Invertible(
                    -- * Creating and executing 'Plan's
                    run,
                    plan,
                    execute,
                    -- * Complex-to-complex transforms
                    U.dft,
                    idft,
                    -- * Real-to-complex transforms
                    U.dftR2C,
                    dftC2R,
                    -- * Real-to-real transforms
                    -- $dct_size
                    -- ** Discrete cosine transforms
                    U.dct1,
                    idct1,
                    U.dct2,
                    idct2,
                    U.dct3,
                    idct3,
                    U.dct4,
                    idct4,
                    -- ** Discrete sine transforms
                    U.dst1,
                    idst1,
                    U.dst2,
                    idst2,
                    U.dst3,
                    idst3,
                    U.dst4,
                    idst4,
                    ) where

import Numeric.FFT.Vector.Base
import qualified Numeric.FFT.Vector.Unnormalized as U
import Data.Complex

-- | A backward discrete Fourier transform which is the inverse of 'U.dft'.  The output and input sizes are the same (@n@).
--
-- @y_k = (1\/n) sum_(j=0)^(n-1) x_j e^(2pi i j k/n)@
idft :: Transform (Complex Double) (Complex Double)
idft = U.idft {normalization = \n -> constMultOutput $ 1 / toEnum n}

-- | 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 :: Transform (Complex Double) Double
dftC2R = U.dftC2R {normalization = \n -> constMultOutput $ 1 / toEnum n}

-- OK, the inverse of each unnormalized operation.

-- $dct_size
-- The real-even (DCT) and real-odd (DST) transforms.  The input and output sizes
-- are the same (@n@).

-- | A type-1 discrete cosine transform which is the inverse of 'U.dct1'.
--
-- @y_k = (1\/(2(n-1)) [x_0 + (-1)^k x_(n-1) + 2 sum_(j=1)^(n-2) x_j cos(pi j k\/(n-1))]@
idct1 :: Transform Double Double
idct1 = U.dct1 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * (n-1))}

-- | A type-3 discrete cosine transform which is the inverse of 'U.dct2'.
--
-- @y_k = (1\/(2n)) [x_0 + 2 sum_(j=1)^(n-1) x_j cos(pi j(k+1\/2)\/n)]@
idct2 :: Transform Double Double
idct2 = U.dct3 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}

-- | A type-2 discrete cosine transform which is the inverse of 'U.dct3'.
--
-- @y_k = (1\/n) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)k\/n)@
idct3 :: Transform Double Double
idct3 = U.dct2 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}

-- | A type-4 discrete cosine transform which is the inverse of 'U.dct4'.
--
-- @y_k = (1\/n) sum_(j=0)^(n-1) x_j cos(pi(j+1\/2)(k+1\/2)\/n)@
idct4 :: Transform Double Double
idct4 = U.dct4 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}

-- | A type-1 discrete sine transform which is the inverse of 'U.dst1'.
--
-- @y_k = (1\/(n+1)) sum_(j=0)^(n-1) x_j sin(pi(j+1)(k+1)\/(n+1))@
idst1 :: Transform Double Double
idst1 = U.dst1 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * (n+1))}

-- | A type-3 discrete sine transform which is the inverse of 'U.dst2'.
--
-- @y_k = (1\/(2n)) [(-1)^k x_(n-1) + 2 sum_(j=0)^(n-2) x_j sin(pi(j+1)(k+1\/2)/n)]@
idst2 :: Transform Double Double
idst2 = U.dst3 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}

-- | A type-2 discrete sine transform which is the inverse of 'U.dst3'.
--
-- @y_k = (1\/n) sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1)\/n)@
idst3 :: Transform Double Double
idst3 = U.dst2 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}

-- | A type-4 discrete sine transform which is the inverse of 'U.dst4'.
--
-- @y_k = (1\/(2n)) sum_(j=0)^(n-1) x_j sin(pi(j+1\/2)(k+1\/2)\/n)@
idst4 :: Transform Double Double
idst4 = U.dst4 {normalization = \n -> constMultOutput $ 1 / toEnum (2 * n)}


{-# LANGUAGE Arrows #-}

-- Homogeneous (or depth-preserving) functions over perfectly balanced trees.

module Main where

import Control.Arrow
import Control.Category
import Data.Complex
import Prelude hiding (id, (.))

infixr 4 :&:

-- Consider the following non-regular type of perfectly balanced trees,
-- or `powertrees' (cf Jayadev Misra's powerlists):

data Pow a = Zero a | Succ (Pow (Pair a))
	deriving Show

type Pair a = (a, a)

-- Here are some example elements:

tree0 = Zero 1
tree1 = Succ (Zero (1, 2))
tree2 = Succ (Succ (Zero ((1, 2), (3, 4))))
tree3 = Succ (Succ (Succ (Zero (((1, 2), (3, 4)), ((5, 6), (7, 8))))))

-- The elements of this type have a string of constructors expressing
-- a depth n as a Peano numeral, enclosing a nested pair tree of 2^n
-- elements.  The type definition ensures that all elements of this type
-- are perfectly balanced binary trees of this form.  (Such things arise
-- in circuit design, eg Ruby, and descriptions of parallel algorithms.)
-- And the type system will ensure that all legal programs preserve
-- this structural invariant.
--
-- The only problem is that the type constraint is too restrictive, rejecting
-- many of the standard operations on these trees.  Typically you want to
-- split a tree into two subtrees, do some processing on the subtrees and
-- combine the results.  But the type system cannot discover that the two
-- results are of the same depth (and thus combinable).  We need a type
-- that says a function preserves depth.  Here it is:

data Hom a b = (a -> b) :&: Hom (Pair a) (Pair b)

-- A homogeneous (or depth-preserving) function is an infinite sequence of
-- functions of type Pair^n a -> Pair^n b, one for each depth n.  We can
-- apply a homogeneous function to a powertree by selecting the function
-- for the required depth:

apply :: Hom a b -> Pow a -> Pow b
apply (f :&: fs) (Zero x) = Zero (f x)
apply (f :&: fs) (Succ t) = Succ (apply fs t)

-- Having defined apply, we can forget about powertrees and do all our
-- programming with Hom's.  Firstly, Hom is an arrow:

instance Category Hom where
	id = id :&: id
	(f :&: fs) . (g :&: gs) = (f . g) :&: (fs . gs)

instance Arrow Hom where
	arr f = f :&: arr (f *** f)
	first (f :&: fs) =
		first f :&: (arr transpose >>> first fs >>> arr transpose)

transpose :: ((a,b), (c,d)) -> ((a,c), (b,d))
transpose ((a,b), (c,d)) = ((a,c), (b,d))

-- arr maps f over the leaves of a powertree.

-- The composition >>> composes sequences of functions pairwise.
--
-- The *** operator unriffles a powertree of pairs into a pair of powertrees,
-- applies the appropriate function to each and riffles the results.
-- It defines a categorical product for this arrow category.

-- When describing algorithms, one often provides a pure function for the
-- base case (trees of one element) and a (usually recursive) expression
-- for trees of pairs.

-- For example, a common divide-and-conquer pattern is the butterfly, where
-- one recursive call processes the odd-numbered elements and the other
-- processes the even ones (cf Geraint Jones and Mary Sheeran's Ruby papers):

butterfly :: (Pair a -> Pair a) -> Hom a a
butterfly f = id :&: proc (x, y) -> do
		x' <- butterfly f -< x
		y' <- butterfly f -< y
		returnA -< f (x', y')

-- The recursive calls operate on halves of the original tree, so the
-- recursion is well-defined.

-- Some examples of butterflies:

rev :: Hom a a
rev = butterfly swap
	where	swap (x, y) = (y, x)

unriffle :: Hom (Pair a) (Pair a)
unriffle = butterfly transpose

-- Batcher's sorter for bitonic sequences:

bisort :: Ord a => Hom a a
bisort = butterfly cmp
	where	cmp (x, y) = (min x y, max x y)

-- This can be used (with rev) as the merge phase of a merge sort.
--
sort :: Ord a => Hom a a
sort = id :&: proc (x, y) -> do
		x' <- sort -< x
		y' <- sort -< y
		yr <- rev -< y'
		p <- unriffle -< (x', yr)
		bisort2 -< p
	where _ :&: bisort2 = bisort

-- Here is the scan operation, using the algorithm of Ladner and Fischer:

scan :: (a -> a -> a) -> a -> Hom a a
scan op b = id :&: proc (x, y) -> do
		y' <- scan op b -< op x y
		l <- rsh b -< y'
		returnA -< (op l x, y')

-- The auxiliary function rsh b shifts each element in the tree one place to
-- the right, placing b in the now-vacant leftmost position, and discarding
-- the old rightmost element:

rsh :: a -> Hom a a
rsh b = const b :&: proc (x, y) -> do
		w <- rsh b -< y
		returnA -< (w, x)

-- Finally, here is the Fast Fourier Transform:

type C = Complex Double

fft :: Hom C C
fft = id :&: proc (x, y) -> do
		x' <- fft -< x
		y' <- fft -< y
		r <- roots (-1) -< ()
		let z = r*y'
		unriffle -< (x' + z, x' - z)

-- The auxiliary function roots r (where r is typically a root of unity)
-- populates a tree of size n (necessarily a power of 2) with the values
-- 1, w, w^2, ..., w^(n-1), where w^n = r.

roots :: C -> Hom () C
roots r = const 1 :&: proc _ -> do
		x <- roots r' -< ()
		unriffle -< (x, x*r')
	where	r' = if imagPart s >= 0 then -s else s
		s = sqrt r

-- Miscellaneous functions:

rrot :: Hom a a
rrot = id :&: proc (x, y) -> do
		w <- rrot -< y
		returnA -< (w, x)

ilv :: Hom a a -> Hom (Pair a) (Pair a)
ilv f = proc (x, y) -> do
		x' <- f -< x
		y' <- f -< y
		returnA -< (x', y')

scan' :: (a -> a -> a) -> a -> Hom a a
scan' op b = proc x -> do
		l <- rsh b -< x
		(id :&: ilv (scan' op b)) -< op l x

riffle :: Hom (Pair a) (Pair a)
riffle = id :&: proc ((x1, y1), (x2, y2)) -> do
		x <- riffle -< (x1, x2)
		y <- riffle -< (y1, y2)
		returnA -< (x, y)

invert :: Hom a a
invert = id :&: proc (x, y) -> do
		x' <- invert -< x
		y' <- invert -< y
		unriffle -< (x', y')

carryLookaheadAdder :: Hom (Bool, Bool) Bool
carryLookaheadAdder = proc (x, y) -> do
		carryOut <- rsh (Just False) -<
			if x == y then Just x else Nothing
		Just carryIn <- scan plusMaybe Nothing -< carryOut
		returnA -< x `xor` y `xor` carryIn
	where	plusMaybe x Nothing = x
		plusMaybe x (Just y) = Just y
		False `xor` b = b
		True `xor` b = not b

-- Global conditional for SIMD

ifAll :: Hom a b -> Hom a b -> Hom (a, Bool) b
ifAll fs gs = ifAllAux snd (arr fst >>> fs) (arr fst >>> gs)
	where	ifAllAux :: (a -> Bool) -> Hom a b -> Hom a b -> Hom a b
		ifAllAux p (f :&: fs) (g :&: gs) =
			liftIf p f g :&: ifAllAux (liftAnd p) fs gs
		liftIf p f g x = if p x then f x else g x
		liftAnd p (x, y) = p x && p y

maybeAll :: Hom a c -> Hom (a, b) c -> Hom (a, Maybe b) c
maybeAll (n :&: ns) (j :&: js) =
	choose :&: (arr dist >>> maybeAll ns (arr transpose >>> js))
	where	choose (a, Nothing) = n a
		choose (a, Just b) = j (a, b)
		dist ((a1, b1), (a2, b2)) = ((a1, a2), zipMaybe b1 b2)
		zipMaybe (Just x) (Just y) = Just (x, y)
		zipMaybe _ _ = Nothing

main = do
	print (apply rev tree3)
	print (apply invert tree3)
	print (apply (invert >>> sort) tree3)
	print (apply (scan (+) 0) tree3)


module Numeric.Physics.Particle where

import Numeric.LinearAlgebra.Vector

data Particle a = Particle
  { pPos     :: !(Vec3 a) -- ^ Position of particle in the world
  , pVel     :: !(Vec3 a) -- ^ Velocity of particle
  , pAcc     :: !(Vec3 a) -- ^ Acceleration of particle
  , pDamp    :: !a        -- ^ damping factor for this particle
  , pInvMass :: !a        -- ^ inverse mass (1/m) of this particle
  } deriving (Read,Show,Ord,Eq)

-- | Takes a particle, a force, and a time duration
-- and gives the particle updated for time slice
integrate :: Floating a => Particle a -> Vec3 a -> a -> Particle a
integrate p f dt =
  let pos  = (dt         *> pVel p) <+> pPos p
      acc' = (pInvMass p *> f)      <+> pAcc p
      vel  = (dt         *> acc')   <+> pVel p
  in p { pPos = pos
       , pVel = ((pDamp p)**dt) *> vel
       }
{-# SPECIALIZE INLINE integrate :: Particle Double -> Vec3 Double -> Double -> Particle Double #-}
{-# SPECIALIZE INLINE integrate :: Particle Float  -> Vec3 Float  -> Float  -> Particle Float  #-}


#!/usr/bin/env stack
-- stack --resolver lts-5.15 --install-ghc runghc --package hmatrix --package MonadRandom

-- FROM: https://github.com/mstksg/inCode/blob/master/code-samples/dependent-haskell/NetworkUntyped.hs

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad
import Control.Monad.Random
import Data.List
import Data.Maybe
import Numeric.LinearAlgebra
import System.Environment
import Text.Read

data Weights = W { wBiases :: !(Vector Double)  -- n
                 , wNodes  :: !(Matrix Double)  -- n x m
                 }                              -- "m to n" layer

data Network = Last !Weights
             | Link !Weights Network

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

logistic' :: Floating a => a -> a
logistic' x = logix * (1 - logix)
  where
    logix = logistic x

runLayer :: Weights -> Vector Double -> Vector Double
runLayer (W wB wN) v = wB + wN #> v

runNet :: Network -> Vector Double -> Vector Double
runNet (O w)      !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in  runNet n' v'

randomWeights :: MonadRandom m => Int -> Int -> m Weights
randomWeights i o = do
    seed1 :: Int <- getRandom
    seed2 :: Int <- getRandom
    let wB = randomVector  seed1 Uniform o * 2 - 1
        wN = uniformSample seed2 o (replicate i (-1, 1))
    return $ W wB wN

randomNet :: MonadRandom m => Int -> [Int] -> Int -> m Network
randomNet i []     o =     O <$> randomWeights i o
randomNet i (h:hs) o = (:&~) <$> randomWeights i h <*> randomNet h hs o

train :: Double           -- ^ learning rate
      -> Vector Double    -- ^ input vector
      -> Vector Double    -- ^ target vector
      -> Network          -- ^ network to train
      -> Network
train rate x0 target = fst . go x0
  where
    go :: Vector Double    -- ^ input vector
       -> Network          -- ^ network to train
       -> (Network, Vector Double)
    -- handle the output layer
    go !x (O w@(W wB wN))
        = let y    = runLayer w x
              o    = logistic y
              -- the gradient (how much y affects the error)
              --   (logistic' is the derivative of logistic)
              dEdy = logistic' y * (o - target)
              -- new bias weights and node weights
              wB'  = wB - scale rate dEdy
              wN'  = wN - scale rate (dEdy `outer` x)
              w'   = W wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (O w', dWs)
    -- handle the inner layers
    go !x (w@(W wB wN) :&~ n)
        = let y          = runLayer w x
              o          = logistic y
              -- get dWs', bundle of derivatives from rest of the net
              (n', dWs') = go o n
              -- the gradient (how much y affects the error)
              dEdy       = logistic' y * dWs'
              -- new bias weights and node weights
              wB'  = wB - scale rate dEdy
              wN'  = wN - scale rate (dEdy `outer` x)
              w'   = W wB' wN'
              -- bundle of derivatives for next step
              dWs  = tr wN #> dEdy
          in  (w' :&~ n', dWs)

netTest :: MonadRandom m => Double -> Int -> m String
netTest rate n = do
    inps <- replicateM n $ do
      s <- getRandom
      return $ randomVector s Uniform 2 * 2 - 1
    let outs = flip map inps $ \v ->
                 if v `inCircle` (fromRational 0.33, 0.33)
                      || v `inCircle` (fromRational (-0.33), 0.33)
                   then fromRational 1
                   else fromRational 0
    net0 <- randomNet 2 [16,8] 1
    let trained = foldl' trainEach net0 (zip inps outs)
          where
            trainEach :: Network -> (Vector Double, Vector Double) -> Network
            trainEach nt (i, o) = train rate i o nt

        outMat = [ [ render (norm_2 (runNet trained (vector [x / 25 - 1,y / 10 - 1])))
                   | x <- [0..50] ]
                 | y <- [0..20] ]
        render r | r <= 0.2  = ' '
                 | r <= 0.4  = '.'
                 | r <= 0.6  = '-'
                 | r <= 0.8  = '='
                 | otherwise = '#'

    return $ unlines outMat
  where
    inCircle :: Vector Double -> (Vector Double, Double) -> Bool
    v `inCircle` (o, r) = norm_2 (v - o) <= r

main :: IO ()
main = do
    args <- getArgs
    let n    = readMaybe =<< (args !!? 0)
        rate = readMaybe =<< (args !!? 1)
    putStrLn "Training network..."
    putStrLn =<< evalRandIO (netTest (fromMaybe 0.25   rate)
                                     (fromMaybe 500000 n   )
                            )

(!!?) :: [a] -> Int -> Maybe a
xs !!? i = listToMaybe (drop i xs)


module Solver (coefficients, approxU, getLMatrix, getBMatrix, domainStart, domainEnd) where

import Data.Maybe (fromJust)
import Math.GaussianQuadratureIntegration (nIntegrate256)
import Numeric.LinearAlgebra (flatten, fromLists, linearSolve, luSolve, toList, toLists, (><))

approxU :: Int -> Double -> Double
approxU n x =
  sum . zipWith (\i c_i -> c_i * basisF i x) [0 ..] . coefficients $ n
  where
    basisF = basisFWithDomain n

coefficients :: Int -> [Double]
coefficients n = 0 : (toList . flatten $ fromJust (linearSolve b_matrix l_vector))
  where
    b_matrix = fromLists $ getBMatrix n
    l_vector = (n >< 1) $ getLMatrix n

getLMatrix :: Int -> [Double]
getLMatrix n =
  [l_ej j | j <- [1 .. n]]
  where
    -- L(v) = - int[0, 2] v*sindx
    -- L(e_j) = - int[0, 2] e_j*sindx
    -- u(0) => e_j = 0, so we skip it
    l_ej j = - integrate domainStart domainEnd (j > 0) (\x -> basisF j x * sin x)
    basisF = basisFWithDomain n

getBMatrix :: Int -> [[Double]]
getBMatrix n =
  [[b_ei_ej i j | j <- [1 .. n]] | i <- [1 .. n]]
  where
    -- B(u,v) = u(2)v(2) + int[0, 2] (-u'v' + uv) dx
    -- B(e_i, e_j) = e_i(2)*e_j(2) + int[0, 2] (- e_i'*e_j' + e_i*e_j) dx
    b_ei_ej i j = u2v2 i j + ʃsub_u'v'_add_uvdx i j
    u2v2 i j = basisF i domainEnd * basisF j domainEnd
    -- we only need to integrate splines with non-zero intersection
    ʃsub_u'v'_add_uvdx i j = integrate lowerBound upperBound (near i j) (\x -> - basisF' i x * basisF' j x + basisF i x * basisF j x)
      where
        -- finding intersection of tent functions
        lowerBound = dx * max (bigger_idx - 1) domainStart
        upperBound = dx * min (smaller_idx + 1) (fromIntegral n)
        dx = domainEnd / fromIntegral n
        smaller_idx = min (fromIntegral i) (fromIntegral j)
        bigger_idx = max (fromIntegral i) (fromIntegral j)
    near i j = abs (i - j) <= 1
    basisF = basisFWithDomain n
    basisF' = basisF'WithDomain n

integrate :: (Fractional a, Ord a) => a -> a -> Bool -> (a -> a) -> a
integrate start end nonZeroCond f
  | nonZeroCond = nIntegrate256 f start end
  | otherwise = 0

-- tent function used for generating vector space
-- https://en.wikipedia.org/wiki/Triangular_function
basisFWithDomain :: Int -> Int -> Double -> Double
basisFWithDomain n i x
  | x < left || right < x = 0
  | x <= center = (x - left) * hInverse
  | otherwise = (right - x) * hInverse
  where
    left = center - h
    right = center + h
    center = domainEnd * fromIntegral i / fromIntegral n
    h = domainEnd / fromIntegral n
    hInverse = fromIntegral n / domainEnd

-- derivative of tent function
basisF'WithDomain :: Int -> Int -> Double -> Double
basisF'WithDomain n i x
  | x < left || right < x = 0
  | x <= center = hInverse
  | otherwise = - hInverse
  where
    left = center - h
    right = center + h
    center = domainEnd * fromIntegral i / fromIntegral n
    h = domainEnd / fromIntegral n
    hInverse = fromIntegral n / domainEnd

domainEnd :: Double
domainEnd = 2

domainStart :: Double
domainStart = 0


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Laplace
-- Copyright : (c) 2015 Mihai Maruseac
-- License   : BSD3
--
-- Maintainer  : mihai.maruseac@maruseac.com
-- Stability   : experimental
-- Portability : portable
--
-- The Laplace distribution.  This is the continuous probability
-- defined as the difference of two iid exponential random variables
-- or a Brownian motion evaluated as exponentially distributed times.
-- It is used in differential privacy (Laplace Method), speech
-- recognition and least absolute deviations method (Laplace's first
-- law of errors, giving a robust regression method)
--
module Statistics.Distribution.Laplace
    (
      LaplaceDistribution
    -- * Constructors
    , laplace
    , laplaceE
    -- * Accessors
    , ldLocation
    , ldScale
    ) where

import Control.Applicative
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import qualified Data.Vector.Generic             as G
import qualified Statistics.Distribution         as D
import qualified Statistics.Quantile             as Q
import qualified Statistics.Sample               as S
import Statistics.Internal

data LaplaceDistribution = LD {
      ldLocation :: {-# UNPACK #-} !Double
    -- ^ Location.
    , ldScale    :: {-# UNPACK #-} !Double
    -- ^ Scale.
    } deriving (Eq, Typeable, Data, Generic)

instance Show LaplaceDistribution where
  showsPrec i (LD l s) = defaultShow2 "laplace" l s i
instance Read LaplaceDistribution where
  readPrec = defaultReadPrecM2 "laplace" laplaceE

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

instance D.ContDistr LaplaceDistribution where
    density    (LD l s) x = exp (- abs (x - l) / s) / (2 * s)
    logDensity (LD l s) x = - abs (x - l) / s - log 2 - log s
    quantile      = quantile
    complQuantile = complQuantile

instance D.Mean LaplaceDistribution where
    mean (LD l _) = l

instance D.Variance LaplaceDistribution where
    variance (LD _ s) = 2 * s * s

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

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

instance D.Entropy LaplaceDistribution where
  entropy (LD _ s) = 1 + log (2 * s)

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

instance D.ContGen LaplaceDistribution where
  genContVar = D.genContinuous

cumulative :: LaplaceDistribution -> Double -> Double
cumulative (LD l s) x
  | x <= l    = 0.5 * exp ( (x - l) / s)
  | otherwise = 1 - 0.5 * exp ( - (x - l) / s )

complCumulative :: LaplaceDistribution -> Double -> Double
complCumulative (LD l s) x
  | x <= l    = 1 - 0.5 * exp ( (x - l) / s)
  | otherwise = 0.5 * exp ( - (x - l) / s )

quantile :: LaplaceDistribution -> Double -> Double
quantile (LD l s) p
  | p == 0             = -inf
  | p == 1             = inf
  | p == 0.5           = l
  | p > 0   && p < 0.5 = l + s * log (2 * p)
  | p > 0.5 && p < 1   = l - s * log (2 - 2 * p)
  | otherwise          =
    error $ "Statistics.Distribution.Laplace.quantile: p must be in [0,1] range. Got: "++show p
  where
    inf = 1 / 0

complQuantile :: LaplaceDistribution -> Double -> Double
complQuantile (LD l s) p
  | p == 0             = inf
  | p == 1             = -inf
  | p == 0.5           = l
  | p > 0   && p < 0.5 = l - s * log (2 * p)
  | p > 0.5 && p < 1   = l + s * log (2 - 2 * p)
  | otherwise          =
    error $ "Statistics.Distribution.Laplace.quantile: p must be in [0,1] range. Got: "++show p
  where
    inf = 1 / 0

-- | Create an Laplace distribution.
laplace :: Double         -- ^ Location
        -> Double        -- ^ Scale
        -> LaplaceDistribution
laplace l s = maybe (error $ errMsg l s) id $ laplaceE l s

-- | Create an Laplace distribution.
laplaceE :: Double         -- ^ Location
         -> Double        -- ^ Scale
         -> Maybe LaplaceDistribution
laplaceE l s
  | s >= 0    = Just (LD l s)
  | otherwise = Nothing

errMsg :: Double -> Double -> String
errMsg _ s = "Statistics.Distribution.Laplace.laplace: scale parameter must be positive. Got " ++ show s

-- | Create Laplace distribution from sample. No tests are made to
--   check whether it truly is Laplace. Location of distribution
--   estimated as median of sample.
instance D.FromSample LaplaceDistribution Double where
  fromSample xs
    | G.null xs = Nothing
    | otherwise = Just $! LD s l
    where
      s = Q.median Q.medianUnbiased xs
      l = S.mean $ G.map (\x -> abs $ x - s) xs


{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE TemplateHaskell  #-}
{-# LANGUAGE ViewPatterns     #-}

module Main where

-- Imports for network logic

import           Control.Lens hiding                   ((<.>))
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop

-- Imports for running network on data

import           Control.DeepSeq
import           Control.Exception
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State
import           Data.Bitraversable
import           Data.Foldable
import           Data.IDX
-- import           Data.List
import           Data.List.Split
import           Data.Traversable
import           Data.Tuple
import           GHC.Generics                    (Generic)
import           GHC.TypeLits
import           Numeric.OneLiner
import           System.Environment
import           System.FilePath hiding          ((<.>))
import           Text.Printf
import qualified Data.Vector                     as V
import qualified Data.Vector.Generic             as VG
import qualified Numeric.LinearAlgebra           as HM
import qualified Numeric.LinearAlgebra.Static    as H
import qualified System.Random.MWC               as MWC
import qualified System.Random.MWC.Distributions as MWC

data Net = N { _weights1 :: L 250 784
             , _bias1    :: R 250
             , _weights2 :: L 10 250
             , _bias2    :: R 10
             }
  deriving (Generic)
makeLenses ''Net

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

softMax
    :: Reifies s W
    => BVar s (R 10)
    -> BVar s (R 10)
softMax x = expx / konst (norm_1V expx)
  where
    expx = exp x

crossEntropy
    :: Reifies s W
    => BVar s (R 10)
    -> BVar s (R 10)
    -> BVar s (Double)
crossEntropy targ res = -(log res <.> targ)

runNet
    :: Reifies s W
    => BVar s Net
    -> BVar s (R 784)
    -> BVar s (R 10)
runNet n x = z
  where
    y = logistic $ (n ^^. weights1) #> x + (n ^^. bias1)
    z = softMax  $ (n ^^. weights2) #> y + (n ^^. bias2)

netErr
    :: Reifies s W
    => BVar s (R 784)
    -> BVar s (R 10)
    -> BVar s Net
    -> BVar s Double
netErr x targ n = crossEntropy targ (runNet n x)

stepNet :: R 784 -> R 10 -> Net -> Net
stepNet x targ net0 = net0 - 0.02 * gr
  where
    gr :: Net
    gr = gradBP (netErr (constVar x) (constVar targ)) net0

-- *********************************************
-- Plumbing for running the network on real data
-- *********************************************

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
    datadir:_ <- getArgs
    Just train <- loadMNIST (datadir </> "train-images-idx3-ubyte")
                            (datadir </> "train-labels-idx1-ubyte")
    Just test  <- loadMNIST (datadir </> "t10k-images-idx3-ubyte")
                            (datadir </> "t10k-labels-idx1-ubyte")
    putStrLn "Loaded data."
    net0 <- MWC.uniformR (-0.5, 0.5) g
    flip evalStateT net0 . forM_ [1..] $ \e -> do
      train' <- liftIO . fmap V.toList $ MWC.uniformShuffle (V.fromList train) g
      liftIO $ printf "[Epoch %d]\n" (e :: Int)

      forM_ ([1..] `zip` chunksOf 5000 train') $ \(b, chnk) -> StateT $ \n0 -> do
        printf "(Batch %d)\n" (b :: Int)

        n' <- evaluate . force $ trainList chnk n0
        printf "Trained on %d points.\n" (length chnk)

        let trainScore = testNet chnk n'
            testScore  = testNet test n'
        printf "Training error:   %.2f%%\n" ((1 - trainScore) * 100)
        printf "Validation error: %.2f%%\n" ((1 - testScore ) * 100)

        return ((), n')

trainList :: [(R 784, R 10)] -> Net -> Net
trainList = flip $ foldl' (\n (x,y) -> stepNet x y n)

testNet :: [(R 784, R 10)] -> Net -> Double
testNet xs n = sum (map (uncurry test) xs) / fromIntegral (length xs)
  where
    test x (H.extract->t)
        | HM.maxIndex t == HM.maxIndex (H.extract r) = 1
        | otherwise                                  = 0
      where
        r = evalBP (`runNet` constVar x) n

loadMNIST
    :: FilePath
    -> FilePath
    -> IO (Maybe [(R 784, R 10)])
loadMNIST fpI fpL = runMaybeT $ do
    i <- MaybeT          $ decodeIDXFile       fpI
    l <- MaybeT          $ decodeIDXLabelsFile fpL
    d <- MaybeT . return $ labeledIntData l i
    MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
  where
    mkImage = H.create . VG.convert . VG.map (\i -> fromIntegral i / 255)
    mkLabel n = H.create $ HM.build 10 (\i -> if round i == n then 1 else 0)

instance Num Net where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance Fractional Net where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational

instance KnownNat n => MWC.Variate (R n) where
    uniform g = H.randomVector <$> MWC.uniform g <*> pure H.Uniform
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where
    uniform g = H.uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance MWC.Variate Net where
    uniform g = N <$> MWC.uniform g
                  <*> MWC.uniform g
                  <*> MWC.uniform g
                  <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance NFData Net


{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Util
Copyright   :  (c) Alberto Ruiz 2013
License     :  GPL

Maintainer  :  Alberto Ruiz (aruiz at um dot es)
Stability   :  provisional

-}
-----------------------------------------------------------------------------

module Numeric.LinearAlgebra.Util(

    -- * Convenience functions
    size, disp,
    zeros, ones,
    diagl,
    row,
    col,
    (&),(!), (¦), (#),
    (?),(¿),
    rand, randn,
    cross,
    norm,
    unitary,
    mt,
    pairwiseD2,
    rowOuters,
    null1,
    null1sym,
    -- * Convolution
    -- ** 1D
    corr, conv, corrMin,
    -- ** 2D
    corr2, conv2, separable,
    -- * Tools for the Kronecker product
    --
    -- | (see A. Fusiello, A matter of notation: Several uses of the Kronecker product in
    --  3d computer vision, Pattern Recognition Letters 28 (15) (2007) 2127-2132)

    --
    -- | @`vec` (a \<> x \<> b) == ('trans' b ` 'kronecker' ` a) \<> 'vec' x@
    vec,
    vech,
    dup,
    vtrans
) where

import Numeric.Container
import Numeric.LinearAlgebra.Algorithms hiding (i)
import Numeric.Matrix()
import Numeric.Vector()

import System.Random(randomIO)
import Numeric.LinearAlgebra.Util.Convolution

disp :: Int -> Matrix Double -> IO ()
-- ^ show a matrix with given number of digits after the decimal point
disp n = putStrLn . dispf n

-- | pseudorandom matrix with uniform elements between 0 and 1
randm :: RandDist
     -> Int -- ^ rows
     -> Int -- ^ columns
     -> IO (Matrix Double)
randm d r c = do
    seed <- randomIO
    return (reshape c $ randomVector seed d (r*c))

-- | pseudorandom matrix with uniform elements between 0 and 1
rand :: Int -> Int -> IO (Matrix Double)
rand = randm Uniform

-- | pseudorandom matrix with normal elements
randn :: Int -> Int -> IO (Matrix Double)
randn = randm Gaussian

-- | create a real diagonal matrix from a list
diagl :: [Double] -> Matrix Double
diagl = diag . fromList

-- | a real matrix of zeros
zeros :: Int -- ^ rows
      -> Int -- ^ columns
      -> Matrix Double
zeros r c = konst 0 (r,c)

-- | a real matrix of ones
ones :: Int -- ^ rows
     -> Int -- ^ columns
     -> Matrix Double
ones r c = konst 1 (r,c)

-- | concatenation of real vectors
infixl 3 &
(&) :: Vector Double -> Vector Double -> Vector Double
a & b = join [a,b]

-- | horizontal concatenation of real matrices
infixl 3 !
(!) :: Matrix Double -> Matrix Double -> Matrix Double
a ! b = fromBlocks [[a,b]]

-- | (00A6) horizontal concatenation of real matrices
infixl 3 ¦
(¦) :: Matrix Double -> Matrix Double -> Matrix Double
a ¦ b = fromBlocks [[a,b]]

-- | vertical concatenation of real matrices
(#) :: Matrix Double -> Matrix Double -> Matrix Double
infixl 2 #
a # b = fromBlocks [[a],[b]]

-- | create a single row real matrix from a list
row :: [Double] -> Matrix Double
row = asRow . fromList

-- | create a single column real matrix from a list
col :: [Double] -> Matrix Double
col = asColumn . fromList

-- | extract selected rows
infixl 9 ?
(?) :: Element t => Matrix t -> [Int] -> Matrix t
(?) = flip extractRows

-- | (00BF) extract selected columns
infixl 9 ¿
(¿) :: Element t => Matrix t -> [Int] -> Matrix t
m ¿ ks = trans . extractRows ks . trans $ m

cross :: Vector Double -> Vector Double -> Vector Double
-- ^ cross product (for three-element real vectors)
cross x y | dim x == 3 && dim y == 3 = fromList [z1,z2,z3]
          | otherwise = error $ "cross ("++show x++") ("++show y++")"
  where
    [x1,x2,x3] = toList x
    [y1,y2,y3] = toList y
    z1 = x2*y3-x3*y2
    z2 = x3*y1-x1*y3
    z3 = x1*y2-x2*y1

norm :: Vector Double -> Double
-- ^ 2-norm of real vector
norm = pnorm PNorm2

-- | Obtains a vector in the same direction with 2-norm=1
unitary :: Vector Double -> Vector Double
unitary v = v / scalar (norm v)

-- | (rows &&& cols)
size :: Matrix t -> (Int, Int)
size m = (rows m, cols m)

-- | trans . inv
mt :: Matrix Double -> Matrix Double
mt = trans . inv

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

-- | Matrix of pairwise squared distances of row vectors
-- (using the matrix product trick in blog.smola.org)
pairwiseD2 :: Matrix Double -> Matrix Double -> Matrix Double
pairwiseD2 x y | ok = x2 `outer` oy + ox `outer` y2 - 2* x <> trans y
               | otherwise = error $ "pairwiseD2 with different number of columns: "
                                   ++ show (size x) ++ ", " ++ show (size y)
  where
    ox = one (rows x)
    oy = one (rows y)
    oc = one (cols x)
    one k = constant 1 k
    x2 = x * x <> oc
    y2 = y * y <> oc
    ok = cols x == cols y

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

-- | outer products of rows
rowOuters :: Matrix Double -> Matrix Double -> Matrix Double
rowOuters a b = a' * b'
  where
    a' = kronecker a (ones 1 (cols b))
    b' = kronecker (ones 1 (cols a)) b

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

-- | solution of overconstrained homogeneous linear system
null1 :: Matrix Double -> Vector Double
null1 = last . toColumns . snd . rightSV

-- | solution of overconstrained homogeneous symmetric linear system
null1sym :: Matrix Double -> Vector Double
null1sym = last . toColumns . snd . eigSH'

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

vec :: Element t => Matrix t -> Vector t
-- ^ stacking of columns
vec = flatten . trans

vech :: Element t => Matrix t -> Vector t
-- ^ half-vectorization (of the lower triangular part)
vech m = join . zipWith f [0..] . toColumns $ m
  where
    f k v = subVector k (dim v - k) v

dup :: (Num t, Num (Vector t), Element t) => Int -> Matrix t
-- ^ duplication matrix (@'dup' k \<> 'vech' m == 'vec' m@, for symmetric m of 'dim' k)
dup k = trans $ fromRows $ map f es
  where
    rs = zip [0..] (toRows (ident (k^(2::Int))))
    es = [(i,j) | j <- [0..k-1], i <- [0..k-1], i>=j ]
    f (i,j) | i == j = g (k*j + i)
            | otherwise = g (k*j + i) + g (k*i + j)
    g j = v
      where
        Just v = lookup j rs

vtrans :: Element t => Int -> Matrix t -> Matrix t
-- ^ generalized \"vector\" transposition: @'vtrans' 1 == 'trans'@, and @'vtrans' ('rows' m) m == 'asColumn' ('vec' m)@
vtrans p m | r == 0 = fromBlocks . map (map asColumn . takesV (replicate q p)) . toColumns $ m
           | otherwise = error $ "vtrans " ++ show p ++ " of matrix with " ++ show (rows m) ++ " rows"
  where
    (q,r) = divMod (rows m) p


-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Kalman
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- =The Theory
-- The model for the extended Kalman filter is given by
--
-- \[
-- \begin{aligned}
-- \boldsymbol{x}_i &= \boldsymbol{a}_i(\boldsymbol{x}_{i-1}, \boldsymbol{u}_{i-1}) + \boldsymbol{\psi}_{i-1} \\
-- \boldsymbol{y}_i &= \boldsymbol{h}_i(\boldsymbol{x}_i) + \boldsymbol{\upsilon}_i
-- \end{aligned}
-- \]
--
-- where
--
-- * \(\boldsymbol{a_i}\)\ is some non-linear vector-valued possibly
-- time-varying state update function.
--
-- * \(\boldsymbol{x_i}\) is the state at time \(i\).
--
-- * \(\boldsymbol{u_i}\) is the control at time \(i\).
--
-- * \(\boldsymbol{\psi}_{i}\) are independent normally
-- distributed random variables with mean 0 representing the fact that
-- the state update is noisy: \(\boldsymbol{\psi}_{i} \sim {\cal{N}}(0,Q_i)\).
--
-- * \(\boldsymbol{h}_i\)\ is some non-linear vector-valued possibly time-varying
-- function describing how we observe the hidden state process.
--
-- * \(\boldsymbol{y_i}\) is the observable at time \(i\).
--
-- * \(\boldsymbol{\upsilon}_i\) are independent normally
-- distributed random variables with mean 0 represent the fact that
-- the observations are noisy: \(\boldsymbol{\upsilon}_{i} \sim {\cal{N}}(0,R_i)\).
--
-- Note that in most presentations of the Kalman filter (the
-- [wikipedia](https://en.wikipedia.org/wiki/Kalman_filter)
-- presentation being an exception), the state update function and the
-- observation function are taken to be constant over time as is the
-- noise for the state and the noise for
-- the observation. In symbols, \(\forall i \, \boldsymbol{a}_i = \boldsymbol{a}\)
-- for some \(\boldsymbol{a}\), \(\forall i \, \boldsymbol{h}_i = \boldsymbol{h}\)
-- for some \(\boldsymbol{h}\),
-- \(\forall i \, Q_i = Q\) for some \(Q\) and \(\forall i \, \boldsymbol{a}_i = R\)
-- for some \(R\).
--
-- We assume the whole process starts at 0 with our belief of the state (aka the
-- prior state) being given by
-- \(\boldsymbol{x}_0 \sim {\cal{N}}(\boldsymbol{\mu}_0, \Sigma_0)\)
--
-- The Kalman filtering process is a recursive procedure as follows:
--
-- (1) Make a prediction of the current system state, given our previous
-- estimation of the system state.
--
-- (2) Update our prediction, given a newly acquired measurement.
--
-- The prediction and update steps depend on our system and measurement
-- models, as well as our current estimate of the system state.
--
-- = An Extended Example
--
-- The equation of motion for a pendulum of unit length subject to
-- [Gaussian white
-- noise](https://en.wikipedia.org/wiki/White_noise#Mathematical_definitions)
-- is
--
-- \[
-- \frac{\mathrm{d}^2\alpha}{\mathrm{d}t^2} = -g\sin\alpha + w(t)
-- \]
--
-- We can discretize this via the usual [Euler method](https://en.wikipedia.org/wiki/Euler_method)
--
-- \[
-- \begin{bmatrix}
-- x_{1,i} \\
-- x_{2,i}
-- \end{bmatrix}
-- =
-- \begin{bmatrix}
-- x_{1,i-1} + x_{2,i-1}\Delta t \\
-- x_{2,i-1} - g\sin x_{1,i-1}\Delta t
-- \end{bmatrix}
-- +
-- \mathbf{q}_{i-1}
-- \]
--
-- where \(q_i \sim {\mathcal{N}}(0,Q)\) and
--
-- \[
-- Q
-- =
-- \begin{bmatrix}
-- \frac{q^c \Delta t^3}{3} & \frac{q^c \Delta t^2}{2} \\
-- \frac{q^c \Delta t^2}{2} & {q^c \Delta t}
-- \end{bmatrix}
-- \]
--
-- Assume that we can only measure the horizontal position of the
-- pendulum and further that this measurement is subject to error so that
--
-- \[
-- y_i = \sin x_i + r_i
-- \]
--
-- where \(r_i \sim {\mathcal{N}}(0,R)\).
--
-- First let's set the time step and the acceleration caused by earth's gravity.
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > import Numeric.Kalman
-- >
-- > deltaT, g :: Double
-- > deltaT = 0.01
-- > g  = 9.81
-- > bigQ :: Sym 2
-- > bigQ = sym $ matrix bigQl
-- >
-- > qc :: Double
-- > qc = 0.01
-- >
-- > bigQl :: [Double]
-- > bigQl = [ qc * deltaT^3 / 3, qc * deltaT^2 / 2,
-- >           qc * deltaT^2 / 2,       qc * deltaT
-- >         ]
-- >
-- > bigR :: Sym 1
-- > bigR  = sym $ matrix [0.1]
-- > stateUpdate :: R 2 -> R 2
-- > stateUpdate u =  vector [x1 + x2 * deltaT, x2 - g * (sin x1) * deltaT]
-- >   where
-- >     (x1, w) = headTail u
-- >     (x2, _) = headTail w
-- >
-- > observe :: R 2 -> R 1
-- > observe a = vector [sin x] where x = fst $ headTail a
-- > linearizedObserve :: R 2 -> L 1 2
-- > linearizedObserve a = matrix [cos x, 0.0] where x = fst $ headTail a
-- >
-- > linearizedStateUpdate :: R 2 -> Sq 2
-- > linearizedStateUpdate u = matrix [1.0,                    deltaT,
-- >                                   -g * (cos x1) * deltaT,    1.0]
-- >   where
-- >     (x1, _) = headTail u
--
-- Now we can create extended and unscented filters which consume a
-- single observation.
--
-- > singleEKF = runEKF (const observe) (const linearizedObserve) (const bigR)
-- >              (const stateUpdate) (const linearizedStateUpdate) (const bigQ)
-- >              undefined
--
-- > singleUKF = runUKF (const observe) (const bigR) (const stateUpdate) (const bigQ)
-- >              undefined
--
-- We start off not too far from the actual value.
--
-- > initialDist =  (vector [1.6, 0.0],
-- >                 sym $ matrix [0.1, 0.0,
-- >                               0.0, 0.1])
--
-- Using some data generated using code made available with Simo
-- Särkkä's
-- [book](http://www.cambridge.org/gb/academic/subjects/statistics-probability/applied-probability-and-stochastic-networks/bayesian-filtering-and-smoothing),
-- we can track the pendulum using the extended Kalman filter.
--
-- > multiEKF obs = scanl singleEKF initialDist (map (vector . pure) obs)
--
-- And then plot the results.
--
-- <<diagrams/src_Numeric_Kalman_diagE.svg#diagram=diagE&height=600&width=500>>
--
-- And also track it using the unscented Kalman filter.
--
-- > multiUKF obs = scanl singleUKF initialDist (map (vector . pure) obs)
--
-- And also plot the results
--
-- <<diagrams/src_Numeric_Kalman_diagU.svg#diagram=diagU&height=600&width=500>>
--
-- === Code for Plotting
--
-- The full code for plotting the results:
--
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- > import Data.Colour
-- > import Data.Colour.Names
-- > import Data.Default.Class
-- > import Control.Lens
-- >
-- > import Data.Csv
-- > import System.IO hiding ( hGetContents )
-- > import Data.ByteString.Lazy ( hGetContents )
-- > import qualified Data.Vector as V
-- >
-- > import Numeric.LinearAlgebra.Static
-- >
-- > chartEstimated :: String ->
-- >               [(Double, Double)] ->
-- >               [(Double, Double)] ->
-- >               [(Double, Double)] ->
-- >               C.Renderable ()
-- > chartEstimated title acts obs ests = C.toRenderable layout
-- >   where
-- >
-- >     actuals = C.plot_lines_values .~ [acts]
-- >             $ C.plot_lines_style  . C.line_color .~ opaque red
-- >             $ C.plot_lines_title .~ "Actual Trajectory"
-- >             $ C.plot_lines_style  . C.line_width .~ 1.0
-- >             $ def
-- >
-- >     measurements = C.plot_points_values .~ obs
-- >                  $ C.plot_points_style  . C.point_color .~ opaque blue
-- >                  $ C.plot_points_title .~ "Measurements"
-- >                  $ def
-- >
-- >     estimas = C.plot_lines_values .~ [ests]
-- >             $ C.plot_lines_style  . C.line_color .~ opaque black
-- >             $ C.plot_lines_title .~ "Inferred Trajectory"
-- >             $ C.plot_lines_style  . C.line_width .~ 1.0
-- >             $ def
-- >
-- >     layout = C.layout_title .~ title
-- >            $ C.layout_plots .~ [C.toPlot actuals, C.toPlot measurements, C.toPlot estimas]
-- >            $ C.layout_y_axis . C.laxis_title .~ "Angle / Horizontal Displacement"
-- >            $ C.layout_y_axis . C.laxis_override .~ C.axisGridHide
-- >            $ C.layout_x_axis . C.laxis_title .~ "Time"
-- >            $ C.layout_x_axis . C.laxis_override .~ C.axisGridHide
-- >            $ def
-- >
-- > diagE = do
-- >   h <- openFile "data/matlabRNGs.csv" ReadMode
-- >   cs <- hGetContents h
-- >   let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
-- >   case df of
-- >     Left _ -> error "Whatever"
-- >     Right generatedSamples -> do
-- >       let xs = take 500 (multiEKF $ V.toList $ V.map fst generatedSamples)
-- >       let mus = map (fst . headTail . fst) xs
-- >       let obs = V.toList $ V.map fst generatedSamples
-- >       let acts = V.toList $ V.map snd generatedSamples
-- >       denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >       let charte = chartEstimated "Extended Kalman Filter"
-- >                                   (zip [0,1..] acts)
-- >                                   (zip [0,1..] obs)
-- >                                   (zip [0,1..] mus)
-- >       return $ fst $ runBackend denv (C.render charte (600, 500))
-- >
-- > diagU = do
-- >   h <- openFile "data/matlabRNGs.csv" ReadMode
-- >   cs <- hGetContents h
-- >   let df = (decode NoHeader cs) :: Either String (V.Vector (Double, Double))
-- >   case df of
-- >     Left _ -> error "Whatever"
-- >     Right generatedSamples -> do
-- >       let ys = take 500 (multiUKF $ V.toList $ V.map fst generatedSamples)
-- >       let nus = map (fst . headTail . fst) ys
-- >       let obs = V.toList $ V.map fst generatedSamples
-- >       let acts = V.toList $ V.map snd generatedSamples
-- >       denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >       let charte = chartEstimated "Unscented Kalman Filter"
-- >                                   (zip [0,1..] acts)
-- >                                   (zip [0,1..] obs)
-- >                                   (zip [0,1..] nus)
-- >       return $ fst $ runBackend denv (C.render charte (600, 500))
--
-----------------------------------------------------------------------------

module Numeric.Kalman
       (
         runKF, runKFPrediction, runKFUpdate
       , runEKF, runEKFPrediction, runEKFUpdate
       , runUKF, runUKFPrediction, runUKFUpdate
       , runKS, runEKS, runUKS
       )
       where

import GHC.TypeLits
import Numeric.LinearAlgebra.Static
import Data.Maybe ( fromJust )
import Prelude hiding ((<>))

-- | Given our system model and our previous estimate of the system state,
-- we generate a prediction of the current system state by taking the
-- mean of our estimated state as our point estimate and evolving it one
-- step forward with the system evolution function.  In other words, we
-- predict that the current system state is the minimum mean squared
-- error (MMSE) state, which corresponds to the maximum a posteriori
-- (MAP) estimate, from the previous estimate evolved one step by the
-- system evolution function of our system model.
--
-- We also generate a predicted covariance matrix by updating the current
-- covariance and adding any noise estimate.  Updating the current
-- covariance requires the linearized version of the system evolution,
-- consequently the deviation of the actual system evolution from its
-- linearization is manifested in the newly predicted covariance.
--
-- Taken together, the prediction step generates a new Gaussian
-- distribution of predicted system states with mean the evolved MAP
-- state from the previous estimate.
runEKFPrediction
  :: KnownNat n
  => (a -> R n -> R n)  -- ^ System evolution function
                        -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                        -- note the order of the control and the
                        -- input
  -> (a -> R n -> Sq n) -- ^ Linearization of the system evolution at a point \(\frac{\partial \boldsymbol{a}_i}{\partial \boldsymbol{x}}\big|_{{\boldsymbol{u}}, \, {\boldsymbol{x}}}\)
  -> (a -> Sym n)       -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                  -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)       -- ^ Current estimate \((\hat{\boldsymbol{x}}_{n-1}, \hat{\boldsymbol{\Sigma}}_{n-1})\)
  -> (R n, Sym n)       -- ^ New prediction \(({\boldsymbol{x}_i^\flat}, {\boldsymbol{\Sigma}_i^\flat})\)

runEKFPrediction evolve linEvol sysCov input (estMu, estCov) =
  (predMu, predCov)
  where
    predMu  = evolve input estMu
    predCov = sym $ lin <> unSym estCov <> tr lin + unSym (sysCov input)
    lin     = linEvol input estMu

runKFPrediction
  :: (KnownNat n)
  => (a -> R n -> Sq n) -- ^ Linear system evolution at a point
                        -- \(\boldsymbol{a}_i(\boldsymbol{u}_{i-1}, \boldsymbol{x}_{i-1})\),
                        -- note the order of the control and the
                        -- input
  -> (a -> Sym n)       -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                  -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)       -- ^ Current estimate
                        -- \((\hat{\boldsymbol{x}}_{i-1}, \hat{\boldsymbol{\Sigma}}_{i-1})\)
  -> (R n, Sym n)       -- ^ New prediction
                        -- \(({\boldsymbol{x}_i^\flat}, {\boldsymbol{\Sigma}_i^\flat})\)
runKFPrediction linEvol =
  runEKFPrediction (\inp sys -> linEvol inp sys #> sys) linEvol

-- | After a new measurement has been taken, we update our original
-- prediction of the current system state using the result of this
-- measurement.  This is accomplished by looking at how much the
-- measurement deviated from our prediction, and then updating our
-- estimated state accordingly.
--
-- This step requires the linearized measurement transformation.
runEKFUpdate
  :: (KnownNat m, KnownNat n)
  => (a -> R n -> R m)   -- ^ System measurement function \(\boldsymbol{h}_i\)
  -> (a -> R n -> L m n) -- ^ Linearization of the observation at a point
                         -- \(\frac{\partial \boldsymbol{h}_i}{\partial \boldsymbol{x}}\big|_{\boldsymbol{x}}\)
  -> (a -> Sym m)        -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> a                   -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)        -- ^ Current prediction \(({\boldsymbol{x}}_i^\flat, {\boldsymbol{\Sigma}}_i^\flat)\)
  -> R m                 -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)        -- ^ Updated prediction \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)

runEKFUpdate measure linMeas measCov input (predMu, predCov') newMeas =
  (newMu, newCov)
  where
    newMu  = predMu + kkMat #> voff
    newCov = sym $ predCov - kkMat <> skMat <> tr kkMat

    predCov = unSym predCov'

    lin   = linMeas input predMu
    voff  = newMeas - measure input predMu
    skMat = lin <> predCov <> tr lin + unSym (measCov input)
    kkMat = predCov <> tr lin <> unsafeInv skMat

runKFUpdate
  :: (KnownNat m, KnownNat n)
  => (a -> R n -> L m n) -- ^ Linear measurement operator at a point \(\boldsymbol{h}_i\)
  -> (a -> Sym m)        -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> a                   -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)        -- ^ Current prediction \(({\boldsymbol{x}}_i^\flat, {\boldsymbol{\Sigma}}_i^\flat)\)
  -> R m                 -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)        -- ^ Updated prediction \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)
runKFUpdate linMeas =
  runEKFUpdate (\inp sys -> linMeas inp sys #> sys) linMeas

-- | Here we combine the prediction and update setps applied to a new
-- measurement, thereby creating a single step of the (extended) Kalman
-- filter.
runEKF
  :: (KnownNat m, KnownNat n)
  => (a -> R n -> R m)   -- ^ System measurement function \(\boldsymbol{h}_i\)
  -> (a -> R n -> L m n) -- ^ Linearization of the measurement at a point
                         -- \(\frac{\partial \boldsymbol{h}_i}{\partial \boldsymbol{x}}\big|_{\boldsymbol{x}}\)
  -> (a -> Sym m)        -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> (a -> R n -> R n)   -- ^ System evolution function
                         -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                         -- note the order of the control and the
                         -- input
  -> (a -> R n -> Sq n)  -- ^ Linearization of the system evolution at a point  \(\frac{\partial \boldsymbol{a}_i}{\partial \boldsymbol{x}}\big|_{{\boldsymbol{u}}, \, {\boldsymbol{x}}}\)
  -> (a -> Sym n)        -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                   -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)        -- ^ Current estimate
                         -- \((\hat{\boldsymbol{x}}_{i-1}, \hat{\boldsymbol{\Sigma}}_{i-1})\)
  -> R m                 -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)        -- ^ New (filtered) estimate \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)
runEKF measure linMeas measCov
  evolve linEvol sysCov
  input estSys newMeas = updatedEstimate
  where
    predictedSystem = runEKFPrediction evolve linEvol sysCov input estSys
    updatedEstimate = runEKFUpdate measure linMeas measCov input predictedSystem newMeas

-- | The ordinary Kalman filter is a special case of the extended Kalman
-- filter, when the state update and measurement transformations are
-- already linear.
runKF
  :: (KnownNat m, KnownNat n)
  => (a -> R n -> L m n) -- ^ Linear measurement operator at a point \(\boldsymbol{h}_i\)
  -> (a -> Sym m)        -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> (a -> R n -> Sq n)  -- ^ Linear system evolution at a point
                         -- \(\boldsymbol{a}_i(\boldsymbol{u}_{i-1}, \boldsymbol{x}_{i-1})\),
                         -- note the order of the control and the
                         -- input
  -> (a -> Sym n)        -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                   -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)        -- ^ Current estimate
                         -- \((\hat{\boldsymbol{x}}_{i-1}, \hat{\boldsymbol{\Sigma}}_{i-1})\)
  -> R m                 -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)        -- ^ New (filtered) estimate \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)
runKF linMeas measCov
  linEvol sysCov
  input estSys newMeas = updatedEstimate
  where
    predictedSystem = runKFPrediction linEvol sysCov input estSys
    updatedEstimate = runKFUpdate linMeas measCov input predictedSystem newMeas

runUKFPrediction
  :: KnownNat n
  => (a -> R n -> R n) -- ^ System evolution function at a point
                       -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                       -- note the order of the control and the
                       -- input
  -> (a -> Sym n)      -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                 -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)      -- ^ Current estimate \((\hat{\boldsymbol{x}}_{n-1}, \hat{\boldsymbol{\Sigma}}_{n-1})\)
  -> (R n, Sym n)      -- ^ Prediction \(({\boldsymbol{x}_i^\flat}, {\boldsymbol{\Sigma}_i^\flat})\)

runUKFPrediction evolve sysCov input (estMu, estCov) =
  (predMu, predCov)
  where
    predMu  = weightM0 * estMu' +
              sum (map (weightCM *) sigmaPoints')

    predCov = sym $ col (weightC0 * (estMu' - predMu)) <> row (estMu' - predMu) +
              sum (map (\sig ->
                          col (weightCM * (sig - predMu)) <> row (sig - predMu)
                       )
                   sigmaPoints') +
              unSym (sysCov input)

    estMu' = evolve input estMu

    sqCov  = chol estCov
    sqRows = map (* sqrt 3) $ toRows sqCov
    sigmaPoints = map (estMu +) sqRows ++ map (estMu -) sqRows -- 2n points
    sigmaPoints' = map (evolve input) sigmaPoints

    -- hand-tuned weights, more explanation required
    n = fromIntegral $ size estMu
    weightM0 = 1 - n / 3
    weightCM = 1 / 6
    weightC0 = 4 - n / 3 - 3 / n

runUKFUpdate
  :: (KnownNat n, KnownNat m)
  => (a -> R n -> R m) -- ^ Measurement transformation \(\boldsymbol{h}_i\)
  -> (a -> Sym m)      -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> a                 -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)      -- ^ Current prediction \(({\boldsymbol{x}}_i^\flat, {\boldsymbol{\Sigma}}_i^\flat)\)
  -> R m               -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)      -- ^ Updated prediction \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)

runUKFUpdate measure measCov input (predMu, predCov) newMeas =
  (newMu, newCov)
  where
    newMu  = predMu + kkMat #> (newMeas - upMu)
    newCov = sym $ (unSym predCov) - kkMat <> skMat <> tr kkMat

    predMu' = measure input predMu

    kkMat = ckMat <> unsafeInv skMat
    upMu  = weightM0 * predMu' +
            sum (map (weightCM *) sigmaPoints')

    skMat = (unSym $ measCov input) +
            (col $ weightC0 * (predMu' - upMu)) <> (row $ predMu' - upMu) +
            sum (map (\sig ->
                        (col $ weightCM * (sig - upMu)) <> (row $ sig - upMu)
                     )
                 sigmaPoints')

    ckMat = sum $
            zipWith (\preds meas ->
                       (col $ weightCM' * (preds - predMu)) <> (row $ meas - upMu)
                    )
            sigmaPoints sigmaPoints'

    sqCov   = chol predCov
    sqRows = map (* sqrt 3) $ toRows sqCov
    sigmaPoints = map (predMu +) sqRows ++ map (predMu -) sqRows -- 2n points
    sigmaPoints' = map (measure input) sigmaPoints

    -- hand tuned weights, more exlanation required
    n = fromIntegral $ size predMu
    weightM0 = 1 - n / 3
    weightCM = 1 / 6
    weightCM' = 1 / 6
    weightC0 = 4 - n / 3 - 3 / n

runUKF
  :: (KnownNat m, KnownNat n)
  => (a -> R n -> R m) -- ^ System measurement function \(\boldsymbol{h}_i\)
  -> (a -> Sym m)      -- ^ Covariance matrix encoding measurement noise \(R_i\)
  -> (a -> R n -> R n) -- ^ System evolution function
                       -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                       -- note the order of the control and the
                       -- input
  -> (a -> Sym n)      -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                 -- ^ Dynamical input \(\boldsymbol{u}_{i-1}\)
  -> (R n, Sym n)      -- ^ Current estimate
                       -- \((\hat{\boldsymbol{x}}_{i-1}, \hat{\boldsymbol{\Sigma}}_{i-1})\)
  -> R m               -- ^ New measurement \(\boldsymbol{y}_i\)
  -> (R n, Sym n)      -- ^ New (filtered) estimate \((\hat{\boldsymbol{x}}_i, \hat{\boldsymbol{\Sigma}}_i)\)
runUKF measure measCov
  evolve sysCov
  input estSys newMeas = updatedEstimate
  where
    predictedSystem = runUKFPrediction evolve sysCov input estSys
    updatedEstimate = runUKFUpdate measure measCov input predictedSystem newMeas

-- | The Kalman smoothing process (sometimes also called the
-- Rauch-Tung-Striebel smoother or RTSS) is a recursive procedure for
-- improving previous estimates of the system state in light of more
-- recently obtained measurements.  That is, for a given state estimate
-- that was produced (by a Kalman filter) using only the available
-- historical data, we improve the estimate (smoothing it) using
-- information that only became available later, via additional
-- measurements after the estimate was originally made.  The result is an
-- estimate of the state of the system at that moment in its evolution,
-- taking full advantage of hindsight obtained from observations of the
-- system in the future of the current step.
--
-- Consequently, the recursive smoothing procedure progresses backwards
-- through the state evolution, beginning with the most recent
-- observation and updating past observations given the newer
-- information.  The update is made by measuring the deviation between
-- what the system actually did (via observation) and what our best
-- estimate of the system at that time predicted it would do, then
-- adjusting the current estimate in view of this deviation.
runEKS
  :: (KnownNat n)
  => (a -> R n -> R n)  -- ^ System evolution function
                        -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                        -- note the order of the control and the
                        -- input
  -> (a -> R n -> Sq n) -- ^ Linearization of the system evolution at a point  \(\frac{\partial \boldsymbol{a}_i}{\partial \boldsymbol{x}}\big|_{{\boldsymbol{u}}, \, {\boldsymbol{x}}}\)
  -> (a -> Sym n)       -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                  -- ^ Dynamical input \(\boldsymbol{u}_{i}\)
  -> (R n, Sym n)       -- ^ Future smoothed estimate
                        -- \((\bar{\boldsymbol{x}}_{i+1}, \bar{\boldsymbol{\Sigma}}_{i+1})\)
  -> (R n, Sym n)       -- ^ Present filtered estimate
                        -- \((\hat{\boldsymbol{x}}_{i}, \hat{\boldsymbol{\Sigma}}_{i})\)
  -> (R n, Sym n)       -- ^ Present smoothed estimate
                        -- \((\bar{\boldsymbol{x}}_{i}, \bar{\boldsymbol{\Sigma}}_{i})\)

runEKS sysEvol linEvol sysCov input (futMu, futCov') (curMu, curCov') =
  (smMu, smCov)
  where
    smMu  = curMu + gkMat #> (futMu - predMu)
    smCov = sym $ curCov + gkMat <> (futCov - predCov) <> tr gkMat

    futCov = unSym futCov'

    curCov = unSym curCov'

    (predMu, predCov') = runEKFPrediction sysEvol linEvol sysCov input (curMu, curCov')
    predCov  = unSym predCov'

    gkMat = curCov <> tr lin <> unsafeInv predCov
    lin   = linEvol input curMu

runKS
  :: (KnownNat n)
  => (a -> R n -> Sq n) -- ^ Linear system evolution at a point
                        -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
                        -- note the order of the control and the
                        -- input
  -> (a -> Sym n)       -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                  -- ^ Dynamical input \(\boldsymbol{u}_{i}\)
  -> (R n, Sym n)       -- ^ Future smoothed estimate
                        -- \((\bar{\boldsymbol{x}}_{i+1}, \bar{\boldsymbol{\Sigma}}_{i+1})\)
  -> (R n, Sym n)       -- ^ Present filtered estimate
                        -- \((\hat{\boldsymbol{x}}_{i}, \hat{\boldsymbol{\Sigma}}_{i})\)
  -> (R n, Sym n)       -- ^ Present smoothed estimate
                        -- \((\bar{\boldsymbol{x}}_{i}, \bar{\boldsymbol{\Sigma}}_{i})\)
runKS linEvol =
  runEKS (\inp sys -> linEvol inp sys #> sys) linEvol

-- | Unscented Kalman smoother
runUKS
  :: (KnownNat n)
  => (a -> R n -> R n) -- ^ System evolution function
                       -- \(\boldsymbol{a}_i(\boldsymbol{u}, \boldsymbol{x})\),
  -> (a -> Sym n)      -- ^ Covariance matrix encoding system evolution noise \(Q_i\)
  -> a                 -- ^ Dynamical input \(\boldsymbol{u}_{i}\)
  -> (R n, Sym n)      -- ^ Future smoothed estimate
                        -- \((\bar{\boldsymbol{x}}_{i+1}, \bar{\boldsymbol{\Sigma}}_{i+1})\)
  -> (R n, Sym n)      -- ^ Present filtered estimate
                        -- \((\hat{\boldsymbol{x}}_{i}, \hat{\boldsymbol{\Sigma}}_{i})\)
  -> (R n, Sym n)      -- ^ Present smoothed estimate
                       -- \((\bar{\boldsymbol{x}}_{i}, \bar{\boldsymbol{\Sigma}}_{i})\)
runUKS evolve sysCov input (futMu, futCov') (curMu, curCov') =
  (smMu, smCov)
  where
    smMu  = curMu + gkMat #> (futMu - predMu)
    smCov = sym $ curCov + gkMat <> (futCov - predCov) <> tr gkMat

    futCov = unSym futCov'

    curCov = unSym curCov'

    (predMu, predCov') = runUKFPrediction evolve sysCov input (curMu, curCov')
    predCov  = unSym predCov'

    gkMat = dkMat <> unsafeInv predCov

    dkMat = sum $
            zipWith (\pres fut ->
                       (col $ weightCM * (pres - curMu)) <> (row $ fut - predMu)
                    )
            sigmaPoints sigmaPoints'

    sqCov   = chol predCov'
    sqRows = map (* sqrt 3) $ toRows sqCov
    sigmaPoints = map (curMu +) sqRows ++ map (curMu -) sqRows -- 2n points
    sigmaPoints' = map (evolve input) sigmaPoints

    -- hand tuned weights, more exlanation required
    weightCM = 1 / 6

unsafeInv :: KnownNat n => Sq n -> Sq n
unsafeInv m = fromJust $ linSolve m eye


{-# LANGUAGE MultiParamTypeClasses #-}

module Q.Options.Bachelier
  (
    Bachelier(..)
  , euOption
  , eucall
  , euput
  ) where

import           Control.Monad.State
import           Data.Random (RVar, stdNormal)
import           Q.MonteCarlo
import           Q.Options
import           Statistics.Distribution (cumulative)
import           Statistics.Distribution.Normal (standard)

data Bachelier = Bachelier YearFrac Forward Rate Vol deriving stock Show

-- | European option valuation with bachelier model.
euOption ::  Bachelier -> OptionType -> Strike -> Valuation
euOption (Bachelier (YearFrac t ) (Forward f) (Rate r) (Vol sigma)) cp (Strike k)
  = Valuation premium delta vega gamma where
    premium = Premium $ df * (q*(f - k)*n(q*d1) + sigma*sqrt(t)/sqrt2Pi * (exp(-0.5 *d1 * d1)))
    delta   = Delta   $ df * n (q * d1)
    vega    = Vega    $ df * (sqrt t) / sqrt2Pi * (exp (-0.5 * d1 * d1))
    gamma   = Gamma   $ (df/(sigma * (sqrt t)))*(recip sqrt2Pi)*(exp(-0.5 *d1 * d1))
    d1 = (f - k) / (sigma * sqrt(t))
    q = cpi cp
    sqrt2Pi = sqrt (2*pi)
    df =  exp $ (-r) * t
    n = cumulative standard

-- | see 'euOption'
euput :: Bachelier -> Strike -> Valuation
euput b =  euOption b Put

-- | see 'euOption'
eucall :: Bachelier -> Strike -> Valuation
eucall b = euOption b Call

instance Model Bachelier Bachelier where
  discountFactor (Bachelier _ _ r _) t1 t2 = return $ exp (scale dt r)
    where dt = t2 - t1

  evolve _ (YearFrac t) = do
    (YearFrac t0, (Bachelier _ (Forward f0) (Rate r) (Vol sigma))) <- get
    let dt = t - t0
    dW <- lift stdNormal::StateT (YearFrac, Bachelier) RVar Double
    let ft = f0 * exp (r * dt) + sqrt(sigma*sigma/2*r * (exp (2 * r * dt) - 1)) * dW
    put (YearFrac t, Bachelier (YearFrac t) (Forward ft) (Rate r) (Vol sigma))
    return ft


-- | The method of moments can be used to estimate a number of commonly used distributions.  This module is still under construction as I work out the best way to handle morphisms from the Moments3 type to types of other distributions.  For more information, see the wikipedia entry: <https://en.wikipedia.org/wiki/Method_of_moments_(statistics)>

module HLearn.Models.Distributions.Univariate.Geometric
    ( Geometric
    )
    where

import Control.DeepSeq
import GHC.TypeLits
import qualified Data.Vector.Unboxed as U
import Data.Vector.Unboxed.Deriving

import qualified Statistics.Distribution as S
import Statistics.Distribution.Geometric

import HLearn.Algebra
import HLearn.Models.Distributions.Common
import HLearn.Models.Distributions.Univariate.Internal.Moments
import HLearn.Models.Distributions.Visualization.Gnuplot

-------------------------------------------------------------------------------
-- Geometric

newtype Geometric sample prob = Geometric {  moments :: (Moments3 sample) }
    deriving (Read,Show,Eq,Ord,Monoid,Group)

instance (Num sample) => HomTrainer (Geometric sample prob) where
    type Datapoint (Geometric sample prob) = sample
    train1dp dp = Geometric $ train1dp dp

instance (Num sample) => Probabilistic (Geometric sample prob) where
    type Probability (Geometric sample prob) = prob

instance (Integral sample, Floating prob) => PDF (Geometric sample prob) where
    pdf dist dp = p*(1-p)^^dp
        where
            p = geo_p dist

instance
    ( PDF (Geometric sample prob)
    , Show prob
    , Show sample
    , Ord sample
    , Ord prob
    , Fractional prob
    , RealFrac prob
    , Integral sample
    ) => PlottableDistribution (Geometric sample prob)
        where

    plotType _ = Points

    samplePoints dist = [min..max]
        where
            min = 0
            max = maximum [20,round $ 3*(fromIntegral $ mean dist)]

geo_p :: (Fractional prob, Integral sample) => Geometric sample prob -> prob
geo_p (Geometric dist) = 1/((fromIntegral $ m1 dist)/(fromIntegral $ m0 dist) +1)

instance (Integral sample, RealFrac prob) => Mean (Geometric sample prob) where
    mean dist = round $ 1/(geo_p dist)

instance (Integral sample, Fractional prob) => Variance (Geometric sample prob) where
    variance dist = (1-p)/(p*p)
        where
            p = geo_p dist


{-# OPTIONS_GHC -Wall #-}

module Numeric.FFT.HMatrix.Matrix (
  -- * Complex-to-complex transforms
  dft
  , idft
    -- * Real-to-complex transforms
  , dftR2C
  , dftC2R
    -- * Real-to-real transforms
    -- ** Discrete cosine transforms
  , dct1
  , dct2
  , dct3
  , dct4
    -- ** Discrete sine transforms
  , dst1
  , dst2
  , dst3
  , dst4
  ) where

import Foreign.Storable ( Storable )
import Numeric.LinearAlgebra.Data ( Matrix, Complex )
import Numeric.LinearAlgebra ( Element )
import qualified Numeric.LinearAlgebra.Data as M
import Numeric.FFT.Vector.Unnormalized ( Transform )
import qualified Numeric.FFT.Vector.Unnormalized as FFT

transResize :: (Storable a, Element b) =>
               (Matrix b -> Int)
            -> ((Int, Int) -> Transform b a)
            -> Matrix b
            -> Matrix a
transResize calcCols f m = M.reshape newcols . FFT.run (f (r, c)) $ M.flatten m
  where
    c = M.cols m
    r = M.rows m
    newcols = calcCols m

trans :: (Storable a, Element b) =>
         ((Int, Int) -> Transform b a)
      -> Matrix b
      -> Matrix a
trans = transResize (M.cols)

-- |
-- >>> M.toLists . dft $ M.fromLists [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
-- [[45.0 :+ 0.0,(-4.5) :+ 2.598076211353316,(-4.5) :+ (-2.598076211353316)],[(-13.5) :+ 7.794228634059947,0.0 :+ 0.0,0.0 :+ 0.0],[(-13.5) :+ (-7.794228634059947),0.0 :+ 0.0,0.0 :+ 0.0]]
dft :: Matrix (Complex Double) -> Matrix (Complex Double)
dft = trans FFT.dft2d

-- |
-- >>> M.toLists . idft $ M.fromLists [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
-- [[45.0 :+ 0.0,(-4.5) :+ (-2.598076211353316),(-4.5) :+ 2.598076211353316],[(-13.5) :+ (-7.794228634059947),0.0 :+ 0.0,0.0 :+ 0.0],[(-13.5) :+ 7.794228634059947,0.0 :+ 0.0,0.0 :+ 0.0]]
idft :: Matrix (Complex Double) -> Matrix (Complex Double)
idft = trans FFT.idft2d

-- |
-- >>> M.toLists . dftR2C $ M.fromLists [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
-- [[45.0 :+ 0.0,(-4.5) :+ 2.598076211353316],[(-13.5) :+ 7.794228634059947,0.0 :+ 0.0],[(-13.5) :+ (-7.794228634059947),0.0 :+ 0.0]]
dftR2C :: Matrix Double -> Matrix (Complex Double)
dftR2C = transResize (\m -> M.cols m `div` 2 + 1) FFT.dftR2C2D

-- |
--
-- >>> let m = M.fromLists [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
-- >>> let mtest = M.fromLists [[33,-3,-3,-9],[3,-3,-9,-3],[3,0,0,0]]
-- >>> M.maxElement (mtest - dftC2R m) < 1e-12
-- True
dftC2R :: Matrix (Complex Double) -> Matrix Double
dftC2R = transResize (\m -> 2 * (M.cols m - 1)) FFT.dftC2R2D

dct1 :: Matrix Double -> Matrix Double
dct1 = trans FFT.dct1_2D

dct2 :: Matrix Double -> Matrix Double
dct2 = trans FFT.dct2_2D

dct3 :: Matrix Double -> Matrix Double
dct3 = trans FFT.dct3_2D

dct4 :: Matrix Double -> Matrix Double
dct4 = trans FFT.dct4_2D

dst1 :: Matrix Double -> Matrix Double
dst1 = trans FFT.dst1_2D

dst2 :: Matrix Double -> Matrix Double
dst2 = trans FFT.dst2_2D

dst3 :: Matrix Double -> Matrix Double
dst3 = trans FFT.dst3_2D

dst4 :: Matrix Double -> Matrix Double
dst4 = trans FFT.dst4_2D


-- This code was automatically generated by lhs2tex --code, from the file
-- HSoM/SpectrumAnalysis.lhs.  (See HSoM/MakeCode.bat.)

module Euterpea.Music.Signal.SpectrumAnalysis where

import Euterpea

import Data.Complex
import Data.List
-- import Math.FFT
dft :: RealFloat a => [Complex a] -> [Complex a]
dft xs =
  let  lenI = length xs
       lenR = fromIntegral lenI
       lenC = lenR :+ 0
  in [  let i = -2 * pi * fromIntegral k / lenR
        in (1/lenC) * sum [  (xs!!n) * exp (0 :+ i * fromIntegral n)
                             | n <- [0,1..lenI-1] ]
        | k <- [0,1..lenI-1] ]
mkTerm :: Int -> Double -> [Complex Double]
mkTerm num n = let f = 2 * pi / fromIntegral num
               in map (:+ 0) [  sin (n * f * fromIntegral i) / n
                                | i <- [0,1..num-1] ]

mkxa, mkxb, mkxc :: Int-> [Complex Double]
mkxa num = mkTerm num 1
mkxb num = zipWith (+) (mkxa num) (mkTerm num 3)
mkxc num = zipWith (+) (mkxb num) (mkTerm num 5)
printComplexL :: [Complex Double] -> IO ()
printComplexL xs  =
  let  f (i,rl:+im) =
            do  putStr (spaces (3 - length (show i))  )
                putStr (show i       ++ ":  ("        )
                putStr (niceNum rl  ++ ", "           )
                putStr (niceNum im  ++ ")\n"          )
  in mapM_ f (zip [0..length xs - 1] xs)

niceNum :: Double -> String
niceNum d =
  let  d' = fromIntegral (round (1e10 * d)) / 1e10
       (dec, fra)  = break (== '.') (show d')
       (fra',exp)  = break (== 'e') fra
  in  spaces (3  - length dec) ++ dec ++ take 11 fra'
      ++ exp ++ spaces (12 - length fra' - length exp)

spaces :: Int -> String
spaces  n = take n (repeat ' ')
mkPulse :: Int -> [Complex Double]
mkPulse n = 100 : take (n-1) (repeat 0)
x1 num = let f = pi * 2 * pi / fromIntegral num
         in map (:+ 0) [  sin (f * fromIntegral i)
                          | i <- [0,1..num-1] ]
mkPolars :: [Complex Double] -> [Complex Double]
mkPolars = map ((\(m,p)-> m:+p) . polar)


{-# LANGUAGE FlexibleContexts #-}
module Workflow.Toolset
    ( Array1D
    , mean, var, characteristicF, distributionF
    , fft, ifft, distToChar, charToDist, linSpace, asort, azipWith
    ) where

import Data.Array.Unboxed
import Data.Complex
import Data.Bits as Bit
import Data.Array.ST
import Control.Monad.ST
import Data.Array.Base

-- | Array indexed by integers
type Array1D = UArray Int

-- | mean estimation.
-- Complexity O(n).
mean :: (IArray UArray a, Fractional a) => Array1D a -> a
mean arr = s / size where
    (i1,i2) = bounds arr
    size = fromIntegral (i2 - i1 + 1)
    s = afoldl (+) 0 arr

-- | Unbiased variance estimation.
-- Complexity O(n).
var :: (IArray UArray a, Fractional a) => Array1D a -> a
var arr = s / size where
    (i1,i2) = bounds arr
    size = fromIntegral (i2 - i1)
    me = mean arr
    s = afoldl (\acc x -> acc + (x - me)^(2::Int)) 0 arr

-- |
-- Empirical characteristic function of the distribuition represented by array (sample).
-- Complexity O(n).
characteristicF :: (IArray UArray a, RealFloat a) => Array1D a -> [a] -> [Complex a]
characteristicF _ [] = []
characteristicF arr ts = size `seq` ts `seq` map ((/size) . f) ts
  where
    (i1,i2) = bounds arr
    size = fromIntegral (i2 - i1 + 1)
    f t = t `seq` afoldl g 0 arr where
        g acc x = x `seq` acc `seq` acc + (cis $ x*t)

-- |
-- Empirical cumulative distribuition function estimation
-- of the distribuition represented by !sorted! array (sample).
-- Uses linear approximation between points.
-- Complexity O(n).
distributionF :: (IArray UArray a, RealFloat a) => Array1D a -> [a] -> [a]
distributionF _ [] = []
distributionF arr tt@(t:ts) = if t <= (arr ! i1)
        then 0:(distributionF arr ts)
        else go (i1+1) tt
  where
    (i1,i2) = bounds arr
    n = fromIntegral (i2 - i1 + 1)
    go _ [] = []
    go i xt@(x:xs) | i == i2     = take (length xt) . repeat $ 1
                   | arr ! i > x = ( (fromIntegral (i - i1) +
                     (x - (arr ! (i-1))) / ((arr ! i) - (arr ! (i-1)))
                     )/ n):(go i xs)
                   | otherwise   = go (i+1) tt

-- | Converts CDF into Char Func., given grid (x,t)
-- Note, length of t is n/2, and length of x is n
-- resulting array also has length of n/2, but CDF array has length of n
distToChar :: ([Double],[Double]) -> [Double] -> [Complex Double]
distToChar (x,t) f = zipWith (\t' f' -> f' * (0 :+ t'*c) * (cis $ t'*a)) t phi'
  where
    (a,b) = (head x, last x)
    c = (a-b)/(n-1)
    n = fromIntegral . flip Bit.shiftL 1 . length $ t  :: Double
    phi' = reverse . fft True $ zipWith (\f' x' -> (f' - (x'-a)/(b-a)) :+ 0) f x

-- | Converts Char Func. into CDF, given grid (x,t)
-- Note, length of t is n/2, and length of x is n
-- resulting array also has length of n, but Char Func. array has length of n/2
charToDist :: ([Double],[Double]) -> [Complex Double] -> [Double]
charToDist (x,t) phi = postPass
    . shiftUp
    $ zipWith (\x' ps -> (realPart ps) + (x'-a)/(b-a)) x $ ifft False psi
  where
    (a,b) = (head x, last x)
    c = (n-1)/(a-b)
    n' = length $ t  :: Int
    n = fromIntegral . flip Bit.shiftL 1 $ n' :: Double
    phi' = zipWith (\t' ph -> conjugate ph * (0 :+ c/t') * (cis $ t'*a)) t phi
    psi = (:) 0 $ (++) phi' . tail . reverse . map conjugate $ phi'
    shiftUp xxs'@(x':_) = map (\z -> z-x') xxs'
    shiftUp [] = []
    postPass xs = map (min 1 . max 0) xs

-- | Fast Fourier Transformation.
-- Only accepts sizes n = 2^k.
-- Boolean parameter is whether to return full array, or only second (right) part of it.
fft :: RealFloat a => Bool -> [Complex a] -> [Complex a]
fft halve xs = fft' xs w halve
  where
    n = length $ xs
    w = cis $ -2 * pi / fromIntegral n

-- | Inverse Fast Fourier Transformation.
-- Only accepts sizes n = 2^k.
-- Boolean parameter is whether to return full array, or only second (right) part of it.
ifft :: RealFloat a => Bool -> [Complex a] -> [Complex a]
ifft halve xs = map (/ fromIntegral n) $ fft' xs w halve
  where
    n = length $ xs
    w = cis $ 2 * pi / fromIntegral n

-- Fourier transform (both, forward and inverse)
fft' :: Num t => [t] -> t -> Bool -> [t]
fft' []  _ _ = []
fft' [x] _ _ = [x]
fft' xs  w halve = xs `seq` ys `seq` zs `seq` w `seq` halve `seq`
    if halve then zipWith (-) ys ts
             else zipWith (+) ys ts ++ zipWith (-) ys ts
    where wx  = w*w
          ys = fft' evens wx False
          zs = fft' odds wx False
          (evens, odds)  = split xs
          split []       = ([], [])
          split [x]       = ([x], [])
          split (x:y:ss) = (x:xt, y:yt) where (xt, yt) = split ss
          ts = zipWith (*) zs . iterate (w*) $ 1

-- |
-- create x and t grids from bounds on x;
-- Arguments: k power of 2 - number of points: n = 2^k; bounds
linSpace :: Int -> (Double,Double) -> ([Double],[Double])
linSpace k (a,b) = (map (\i' -> a + i' * (b-a) / (n-1)) i, map (\j' -> 2*pi*j'*(n-1)/n/(b-a)) j)
  where
    n = fromIntegral $ (Bit.shiftL 1 k :: Int ):: Double
    i = [0..n-1] :: [Double]
    j = [1..n/2] :: [Double]

-- import Data.Foldable
-- | Left fold on an array
afoldl :: (IArray UArray a) => (b -> a -> b) -> b -> Array1D a -> b
afoldl f x0 arr = f `seq` x0 `seq` arr `seq` go x0 start where
    (start, end) = bounds arr
    g x y = x `seq` y `seq` f x y
    go x i | i > end   = x
           | otherwise = x `seq` go (g x (arr ! i)) (i+1)

---- | Right fold on an array
--afoldr :: (IArray UArray a) => (a -> b -> b) -> b -> Array1D a -> b
--afoldr f x0 arr = f `seq` x0 `seq` arr `seq` go x0 end where
--    (start, end) = bounds arr
--    go x i | i < start = x
--           | otherwise = x `seq` go (f (arr ! i) x) (i-1)

-- | zipWith for arrays
azipWith :: (IArray UArray a, IArray UArray b, IArray UArray c) => (a -> b -> c) -> Array1D a -> Array1D b -> Array1D c
azipWith f a b = a `seq` b `seq` listArray bnds [f (a ! i) (b ! i) | i <- range bnds] where
    bnds = bounds a

-- | working implementation of mergesort; uses the same array of additional memory, slow just as original Data.List.sort
asort :: Array1D Double -> Array1D Double
asort arr = if size < 2 then arr
                        else runSTUArray result
  where
  (start,end) = bounds arr :: (Int, Int)
  size = end-start+1 :: Int
  result = do
    xRight <- (thaw :: Array1D Double -> ST s (STUArray s Int Double)) arr
    xLeft <- (thaw :: Array1D Double -> ST s (STUArray s Int Double)) arr
    let f _ 1 _ = return ()
        f i 2 right = let x = if right then xLeft else xRight
                      in do
          x1 <- unsafeRead x i
          x2 <- unsafeRead x (i+1)
          if x1 > x2 then do
            unsafeWrite x i x2
            unsafeWrite x (i+1) x1
          else return ()
        f i 3 right = let x = if right then xLeft else xRight in do
          x1 <- unsafeRead x i
          x2 <- unsafeRead x (i+1)
          x3 <- unsafeRead x (i+2)
          case (x1 > x2, x1 > x3, x2 > x3) of
            (True,False,False) -> do
                unsafeWrite x i x2
                unsafeWrite x (i+1) x1
            (True,True,True) -> do
                unsafeWrite x i x3
                unsafeWrite x (i+2) x1
            (False,False,True) -> do
                unsafeWrite x (i+1) x3
                unsafeWrite x (i+2) x2
            (False,True,True) -> do
                unsafeWrite x i x3
                unsafeWrite x (i+1) x1
                unsafeWrite x (i+2) x2
            (True,True,False) -> do
                unsafeWrite x i x2
                unsafeWrite x (i+1) x3
                unsafeWrite x (i+2) x1
            _ -> return ()
        f i n right = do
            f i n1 (not right)
            f (i+n1) n2 (not right)
            merge 0 0
          where
            x = if right then xRight else xLeft
            y = if right then xLeft else xRight
            n1 = shiftR n 1
            n2 = n - n1
            i2 = i+n1
            merge k j | k == n1 && j < n2  = sequence_
                          [unsafeRead x (i2+m) >>= unsafeWrite y (i+k+m) | m <- range (j,n2-1)]
                      | k < n1 && j == n2  = sequence_
                          [unsafeRead x (i +m) >>= unsafeWrite y (i+m+j) | m <- range (k,n1-1)]
                      | k < n1 && j < n2 = do
               x1 <- unsafeRead x (i+k)
               x2 <- unsafeRead x (i2+j)
               if x1 > x2 then unsafeWrite y (i+k+j) x2 >> merge k (j+1)
                          else unsafeWrite y (i+k+j) x1 >> merge (k+1) j
                      | otherwise = return ()
    f 0 size True
    return xLeft


import           Control.Applicative            ( liftA2 )

import           Math.Combinat.Partitions.Set
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data

(∨) = min
(f .: g) x y = f $ g x y

-- Partitions [l]→[n] factoring through a 2-fibration [2l]↠[l]
isTwoPartition = all even . map length . fromSetPartition

-- l≤n: n

twoPartitions n l =
  filter isTwoPartition . setPartitionsWithKParts (n ∨ (l - n))

-- Constraint: lists must be finite (forcing `length` on both)
relationToMatrix :: (r -> c -> Bool) -> [r] -> [c] -> Matrix R
relationToMatrix p rs cs = (length rs >< length cs) $ liftA2 ((↯) .: p) rs cs
  where (↯) b = if b then 1 else 0


import Control.Parallel

import Numeric.Container hiding (linspace)
import Numeric.LinearAlgebra.Util
import Numeric.LinearAlgebra.Data hiding (linspace)
import Numeric.LinearAlgebra.Algorithms
import Numeric.LinearAlgebra.Devel
import Data.List hiding (find)
import DSP.Basic

--For Charts
--import Graphics.Rendering.Chart hiding (Matrix, Vector)
--import Data.Colour
--import Data.Colour.Names
--import Data.Default.Class
--import Graphics.Rendering.Chart.Backend.Cairo
--import Control.Lens

--gridsize
n = 60

--index function
index :: (Int, Int) -> Int -> Double
index (i,j) n
	| i == j       =  4
	| i == j + n   = -1
	| i == j - n   = -1
	| i == j + 1   = -1
	| i == j - 1   = -1
	| otherwise    =  0

-- Laplace Matrix in 2 dimensions
t :: Matrix Double
t = buildMatrix (n*n) (n*n) (\(i,j) -> index (i,j) n)

-- Compute eigenvalues and corresponding eigenvectors
e = sort(toList(mapVector realPart (fst(eig(t)))))
psi = toColumns (mapMatrix realPart (snd(eig(t))))

--For further image processing
grid :: Vector Double
grid = fromList (linspace 0.0 1.0 n)

ygrid = fromColumns $ replicate n grid
xgrid = fromRows    $ replicate n grid
-- eigenfunction matrix of n=12 for plotting
z1grid = reshape n $ mapVector (\x -> abs(x)^2) (psi !! 12)

main :: IO ()
main = do
-- write out matrices to visualize with python
saveMatrix "xmat.txt" "%f" xgrid
saveMatrix "ymat.txt" "%f" ygrid
saveMatrix "zmat.txt" "%f" z1grid


-----------------------------------------------------------------------------
-- |
-- Module  :  ForSyDe.DFT
-- Copyright   :  (c) ForSyDe Group, KTH 2007-2008
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  forsyde-dev@ict.kth.se
-- Stability   :  experimental
-- Portability :  portable
--
-- This module includes the standard Discrete Fourier Transform (DFT)
-- function, and a fast Fourier transform (FFT) algorithm, for
-- computing the DFT, when the input vectors' length is a power of 2.
-----------------------------------------------------------------------------
module ForSyDe.Shallow.Utility.DFT(dft, fft) where

import ForSyDe.Shallow.Core.Vector
import Data.Complex

-- | The function 'dft' performs a standard Discrete Fourier Transformation
dft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
dft bigN x | bigN == (lengthV x) = mapV (bigX_k bigN x) (nVector x)
           | otherwise = error "DFT: Vector has not the right size!"
  where
    nVector x'        = iterateV (lengthV x') (+1) 0
    bigX_k bigN' x' k = sumV (zipWithV (*) x' (bigW' k bigN'))
    bigW' k' bigN'    = mapV (** k') (mapV cis (fullcircle bigN'))
    sumV              = foldlV (+) (0:+ 0)

fullcircle :: Int -> Vector Double
fullcircle n = fullcircle1 0 (fromIntegral n) n
  where
    fullcircle1 l m n'
      | l == m    = NullV
      | otherwise = -2*pi*l/(fromIntegral n')
                    :> fullcircle1 (l+1) m n'

-- | The function 'fft' implements a fast Fourier transform (FFT)
-- algorithm, for computing the DFT, when the size N is a power of 2.
fft :: Int -> Vector (Complex Double) -> Vector (Complex Double)
fft bigN xv | bigN == (lengthV xv) = mapV (bigX xv) (kVector bigN)
            | otherwise = error "FFT: Vector has not the right size!"

kVector :: (Num b, Num a, Eq a) => a -> Vector b
kVector bigN = iterateV bigN (+1) 0

bigX :: Vector (Complex Double) -> Int -> Complex Double
bigX (x0:>x1:>NullV) k | even k = x0 + x1 * bigW 2 0
                       | odd k  = x0 - x1 * bigW 2 0
bigX xv k = bigF_even k + bigF_odd k * bigW bigN (fromIntegral k)
  where bigF_even k' = bigX (evens xv) k'
        bigF_odd k' = bigX (odds xv) k'
        bigN = lengthV xv

bigW :: Int -> Int -> Complex Double
bigW bigN k = cis (-2 * pi * (fromIntegral k) / (fromIntegral bigN))

evens :: Vector a -> Vector a
evens NullV   = NullV
evens (v1:>NullV) = v1 :> NullV
evens (v1:>_:>v)  = v1 :> evens v

odds :: Vector a -> Vector a
odds NullV    = NullV
odds (_:>NullV)   = NullV
odds (_:>v2:>v)   = v2 :> odds v


{-# LINE 8 "SpectrumAnalysis.lhs" #-}
--  This code was automatically generated by lhs2tex --code, from the file
--  HSoM/SpectrumAnalysis.lhs.  (See HSoM/MakeCode.bat.)
{-# LINE 18 "SpectrumAnalysis.lhs" #-}
{-#  LANGUAGE Arrows  #-}

module Euterpea.Music.Signal.SpectrumAnalysis where

import Euterpea
import Euterpea.Experimental (fftA)

import Data.Complex (Complex ((:+)), polar)
import Data.Maybe (listToMaybe, catMaybes)
{-# LINE 560 "SpectrumAnalysis.lhs" #-}
dft :: RealFloat a => [Complex a] -> [Complex a]
dft xs =
  let  lenI = length xs
       lenR = fromIntegral lenI
       lenC = lenR :+ 0
  in [  let i = -2 * pi * fromIntegral k / lenR
        in (1/lenC) * sum [  (xs!!n) * exp (0 :+ i * fromIntegral n)
                             | n <- [0,1..lenI-1] ]
        | k <- [0,1..lenI-1] ]
{-# LINE 583 "SpectrumAnalysis.lhs" #-}
mkTerm :: Int -> Double -> [Complex Double]
mkTerm num n = let f = 2 * pi / fromIntegral num
               in [  sin (n * f * fromIntegral i) / n :+ 0
                     | i <- [0,1..num-1] ]

mkxa, mkxb, mkxc :: Int-> [Complex Double]
mkxa num = mkTerm num 1
mkxb num = zipWith (+) (mkxa num) (mkTerm num 3)
mkxc num = zipWith (+) (mkxb num) (mkTerm num 5)
{-# LINE 610 "SpectrumAnalysis.lhs" #-}
printComplexL :: [Complex Double] -> IO ()
printComplexL xs  =
  let  f (i,rl:+im) =
            do  putStr (spaces (3 - length (show i))  )
                putStr (show i       ++ ":  ("        )
                putStr (niceNum rl  ++ ", "           )
                putStr (niceNum im  ++ ")\n"          )
  in mapM_ f (zip [0..length xs - 1] xs)

niceNum :: Double -> String
niceNum d =
  let  d' = fromIntegral (round (1e10 * d)) / 1e10
       (dec, fra)  = break (== '.') (show d')
       (fra',exp)  = break (== 'e') fra
  in  spaces (3  - length dec) ++ dec ++ take 11 fra'
      ++ exp ++ spaces (12 - length fra' - length exp)

spaces :: Int -> String
spaces  n = take n (repeat ' ')
{-# LINE 679 "SpectrumAnalysis.lhs" #-}
mkPulse :: Int -> [Complex Double]
mkPulse n = 100 : take (n-1) (repeat 0)
{-# LINE 721 "SpectrumAnalysis.lhs" #-}
x1 num = let f = pi * 2 * pi / fromIntegral num
         in map (:+ 0) [  sin (f * fromIntegral i)
                          | i <- [0,1..num-1] ]
{-# LINE 757 "SpectrumAnalysis.lhs" #-}
mkPolars :: [Complex Double] -> [Complex Double]
mkPolars = map ((\(m,p)-> m:+p) . polar)