module Main where

import Math.Probably.Sampler
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import qualified Math.Probably.PDF as PDF

import Numeric.LinearAlgebra

import Math.Probably.IterLap

import Graphics.Gnewplot.Exec
import Graphics.Gnewplot.Types
import Graphics.Gnewplot.Style
import Graphics.Gnewplot.Panels
import Graphics.Gnewplot.Instances
import Graphics.Gnewplot.Histogram
import Math.Probably.MCMC

--cov = (2><2) [0.1,0.08,0.08,0.1]

--mn = (2 |> [1,1])

cov3 = (3><3) [0.1,  0.02, 0.08,
               0.02, 0.1,  0.05,
               0.08, 0.05, 0.1]

mn3' = (3 |> [1,1,1])

mn3 = (3 |> [2,2,2])

main = runRIO $ do
  (mn', cov') <- sample $ iterLap [100, 100, 100] (PDF.multiNormal mn3 cov3) (mn3', cov3)
  io $ print mn'
  io $ print cov'

{-  iniampar <- sample $ initialAdaMetWithCov 500 (PDF.multiNormal mn3 cov3) cov' mn3'
  vsamples <- runAdaMetRIO 1500 True iniampar $ (PDF.multiNormal mn3 cov3)
  let (mean,sd) =  (both meanF stdDevF) `runStat` vsamples
  io $ print mean
  io $ print sd -}

--  io $ print ws


{-# 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 #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Network.Activation
  (ActivationFunction, unAct, act, name, activate, derivate, logistic, softmax, softmaxMat, relu) where

import Numeric.AD.Mode.Kahn
import Numeric.LinearAlgebra
import Data.Network.Utils(Differentiable, squeeze, derivate, (.*))

data ActivationFunction a = Act {unAct :: (a -> (a,a)), name :: String}

act :: Floating a => (forall s. AD s (Kahn a) -> (AD s (Kahn a))) -> String -> ActivationFunction a
act f = Act (diff' f)

activate :: ActivationFunction a -> a -> a
activate af = fst . (unAct af)

instance Differentiable a a (ActivationFunction a) where
  derivate af = snd . (unAct af)

-- derivate :: ActivationFunction a -> a -> a
-- derivate af = snd . (unAct af)

logistic :: Floating a => ActivationFunction a
logistic = let f x =  1 / (1 + exp (-x))
           in  act f "logistic"

deltaf :: Floating a => Int -> Int -> a
deltaf i j = if (i == j) then 1 else 0

softmax :: (Fractional e, Linear e Vector, Floating (Vector e), Container Vector e) => ActivationFunction (Vector e)
softmax = let f x = let n = sumElements (exp x)
                    in (recip n) .* (exp x)
              f' x = 1 -- Softmax is used only for last layer, so derivative is not needed.
          in  Act (\x -> (f x, f' x)) "softmax"

-- Matrix version for batch computing
softmaxMat :: (Floating e, Container Matrix e, Num (Matrix e), Numeric e, Floating (Matrix e), Floating (Vector e)) =>
  ActivationFunction (Matrix e)
softmaxMat = let f x = let n = fromColumns (replicate (cols x) $ squeeze (exp x))
                       in (recip n) * (exp x)
                 f' x = 1 -- Softmax is used only for last layer, so derivative is not needed.
          in  Act (\x -> (f x, f' x)) "softmax"

relu = let f x = max x 0
       in  act f "relu"

-- We can add activation functions and their derivative using the "diff" function from Numeric.AD


{-# LANGUAGE ScopedTypeVariables #-}

module DecisionTree where

import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.List as List
import qualified Data.List.Extras.Argmax as Argmax
import qualified Statistics.Sample as S

-- Returns the list of unique elements of x --
uniqueElems :: (Eq a) => V.Vector a -> [a]
uniqueElems x = List.nub $ V.toList x

-- Divide two Integral values --
intDiv :: (Integral a, Fractional b) => a -> a -> b
intDiv x y = fromIntegral x / fromIntegral y

-- Returns the sample probability mass function of `samples` P(X = x)
samplePMF :: (Eq a) => V.Vector a -> a -> Double
samplePMF samples x = V.length (V.filter (== x) samples) `intDiv` V.length samples

-- Returns the entropy from the samples `y`. In terms of the sample
-- frequencies p_j (for the j-the value) it is defined by
--                 H(Y) =  -\sum_j p_j * log_2 p_j
entropy :: forall a. (Eq a) => V.Vector a -> Double
entropy y = (-1.0) * sum (map partialEntropy (uniqueElems y))
    where partialEntropy :: a -> Double
          partialEntropy i = pr i * logBase 2 (pr i)

          pr :: a -> Double
          pr = samplePMF y

-- Calculates the "information gain" of the sample `y` when we subdivide
-- it into groups according to the labels given by `lab`
infoGain :: forall a b. (Eq a, Eq b) => V.Vector a -> V.Vector b -> Double
infoGain y lab
    | V.length y /= V.length lab = error $ "DecisionTree.hs:infoGain " ++
                                           "Dimensions do not match."
    | otherwise = (-1.0) * sum (map partialInfoGain subsets)
        where partialInfoGain :: V.Vector a -> Double
              partialInfoGain a = pr_a * entropy a
                  where pr_a = V.length a `intDiv` V.length y

              subsets :: [V.Vector a]
              subsets = map elemsWithLabel (uniqueElems lab)

              elemsWithLabel :: b -> V.Vector a
              elemsWithLabel a = VG.ifilter (\ i _ -> lab V.! i == a) y

fitTree :: forall a b. (Ord a, Eq b) => [V.Vector a] -> V.Vector b -> DNode [a] [Double]
fitTree x y = Decision ((decF bestP) . bestP)  (Result "left") (Result "right")
    where decF :: [a] -> a -> a -> Bool
          decF p z = z < S.mean (p x)

          projectors :: [[a] -> a]
          projectors = [\z -> (z !! n) | n <- [0..(length x - 1)]]
          bestP :: [a] -> a
          bestP = Argmax.argmax (\ p -> (infoGain y $ V.map (decF p) (p x))) projectors

-- a: experiment to decide on
-- b: type of result i.e. classification or classification probability
data DNode a b = Result b
               | Decision (a -> Bool) (DNode a b) (DNode a b)

-- Apply the decision tree to a single experiment
decide :: DNode a b -> a -> b
decide (Result p) = p
decide (Decision df tbranch fbranch) x =
        if df x
            then decide tbranch x
            else decide fbranch x


module Form where

import           Data.Complex
import           Control.Arrow
import           Settings

type Form = Double -> Complex Double

--f(t) = c * exp(n * i * 2pi * t)
form :: Complex Double -> Int -> Double -> Complex Double
form c n t = c * cis (fromIntegral n * 2 * pi * t)

-- exact coordinates
formsToCoordsExact :: [Form] -> [(Double, Double)]
formsToCoordsExact forms = map (realPart &&& imagPart) $ ps timeIntervals forms
  where
    ps []       _  = []
    ps (t : ts) fs = sum (map ($t) fs) : ps ts fs

-- [0, 0.001, ...]
timeIntervals :: [Double]
timeIntervals = ti amount [1]
  where
    ti 0 ds = ds
    ti i ds = ti (i - 1) (interval * i : ds)
    amount   = fromIntegral pointCount
    interval = 1 / (amount + 1)

formsToCoords :: [Form] -> [(Int, Int)]
formsToCoords fs = map (f *** f) $ formsToCoordsExact fs
    where f = (*) scaleFactor . round


{-# LANGUAGE OverloadedLists #-}
module Tasks.Curve where
import qualified Data.Vector.Storable as VS
import           Numeric.LinearAlgebra (Matrix, Vector, flatten,(><))
import           Data.Monoid ((<>))

hermiteMatrix :: Matrix Float
hermiteMatrix = (4><4)
  [ 2,-2, 1, 1
  ,-3, 3,-2,-1
  , 0, 0, 1, 0
  , 1, 0, 0, 0 ]

curvePoints :: Matrix Float
curvePoints = (4><3)
  [ 20, 0, -10
  , 20, 0, -10
  , 100, 100, 0
  , 100, -100, 0]

countPosition :: Float -> Vector Float
countPosition t = flatten $ (1><4) [t^3,t^2,t^1,1] <> hermMat
  where
    hermMat = hermiteMatrix <> curvePoints

lowerPoints :: [Vector Float]
lowerPoints = map countPosition  [0,0.05..0.999]

curveVertices :: Vector Float
curveVertices = VS.concat . map (VS.++ vData) $ quad =<< [0..length lowerPoints - 2]
  where
    lowerTri i = [lVert $ i+1,lVert i,uVert i] :: [Vector Float]
    upperTri i = [lVert $ i+1, uVert i, uVert $ i+1] :: [Vector Float]
    quad i = lowerTri i ++ upperTri i
    lVert i = lowerPoints!!i
    uVert i = (lowerPoints!!i) + ([0,1,0] :: Vector Float)
    vData = [0,0,1,0,0,1,0,1,0,0,0,0,0,0] :: Vector Float


{-|
 Module      : Data.Boltzmann.System.Utils
 Description : Various routines for combinatorial systems.
 Copyright   : (c) Maciej Bendkowski, 2017-2020

 License     : BSD3
 Maintainer  : maciej.bendkowski@tcs.uj.edu.pl
 Stability   : experimental

 The following module implements subroutines for combinatorial systems based
 on the excellent paper of Carine Pivoteau, Bruno Salvy, and Michèle Soria:

 Algorithms for combinatorial structures: Well-founded systems and Newton iterations.
 Journal of Combinatorial Theory, Series A 119 (2012) p. 1711–1773.
 -}
module Data.Boltzmann.System.Utils
    ( isEmptyAtZero
    , zeroCoordinates
    , wellFoundedAtZero
    , polynomial
    ) where

import           Prelude hiding ((<>))

import           Data.Map (Map)
import qualified Data.Map.Strict as M
import           Numeric.LinearAlgebra hiding (size)

import           Data.Boltzmann.System

-- | Non-strict multiplication.
(.*) :: (Num a, Eq a) => a -> a -> a
0 .* _ = 0
_ .* 0 = 0
1 .* x = x
x .* 1 = x
x .* y = x * y

-- | Evaluates the given system Y = H(Z, Y) at coordinates
--   Z = 0 and Y = 0 and checks if the outcome system is empty,
--   i.e. H(0, 0) = 0 or not.
isEmptyAtZero :: (Eq a, Num a) => System a -> Bool
isEmptyAtZero sys = all isZeroType ts
    where ts = M.toList $ defs sys

          isZeroType (_, consL) = all isZeroCons consL
          -- Note: A given constructor can give a non-zero coordinate
          --       iff it has no arguments and no positive weight. List
          --       arguments are irrelevant as they yield powers of unity.
          isZeroCons cons = not $ weight cons == 0 && null (args cons)

evalAtZero :: Num a => Arg -> a
evalAtZero (Type _) = 0
evalAtZero (List _) = 1

evalAtZeroL :: (Eq a, Num a) => [Arg] -> a
evalAtZeroL xs = foldl (.*) 1 $ map evalAtZero xs

derivTypeAtZero :: (Eq a, Eq b, Num a, Num b)
                => String -> [Cons b] -> a

derivTypeAtZero dt consL =
    sum $ map (derivConsAtZero dt) consL

derivConsAtZero :: (Eq a, Eq b, Num a, Num b)
                => String -> Cons a -> b

derivConsAtZero dt cons
    | weight cons /= 0 = 0
    | otherwise = derivConsAtZeroL dt (args cons)

derivConsAtZeroL :: (Eq a, Num a)
                 => String -> [Arg] -> a

derivConsAtZeroL _ [] = 1
derivConsAtZeroL dt (x:xs) = lhs + rhs
    where lhs = evalAtZero x .* derivConsAtZeroL dt xs
          rhs = derivAtZero dt x .* evalAtZeroL xs

derivAtZero :: Num a
            => String -> Arg -> a

derivAtZero dt arg
    | argName arg /= dt = 0
    | otherwise = 1

-- | Computes the Jacobian matrix of the given system
--   and evaluates it at zero-valued coordinates.
jacobian :: (Eq a, Num a) => System a -> Matrix Double
jacobian sys = (n><n) [idx i j | i <- [0..n-1], j <- [0..n-1]]
    where n = size sys
          ts = defs sys
          idx i j =
              let typ = snd $ M.elemAt i ts
                  dt  = fst $ M.elemAt j ts
                  in derivTypeAtZero dt typ

-- | Checks whether the given matrix is zero or not.
isZero :: Matrix Double -> Bool
isZero mat = all (== 0) $ concat (toLists mat)

-- | Checks whether the given Jacobian matrix nilpotent is or not.
isNilpotent :: Matrix Double -> Bool
isNilpotent mat = isZero $ power mat m
    where m = rows mat

-- | Squares the given matrix.
square :: Matrix Double -> Matrix Double
square m = m <> m

-- | Fast matrix exponentiation.
power :: Integral a
      => Matrix Double -> a -> Matrix Double

power m 1 = m
power m n
  | odd n = m <> square (power m $ n-1)
  | otherwise = square (power m $ n `div` 2)

data Coordinate = Zero | Z deriving (Eq)

-- | Maps coordinates into booleans.
coord :: Coordinate -> Bool
coord Zero = True
coord Z    = False

-- | Gives the constructor list of a given type.
constrList :: System a -> String -> [Cons a]
constrList sys = (M.!) (defs sys)

-- | Detects whether the given system, satisfying H(0,0) = 0
--   and having a nilpotent Jacobian matrix, admits zero coordinates
--   in its solution. See also the 0-coord subroutine of Pivoteau et al.
zeroCoordinates :: System a -> Bool
zeroCoordinates sys = zeroCoordinates' sys m f
    where f = M.fromSet (const Zero) $ types sys
          m = size sys

zeroCoordinates' :: (Eq b, Num b)
                 => System a -> b -> Map String Coordinate -> Bool

zeroCoordinates' _ 0 f = elem Zero $ M.elems f
zeroCoordinates' sys m f = zeroCoordinates' sys (m-1) f'
    where xs = map (\k -> (k, mapF $ constrList sys k)) (M.keys f)
          f' = M.fromAscList xs

          mapF consL
              | all (coord . mapF') consL = Zero
              | otherwise                 = Z

          mapF' cons
              | any (coord . mapF'') (args cons) = Zero
              | otherwise                        = Z

          mapF'' (List _) = Z
          mapF'' (Type t) =
              case t `M.lookup` f of
                  Just Zero -> Zero
                  _         -> Z

-- | Checks whether the given system is well-founded at zero. Note that
--   the input system is assumed to satisfy H(0,0) = 0. See also the
--   isWellFoundedAt0 subroutine of Pivoteau et al.
wellFoundedAtZero :: (Eq a, Num a)
                  => System a -> Bool

wellFoundedAtZero sys
    | isNilpotent (jacobian sys) = not $ zeroCoordinates sys
    | otherwise                  = False

-- | Yields an infinite stream of successive evaluation iterations
--   in form of Y[m+1] = H(Z, Y[m]); starting with the given Y[0] and z.
evals :: System Int -> Vector Double -> Double -> [Vector Double]
evals sys ys z = ys : evals sys (eval sys ys z) z

-- | Checks whether the given system, assumed to satisfy H(0,0) = 0
--   and having a nilpotent Jacobian, encodes an implicit polynomial
--   species. See also the isPolynomial subroutine of Pivoteau et al.
polynomial :: System Int -> Bool
polynomial sys = hs !! m  == hs !! (m+1)
    where m   = size sys
          hs  = evals sys vec 1 -- note: numerical evaluation.
          vec = vector $ replicate m 0


{-# LANGUAGE RankNTypes, FlexibleContexts #-}
module Numeric.ADLens.HMatrix
    ( mdot'
    , add'
    , dot'
    , sumElements'
    , cmap'
    ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel (zipVectorWith)
import Numeric.ADLens.Lens
-- import Data.Vector as V

dot' :: (Container Vector t, Numeric t) => Lens' (Vector t, Vector t) t
dot' = lens'' $ \(v1,v2) -> (v1 <.> v2, \ds -> (scale ds v2, scale ds v1))

mdot' :: (Product t, Numeric t) => Lens' (Matrix t, Vector t) (Vector t)
mdot' = lens'' $ \(a,v) -> (a #> v, \dv -> (outer dv v, dv <# a))

add' :: Additive c => Lens' (c, c) c
add' = lens'' $ \(v1,v2) -> (add v1 v2, \dv -> (dv, dv))

-- I need konst I think?
sumElements' :: (Container Vector t, Numeric t) => Lens' (Vector t) t
sumElements' = lens'' $ \v -> (sumElements v, \ds -> scalar ds)

reshape' :: Container Vector t => Int -> Lens' (Vector t) (Matrix t)
reshape' n = lens'' $ \v -> (reshape n v,  \dm -> flatten dm)

-- conjugate transpose not trace
tr'' ::  (Transposable m mt, Transposable mt m) => Lens' m mt
tr'' = lens'' $ \x -> (tr x, \dt -> tr dt)

flatten' :: (Num t, Container Vector t) => Lens' (Matrix t) (Vector t)
flatten' = lens'' $ \m -> let s = fst $ size m in
                          (flatten m,  \dm -> reshape s dm)

norm_2' :: (Container c R, Normed (c R), Linear R c) => Lens' (c R) R
norm_2' = lens'' $ \v -> let nv = norm_2 v in (nv, \dnv -> scale (2 * dnv / nv) v )

cmap' :: (Element b, Container Vector e) => (Lens' e b) -> Lens' (Vector e) (Vector b)
cmap' l = lens'' $ \c -> (cmap f c, \dc -> zipVectorWith f' c dc) where
        f = view l
        f' = set l

{-
maxElement' :: Container c e => Lens' (c e) e
maxElement' = lens'' $ \v -> let i = maxIndex v in (v ! i, dv -> scalar 0)
-}

det' :: Field t => Lens' (Matrix t) t
det' = lens'' $ \m -> let (minv, (lndet, phase)) = invlndet m in
                    let detm = phase * exp detm in
                    (detm, \ds -> (scale (ds * detm) minv))

diag' :: (Num a, Element a) => Lens' (Vector a) (Matrix a)
diag' = lens'' $ \v -> (diag v, takeDiag)

takeDiag' :: (Num a, Element a) => Lens' (Matrix a) (Vector a)
takeDiag' = lens'' $ \m -> (takeDiag m, diag)


{-# 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)


module Main where
import qualified Data.ByteString.Lazy as BS hiding (putStrLn)
import qualified Data.ByteString.Lazy as BS (putStrLn)
import Codec.Compression.GZip (decompress)
import Numeric.LinearAlgebra hiding (split)
import Prelude hiding (readFile)
import System.Random ( StdGen, mkStdGen )
import System.Random.Shuffle
import Control.Monad
import Data.List.Split
import Data.List

import Control.Monad.State
import Control.Monad (join)
import Data.Bifunctor (bimap)
import Data.ByteString.Conversion.From
import Data.ByteString.Char8 (readInt)
import Data.Int

data Network = Network { layers  :: [Int]
                       , biases  :: [Vector R]
                       , weights :: [Matrix R]
                       , rate    :: Double
                       } deriving (Show)

main :: IO ()
main = do
  let nInputs  = 784
  let nHidden  = 30
  let nOutputs = 10
  let seed = 1
  let layers = [nInputs, nHidden, nOutputs]
  let biases = fmap (\n -> randomVector seed Gaussian n) $ drop 1 layers
  let weights = fmap (\(x, y) -> randomMatrix x y seed + 10)  $ zip (drop 1 layers) layers
  let rate = 1.5
  let net = Network layers biases weights rate

  trainingData <- readImages "images.gz" "labels.gz"
  let vec = snd $ trainingData !! 0
  --print . show . feedForward net $ vec
  print . show . fst $ trainingData !! 1
  --let input = (take 20 trainingData)
  --t <- forM input $ \(digit, input) ->
      --return (digit, train net 20 g input)

------------------- Reading MNIST -------------------------

readImages :: String -> String -> IO ([(Vector R, Vector R)])
readImages iPath lPath = do
  iBytes <- decompress <$> BS.readFile iPath
  lBytes <- decompress <$> BS.readFile lPath
  let is = readImage iBytes <$> [1..50000]
  let ls = readLabel lBytes <$> [1..50000]
  return $ zip ls is

readImage :: BS.ByteString -> Int64 -> Vector R
readImage bytes n = vector $ fmap (fromIntegral . BS.index bytes . (n*28^2 + 16 +)) [0..783]

readLabel :: BS.ByteString -> Int64 -> Vector R
readLabel bytes n = vector [if x == i then 1 else 0 | x <- [0..9]] where
                               i = fromIntegral $ BS.index bytes (n + 8)

---------------------- Training ------------------------

train :: Network-> Int -> StdGen-> [(Vector R, Vector R)] -> Network
train net batchSize gen input = net' where
  net' = foldl' (foldl' processTrainingData) net batches
  batches = chunksOf batchSize $ shuffledInput
  shuffledInput = shuffle' input (length input) gen

processTrainingData :: Network -> (Vector R, Vector R) -> Network
processTrainingData net input = Network (layers net) bs ws (rate net) where
  ws = zipWith add (weights net) $ fmap (scale (-eta / l)) dw
  bs = zipWith add (biases net) $ fmap (scale (-eta / l)) db
  (dw, db) = unzip $ backprop net input
  l = fromIntegral $ length input
  eta = rate net

-- Returns nablaWeights and nableBiases in a tuple
backprop :: Network -> (Vector R, Vector R) -> [(Matrix R, Vector R)]
backprop net (y, x) = undefined where

  -- Propagate backwards calculating
  -- feltas thourhgout network, with input layer as head
  ds = foldl' calcDelta [outputDelta] list where
    outputDelta = cost' (head as) y * sigmoid' (head zs)
    list = init $ zip3 ws zs $ drop 1 as
    ws = reverse $ weights net
    calcDelta :: [Vector R] -> (Matrix R, Vector R, Vector R) -> [Vector R]
    calcDelta (d:ds) (w, z, a) = d':d:ds where
      d' = (tr' w #> d * sigmoid' z)

  -- Feed forward through the network
  -- calculating all zs and as, with the output layer as head
  (zs, as) = feedForward net x

  {-
  [(dw, db)] = foldl' calcDeltas [(calcDw d as, d)] $ init $ zip3 ws zs $ drop 1 as
ws = reverse $ weights net

calcDw :: Vector R -> [Vector R] -> R
  calcDw d (a:a':_) = d `dot` a'
  calcD (a:_) y = sigmoid' a * cost' a y

  calcDeltas :: [Vector R] -> (Matrix R, Vector R, Vector R) -> [Vector R]
  calcDeltas (d:ds) (w, z, a) = d':d:ds where
    d' = (w #> d * sigmoid' z)

  -- Calculate the deltas, with the input layer at index 0
  ds = foldl' calcDelta [(calcDw d as, delta as y)] $ init $ zip3 ws zs $ drop 1 as where
    calcDelta :: [Vector R] -> (Matrix R, Vector R, Vector R) -> [Vector R]
    calcDelta (d:ds) (w, z, a) = (w #> d * sigmoid' z):d:ds
    delta as y = sigmoid' a * cost' a y
    ws = reverse $ weights net
-}

-- Returns a tuple of zs and as throughout the network
-- with the last activation first (like a stack)
feedForward :: Network -> Vector R -> ([Vector R], [Vector R])
feedForward net a = foldl' ffd ([a], []) $ zip (weights net) (biases net)  where
  ffd :: ([Vector R], [Vector R]) -> (Matrix R, Vector R) -> ([Vector R], [Vector R])
  ffd (a:as, zs) (w, b) = (a':a:as, z:zs) where
    a' = sigmoid z
    z  = w #> a + b

{-
-- Perform stochastic gradient descent on the batch
-- and return a new network with updated weights and biases
processBatch :: Network -> [(Vector R, Vector R)] -> Network
processBatch net batch = foldl' processTrainingData net batch where
  processTrainingData :: Network -> (Vector R, Vector R) -> Network
  processTrainingData = \net input ->
    Network (layers net) bs ws (rate net) where
      ws = zipWith add (weights net) $ fmap (scale (-eta / l)) dw
      bs = zipWith add (biases net) $ fmap (scale (-eta / l)) db
      (dw, db) = unzip $ backprop net input
      l = fromIntegral $ length batch
      eta = rate net
-}

-- Returns a tuple of zs and as throughout the network
-- with the last activation first (like a stack)
{-
oldfeedForward :: Network -> Vector R -> ([Vector R], [Vector R])
oldfeedForward net a = ffd (weights net) (biases net) ([a], []) where
  ffd :: [Matrix R] -> [Vector R] -> ([Vector R], [Vector R]) -> ([Vector R], [Vector R])
  ffd [] [] a = a
  ffd (w:ws) (b:bs) (a:as, zs) = ffd ws bs (a':a:as, z:zs) where
    a' = sigmoid z
    z  = w #> a + b
-}
  -- Do the backpropagation
--forM ds $ \(lbl, img) ->
--                                    let a = feedForward n img

-- Cost function
cost w b n xs a = 1/(2*n) * (sum [(norm_2 (y_ x - a))^2 | x <- xs])

-- Derivative of cost
cost' :: Vector R -> Vector R -> Vector R
cost' a y = a - y

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

sigmoid' :: Floating a => a -> a
sigmoid' z = sigmoid(z)*(1 - sigmoid(z))

--------------------  Utility   ---------------------------

randomMatrix :: Int -> Int -> Int -> Matrix R
randomMatrix x y seed = (x><y) . toList . randomVector seed Gaussian $ x * y

y_ :: Vector R -> Vector R
y_ = undefined

-- Test stuff
{-
getImage s n = fromIntegral . BS.index s . (n*28^2 + 16 +) <$> [0..28^2 - 1]

render n = let s = " .:oO@" in s !! (fromIntegral n * length s `div` 256)
pasta = do
  s <- decompress <$> BS.readFile "images.gz"
  l <- decompress <$> BS.readFile "labels.gz"
  let n = 1

  -- Labels 8 bytes
  -- Images has offset 16 8 * 28  28?

  --n <- (`mod` 60000) <$> randomIO
  let t =  (getImage s n :: [Double])
  putStrLn $ show t
  {-
  putStr . unlines $
    [(render . BS.index s . (n*28^2 + 16 + r*28 +)) <$> [0..27] | r <- [0..27]]
  print $ BS.index l (n + 8)
-}
-}


{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}

module Vector where

import Tensor
import Data.Monoid
import Data.Maybe
import Data.Complex
import Helper
import Data.Either
import System.Random

data Vector a = Scalar a | Row [Vector a] | Column [Vector a]

-- This Eq implementation ignores structural differences
instance (Eq a) => Eq (Vector a) where
    x == y = wx == wy && hx == hy && go 0 0
        where
            wx = width x
            wy = width y
            hx = height x
            hy = height y
            go i j
                | i == hx = True
                | j == wx = go (i + 1) 0
                | otherwise = atM x (i, j) == atM y (i, j) && go i (j + 1)

instance Functor Vector where
    fmap f (Scalar s) = Scalar $ f s
    fmap f (Row t) = Row $ map (fmap f) t
    fmap f (Column t) = Column $ map (fmap f) t

instance Applicative Vector where
    pure = Scalar
    Scalar s <*> t = fmap s t
    t <*> (Scalar s) = fmap ($s) t
    Row r <*> t = Row $ map (<*> t) r
    Column c <*> t = Column $ map (<*> t) c

instance Monad Vector where
    Scalar s >>= f = f s
    Row r >>= f = Row $ map (>>= f) r
    Column c >>= f = Column $ map (>>= f) c

instance Num a => Num (Vector a) where
    Row x + (Row y) = Row $ zipWith (+) x y
    Column x + (Column y) = Column $ zipWith (+) x y
    Scalar x + y = fmap (+ x) y
    x + (Scalar y) = fmap (+ y) x
    _+_ = error "Vector Addition: Dimension mismatch"
    Scalar x * y = fmap (* x) y
    x * Scalar y = fmap (* y) x
    Row x * (Column y) = getSum $ foldMap Sum $ zipWith (*) x y
    Column x * (Column y) = Column $ map (* Column y) x
    x * (Row y) = Row $ map (x *) y
    abs = fmap abs
    signum = fmap signum
    fromInteger = Scalar . fromInteger
    negate = fmap negate

instance Foldable Vector where
  foldr f acc (Scalar s) = f s acc
  foldr f acc (Row r) = foldr (flip (foldr f)) acc r
  foldr f acc (Column c) = foldr (flip (foldr f)) acc c

instance Traversable Vector where
    traverse f (Scalar s) = Scalar <$> f s
    traverse f (Row r) = Row <$> traverse (traverse f) r
    traverse f (Column c) = Column <$> traverse (traverse f) c

instance (Num a) => Semigroup (Vector a) where
    x <> y = (*) <$> x <*> y

instance (Num a) => Monoid (Vector a) where
    mempty = Scalar 1

instance (Cmplx a) => Cmplx (Vector a) where
    conj = fmap conj . transposeStruct

instance (Num a, Cmplx a) => Tensor Vector a where
    inner = dot
    -- TODO: norm for matrices
    norm t = (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t
    fromList xs dim = go (foldr (con . Scalar) (Column []) xs) dim
        where
            go (Column c) n
                | n == 1 = Column [go (Row c) dim]
                | l <= fromIntegral dim = Column c
                | otherwise = go (Row $ take (l `div` fromIntegral n) c) dim
                    `con` go (Column $ drop (l `div` fromIntegral n) c) (n - 1)
                where l = length c

            go (Row r) n
                | n == 1 = Row [go (Column r) dim]
                | l <= fromIntegral dim = Row r
                | otherwise = go (Column $ take (l `div` fromIntegral n) r) dim
                    `con` go (Row $ drop (l `div` fromIntegral n) r) (n - 1)
                where l = length r

            go s n = error "error when adding tensor rank"
    rank (Scalar s) = 0
    rank (Row []) = 1
    rank (Column []) = 1
    rank (Row (r:rs)) = rank r + 1
    rank (Column (c:cs)) = rank c + 1

    isSquare t
        | (Scalar _) <- t = True
        | (Column c) <- t = getAll $ go (Column c) (length c)
        | (Row r) <- t = getAll $ go (Row r) (length r)
        where
            go (Column c) l
                | length c == l = foldMap (`go` l) c
                | otherwise = All False
            go (Row r) l
                | length r == l = foldMap (`go` l) r
                | otherwise = All False
            go (Scalar _) _ = All True

    conjTr = conj

instance (Floating a, Eq a, Show a, Ord a, UniformRange a, Cmplx a) => QuantumRegister Vector a where

    densityOp x = transposeStruct x * x

    initQubit zero one
        | abs zero ** 2 + abs one ** 2 == 1 = Right $ fromList [zero, one] 2
        | otherwise = Left "Sum of probability of |0> and |1> state does not equal to 1"

    initQubit0 = fromList [1, 0] 2

    initNumQubit0 n
        | n < 1 = Left "Invalid qubit number"
        | otherwise = Right $ foldr (|*|) initQubit0 $ replicate (fromIntegral n - 1) initQubit0

    toQuantumRegister q
        | isLeft num = Left "Input list length is not a power of 2"
        | otherwise = Right . norm $ fromList q 2
        where num = integralLog 2 (length q)

    illegalPeek (Scalar s) = show s
    illegalPeek q = go 0 0
        where
            w = width q
            h = height q
            go i j
                | i == h = mempty
                | j == w = "\t|\n" ++ go (i + 1) 0
                | j == 0 = "|\t" <> current <> "\t" <> go i (j + 1)
                | otherwise = current <> "\t" <> go i (j + 1)
                where
                    current = maybe "null" show (atM q (i, j))

    subSystem lst q
        | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid Quantum Register"
        | otherwise = Right . fmap sqrt $ go lst (fmap (\x -> conj x * x) q)
        where
            go [] (Scalar s) = Scalar s
            go is (Row r)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Row $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) r
                where
                    preSum = map (go $ fmap (flip (-) 1) is) r
            go is (Column c)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Column $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) c
                where
                    preSum = map (go $ fmap (flip (-) 1) is) c
            go _ _ = error "Exceptions not handled by subSystem function"

    isEntangled lst q = do
            reduced <- reducedDensity lst q
            return $ abs (trace (reduced * reduced) - 1) > 1e-13

    mask0 = Row [
                Column[1, 0],
                Column[0, 0]
            ]

    mask1 = Row [
                Column[0, 0],
                Column[0, 1]
            ]

    pauliId = Row [
                Column [1, 0],
                Column [0, 1]
            ]

instance (RealFloat a, Show a, UniformRange a) => Gates Vector (Complex a) where
    pauliX = Row [
            Column [Scalar (0 :+ 0), Scalar (1 :+ 0)],
            Column [Scalar (1 :+ 0), Scalar (0 :+ 0)]
        ]

    pauliY = Row [
                Column [Scalar (0 :+ 0), Scalar (0 :+ 1)],
                Column [Scalar (0 :+ (-1)), Scalar (0 :+ 0)]
            ]

    pauliZ = Row [
                Column [Scalar (1 :+ 0), Scalar (0 :+ 0)],
                Column [Scalar (0 :+ 0), Scalar ((-1) :+ 0)]
            ]

    hadamard = Row [
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar (1 / sqrt 2 :+ 0)],
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar ((-1) / sqrt 2 :+ 0)]
            ]

at :: Integral b => a -> Vector a -> (b, b) -> a
at x y z = fromMaybe x (atM y z)

atM :: Integral b => Vector a -> (b, b) -> Maybe a
atM (Scalar s) (0, 0) = Just s
atM (Column (c:cs)) (i, j)
    | height (Column (c:cs)) > i = let offset = height c in
                                    if offset > i
                                    then atM c (i, j)
                                    else atM (Column cs) (i - offset, j)
    | otherwise = Nothing
atM (Row (r:rs)) (i, j)
    | width (Row (r:rs)) > j = let offset = width r in
                                    if offset > j
                                    then atM r (i, j)
                                    else atM (Row rs) (i, j - offset)
    | otherwise = Nothing
atM _ _ = Nothing

con :: Vector a -> Vector a -> Vector a
con x (Row y) = Row $ x : y
con x (Column y) = Column $ x : y
con _ _ = error "Appending to Scalar is prohibited"

append :: Vector a -> Vector a -> Vector a
append (Row x) (Row y) = Row $ x ++ y
append (Column x) (Column y) = Column $ x ++ y
append _ _ = error "Vector direction mismatch"

width :: Integral b => Vector a -> b
width (Scalar s) = 1
width (Row r) = getSum $ foldMap (Sum . width) r
width (Column c) = width $ Prelude.head c

height :: Integral b => Vector a -> b
height (Scalar s) = 1
height (Column c) = getSum $ foldMap (Sum . height) c
height (Row r) = height $ Prelude.head r

-- Transform structure retaining the original rank and shape
transformStruct :: Vector a -> Vector a
transformStruct (Row (Column c:cs)) = Column . go $ Column c:cs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go cs = Row (map (transformStruct . Vector.head) cs)
            : go (map Vector.tail cs)

transformStruct (Column (Row r:rs)) = Row . go $ Row r:rs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go rs = Column (map (transformStruct . Vector.head) rs)
            : go (map Vector.tail rs)

transformStruct (Column cs) = Column $ map transformStruct cs
transformStruct (Row rs) = Row $ map transformStruct rs
transformStruct t = t

-- transpose via structure
transposeStruct :: Vector a -> Vector a
transposeStruct (Scalar s) = Scalar s
transposeStruct (Row r) = Column $ map transposeStruct r
transposeStruct (Column c) = Row $ map transposeStruct c

-- transpose retaining structure
transpose :: Vector a -> Vector a
transpose = transformStruct . transposeStruct

head :: Vector a -> Vector a
head (Scalar s) = Scalar s
head (Row r) = Prelude.head r
head (Column c) = Prelude.head c

last :: Vector a -> Vector a
last (Scalar s) = Scalar s
last (Row r) = Prelude.last r
last (Column c) = Prelude.last c

tail :: Vector a -> Vector a
tail (Scalar s) = Scalar s
tail (Row r) = Row $ Prelude.tail r
tail (Column c) = Column $ Prelude.tail c

init :: Vector a -> Vector a
init (Scalar s) = Scalar s
init (Row r) = Row $ Prelude.init r
init (Column c) = Column $ Prelude.init c

dot :: Num a => Vector a -> Vector a -> Vector a
dot x y = transposeStruct x * y

reducedDensity :: (Num a, Cmplx a, Floating a) => [Int] -> Vector a -> Either String (Vector a)
reducedDensity lst q
    | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
    | not (isSquare q) = Left "Invalid Quantum Register"
    | otherwise = Right $ go lst q (transposeStruct q)
    where
        go [] (Scalar s) (Scalar t) = Scalar (s * t)
        go is (Column c) (Row r)
            | null is || Prelude.head is /= 0 = (getSum . foldMap Sum) $ zipWith (go (flip (-) 1 <$> is)) c r
            | otherwise = Row $ map (Column . ($ c) . (\ b a -> map (\ x -> go (flip (-) 1 <$> Prelude.tail is) x b) a)) r
        go _ _ _ = error "Input is not a valid statevector"

isColumn :: Vector a -> Bool
isColumn (Column _) = True
isColumn _ = False

isRow :: Vector a -> Bool
isRow (Row _) = True
isRow _ = False

isScalar :: Vector a -> Bool
isScalar (Scalar _) = True
isScalar _ = False

trace :: (Num a) => Vector a -> a
trace (Scalar s) = s
trace (Row r) = (getSum . foldMap Sum) $ zipWith go [0, 1 ..] r
    where
        go i (Row r) = trace (r !! i)
        go i (Column c) = trace (c !! i)
        go i (Scalar s) = s
{-
    Tensor type:
    All one-dimensional tensors are column vectors. This means matrices,
    which are rank 2 tensors, compose of multiple column tensors, which
    may not follow conventional representation in other traversable types.
-}

-- data Tensor a = Scalar a | Tensor [Tensor a]
--     deriving(Eq)

-- instance (Show a) => Show (Tensor a) where
--     show (Scalar s) = show s
--     show (Tensor ((Tensor t):ts)) = "[\n" ++ unlines (go (Tensor t:ts)) ++ "]"
--         where go = fmap ("  " ++).concatMap(lines . show)
--     show (Tensor s) = show s ++ "\n"

-- instance Functor Tensor where
--     fmap f (Scalar s) = Scalar $ f s
--     fmap f (Tensor t) = Tensor $ map (fmap f) t

-- instance Applicative Tensor where
--     pure = Scalar
--     (Scalar s) <*> t = fmap s t
--     t <*> (Scalar s) = fmap ($ s) t
--     (Tensor []) <*> _ = Tensor []
--     (Tensor (Scalar s : ss)) <*> t = Tensor $ fmap s t : tensorList
--         where Tensor tensorList = Tensor ss <*> t
--     (Tensor (Tensor t : ts)) <*> t2 = Tensor $ (Tensor t <*> t2) : tensorList
--         where Tensor tensorList = Tensor ts <*> t2

-- instance Monad Tensor where
--     (Scalar s) >>= f = f s
--     (Tensor t) >>= f = Tensor $ map (>>= f) t

-- instance (Num a) => Num (Tensor a) where
--     (Scalar s) + t = fmap (+ s) t
--     t + (Scalar s) = fmap (+ s) t
--     (Tensor xs) + (Tensor ys) = Tensor $ zipWith (+) xs ys
--     (Scalar s) * t = fmap (* s) t
--     t * (Scalar s) = fmap (* s) t
--     _ * Tensor[] = Tensor []
--     (Tensor x) * (Tensor y) = getSum $ foldMap Sum (zipWith go x y)
--         where
--             go (Tensor a) b = Tensor $ map (* b) a
--             go a b = a * b
--     abs = fmap abs
--     signum = fmap signum
--     negate = fmap negate
--     fromInteger i = Scalar $ fromInteger i

-- instance Foldable Tensor where
--     foldr f acc (Tensor t) = foldr (flip $ foldr f) acc t
--     foldr f acc (Scalar s) = f s acc

-- instance Traversable Tensor where
--     traverse f (Scalar s) = Scalar <$> f s
--     traverse f (Tensor t) = Tensor <$> traverse (traverse f) t

-- -- TODO: revise function for higher-dimension tensors
-- fromList :: [a] -> Int -> Tensor a
-- fromList xs dim = go (foldr (con . Scalar) (Tensor[]) xs) dim
--     where
--         go (Tensor t) n
--             | n == 1 = Tensor [go (Tensor t) dim]
--             | length t <= dim = Tensor t
--             | otherwise = go (Tensor $ take (length t `div` n) t) dim
--                 `con` go (Tensor $ drop (length t `div` n) t) (n - 1)
--         go s n = error "error when adding tensor rank"

-- product :: (Num a) => Tensor a -> Tensor a -> Tensor a
-- product (Scalar s) t = Scalar s * t
-- product t (Scalar s) = t * Scalar s
-- product (Tensor []) _ = Tensor []
-- product _ (Tensor []) = Tensor []
-- product (Tensor (x:xs)) y = Tensor $ current : rest
--     where
--         current = Tensor.product x y
--         Tensor rest = Tensor.product (Tensor xs) y

-- norm :: (Floating a) => Tensor a -> Tensor a
-- norm t= (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t

-- append :: Tensor a -> Tensor a -> Tensor a
-- append (Tensor x) (Tensor y) = Tensor $ x ++ y
-- append _ _ = error "Concatenating scalars is prohibited"

-- con :: Tensor a -> Tensor a -> Tensor a
-- con x (Tensor y) = Tensor $ x : y
-- con _ _ = error "Appending to Scalar is prohibited"

-- dot :: Num a => Tensor a -> Tensor a -> a
-- dot (Tensor x) (Tensor y) = getSum $ foldMap Sum (zipWith dot x y)
-- dot (Scalar x) (Scalar y) = x * y
-- dot _ _ = error "Length mismatch"

-- flatten2d :: Tensor a -> Tensor a
-- flatten2d (Tensor ((Tensor ((Tensor (t:ts)):ss)):xs)) = Tensor . go $ Tensor (Tensor (t:ts):ss):xs
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor (Tensor (Scalar s:ss):us):ys) = foldr (\a acc ->
--             let (Tensor result) = Tensor.transpose a
--                 in result ++ acc
--             ) [] (Tensor (Tensor (Scalar s:ss):us):ys)

--         go y = go flattened
--             where flattened = foldr (\a acc ->
--                     let (Tensor next) = a
--                         in Tensor (go next) : acc
--                     ) [] y
-- flatten2d t = t

-- head :: Tensor a -> Tensor a
-- head (Tensor ts) = Prelude.head ts
-- head t = t

-- last :: Tensor a -> Tensor a
-- last (Tensor ts) = Prelude.last ts
-- last t = t

-- tail :: Tensor a -> Tensor a
-- tail (Tensor ts) = Tensor $ Prelude.tail ts
-- tail t = t

-- init :: Tensor a -> Tensor a
-- init (Tensor ts) = Tensor $ Prelude.init ts
-- init t = t

-- transpose :: Tensor a -> Tensor a
-- transpose (Tensor (Tensor t:ts)) = Tensor $ go $ Tensor t:ts
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor [] : _) = []
--         go xs = Tensor (map Tensor.head xs) : go (map Tensor.tail xs)
-- transpose (Tensor (Scalar s:ss)) = Tensor $ map (Tensor.(:[])) (Scalar s:ss)
-- transpose t = t

-- isSquare :: Tensor a -> Bool
-- isSquare (Scalar _) = True
-- isSquare (Tensor []) = True
-- isSquare (Tensor t) = getAll $ go (Tensor t) (length t)
--     where
--         go (Tensor x) l
--             | length x == l = foldMap (`go` l) x
--             | otherwise = All False
--         go (Scalar _) _ = All True

-- rank :: Integral b => Tensor a -> b
-- rank (Scalar s) = 0
-- rank (Tensor []) = 1
-- rank (Tensor (t:ts)) = rank t + 1


{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric,
    FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample.Powers
-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Very fast statistics over simple powers of a sample.  These can all
-- be computed efficiently in just a single pass over a sample, with
-- that pass subject to stream fusion.
--
-- The tradeoff is that some of these functions are less numerically
-- robust than their counterparts in the 'Statistics.Sample' module.
-- Where this is the case, the alternatives are noted.

module Statistics.Sample.Powers
    (
    -- * Types
      Powers

    -- * Constructor
    , powers

    -- * Descriptive functions
    , order
    , count
    , sum

    -- * Statistics of location
    , mean

    -- * Statistics of dispersion
    , variance
    , stdDev
    , varianceUnbiased

    -- * Functions over central moments
    , centralMoment
    , skewness
    , kurtosis

    -- * References
    -- $references
    ) where

import Control.Monad.ST
import Data.Aeson            (FromJSON, ToJSON)
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import Data.Vector.Binary    ()
import Data.Vector.Unboxed   ((!))
import GHC.Generics          (Generic)
import Numeric.SpecFunctions (choose)
import Prelude hiding (sum)
import Statistics.Function   (indexed)
import qualified Data.Vector          as V
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Unboxed.Mutable as MU
import qualified Statistics.Sample.Internal  as S

newtype Powers = Powers (U.Vector Double)
    deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Powers
instance ToJSON Powers

instance Binary Powers where
    put (Powers v) = put v
    get = fmap Powers get

-- | O(/n/) Collect the /n/ simple powers of a sample.
--
-- Functions computed over a sample's simple powers require at least a
-- certain number (or /order/) of powers to be collected.
--
-- * To compute the /k/th 'centralMoment', at least /k/ simple powers
--   must be collected.
--
-- * For the 'variance', at least 2 simple powers are needed.
--
-- * For 'skewness', we need at least 3 simple powers.
--
-- * For 'kurtosis', at least 4 simple powers are required.
--
-- This function is subject to stream fusion.
powers :: G.Vector v Double =>
          Int                   -- ^ /n/, the number of powers, where /n/ >= 2.
       -> v Double
       -> Powers
powers k sample
  | k < 2     = error "Statistics.Sample.powers: too few powers"
  | otherwise = runST $ do
      acc <- MU.replicate l 0
      G.forM_ sample $ \x ->
        let loop !i !xk
              | i == l    = return ()
              | otherwise = do MU.write acc i . (+ xk) =<< MU.read acc i
                               loop (i+1) (xk * x)
        in loop 0 1
      fmap Powers $ U.unsafeFreeze acc
  where
    l = k + 1
{-# SPECIALIZE powers :: Int -> U.Vector  Double -> Powers #-}
{-# SPECIALIZE powers :: Int -> V.Vector  Double -> Powers #-}
{-# SPECIALIZE powers :: Int -> SV.Vector Double -> Powers #-}

-- | The order (number) of simple powers collected from a 'sample'.
order :: Powers -> Int
order (Powers pa) = U.length pa - 1

-- | Compute the /k/th central moment of a sample.  The central
-- moment is also known as the moment about the mean.
centralMoment :: Int -> Powers -> Double
centralMoment k p@(Powers pa)
    | k < 0 || k > order p =
                  error ("Statistics.Sample.Powers.centralMoment: "
                         ++ "invalid argument")
    | k == 0    = 1
    | otherwise = (/n) . S.sum . U.map go . indexed . U.take (k+1) $ pa
  where
    go (i , e) = (k `choose` i) * ((-m) ^ (k-i)) * e
    n = U.head pa
    m = mean p

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.  This is
-- the second central moment of the sample.
--
-- This is less numerically robust than the variance function in the
-- 'Statistics.Sample' module, but the number is essentially free to
-- compute if you have already collected a sample's simple powers.
--
-- Requires 'Powers' with 'order' at least 2.
variance :: Powers -> Double
variance = centralMoment 2

-- | Standard deviation.  This is simply the square root of the
-- maximum likelihood estimate of the variance.
stdDev :: Powers -> Double
stdDev = sqrt . variance

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
--
-- Requires 'Powers' with 'order' at least 2.
varianceUnbiased :: Powers -> Double
varianceUnbiased p@(Powers pa)
    | n > 1     = variance p * n / (n-1)
    | otherwise = 0
  where n = U.head pa

-- | Compute the skewness of a sample. This is a measure of the
-- asymmetry of its distribution.
--
-- A sample with negative skew is said to be /left-skewed/.  Most of
-- its mass is on the right of the distribution, with the tail on the
-- left.
--
-- > skewness . powers 3 $ U.to [1,100,101,102,103]
-- > ==> -1.497681449918257
--
-- A sample with positive skew is said to be /right-skewed/.
--
-- > skewness . powers 3 $ U.to [1,2,3,4,100]
-- > ==> 1.4975367033335198
--
-- A sample's skewness is not defined if its 'variance' is zero.
--
-- Requires 'Powers' with 'order' at least 3.
skewness :: Powers -> Double
skewness p = centralMoment 3 p * variance p ** (-1.5)

-- | Compute the excess kurtosis of a sample.  This is a measure of
-- the \"peakedness\" of its distribution.  A high kurtosis indicates
-- that the sample's variance is due more to infrequent severe
-- deviations than to frequent modest deviations.
--
-- A sample's excess kurtosis is not defined if its 'variance' is
-- zero.
--
-- Requires 'Powers' with 'order' at least 4.
kurtosis :: Powers -> Double
kurtosis p = centralMoment 4 p / (v * v) - 3
    where v = variance p

-- | The number of elements in the original 'Sample'.  This is the
-- sample's zeroth simple power.
count :: Powers -> Int
count (Powers pa) = floor $ U.head pa

-- | The sum of elements in the original 'Sample'.  This is the
-- sample's first simple power.
sum :: Powers -> Double
sum (Powers pa) = pa ! 1

-- | The arithmetic mean of elements in the original 'Sample'.
--
-- This is less numerically robust than the mean function in the
-- 'Statistics.Sample' module, but the number is essentially free to
-- compute if you have already collected a sample's simple powers.
mean :: Powers -> Double
mean p@(Powers pa)
    | n == 0    = 0
    | otherwise = sum p / n
    where n     = U.head pa

-- $references
--
-- * Besset, D.H. (2000) Elements of statistics.
--   /Object-oriented implementation of numerical methods/
--   ch. 9, pp. 311–331.
--   <http://www.elsevier.com/wps/product/cws_home/677916>
--
-- * Anderson, G. (2009) Compute /k/th central moments in one
--   pass. /quantblog/. <http://quantblog.wordpress.com/2009/02/07/compute-kth-central-moments-in-one-pass/>


module GradientDescent where
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util

--  Performs gradient descent to learn theta
--  theta = GRADIENTDESCENT(x, y, theta, alpha, num_iters) updates theta by
--  taking num_iters gradient steps with learning rate alpha
gradientDescent :: (Matrix Double -> Matrix Double) -> Matrix Double -> Double -> Int -> Matrix Double
gradientDescent gradFunc initialTheta alpha num_iters = iterate initialTheta num_iters
            where
                iterate theta 0 = theta
                iterate theta iters = iterate (updateTheta theta) (iters-1)
                updateTheta t = t - (scalar alpha) * (gradFunc t)


module Deanie.Density (
    logDensityBernoulli
  , logDensityBeta
  , logDensityGamma
  , logDensityGaussian
  ) where

import Numeric.SpecFunctions (logGamma)

logDensityBernoulli :: Double -> Bool -> Double
logDensityBernoulli p x
    | p < 0 || p > 1 = log 0
    | x              = log p
    | otherwise      = log (1 - p)
{-# INLINE logDensityBernoulli #-}

logDensityBeta :: Double -> Double -> Double -> Double
logDensityBeta a b x
  | x <= 0 || x >= 1 = log 0
  | a < 0 || b < 0   = log 0
  | otherwise        = (a - 1) * log x + (b - 1) * log (1 - x)
{-# INLINE logDensityBeta #-}

-- | The log-density for the gamma distribution with rate parameter b.
--
--     f(x) = b ^ a * x ^ (a - 1) * x ^ (-b * x) / Gamma(a)
logDensityGamma :: Double -> Double -> Double -> Double
logDensityGamma a b x
  | a < 0 || b < 0 = log 0
  | otherwise      = a * log b + (a - 1) * log x - b * x - logGamma a
{-# INLINE logDensityGamma #-}

-- | The density for the normal distribution with specified mean and standard
--   deviation.
logDensityGaussian :: Double -> Double -> Double -> Double
logDensityGaussian m sd x
  | sd <= 0   = log 0
  | otherwise = negate (log sd) - (x - m) * (x - m) / (2 * sd * sd)
{-# INLINE logDensityGaussian #-}


{-|
Module: MachineLearning.Optimization
Description: Optimization
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

Optimization module.
-}

module MachineLearning.Optimization
(
    MinimizeMethod(..)
  , minimize
  , checkGradient
)

where

import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Model as Model
import MachineLearning.Regularization (Regularization)
import qualified MachineLearning.Optimization.GradientDescent as GD
import qualified MachineLearning.Optimization.MinibatchGradientDescent as MGD
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

import qualified Numeric.GSL.Minimization as Min

data MinimizeMethod = GradientDescent R       -- ^ Gradient descent, takes alpha. Requires feature normalization.
                    | MinibatchGradientDescent Int Int R  -- ^ Minibacth Gradietn Descent, takes seed, batch size and alpha
                    | ConjugateGradientFR R R -- ^ Fletcher-Reeves conjugate gradient algorithm,
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).
                    | ConjugateGradientPR R R -- ^ Polak-Ribiere conjugate gradient algorithm.
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).
                    | BFGS2 R R               -- ^ Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm,
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).

-- | Returns solution vector (theta) and optimization path.
-- Optimization path's row format:
-- [iter number, cost function value, theta values...]
minimize :: Model.Model a =>
            MinimizeMethod
         -> a       -- ^ model (Least Squares, Logistic Regression etc)
         -> R   -- ^ epsilon, desired precision of the solution
         -> Int     -- ^ maximum number of iterations allowed
         -> Regularization   -- ^ regularization parameter
         -> Matrix  -- ^ X
         -> Vector  -- ^ y
         -> Vector  -- ^ initial solution, theta
         -> (Vector, Matrix) -- ^ solution vector and optimization path

minimize (BFGS2 firstStepSize tol) = minimizeVD Min.VectorBFGS2 firstStepSize tol
minimize (ConjugateGradientFR firstStepSize tol) = minimizeVD Min.ConjugateFR firstStepSize tol
minimize (ConjugateGradientPR firstStepSize tol) = minimizeVD Min.ConjugatePR firstStepSize tol
minimize (GradientDescent alpha) = GD.gradientDescent alpha
minimize (MinibatchGradientDescent seed batchSize alpha) = MGD.minibatchGradientDescent seed batchSize alpha

minimizeVD method firstStepSize tol model epsilon niters reg x y initialTheta
  = Min.minimizeVD method epsilon niters firstStepSize tol costf gradientf initialTheta
  where costf = Model.cost model reg x y
        gradientf = Model.gradient model reg x y

-- | Gradient checking function.
-- Approximates the derivates of the Model's cost function
-- and calculates derivatives using the Model's gradient functions.
-- Returns norn_2 between 2 derivatives.
-- Takes model, regularization, X, y, theta and epsilon (used to approximate derivatives, 1e-4 is a good value).
checkGradient :: Model a => a -> Regularization -> Matrix -> Vector -> Vector -> R -> R
checkGradient model reg x y theta eps = LA.norm_2 $ grad1 - grad2
  where theta_m = LA.asColumn theta
        eps_v = V.replicate (V.length theta) eps
        eps_m = LA.diag eps_v
        -- +/- eps_m in case of zero theta
        theta_m1 = theta_m * (1 + eps_m) + eps_m
        theta_m2 = theta_m * (1 - eps_m) - eps_m
        cost1 = LA.vector $ map (cost model reg x y) $ LA.toColumns theta_m1
        cost2 = LA.vector $ map (cost model reg x y) $ LA.toColumns theta_m2
        eps_v1 = LA.takeDiag $ theta_m1 - theta_m
        eps_v2 = LA.takeDiag $ theta_m - theta_m2
        grad1 = (cost1 - cost2) / (eps_v1 + eps_v2)
        grad2 = gradient model reg x y theta


module MachineLearningTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Plus

import MachineLearning

import qualified Numeric.LinearAlgebra as LA

x = LA.matrix 3 [17, 19, 29]
x2 = LA.matrix 9 [17, 19, 29, 289, 323, 493, 361, 551, 841]
x3 = LA.matrix 19 [17, 19, 29, 289, 323, 493, 361, 551, 841, 4913, 5491, 8381, 6137, 9367, 14297, 6859, 10469, 15979, 24389]
eps = 1e-5

tests = [testGroup "mapFeatures" [
            testCase "mapfeatures 1" $ assertMatrix "" eps x (mapFeatures 1 x)
            , testCase "mapfeatures 2" $ assertMatrix "" eps x2 (mapFeatures 2 x)
            , testCase "mapfeatures 3" $ assertMatrix "" eps x3 (mapFeatures 3 x)
            ]
        ]


{-# LANGUAGE ScopedTypeVariables, TypeFamilies #-}

module Fuml.Base.KNN where

import qualified Data.Vector.Storable as VS
import Numeric.LinearAlgebra
import Data.List (nub, sortBy)
import Data.Ord (comparing)

euclideanDistance :: Vector Double -> Vector Double -> Double
euclideanDistance v1 v2 = sqrt $ VS.sum $ VS.map (^2) $ VS.zipWith (-) v1 v2

weightedBoolVote :: [(Double, Bool)] -> Double
weightedBoolVote distBools =
  let wtrue = sum $ map (recip . fst) $ filter snd distBools
      wfalse = sum $ map (recip . fst) $ filter (not . snd) distBools
  in exp wtrue / (exp wtrue + exp wfalse)

majorityVote :: Eq a => [(Double, a)] -> a
majorityVote distXs = let classes = nub $ map snd distXs
                          occurences c = (c,negate $ length $ filter ((==c) . snd) distXs)
                      in fst $ head $ sortBy (comparing snd) $ map occurences classes

weightedMajorityVote :: Eq a => [(Double, a)] -> a
weightedMajorityVote distXs
  = let classes = nub $ map snd distXs
        weight c = (c,negate $ sum $ map (recip . fst) $ filter ((==c) . snd) distXs)
        in fst $ head $ sortBy (comparing snd) $ map weight classes

avgVote :: [(Double, Double)] -> Double
avgVote distXs = let n = realToFrac $ length distXs
                 in (sum $ map snd distXs) / n

weightedAvgVote :: [(Double, Double)] -> Double
weightedAvgVote distXs
  = let wsum = sum $ map (recip . fst) distXs
    in  (sum $ map (uncurry (*)) distXs) / wsum


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Eigen
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Eigenvalue decompositions.
--
module Numeric.LinearAlgebra.Matrix.Eigen (
    hermEigen,
    hermEigenvalues,
    hermEigenM_,
    hermEigenvaluesM_,
    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Foreign( allocaArray, nullPtr )
import Text.Printf( printf )

import Foreign.LAPACK( LAPACK, EigJob(..), EigRange(..) )
import qualified Foreign.LAPACK as LAPACK

import Numeric.LinearAlgebra.Types( Herm(..) )
import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import Numeric.LinearAlgebra.Vector( Vector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V

-- | Compute the eigenvalues and eigenvectors of a Hermitian matrix.
-- Return the eigenvalues are in ascending order in the result vector;
-- store the corresponding eigenvectors are in the columns of the result
-- matrix.
hermEigen :: (LAPACK e) => Herm Matrix e -> (Vector Double, Matrix e)
hermEigen (Herm uplo a) = runST $ do
    (_,n) <- M.getDim a
    ma <- M.newCopy a
    mw <- V.new_ n
    mz <- M.new_ (n,n)
    hermEigenM_ (Herm uplo ma) mw mz
    w <- V.unsafeFreeze mw
    z <- M.unsafeFreeze mz
    return (w,z)

-- | Return the eigenvalues of a Hermitian matrix in ascending order.
hermEigenvalues :: (LAPACK e) => Herm Matrix e -> Vector Double
hermEigenvalues (Herm uplo a) = V.create $ do
    (_,n) <- M.getDim a
    ma <- M.newCopy a
    w <- V.new_ n
    hermEigenvaluesM_ (Herm uplo ma) w
    return w

-- | Compute and copy the eigenvalues and eigenvectors of a mutable
-- Hermitian matrix.  This destroys the original Hermitian matrix.
hermEigenM_ :: (LAPACK e)
                  => Herm (STMatrix s) e
                  -> STVector s Double
                  -> STMatrix s e
                  -> ST s ()
hermEigenM_ (Herm uplo a) w z = do
    (ma,na) <- M.getDim a
    (mz,nz) <- M.getDim z
    nw <- V.getDim w
    let n = na

    when (ma /= na) $ error $
        printf ("hermEigenM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>): nonsquare matrix"
               ) ma na
    when ((not . and) [ (ma,na) == (n,n)
                      , nw == n
                      , (mz,nz) == (n,n)
                      ]) $ error $
        printf ("hermEigenM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>:"
                ++ " <matrix with dim (%d,%d)>"
                ++ " dimension mismatch")
               ma na nw mz nz

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith w $ \pw ->
        M.unsafeWith z $ \pz ldz ->
        allocaArray (2*n) $ \psuppz -> do
            _m <- LAPACK.heevr jobz range uplo n pa lda abstol pw pz ldz
                               psuppz
            return ()
  where
    jobz = EigVectors
    range = AllEigs
    abstol = 0

-- | Compute and copy the eigenvalues of a mutable Hermitian matrix.  This
-- destroys the original Hermitian matrix.
hermEigenvaluesM_ :: (LAPACK e)
                        => Herm (STMatrix s) e
                        -> STVector s Double
                        -> ST s ()
hermEigenvaluesM_ (Herm uplo a) w = do
    (ma,na) <- M.getDim a
    nw <- V.getDim w
    let n = na

    when (ma /= na) $  error $
        printf ("hermEigenvaluesM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>): nonsquare matrix"
               ) ma na

    when ((not . and) [ (ma,na) == (n,n)
                      , nw == n
                      ]) $ error $
        printf ("hermEigenvaluesM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>: dimension mismatch")
               ma na nw

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith w $ \pw -> do
            _m <- LAPACK.heevr jobz range uplo n pa lda abstol pw pz ldz
                               psuppz
            return ()

  where
    jobz = NoEigVectors
    range = AllEigs
    abstol = 0
    pz = nullPtr
    ldz = 1
    psuppz = nullPtr


module Statistics where

import           Control.Monad
import           Control.Monad.Bayes.Class      ( bernoulli
                                                , MonadSample
                                                , discrete
                                                )
import           Numeric.Log
import           Statistics.Distribution        ( logProbability )
import           Statistics.Distribution.Poisson
                                                ( poisson )
import qualified Statistics.Distribution.Binomial
                                               as SB

boolToInt :: Bool -> Int
boolToInt True  = 1
boolToInt False = 0

binomial :: MonadSample m => Int -> Double -> m Int
--binomial n p = Prelude.sum <$> replicateM n (boolToInt <$> bernoulli p)
binomial n p = discrete $ SB.binomial n p

poissonPdf
  :: Double
  ->
  -- | lambda
     Int
  ->
  -- | x
     Log Double
poissonPdf lambda x = Exp $ logProbability (poisson lambda) x


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict #-}
module STC.Reversal where

import           Data.Array.Repa      as R
import           Data.Complex
import           Data.List            as L
import           Data.Vector.Storable as VS
import           Data.Vector.Unboxed  as VU
import           STC.DFTArray
import           Types
import FourierPinwheel.Array

{-# INLINE computeReversalR2Z2T0S0 #-}
computeReversalR2Z2T0S0 ::
     (R.Source r (Complex Double))
  => [Double]
  -> R.Array r DIM6 (Complex Double)
  -> R.Array D DIM6 (Complex Double)
computeReversalR2Z2T0S0 theta0Freqs sourceArr =
  R.traverse2
    sourceArr
    (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
    const $ \f1 f2 idx@(Z :. _ :. _ :. l :. _ :. i :. j) ->
    f1 idx * exp (0 :+ f2 (Z :. l) * pi)

{-# INLINE timeReversal #-}
timeReversal :: DFTArray -> DFTArray
timeReversal arr@(DFTArray _ _ thetaFreqs rFreqs _) =
  repaToDFTArray thetaFreqs rFreqs .
  computeS .
  R.traverse2
    (dftArrayToRepa arr)
    (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
    const $ \fArr fFreq idx@(Z :. _ :. theta :. _ :. _) ->
    fArr idx * (cis $ (fFreq (Z :. theta)) * pi)

{-# INLINE timeReversal' #-}
timeReversal' :: DFTArray -> DFTArray
timeReversal' arr@(DFTArray rows cols thetaFreqs _ _) =
  repaToDFTArray thetaFreqs [] .
  computeS .
  R.traverse2
    (fromUnboxed (Z :. (1 :: Int) :. (L.length thetaFreqs) :. cols :. rows) .
     VS.convert . VS.concat . getDFTArrayVector $
     arr)
    (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
    const $ \fArr fFreq idx@(Z :. _ :. theta :. _ :. _) ->
    fArr idx * (cis $ (fFreq (Z :. theta)) * pi)

{-# INLINE timeReversal1 #-}
timeReversal1 :: Int -> Int -> DFTArray -> DFTArray
timeReversal1 numThetaFreq numRFreq arr@(DFTArray _ _ thetaFreqs rFreqs _) =
  repaToDFTArray thetaFreqs rFreqs .
  computeS .
  R.traverse2
    (dftArrayToRepa1 numThetaFreq numRFreq arr)
    (fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs)
    const $ \fArr fFreq idx@(Z :. _ :. theta :. _ :. _) ->
    fArr idx * (cis $ (fFreq (Z :. theta)) * pi)

{-# INLINE timeReversalRepa #-}
timeReversalRepa ::
     (Num e, RealFloat e, R.Source s (Complex e), Unbox e)
  => [e]
  -> R.Array s DIM4 (Complex e)
  -> R.Array D DIM4 (Complex e)
timeReversalRepa thetaFreqs arr =
  R.traverse2 arr (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs) const $ \fArr fFreq idx@(Z :. _ :. theta :. _ :. _) ->
    fArr idx * cis (fFreq (Z :. theta) * pi)


module Main where

import Control.Applicative
import Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import qualified Graphics.Rendering.Chart.Easy as Plt
import Kalman
import Numeric.LinearAlgebra
import Sim

fA1 :: (Floating a) => [a] -> [a]
fA1 [x, y, xh, yh, ax, ay] =
  [ x + xh * dt + c * ax * dt * dt,
    y + yh * dt + c * ay * dt * dt,
    xh + ax * dt,
    yh + ay * dt,
    ax,
    ay
  ]
  where
    dt = 0.1
    c = 0.5

fH1 :: (Floating a) => [a] -> [a]
fH1 [x, y, xh, yh, ax, ay] = [x, y]

cannonBall = toFile Plt.def "CannonBall.svg" $ do
  -- KF Initial conditions
  let x0 = [0, 0, 10, 10, 0, -9.8]
  let q = diagl [0.1, 0.1, 1e-3, 1e-3, 1e-3, 1e-3]
  let r = diagl (replicate 2 0.5)
  let s0 = Estimate (vector x0) (diagl (replicate 6 5))

  -- Fire the cannon ball
  let plant = Plant fA1 id fH1 id
  let sim = take 50 $ simulate plant (vector x0) 0
  let measurements = map snd sim
  let trajectory = map fst sim

  -- Run KF/EKF/UKF
  let sys = System fA1 fH1 q r
  let ekfs = scanl (ekf sys) s0 measurements
  let ukfs = scanl (ukf sys) s0 measurements

  -- Plot the results
  let vector2Chart = map ((\[x, y] -> (x, y)) . toList . subVector 0 2)
  let state2Chart xs = vector2Chart $ map sX xs

  Plt.plot (Plt.line "Plant" [vector2Chart trajectory])
  Plt.plot (Plt.line "UKF" [state2Chart ukfs])
  Plt.plot (Plt.line "EKF" [state2Chart ekfs])
  Plt.plot (Plt.points "measurements" (vector2Chart measurements))

main = cannonBall


{-# LANGUAGE FlexibleContexts #-}

module Main where

import Graphics.Plot

import Data.Packed.Vector
import Data.Packed.Matrix
import Numeric.Container
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.HMatrix (tr)

import System.Environment (getArgs)

-- TODO (a->b->c) -> Vector a -> Vector b -> Matrix c
outerOp :: (Element a, Element b) => (Matrix a -> Matrix b -> Matrix c) -> Vector a -> Vector b -> Matrix c
outerOp k x x' = k xp xq
    where xp = repmat (asColumn x) 1 (dim x')
          xq= repmat (asRow x') (dim x) 1

k_SE sigma l x x' = (sigma^2*) $ exp $ (-1/(2*l^2))*dx^2
    where dx = outerOp (-) x x'

k = k_SE 3 0.3

inferenceTrivial k x y ySigma2 domain = (mean, variance)
    where mean = k domain x <> invkxx <> y
          variance = k domain domain - (kdomainx <> invkxx <> k x domain)
          invkxx = inv $ k x x + diag ySigma2
          kdomainx = k domain x

inferenceCholesky k x y ySigma2 domain = (mean, variance)
    where mean = (tr kxdomain) <> alpha
          variance = k domain domain - (tr v <> v)
          alpha = tr l <\> (l <\> y)
          l = chol $ k x x + diag ySigma2
          v = l <\> kxdomain
          kxdomain = k x domain

plotProcessDistribution :: Vector Double -> Vector Double -> Vector Double -> IO ()
plotProcessDistribution domain m v = mplot [domain, m, m + v, m - v, m + 2 * v, m - 2 * v]

main :: IO ()
main = do
    datfile:[] <- getArgs
    dat <- fmap readMatrix $ readFile datfile
    let [x, y] = toColumns dat -- [x,y] or [x,f]
    let domain = linspace 128 (0.5, 10.5)

    let n = 1
    let (x',y') = (vjoin $ replicate n x, vjoin $ take n $ repeat y)

    let (m, v') = inferenceTrivial k x' y' (constant 1 $ dim y') domain
    let v = takeDiag v'

    plotProcessDistribution domain m v


{-# LANGUAGE ScopedTypeVariables #-}

module Qubism.AlgebraTests
  ( module Test.QuickCheck.Property.Generic
  , isVectorSpace, isHilbertSpace, isAlgebra)
  where

import Test.Hspec
import Test.QuickCheck
import Test.QuickCheck.Property.Generic
import Test.QuickCheck.Property.Common.Internal
import Data.Complex

import Qubism.Algebra

tupled :: T a -> T (a,a)
tupled _ = T

prop_Distributive
  :: (a -> b -> b) -> (b -> b -> b) -> T a -> T b -> a -> b -> b -> Equal b
prop_Distributive f g _ _ a b c = Equal left right
  where left  = a `f` (b `g` c)
        right = (a `f` b) `g` (a `f` c)

isVectorSpace :: (Show v, Arbitrary v, VectorSpace v) => T v -> Spec
isVectorSpace t = do
  it "+: is associative"          $ property . eq $ prop_Associative (+:) t
  it "+: is commutative"          $ property . eq $ prop_Commutative (+:) (tupled t)
  it "zero is the identity of +:" $ property . eq $ prop_Identity zero (+:) t
  it "neg produces the inverse"   $ property $
    let test (_ :: T v) (v :: v) = neg v +: v == zero
    in  test t
  it ".: distrbutes over +:" $ property . eq $
    prop_Distributive (.:) (+:) (T :: T C) t

prop_InnerLinear :: VectorSpace v => (v -> C) -> T v -> C -> v -> v -> Bool
prop_InnerLinear f _ a u v = magnitude (left - right) < 0.00001
  where left  = f $ (a .: u) +: v
        right = a * f u + f v

isHilbertSpace :: (Show v, Arbitrary v, HilbertSpace v) => T v -> Spec
isHilbertSpace t = do
  it "<.> is linear in the second argument" $ property $
    \w a u v -> prop_InnerLinear (w <.>) t a u v
  it "<a.b> = conj <b.a>" $ property $
    let test (_ :: T v)(a :: v)(b :: v) = a <.> b `shouldBe` conjugate (b <.> a)
    in  test t

prop_AlgLinear :: (Show v, Algebra v) => (v -> v) -> T v -> C -> v -> v -> Expectation
prop_AlgLinear f _ a u v = left `shouldBe` right
  where left  = f $ (a .: u) +: v
        right = (a .: f u) +: f v

isAlgebra :: (Show v, Arbitrary v, Algebra v) => T v -> Spec
isAlgebra t = do
  it "*: is linear in the first argument" $ property $
    \w a u v -> prop_AlgLinear (*: w) t a u v
  it "*: is linear in the second argument" $ property $
    \w a u v -> prop_AlgLinear (w *:) t a u v


module Fractal.Utils.Field(
  generateField,
  generateRow,
  generateRowField,
  getFieldLength
)where

import Data.Complex

type Value = Complex Double
type Point = Value
type Row = [Point]
type Field = [Row]
type Step = Double

generateField :: Point -> Point -> Step -> Field
generateField start end step
  | (imagPart start) <= (imagPart end)     = (generateRow start end step):[]
  | otherwise                              = (generateRow start end step) : (generateField newStart end step)
      where newStart = (realPart start) :+ ((imagPart start) - step)

generateRow :: Point -> Point -> Step -> Row
generateRow start end step
  | (realPart start) >= (realPart end)  = [start]
  | otherwise                           = start : (generateRow newStart end step)
      where newStart = ((realPart start) + step) :+ (imagPart start)

generateRowField :: Point -> Point -> Step -> Row
generateRowField start end step = generateRowField' start start end step

generateRowField' :: Point -> Point -> Point -> Step -> Row
generateRowField' mark start end step
  | (imagPart start) <= (imagPart end) = endroutine
  | otherwise = fieldroutine
    where endroutine = if (realPart start) >= (realPart end) then [start] else start : (generateRowField' mark newStart end step)
          fieldroutine = if (realPart start) >= (realPart end) then start:(generateRowField' mark newRowStart end step) else start : (generateRowField' mark newStart end step)
          newStart = ((realPart start) + step) :+ (imagPart start)
          newRowStart = (realPart mark) :+ ((imagPart start) - step)

getFieldLength :: Point -> Point -> Step -> Int
getFieldLength a b step = count a b 0
  where count x y i = if (realPart x) >= (realPart y) then (i+1) else count (newX x) y (i+1)
        newX x = ((realPart x) + step) :+ (imagPart x)


module Main where

import Criterion.Main (defaultMain, bgroup, bench, nf, nfIO, whnf)
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.Morpheus as M

main :: IO ()
main = do
  !a <- LA.rand 1000 1000
  !b <- LA.rand 1000 1000
  !x <- LA.flatten <$> LA.rand 1000 1000
  !y <- LA.flatten <$> LA.rand 1000 1000
  let !ata = LA.mTm a

  defaultMain [
    bgroup "Linear Algebra" [
        bench "random matrix" $ nfIO (LA.rand 1000 1000)
        , bench "inv" $ nf LA.inv a
        , bench "Cholesky decomposition" $ nf LA.chol ata
        , bench "vector dot product" $ whnf (LA.dot x) y
        , bench "matrix transpose" $ nf LA.tr a
        , bench "matrix multiplication" $ nf (a LA.<>) b
        , bench "sigmoid" $ nf sigmoid a
        , bench "ReLU" $ nf relu a
        , bench "sum by rows" $ nf (reduceByRows V.sum) a
        , bench "sum by columns" $ nf (reduceByColumns V.sum) a
        , bench "max index in rows" $ nf (reduceByRowsV (fromIntegral . LA.maxIndex) ) a
        , bench "max index in columns" $ nf (reduceByColumnsV (fromIntegral . LA.maxIndex) ) a
        , bench "sigmoid (morpheus)" $ nf M.sigmoid a
        , bench "sigmoid gradient (morpheus)" $ nf M.sigmoidGradient a
        , bench "ReLU (morpheus)" $ nf M.relu a
        , bench "sum by rows (morpheus)" $ nf M.rowSum a
        , bench "sum by columns (morpheus)" $ nf M.columnSum a
        , bench "max index in rows (morpheus)" $ nf M.rowMaxIndex a
        , bench "max index in columns (morpheus)" $ nf M.columnMaxIndex a

        ]
    ]

type Matrix = LA.Matrix LA.R
type Vector = LA.Vector LA.R
type R = LA.R

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

relu :: Matrix -> Matrix
relu x = x * (LA.step x)

reduceByRowsV :: (Vector -> R) -> Matrix -> Vector
reduceByRowsV f = LA.vector . map f . LA.toRows

reduceByColumnsV :: (Vector -> R) -> Matrix -> Vector
reduceByColumnsV f = LA.vector . map f . LA.toColumns

reduceByRows :: (Vector -> R) -> Matrix -> Matrix
reduceByRows f = LA.asColumn . reduceByRowsV f

reduceByColumns :: (Vector -> R) -> Matrix -> Matrix
reduceByColumns f = LA.asRow . reduceByColumnsV f


import Data.Map.Strict ((!))
import Numeric.LinearAlgebra hiding ((!))
import LeastSquares

leastSquaresEstimate :: Vector R
leastSquaresEstimate = result ! "x"
  where result = minimize (SumSquares (a*x-b)) [c*x:==:d]
        a = Mat $ (3><2) [1.0, 1.0, 1.0, 4.0, 2.0, -1.0]
        b = Vec $ vector [13.0, 27.0, 1.0]
        c = Mat $ (1><2) [3.0, 7.0]
        d = Vec $ vector [-1.0]
        x = Var "x" 2


module Elliptic
  where
import Carlson
import Data.Complex
import Internal

ellipticF' :: Double -> Cplx -> Cplx -> Cplx
ellipticF' err phi m
  | phi == 0 =
    toCplx 0
  | m == 1 && abs(realPart phi) == pi/2 =
    toCplx (0/0)
  | m == 1 && abs(realPart phi) < pi/2 =
    atanh(sin phi)
  | abs(realPart phi) <= pi/2 =
    if m == 0
      then
        phi
      else
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * carlsonRF' err cosine2 oneminusmsine2 1
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticF' err (pi/2) m + ellipticF' err phi' m

ellipticF :: Cplx -> Cplx -> Cplx
ellipticF = ellipticF' 1e-15

ellipticE' :: Double -> Cplx -> Cplx -> Cplx
ellipticE' err phi m
  | phi == 0 =
    toCplx 0
  | abs(realPart phi) <= pi/2 =
    case m of
      0 -> phi
      1 -> sin phi
      _ ->
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * (carlsonRF' err cosine2 oneminusmsine2 1 -
          m * sine2 / 3 * carlsonRD' err cosine2 oneminusmsine2 1)
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticE' err (pi/2) m + ellipticE' err phi' m

ellipticE :: Cplx -> Cplx -> Cplx
ellipticE = ellipticE' 1e-15

ellipticPI' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
ellipticPI' err phi n m
  | phi == 0 =
    toCplx 0
  | phi == pi/2 && n == 1 =
    0/0
  | phi == pi/2 && m == 0 =
    pi/2/sqrt(1-n)
  | phi == pi/2 && m == n =
    ellipticE' err (pi/2) m / (1-m)
  | phi == pi/2 && n == 0 =
    ellipticF' err (pi/2) m
  | abs(realPart phi) <= pi/2 =
    let sine = sin phi in
    let sine2 = sine*sine in
    let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
    sine * (carlsonRF' err cosine2 oneminusmsine2 1 +
      n * sine2 / 3 * carlsonRJ' err cosine2 oneminusmsine2 1 (1-n*sine2))
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticPI' err (pi/2) n m + ellipticPI' err phi' n m

ellipticPI :: Cplx -> Cplx -> Cplx -> Cplx
ellipticPI = ellipticPI' 1e-15

jacobiZeta' :: Double -> Cplx -> Cplx -> Cplx
jacobiZeta' err phi m =
  if m == 1
    then
      if abs(realPart phi) <= pi/2
        then sin phi
        else let (phi',_) = getPhiK phi in sin phi'
    else
      ellipticE' err phi m -
        ellipticE' err (pi/2) m / ellipticF' err (pi/2) m *
        ellipticF' err phi m

jacobiZeta :: Cplx -> Cplx -> Cplx
jacobiZeta = jacobiZeta' 1e-15