{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}

module Advection1D
where

import qualified Basis
import           Data.List
import qualified Element
import           Flux
import           Mesh
import           Numeric.LinearAlgebra
--import           Numeric.LinearAlgebra.HMatrix
import           Quadrature            (integrate)
import qualified ShapeFcns

-- Synonym for numeric/fractional constraints
type FrElNuFi a = (Fractional a,Element a,Numeric a,Field a)

-- Element stiffness integrand
stiffnessIntegrand :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> [a] -> Matrix a
--stiffnessIntegrand elem shpFcn xi = fromLists [[
--  ShapeFcns.n shpFcn xi i [0] * atIndex (Element.dndx elem shpFcn xi j) 0 * Element.computeJacobianDet elem shpFcn xi
--  | i <- idxRange] | j <- idxRange]
stiffnessIntegrand elem shpFcn xi = fromLists [[
  (atIndex (Element.dndx elem shpFcn xi i) 0 * atIndex (Element.dndx elem shpFcn xi j) 0 + ShapeFcns.n shpFcn xi i [0] * ShapeFcns.n shpFcn xi j [0]) * Element.computeJacobianDet elem shpFcn xi
  | i <- idxRange] | j <- idxRange]
  where
    idxRange = [0..(Element.getNumNodes elem - 1)]

-- Element mass integrand
massIntegrand :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> [a] -> Matrix a
massIntegrand elem shpFcn xi = fromLists [[
  ShapeFcns.n shpFcn xi i [0] * ShapeFcns.n shpFcn xi j [0] * Element.computeJacobianDet elem shpFcn xi
  | i <- idxRange] | j <- idxRange]
  where
    idxRange = [0..(ShapeFcns.getShapeFcnOrder shpFcn)]

-- Computes the element flux matrix assuming an upwind flux.
-- Assumes that advSpeed is positive.
elemFluxMatrix :: FrElNuFi a => a -> Matrix a
elemFluxMatrix advSpeed = scale advSpeed $ fromLists [[-1.0, 0.0, 0.0],[0.0, 0.0, 1.0]]

-- Element mass and stiffness matrices
-- Must multiply by the determinant of the Jacobian here!!!
elemMatrices :: (Element.Element e,ShapeFcns.ShapeFcn s,Basis.Basis b,FrElNuFi a) => e a -> s b -> Int -> (Matrix a,Matrix a)
elemMatrices elem shpFcn ngpts = (stiffnessMat, massMat)
  where
    stiffnessMat = integrate 1 ngpts (stiffnessIntegrand elem shpFcn)
    massMat      = integrate 1 ngpts (massIntegrand elem shpFcn)

-- Function which computes the dimensions of the global matrices
globalMatrixDimension :: Int -> Int
globalMatrixDimension numNodes = 2 * numNodes - 2

-- Function which takes a local element matrix along with an element and transforms it to global
localMatToGlobal :: (Element.Element e,FrElNuFi a) => e a -> Matrix a -> [((Int, Int), a)]
localMatToGlobal elem elemMat = concat [[((Mesh.globalNodeNum en i, Mesh.globalNodeNum en j), atIndex elemMat (i,j) ) | i <- [0..dim-1]] | j <- [0..dim-1]]
  where
    en = Element.getElementNumber elem
    dim = (fst . size) elemMat

-- Function which does the same as the above, but with contributions due to fluxes.
localFluxMatToGlobal :: (Element.Element e,FrElNuFi a) => e a -> Matrix a -> [((Int, Int), a)]
localFluxMatToGlobal elem elemMat = concat [[((Mesh.globalNodeNum en i, Mesh.globalNodeNum en j - 1), atIndex elemMat (i,j) ) | i <- [0..dim-1]] | j <- [0..dim]]
  where
    en = Element.getElementNumber elem
    dim = (fst . size) elemMat

-- Function which assembles global stiffness and mass matrices
-- This function takes a mesh and returns a tuple which corresponds to the
-- global stiffness and global mass matrices, respectively.
-- THIS ONLY WORKS WITH THE DISCONTINUOUS MESH...  IT DOES NOT SUM DUPLICATE ENTRIES.
assembleGlobalMatrices :: (Element.Element e, ShapeFcns.ShapeFcn s,Basis.Basis b, FrElNuFi a) => Mesh e a -> s b -> Int -> a -> [[((Int,Int),a)]]
assembleGlobalMatrices grid shpFcn ngpts advSpd = [globalK, globalM, globalF]
  where
    elems = Mesh.getMeshElements grid
    nnodes = length (concatMap Element.getElementNodes elems)
    elemMats = map (\elem -> elemMatrices elem shpFcn ngpts) elems
    elemK = map fst elemMats
    elemM = map snd elemMats
    elemF = elemFluxMatrix advSpd
    globalK = concat (zipWith localMatToGlobal elems elemK)
    globalM = concat (zipWith localMatToGlobal elems elemM)
    globalF = filter (\entry -> (snd . fst) entry >= 0 && (snd . fst) entry < nnodes) (concatMap (`localFluxMatToGlobal` elemF) elems)

-- This function computes the right hand side of the semi-discrete statement.
-- udot_j = (M_ij)^(-1) ( a K_ij - F_ij ) u_j.
-- This is what needs to be integrated point-by-point.
uDot :: (Floating a,Field a,Num (Vector a)) => Matrix a -> Matrix a -> Matrix a -> a -> Vector a -> Vector a
uDot globalM globalK globalF advSpd u = inv globalM #> ((scale advSpd globalK - globalF) #> u)

-- This function advances the solution one step in time.
advance :: (Floating a,Field a,Num (Vector a)) => Matrix a -> Matrix a -> Matrix a -> a -> Vector a -> a -> Vector a
advance globalM globalK globalF advSpd u dt = scale dt (uDot globalM globalK globalF advSpd u - u)

-- This function applies boundary conditions to the semi-discrete system.
applyBCs _ = []

-- This function computes the solution as a function of time.
{-
solveAdvection numElements icFunc tf = do {
  linBasis <- Lagrange 1;
  shpFcn   <- TensorProduct linBasis;
  grid     <- genDMesh 0.0 1.0 numElements StructElem;
  }
-}


import Criterion.Main
import FFT.Orig
import FFT.FFTControl
import FFT.FFTData
import FFT.DFTData
import FFT.DFTControl
import FFT.Samples
import Data.Complex
import System.Environment

main = defaultMain [
                bgroup "dft-data" [
                    bench "orig" $ nf sum (dft (samples 1 1000)),
                    bench "foldMap" $ nf sum (dftFoldMap (samples 1 1000))
                ],
                bgroup "dft-control" [

                ],
                bgroup "fft-control" [
                    bench "orig" $ nf sum (fft (samples 1 100000)),
                    bench "bflySPipeline" $ nf sum (fftbflySPipeline (samples 1 100000))
                ],
                bgroup "fft-data" [
                ]
    ]


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Marvin.API.Algorithms.KMeans (
  KMeans (..)

  , CentroidModel
  , distancesFromNearestCentroids
  , centroids
) where

import Control.Monad.Reader

import Data.Foldable
import Data.Array.IArray
import Data.Ord
import Control.Arrow
import Control.Monad.Except

import Marvin.API.Table.DataType
import Marvin.API.Fallible
import Marvin.API.Meta.Model
import Marvin.API.Table.DataType
import Marvin.API.Fallible
import Marvin.API.Table.Internal

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

import qualified Numeric.LinearAlgebra as LA

-- * Parameters

-- | Parameters.
data KMeans = KMeans {
  k :: Int -- ^ Number of clusters (centroids).
  , numberOfIterations :: Int -- ^ Number of iterations.
}

-- | Default parameters.
defaultKMeans = KMeans {
  k = 5
  , numberOfIterations = 10
}

-- * Model

-- | Resulting model.
data CentroidModel = CentroidModel {
  centroids_ :: Array Label Centroid
  , dimension :: Int
} deriving Show

-- | Retrieves a list of the centroids.
centroids :: CentroidModel -> [[Double]]
centroids model = fmap LA.toList $ fmap snd $ assocs $ centroids_ model

-- | Retrieves the distances from the nearest centroids for some points
distancesFromNearestCentroids :: CentroidModel -> NumericTable -> Fallible NumericColumn
distancesFromNearestCentroids model tab = do
  if dimension model /= numberOfColumns tab
    then throwError $ RowLengthMismatch $ "While trying the calculate nearest distances " ++
      "with CentroidModel."
    else return ()
  let ps = tableToRows tab
  let distances = fmap (distanceFromNearest model) ps
  fromList $ Vec.toList $ distances

-- | Type aliases.
type Point = LA.Vector Double
type Centroid = LA.Vector Double
type Label = Int

-- | Distance from nearest centroid.
distanceFromNearest :: CentroidModel -> Point -> Double
distanceFromNearest model p = minimum $ fmap (dist p) cs
  where
    cs = fmap snd $ assocs $ centroids_ model

instance Estimator KMeans NumericTable where
  type ResultingModel KMeans NumericTable = CentroidModel
  fit = fit'

fit' :: KMeans -> NumericTable -> Fallible CentroidModel
fit' kMeans table = do
  ensurePositive
    (InvalidAlgorithmParameter "KMeans" "numberOfIterations must be positive") numIter
  ensurePositive
    (InvalidAlgorithmParameter "KMeans" "number of centroids (k) must be positive") numCentroids
  ensureNonEmptyTable "While fitting KMeans." table
  return $ CentroidModel {
    centroids_ = centroidArray
    , dimension = numberOfColumns table
  }
  where
    centroidArray = fitKMeans params numIter
    numIter = numberOfIterations kMeans
    numCentroids = k kMeans
    pts = tableToRows table
    params = KMeansParams {
      numOfCents = numCentroids
      , points = Vec.toList pts
    }

instance Predictor CentroidModel where
  type Testing CentroidModel = NumericTable
  type Prediction CentroidModel = NaturalColumn
  predict model numericTable = if numberOfColumns numericTable /= dimension model
      then throwError $ RowLengthMismatch "When predicting with a CentroidModel."
      else fromList $ Vec.toList labels
    where
      pts = tableToRows numericTable
      labels = fmap (nearestLabel (centroids_ model)) pts

tableToRows =
  Vec.map (LA.fromList . UVec.toList) . transposeVectorVectorU . Vec.map values . columnsVec

data InitStrategy = TakeFirstK
data KMeansParams = KMeansParams {
  points :: [Point]
  , numOfCents :: Int
}

dimensions :: KMeansParams -> Int
dimensions params = LA.size $ head $ points params

initCents :: Reader KMeansParams (Array Label Centroid)
initCents = do
  ps <- asks points
  k <- asks numOfCents
  return $ listArray (1,k) (take k (ps))

dist :: LA.Vector Double -> LA.Vector Double -> Double
dist xs ys = LA.norm_2 $ xs - ys

nearestLabels :: Array Label Centroid -> Reader KMeansParams [(Label, Point)]
nearestLabels cs = do
  ps <- asks points
  return $
    fmap (\p -> (nearestLabel cs p, p)) ps

nearestLabel :: Array Label Centroid -> Point -> Label
nearestLabel cs p = fst $ minimumBy (comparing (dist p . snd)) (assocs cs)

mean :: [Vector Double] -> Vector Double
mean xs = UVec.map (\x -> x / n) $ sumVecs xs
  where
    n = fromIntegral (length xs)
    sumVecs = foldr1 addVecs
    addVecs x y = UVec.zipWith (+) x y

accumArrWithLabs :: Ix i => Array i (Point, Int) -> [(i, Point)] -> Array i (Point, Int)
accumArrWithLabs = accum meanAccum

meanAccum :: (Point, Int) -> Point -> (Point, Int)
meanAccum (sum, cnt) p = (sum + p, cnt + 1)

nextCentroids :: Array Label Centroid -> Reader KMeansParams (Array Label Centroid)
nextCentroids cs = do
  ps <- asks points
  k <- asks numOfCents
  n <- asks dimensions
  let initArr = listArray (1,k) (repeat (n LA.|> [0,0..], 0))
  labeledPoints <- nearestLabels cs
  let means = accum meanAccum initArr labeledPoints
  return $ amap (\(sum, cnt) -> LA.scale (1 / fromIntegral cnt) sum) means

fitKMeans :: KMeansParams -> Int
  -> Array Label Centroid
fitKMeans params numIter = runReader readerIter params
  where
    readerIter = do
      init <- initCents
      runIterationM numIter nextCentroids init

runIterationM :: Monad m => Int -> (a -> m a) -> a -> m a
runIterationM n f init = iterate (>>= f) (return init) !! n


-- Exercisee/Ch2/Pt1/Ex7.hs
module Ex7 where

-- Exercise 7
-- Add data types and parsers to support the full numeric tower of Scheme
-- numeric types. Haskell has built-in types to represent many of these; check
-- the Prelude. For others, you can define compound types that represent eg. a
-- `Rational` as a numerator and denominator, or a `Complex` as a real and an
-- imaginary part (each itself a `Real`).

import Data.Char                            (digitToInt, isSpace)
import Data.Complex
import Data.Ratio
import Control.Monad
import Numeric
import System.Environment

import Text.ParserCombinators.Parsec hiding (spaces)

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal -- (a b c ... . z)
             | Number LispNum
             | Character String
             | String String
             | Bool Bool
             deriving Show -- debug purposes

data LispNum = Complex (Complex Float)
             | Real Float
             | Rational (Ratio Integer)
             | Integer Integer
             deriving (Eq, Show)

parseCharacter :: Parser LispVal
parseCharacter = do
  string "#\\"
  character <- many $ satisfy $ not . isSpace
  return $ Character character

parseString :: Parser LispVal
parseString = do
  char '"'
  -- order is important here; need to try parsing an escape sequence first because
  -- otherwise we fail when we reach the second character of the escape
  -- sequence
  x <- many $ choice $ escChars ++ [nonQuote]
  char '"'
  return $ String x
  where
    nonQuote = noneOf "\""
    -- taken from here:
    -- https://en.wikipedia.org/wiki/Escape_sequences_in_C#Table_of_escape_sequences
    -- (excluded characters by oct/hex codes)
    escChars = [ char '\\' >> char x | x <- "abfnrtv\\'\"?" ]

parseBool :: Parser LispVal
parseBool = do
  try (char '#')
  v <- oneOf "tf"
  return $ case v of
    't' -> Bool True
    'f' -> Bool False

parseAtom :: Parser LispVal
parseAtom = do
              first <- letter <|> symbol
              rest <- many (letter <|> digit <|> symbol)
              return $ Atom (first:rest)

parseNumber :: Parser LispVal
parseNumber = do
  n <- (try parseBin) <|> (try parseOct) <|> (try parseHex) <|> (try parseFlt)
    <|> (try parseRat) <|> (try parseCmp) <|> parseDec
  return $ Number n
  where
    parseBin = do
      string "#b"
      binStr <- many1 $ oneOf "01"
      -- cribbed from http://stackoverflow.com/a/26961027/1893155
      let binVal = foldl (\acc x -> acc * 2 + digitToInt x) 0 binStr
      return $ Integer (toInteger binVal)
    parseOct = do
      string "#o"
      octStr <- many1 octDigit
      let octVal = fst $ (readOct octStr) !! 0
      return $ Integer octVal
    parseDec = parseDecNoPre <|> parseDecPre
    parseDecNoPre = do
      decStr <- many1 digit
      return $ (Integer . read) decStr
    parseDecPre = do
      string "#d"
      decStr <- many1 digit
      return $ (Integer . read) decStr
    parseHex = do
      string "#x"
      hexStr <- many1 hexDigit
      let hexVal = fst $ (readHex hexStr) !! 0
      return $ Integer hexVal
    parseFlt = do
      whole <- many1 digit
      string "."
      mantissa <- many1 digit
      let fltVal = fst $ (readFloat (whole ++ "." ++ mantissa) !! 0)
      return $ Real $ fltVal
    parseRat = do
      m <- many1 digit
      many space
      string "/"
      many space
      n <- many1 digit
      let (m', n') = (fst $ (readDec m) !! 0, fst $ (readDec n) !! 0)
      return $ Rational $ m' % n'
    parseCmp = do
      (Real a) <- (try parseDec) <|> parseFlt
      many space
      string "+"
      many space
      (Real b) <- (try parseDec) <|> parseFlt
      char 'i'
      return $ Complex $ a :+ b

-- TODO : I don't like all these `try`s -- better way to do this?
parseExpr :: Parser LispVal
parseExpr = (try parseNumber)
         <|> (try parseString)
         <|> (try parseCharacter)
         <|> (try parseBool)
         <|> (try parseAtom)

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

spaces :: Parser ()
spaces = skipMany1 space

readExpr :: String -> String
readExpr input = case parse parseExpr "lisp" input of
  Left err -> "No match: " ++ show err
  Right val -> "Found value"

main :: IO ()
main = do
  (expr:_) <- getArgs
  putStrLn (readExpr expr)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Gelu
Description : Gaussian Error Linear Unit (GELU)
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental

This module implements the Gaussian Error Linear Unit (GELU) activiation function. See

Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).

Available at: https://arxiv.org/pdf/1606.08415.pdf

As in the paper we simply let μ = 0 and σ = 1. Futher, we use the simplified and thus fast representation: x * sigmoid (1.702 * x)

-}
module Grenade.Layers.Gelu (
    Gelu (..)
  , SpecGelu (..)
  , specGelu1D
  , specGelu2D
  , specGelu3D
  , gelu
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Constraint                (Dict (..))
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build

-- | A Gaussion Error Linear Unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
--
--   Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).
data Gelu = Gelu
  deriving (Generic, NFData, Show)

instance UpdateLayer Gelu where
  type Gradient Gelu = ()
  runUpdate _ _ _ = Gelu

instance RandomLayer Gelu where
  createRandomWith _ _ = return Gelu

instance Serialize Gelu where
  put _ = return ()
  get = return Gelu

geluForwardFast :: Floating x => x -> x
geluForwardFast x = x / (e ** (-1.702 * x) + 1) -- = x * sigmoid (1.702 * x)
  where
    e = 2.71828

geluBackwardFast :: Floating x => x -> x
geluBackwardFast x = (e ** (1.702 * x) * (1 + e ** (1.702 * x) + 1.702 * x)) / (1 + e ** (1.702 * x)) ** 2
  where
    e = 2.71828

instance (KnownNat i) => Layer Gelu ('D1 i) ('D1 i) where
  type Tape Gelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (gelu y))
    where
      gelu = LAS.dvmap geluForwardFast
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (gelu' y * dEdy))
    where
      gelu' = LAS.dvmap geluBackwardFast

instance (KnownNat i, KnownNat j) => Layer Gelu ('D2 i j) ('D2 i j) where
  type Tape Gelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Gelu ('D3 i j k) ('D3 i j k) where

  type Tape Gelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

-------------------- DynamicNetwork instance --------------------

instance FromDynamicLayer Gelu where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecGelu (tripleFromSomeShape inp)

instance ToDynamicLayer SpecGelu where
  toDynamicLayer _ _ (SpecGelu (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     case (rows, cols, depth) of
         (_, 1, 1)    -> return $ SpecLayer Gelu (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer Gelu (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer Gelu (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a elu layer.
specGelu1D :: Integer -> SpecNet
specGelu1D i = specGelu3D (i, 1, 1)

-- | Create a specification for a elu layer.
specGelu2D :: (Integer, Integer) -> SpecNet
specGelu2D (i, j) = specGelu3D (i, j, 1)

-- | Create a specification for a elu layer.
specGelu3D :: (Integer, Integer, Integer) -> SpecNet
specGelu3D = SpecNetLayer . SpecGelu

-- | Add a Gelu layer to your build.
gelu :: BuildM ()
gelu = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecGelu

-------------------- GNum instances --------------------

instance GNum Gelu where
  _ |* Gelu = Gelu
  _ |+ Gelu = Gelu


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.ChiSquared
-- Copyright : (c) 2010 Alexey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The chi-squared distribution. This is a continuous probability
-- distribution of sum of squares of k independent standard normal
-- distributions. It's commonly used in statistical tests
module Statistics.Distribution.ChiSquared (
          ChiSquared
        , chiSquaredNDF
        -- * Constructors
        , chiSquared
        , chiSquaredE
        ) where

import Control.Applicative
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.SpecFunctions ( incompleteGamma,invIncompleteGamma,logGamma,digamma)
import Numeric.MathFunctions.Constants (m_neg_inf)
import qualified System.Random.MWC.Distributions as MWC

import qualified Statistics.Distribution         as D
import Statistics.Internal

-- | Chi-squared distribution
newtype ChiSquared = ChiSquared
  { chiSquaredNDF :: Int
    -- ^ Get number of degrees of freedom
  }
  deriving (Eq, Typeable, Data, Generic)

instance Show ChiSquared where
  showsPrec i (ChiSquared n) = defaultShow1 "chiSquared" n i
instance Read ChiSquared where
  readPrec = defaultReadPrecM1 "chiSquared" chiSquaredE

-- | Construct chi-squared distribution. Number of degrees of freedom
--   must be positive.
chiSquared :: Int -> ChiSquared
chiSquared n = maybe (error $ errMsg n) id $ chiSquaredE n

-- | Construct chi-squared distribution. Number of degrees of freedom
--   must be positive.
chiSquaredE :: Int -> Maybe ChiSquared
chiSquaredE n
  | n <= 0    = Nothing
  | otherwise = Just (ChiSquared n)

errMsg :: Int -> String
errMsg n = "Statistics.Distribution.ChiSquared.chiSquared: N.D.F. must be positive. Got " ++ show n

instance D.Distribution ChiSquared where
  cumulative = cumulative

instance D.ContDistr ChiSquared where
  density chi x
    | x <= 0    = 0
    | otherwise = exp $ log x * (ndf2 - 1) - x2 - logGamma ndf2 - log 2 * ndf2
    where
      ndf  = fromIntegral $ chiSquaredNDF chi
      ndf2 = ndf/2
      x2   = x/2

  logDensity chi x
    | x <= 0    = m_neg_inf
    | otherwise = log x * (ndf2 - 1) - x2 - logGamma ndf2 - log 2 * ndf2
    where
      ndf  = fromIntegral $ chiSquaredNDF chi
      ndf2 = ndf/2
      x2   = x/2

  quantile = quantile

instance D.Mean ChiSquared where
    mean (ChiSquared ndf) = fromIntegral ndf

instance D.Variance ChiSquared where
    variance (ChiSquared ndf) = fromIntegral (2*ndf)

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

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

instance D.Entropy ChiSquared where
  entropy (ChiSquared ndf) =
    let kHalf = 0.5 * fromIntegral ndf in
    kHalf
    + log 2
    + logGamma kHalf
    + (1-kHalf) * digamma kHalf

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

instance D.ContGen ChiSquared where
    genContVar (ChiSquared n) = MWC.chiSquare n

cumulative :: ChiSquared -> Double -> Double
cumulative chi x
  | x <= 0    = 0
  | otherwise = incompleteGamma (ndf/2) (x/2)
  where
    ndf = fromIntegral $ chiSquaredNDF chi

quantile :: ChiSquared -> Double -> Double
quantile (ChiSquared ndf) p
  | p == 0         = 0
  | p == 1         = 1/0
  | p > 0 && p < 1 = 2 * invIncompleteGamma (fromIntegral ndf / 2) p
  | otherwise      =
    error $ "Statistics.Distribution.ChiSquared.quantile: p must be in [0,1] range. Got: "++show p


{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

#ifndef MIN_VERSION_hashable
#define MIN_VERSION_hashable(x,y,z) 1
#endif

-----------------------------------------------------------------------------
-- |
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on affine spaces.
-----------------------------------------------------------------------------
module Linear.Affine where

import Control.Applicative
import Control.DeepSeq
import Control.Monad (liftM)
import Control.Lens
import Data.Binary as Binary
import Data.Bytes.Serial
import Data.Coerce
import Data.Complex (Complex)
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind
import Data.Functor.Classes
import Data.Functor.Product
import Data.Functor.Rep as Rep
import Data.HashMap.Lazy (HashMap)
import Data.Hashable
import Data.Hashable.Lifted
import Data.IntMap (IntMap)
import Data.Ix
import Data.Kind
import Data.Map (Map)
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup)
#endif
import Data.Serialize as Cereal
import Data.Vector (Vector)
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed.Base as U
import Foreign.Storable
import GHC.Generics (Generic, Generic1)
import Linear.Epsilon
import Linear.Metric
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector
import System.Random (Random(..))

-- | An affine space is roughly a vector space in which we have
-- forgotten or at least pretend to have forgotten the origin.
--
-- > a .+^ (b .-. a)  =  b@
-- > (a .+^ u) .+^ v  =  a .+^ (u ^+^ v)@
-- > (a .-. b) ^+^ v  =  (a .+^ v) .-. q@
class Additive (Diff p) => Affine p where
  type Diff p :: Type -> Type

  infixl 6 .-.
  -- | Get the difference between two points as a vector offset.
  (.-.) :: Num a => p a -> p a -> Diff p a

  infixl 6 .+^
  -- | Add a vector offset to a point.
  (.+^) :: Num a => p a -> Diff p a -> p a

  infixl 6 .-^
  -- | Subtract a vector offset from a point.
  (.-^) :: Num a => p a -> Diff p a -> p a
  p .-^ v = p .+^ negated v
  {-# INLINE (.-^) #-}

instance (Affine f, Affine g) => Affine (Product f g) where
  type Diff (Product f g) = Product (Diff f) (Diff g)
  Pair a b .-. Pair c d = Pair (a .-. c) (b .-. d)
  Pair a b .+^ Pair c d = Pair (a .+^ c) (b .+^ d)
  Pair a b .-^ Pair c d = Pair (a .+^ c) (b .+^ d)

-- | Compute the quadrance of the difference (the square of the distance)
qdA :: (Affine p, Foldable (Diff p), Num a) => p a -> p a -> a
qdA a b = Foldable.sum (fmap (join (*)) (a .-. b))
{-# INLINE qdA #-}

-- | Distance between two points in an affine space
distanceA :: (Floating a, Foldable (Diff p), Affine p) => p a -> p a -> a
distanceA a b = sqrt (qdA a b)
{-# INLINE distanceA #-}

#define ADDITIVEC(CTX,T) instance CTX => Affine T where type Diff T = T ; \
  (.-.) = (^-^) ; {-# INLINE (.-.) #-} ; (.+^) = (^+^) ; {-# INLINE (.+^) #-} ; \
  (.-^) = (^-^) ; {-# INLINE (.-^) #-}
#define ADDITIVE(T) ADDITIVEC((), T)

ADDITIVE([])
ADDITIVE(Complex)
ADDITIVE(ZipList)
ADDITIVE(Maybe)
ADDITIVE(IntMap)
ADDITIVE(Identity)
ADDITIVE(Vector)
ADDITIVE(V0)
ADDITIVE(V1)
ADDITIVE(V2)
ADDITIVE(V3)
ADDITIVE(V4)
ADDITIVE(Plucker)
ADDITIVE(Quaternion)
ADDITIVE(((->) b))
ADDITIVEC(Ord k, (Map k))
ADDITIVEC((Eq k, Hashable k), (HashMap k))
ADDITIVEC(Dim n, (V n))

-- | A handy wrapper to help distinguish points from vectors at the
-- type level
newtype Point f a = P (f a)
  deriving ( Eq, Ord, Show, Read, Monad, Functor, Applicative, Foldable
           , Eq1, Ord1, Show1, Read1
           , Traversable, Apply, Additive, Metric
           , Fractional , Num, Ix, Storable, Epsilon
           , Semigroup, Monoid
           , Random, Hashable
           , Generic, Generic1, Data
           )

instance Finite f => Finite (Point f) where
  type Size (Point f) = Size f
  toV (P v) = toV v
  fromV v = P (fromV v)

instance NFData (f a) => NFData (Point f a) where
  rnf (P x) = rnf x

instance Serial1 f => Serial1 (Point f) where
  serializeWith f (P p) = serializeWith f p
  deserializeWith m = P `liftM` deserializeWith m

instance Serial (f a) => Serial (Point f a) where
  serialize (P p) = serialize p
  deserialize = P `liftM` deserialize

instance Binary (f a) => Binary (Point f a) where
  put (P p) = Binary.put p
  get = P `liftM` Binary.get

instance Serialize (f a) => Serialize (Point f a) where
  put (P p) = Cereal.put p
  get = P `liftM` Cereal.get

instance Hashable1 f => Hashable1 (Point f) where
  liftHashWithSalt h s (P f) = liftHashWithSalt h s f
  {-# INLINE liftHashWithSalt #-}

lensP :: Lens' (Point g a) (g a)
lensP afb (P a) = P <$> afb a
{-# INLINE lensP #-}

_Point :: Iso' (Point f a) (f a)
_Point = iso (\(P a) -> a) P
{-# INLINE _Point #-}

instance (t ~ Point g b) => Rewrapped (Point f a) t
instance Wrapped (Point f a) where
  type Unwrapped (Point f a) = f a
  _Wrapped' = _Point
  {-# INLINE _Wrapped' #-}

-- These are stolen from Data.Profunctor.Unsafe
(.#) :: Coercible b a => (b -> c) -> (a -> b) -> a -> c
f .# _ = coerce f
{-# INLINE (.#) #-}

(#.) :: Coercible c b => (b -> c) -> (a -> b) -> a -> c
(#.) _ = coerce (\x -> x :: b) :: forall a b. Coercible b a => a -> b
{-# INLINE (#.) #-}

unP :: Point f a -> f a
unP (P x) = x
{-# INLINE unP #-}

-- We can't use GND to derive 'Bind' because 'join' causes
-- role troubles. However, GHC 7.8 and above let us use
-- explicit coercions for (>>-).
instance Bind f => Bind (Point f) where
  (>>-) = ((P .) . (. (unP .))) #. (>>-) .# unP
  join (P m) = P $ m >>- \(P m') -> m'

instance Distributive f => Distributive (Point f) where
  distribute = P . collect (\(P p) -> p)
  collect = (P .) #. collect .# (unP .)

instance Representable f => Representable (Point f) where
  type Rep (Point f) = Rep f
  tabulate = P #. tabulate
  {-# INLINE tabulate #-}
  index = Rep.index .# unP
  {-# INLINE index #-}

type instance Index (Point f a) = Index (f a)
type instance IxValue (Point f a) = IxValue (f a)

instance Ixed (f a) => Ixed (Point f a) where
  ix l = lensP . ix l
  {-# INLINE ix #-}

instance Traversable f => Each (Point f a) (Point f b) a b where
  each = traverse
  {-# INLINE each #-}

instance R1 f => R1 (Point f) where
  _x = lensP . _x
  {-# INLINE _x #-}

instance R2 f => R2 (Point f) where
  _y = lensP . _y
  {-# INLINE _y #-}
  _xy = lensP . _xy
  {-# INLINE _xy #-}

instance R3 f => R3 (Point f) where
  _z = lensP . _z
  {-# INLINE _z #-}
  _xyz = lensP . _xyz
  {-# INLINE _xyz #-}

instance R4 f => R4 (Point f) where
  _w = lensP . _w
  {-# INLINE _w #-}
  _xyzw = lensP . _xyzw
  {-# INLINE _xyzw #-}

instance Additive f => Affine (Point f) where
  type Diff (Point f) = f
  (.-.) = (. unP) #. (^-^) .# unP
  {-# INLINE (.-.) #-}
  (.+^) = (P .) #. (^+^) .# unP
  {-# INLINE (.+^) #-}
  (.-^) = (P .) #. (^-^) .# unP
  {-# INLINE (.-^) #-}

-- | Vector spaces have origins.
origin :: (Additive f, Num a) => Point f a
origin = P zero

-- | An isomorphism between points and vectors, given a reference
--   point.
relative :: (Additive f, Num a) => Point f a -> Iso' (Point f a) (f a)
relative p0 = iso (.-. p0) (p0 .+^)
{-# INLINE relative #-}

newtype instance U.Vector    (Point f a) =  V_P (U.Vector    (f a))
newtype instance U.MVector s (Point f a) = MV_P (U.MVector s (f a))
instance U.Unbox (f a) => U.Unbox (Point f a)

instance U.Unbox (f a) => M.MVector U.MVector (Point f a) where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  basicLength (MV_P v) = M.basicLength v
  basicUnsafeSlice m n (MV_P v) = MV_P (M.basicUnsafeSlice m n v)
  basicOverlaps (MV_P v) (MV_P u) = M.basicOverlaps v u
  basicUnsafeNew n = MV_P `liftM` M.basicUnsafeNew n
  basicUnsafeRead (MV_P v) i = P `liftM` M.basicUnsafeRead v i
  basicUnsafeWrite (MV_P v) i (P x) = M.basicUnsafeWrite v i x
  basicInitialize (MV_P v) = M.basicInitialize v
  {-# INLINE basicInitialize #-}

instance U.Unbox (f a) => G.Vector U.Vector (Point f a) where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw   #-}
  {-# INLINE basicLength       #-}
  {-# INLINE basicUnsafeSlice  #-}
  {-# INLINE basicUnsafeIndexM #-}
  basicUnsafeFreeze (MV_P v) = V_P `liftM` G.basicUnsafeFreeze v
  basicUnsafeThaw   ( V_P v) = MV_P `liftM` G.basicUnsafeThaw   v
  basicLength       ( V_P v) = G.basicLength v
  basicUnsafeSlice m n (V_P v) = V_P (G.basicUnsafeSlice m n v)
  basicUnsafeIndexM (V_P v) i = P `liftM` G.basicUnsafeIndexM v i


{-# LANGUAGE Arrows, ExistentialQuantification, ScopedTypeVariables #-}

module SignalLevel.SpecAnalysis20 where

import Data.Complex
import Data.Maybe (catMaybes,listToMaybe)
import qualified Data.Map as M
import Control.Arrow.ArrowP
import Control.DeepSeq (NFData)
import Control.SF.SF (SF(..))
import FRP.UISF (SEvent,asyncVT,Time,DeltaT)
import FRP.UISF.Asynchrony (Automaton(..))
import FRP.UISF.Graphics.Color (rgbE)

import Euterpea hiding (SF)
import HSoM
import HSoM.MUI (asyncV)
import HSoM.Examples.FFT (fftA,quantize,presentFFT)

import SignalLevel.SigFun19

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]]

-- num = # samples
-- n   = term

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)

mkPulse :: Int -> [Complex Double]
mkPulse n = 100 : take (n-1) (repeat 0)

x1 num = let f = 2*pi*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)

---------------
-- Exercises --
---------------

{- Exercise 20.1©:Write a Haskell function idft that implements the inverse DFT
 - as captured in Equation 20.3. Test your code by applying idft to one of the
 - signals used earlier in this section. In other words, show empirically that,
 - up to round-off errors, idft (dft xs) == xs. -}

idft :: RealFloat a => [Complex a] -> [Complex a]
idft xs = let lenI = length xs
              lenR = fromIntegral lenI
          in [let i = 2*pi* fromIntegral n / lenR
              in sum [(xs !! k) * exp (0 :+ i * fromIntegral k)
                     | k <- [0..lenI-1]]
             | n <- [0..lenI-1]]

{- Exercise 20.2: Use dft to analyze some of the signals generated using signal
 - functions defined in Chapter 19. -}

fft1 :: forall p . Clock p => SigFun p () (SEvent (M.Map Double Double))
fft1 = s1 >>> (fftA 128 256) >>> arr (fmap $ presentFFT $ rate (undefined :: p))

{- Exercise 20.3:© Define a function mkSqWave :: Int -> Int -> [Complex Double]
 - such that mkSqWave num n is the sum of the first n terms of the Fourier
 - Series of a square wave, having num samples in the result. -}

mkSqWave :: Int -> Int -> [Complex Double]
mkSqWave num n =
  foldr (zipWith (+) . (\i -> mkTerm i))
        (replicate num 0)
        (filter odd [1..(2*n-1)])
  where mkTerm x = let f = 2 * pi / fromIntegral num
                       t = fromIntegral x
                   in [sin (t*f*fromIntegral i)/t:+0
                      | i <- [0,1..num-1]]

{- Exercise 20.4: -}

{- a. © Prove mathematically that x and x^ are inverses.

  Let X^ and X be the NxN matrices defined by

  X^(k,n) = (1/N) * w^(kn)
  X(n,k)    = (1/N) * w^(-(kn))

  where w = e^(-j2π/N)
  and (a,b) corresponds to the entry in the ath row and bth column of the matrix

  Prove X^ and X are inverses
  Proof:

    It suffices to prove (X^)X = I

    Thus I show
      if p=q then X^X(p,q) = 1
      otherwise   X^X(p,q) = 0

    X^X(p,q) = X^(p,*)·X(*,q) = (1/N) ∑w^(pn)*w^(-(n*q))

    where X^(p,*) is the pth row of X^ , X(*,q) is the qth column of X and
    ∑ is a sum from 0 to N-1

    if p = q then
      X^X(p,q) = (1/N) ∑ w^(pn)/w^(nq) = (1/N) ∑ 1 =  (1/N)*N = 1

    if p ≠ q then
      X^X(p,q) = (1/N) ∑ w^(n(p-q))
               = (1 - (e^(-j*2π(p-q)/N))^N)/(1 - e^(-j*2π(p-q)/N))
               = (1 - (e^(-j*2π(p-q)))/(1 - e^(-j*2π(p-q)/N))
               = (1 - 1/(cos(2π(p-q)) + jsin(2π(p-q))))/(1 - e^(-j*2π(p-q)/N))
               = (1 - 1)/(1 - e^(-j*2π(p-q)/N))
               = 0

    Thus X^X = I, proving X^ and X are inverses.
-}

{-
e1 =
  \i -> (1/lenC)*sum[(cxs !! n)*exp(0 :+ i*fromIntegral n) | n <- [0,1..lenI-1]]

f = \exp -> let lenI = length cxs
            in let lenR = fromIntegral lenI
               in let lenC = lenR :+ 0
                  in exp

-> fdft = (f (concatMap
               (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
               [0,1..lenI-1]
             )
          )

fix exp = \x -> let x = exp x in x
-}

{- b. Also prove, using equational reasoning, that dft and idft are
 - inverses. (..., you may assume that Haskell numeric types
 - obey the standard axioms of real arithmetic.)

idft(dft(cxs))  = cxs

idft(dft(cxs))
-> idft (let lenI = length cxs
             lenR = fromIntegral lenI
             lenC = lenR :+ 0
         in [let i = -2*pi*fromIntegral k/lenR
             in (1/lenC)*sum[(cxs !! n)*exp (0 :+ i*fromIntegral n)
                             | n <-  [0,1..lenI-1]]
              |k <- [0,1..lenI-1]] )

-> idft (let lenI = length cxs
             lenR = fromIntegral lenI
             lenC = lenR :+ 0
         in [let i = -2*pi*fromIntegral k/lenR
             in e1 i | k <- [0,1..lenI-1]] )

  -> concatMap
       (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
       [0,1..lenI-1]

-> idft (f (concatMap
             (\k -> let i = -2*pi*fromIntegral k/lenR in e1 i)
             [0,1..lenI-1]
           )
        )

-> let lenI2 = length fdft
       lenR2 = fromIntegral lenI2
   in [let i2 = 2*pi* fromIntegral n2 / lenR2
       in sum [(fdft !! k) * exp (0 :+ i2 * fromIntegral k) | k <- [0..lenI2-1]]
        | n2 <- [0..lenI2-1]]

-> let lenI2 = length fdft
       lenR2 = fromIntegral lenI2
   in [let i2 = 2*pi* fromIntegral n2 / lenR2
       in sum [(fdft !! k) * exp (0 :+ i2 * fromIntegral k) | k <- [0..lenI2-1]]
        | n2 <- [0..lenI2-1]]
-}

fftEx :: UISF () ()
fftEx =
  proc _ -> do
    f <- withDisplay $ hSlider (1,2000) 440 -< ()
    d <- clockedSFToUISF (0.1) fftDataSig -< f
    let (ds,ts) = unzip d
    _ <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds))
    _ <- realtimeGraph (makeLayout (Fixed {fixedSize = 500})
                                   (Fixed {fixedSize = 150}))
                                   1 Black -< (zip (fst $ unzip ds) ts)
    m  <- withDisplay $ hSlider (1,2000) 20 -< ()
    d1 <- clockedSFToUISF (0.1) fftDataSig -< m
    let (ds1,ts1) = unzip d1
    _  <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds1))
    _ <- realtimeGraph (makeLayout (Fixed {fixedSize = 500})
                                   (Fixed {fixedSize = 150}))
                                   1 Black -< (zip (fst $ unzip ds1) ts1)
    d3 <- clockedSFToUISF (0.1) fftDataSig -< (f+m)
    let (ds3,ts3) = unzip d3
    _  <- histogram (makeLayout (Fixed {fixedSize = 500})
                               (Fixed {fixedSize = 150}))
                                -< listToMaybe (catMaybes (snd $ unzip ds3))
    outA -< ()
    where
      fftDataSig :: SigFun CtrRate Double (Double, SEvent [Double])
      fftDataSig = proc f -> do
        s <- osc (tableSinesN 4096 [1]) 0 -< f
        fftData <- fftA 128 256 -< s
        outA -< (s,fftData)

t0 = runMUI (defaultMUIParams
               {uiSize       = (500,900),
                uiBackground = rgbE 93 10 13,
                uiTitle      = "FFT Test"})
            fftEx

clockedSFToUISF :: forall a b c . (NFData b,Clock c) =>
                   DeltaT -> SigFun c a b -> UISF a [(b,Time)]
clockedSFToUISF buffer ~(ArrowP sf) = let r = rate (undefined :: c)
  in asyncV r buffer (toAutomaton sf)

toAutomaton :: forall a b . SF a b -> Automaton (->) a b
toAutomaton ~(SF f) = Automaton $ \a -> let (b,sf) = f a in (b,toAutomaton sf)

{- Exercise 20.5: Modify the program in 20.5 (fftEx) in the following ways:
 -
 - 1. Add a second slider, and use it to control the frequency of a second
 - oscillator.©
 - 2. Let s1 and s2 be the names of the signals whose frequencies are controlled
 - by the first and second sliders, respectively. Instead of displaying the FFT
 - of just s1, try a variety of combinations of s1 and s2, such as
 - s1 + s2
 - s1 - s2
 - s1 * s2
 - 1/s1 + 1/s2
 - s1/s2
 - Comment on the results.
 - 3. Use s2 to control the frequency of s1 (as was done with vibrato in Ch.19).
 - Plot the fft of s1 and comment on the result.
 - 4. Instead of using osc to generate a pure sine wave, try using other
 - oscillators and/or table generators to create more complex tones, and plot
 - their FFTs. Comment on the results. -}

----------
-- Util --
----------
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 ' ')


{- |
Module      :  Numeric.GSL.Fitting
Copyright   :  (c) Alberto Ruiz 2010
License     :  GPL

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

Nonlinear Least-Squares Fitting

<http://www.gnu.org/software/gsl/manual/html_node/Nonlinear-Least_002dSquares-Fitting.html>

The example program in the GSL manual (see examples/fitting.hs):

@dat = [
 ([0.0],([6.0133918608118675],0.1)),
 ([1.0],([5.5153769909966535],0.1)),
 ([2.0],([5.261094606015287],0.1)),
 ...
 ([39.0],([1.0619821710802808],0.1))]

expModel [a,lambda,b] [t] = [a * exp (-lambda * t) + b]

expModelDer [a,lambda,b] [t] = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]

(sol,path) = fitModelScaled 1E-4 1E-4 20 (expModel, expModelDer) dat [1,0,0]

\> path
(6><5)
 [ 1.0,  76.45780563978782, 1.6465931240727802, 1.8147715267618197e-2, 0.6465931240727797
 , 2.0, 37.683816318260355,  2.858760367632973,  8.092094813253975e-2, 1.4479636296208662
 , 3.0,    9.5807893736187,  4.948995119561291,   0.11942927999921617, 1.0945766509238248
 , 4.0,  5.630494933603935,  5.021755718065913,   0.10287787128056883, 1.0338835440862608
 , 5.0,  5.443976278682909,  5.045204331329302,   0.10405523433131504,  1.019416067207375
 , 6.0, 5.4439736648994685,  5.045357818922331,   0.10404905846029407, 1.0192487112786812 ]
\> sol
[(5.045357818922331,6.027976702418132e-2),
(0.10404905846029407,3.157045047172834e-3),
(1.0192487112786812,3.782067731353722e-2)]@

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

module Numeric.GSL.Fitting (
    -- * Levenberg-Marquardt
    nlFitting, FittingMethod(..),
    -- * Utilities
    fitModelScaled, fitModel
) where

import Data.Packed.Internal
import Numeric.LinearAlgebra
import Numeric.GSL.Internal

import Foreign.Ptr(FunPtr, freeHaskellFunPtr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

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

data FittingMethod = LevenbergMarquardtScaled -- ^ Interface to gsl_multifit_fdfsolver_lmsder. This is a robust and efficient version of the Levenberg-Marquardt algorithm as implemented in the scaled lmder routine in minpack. Minpack was written by Jorge J. More, Burton S. Garbow and Kenneth E. Hillstrom.
                   | LevenbergMarquardt -- ^ This is an unscaled version of the lmder algorithm. The elements of the diagonal scaling matrix D are set to 1. This algorithm may be useful in circumstances where the scaled version of lmder converges too slowly, or the function is already scaled appropriately.
        deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional least-squares fitting.
nlFitting :: FittingMethod
      -> Double                     -- ^ absolute tolerance
      -> Double                     -- ^ relative tolerance
      -> Int                        -- ^ maximum number of iterations allowed
      -> (Vector Double -> Vector Double)     -- ^ function to be minimized
      -> (Vector Double -> Matrix Double)   -- ^ Jacobian
      -> Vector Double                   -- ^ starting point
      -> (Vector Double, Matrix Double)  -- ^ solution vector and optimization path

nlFitting method epsabs epsrel maxit fun jac xinit = nlFitGen (fi (fromEnum method)) fun jac xinit epsabs epsrel maxit

nlFitGen m f jac xiv epsabs epsrel maxit = unsafePerformIO $ do
    let p   = dim xiv
        n   = dim (f xiv)
    fp <- mkVecVecfun (aux_vTov (checkdim1 n p . f))
    jp <- mkVecMatfun (aux_vTom (checkdim2 n p . jac))
    rawpath <- createMatrix RowMajor maxit (2+p)
    app2 (c_nlfit m fp jp epsabs epsrel (fi maxit) (fi n)) vec xiv mat rawpath "c_nlfit"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toRows $ dropRows (it-1) path
    freeHaskellFunPtr fp
    freeHaskellFunPtr jp
    return (subVector 2 p sol, path)

foreign import ccall safe "nlfit"
    c_nlfit:: CInt -> FunPtr TVV -> FunPtr TVM -> Double -> Double -> CInt -> CInt -> TVM

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

checkdim1 n _p v
    | dim v == n = v
    | otherwise = error $ "Error: "++ show n
                        ++ " components expected in the result of the function supplied to nlFitting"

checkdim2 n p m
    | rows m == n && cols m == p = m
    | otherwise = error $ "Error: "++ show n ++ "x" ++ show p
                        ++ " Jacobian expected in nlFitting"

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

err (model,deriv) dat vsol = zip sol errs where
    sol = toList vsol
    c = max 1 (chi/sqrt (fromIntegral dof))
    dof = length dat - (rows cov)
    chi = norm2 (fromList $ cost (resMs model) dat sol)
    js = fromLists $ jacobian (resDs deriv) dat sol
    cov = inv $ trans js <> js
    errs = toList $ scalar c * sqrt (takeDiag cov)

-- | Higher level interface to 'nlFitting' 'LevenbergMarquardtScaled'. The optimization function and
-- Jacobian are automatically built from a model f vs x = y and its derivatives, and a list of
-- instances (x, (y,sigma)) to be fitted.

fitModelScaled
         :: Double -- ^ absolute tolerance
         -> Double -- ^ relative tolerance
         -> Int    -- ^ maximum number of iterations allowed
         -> ([Double] -> x -> [Double], [Double] -> x -> [[Double]]) -- ^ (model, derivatives)
         -> [(x, ([Double], Double))] -- ^ instances
         -> [Double] -- ^ starting point
         -> ([(Double, Double)], Matrix Double) -- ^ (solution, error) and optimization path
fitModelScaled epsabs epsrel maxit (model,deriv) dt xin = (err (model,deriv) dt sol, path) where
    (sol,path) = nlFitting LevenbergMarquardtScaled epsabs epsrel maxit
                (fromList . cost (resMs model) dt . toList)
                (fromLists . jacobian (resDs deriv) dt . toList)
                (fromList xin)

-- | Higher level interface to 'nlFitting' 'LevenbergMarquardt'. The optimization function and
-- Jacobian are automatically built from a model f vs x = y and its derivatives, and a list of
-- instances (x,y) to be fitted.

fitModel :: Double -- ^ absolute tolerance
         -> Double -- ^ relative tolerance
         -> Int    -- ^ maximum number of iterations allowed
         -> ([Double] -> x -> [Double], [Double] -> x -> [[Double]]) -- ^ (model, derivatives)
         -> [(x, [Double])]  -- ^ instances
         -> [Double] -- ^ starting point
         -> ([Double], Matrix Double) -- ^ solution and optimization path
fitModel epsabs epsrel maxit (model,deriv) dt xin = (toList sol, path) where
    (sol,path) = nlFitting LevenbergMarquardt epsabs epsrel maxit
                (fromList . cost (resM model) dt . toList)
                (fromLists . jacobian (resD deriv) dt . toList)
                (fromList xin)

cost model ds vs = concatMap (model vs) ds

jacobian modelDer ds vs = concatMap (modelDer vs) ds

-- | Model-to-residual for association pairs with sigma, to be used with 'fitModel'.
resMs :: ([Double] -> x -> [Double]) -> [Double] -> (x, ([Double], Double)) -> [Double]
resMs m v = \(x,(ys,s)) -> zipWith (g s) (m v x) ys where g s a b = (a-b)/s

-- | Associated derivative for 'resMs'.
resDs :: ([Double] -> x -> [[Double]]) -> [Double] -> (x, ([Double], Double)) -> [[Double]]
resDs m v = \(x,(_,s)) -> map (map (/s)) (m v x)

-- | Model-to-residual for association pairs, to be used with 'fitModel'. It is equivalent
-- to 'resMs' with all sigmas = 1.
resM :: ([Double] -> x -> [Double]) -> [Double] -> (x, [Double]) -> [Double]
resM m v = \(x,ys) -> zipWith g (m v x) ys where g a b = a-b

-- | Associated derivative for 'resM'.
resD :: ([Double] -> x -> [[Double]]) -> [Double] -> (x, [Double]) -> [[Double]]
resD m v = \(x,_) -> m v x


{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module VecSpec where

import Data.Complex

import Test.QuickCheck
import Test.QuickCheck.Poly

import Category
import Functor
import Unboxed
import Vec

type N = 10

prop_Vec_Functor_id :: FnProp (Vec3 N A)
prop_Vec_Functor_id = checkFnEqual law_Functor_id

prop_Vec_Functor_comp :: Fun B C -> Fun A B -> FnProp (Vec3 N A)
prop_Vec_Functor_comp f g = checkFnEqual (law_Functor_comp f g)

type UA = Int
type UB = Double
type UC = Complex Double

prop_UVec3_Functor_id :: FnProp (UVec3 N UA)
prop_UVec3_Functor_id = checkFnEqual law_Functor_id

prop_UVec3_Functor_comp ::
    (UB -#> UC) -> (UA -#> UB) -> FnProp (UVec3 N UA)
prop_UVec3_Functor_comp f g = checkFnEqual (law_Functor_comp f g)


module NeuralNet.Train (
  nnBackward,
  updateNNParams
) where

import NeuralNet.Net
import NeuralNet.Layer
import NeuralNet.Example
import NeuralNet.Activation
import NeuralNet.Matrix
import Numeric.LinearAlgebra

nnBackward :: NeuralNet -> [ForwardPropStep] -> ExampleSet -> [(Matrix Double, Matrix Double)]
nnBackward nn steps examples
  | numLayers /= stepsSize - 1 = error ("should be one entry (" ++ show stepsSize ++ ") for each layer (" ++ show numLayers ++ ")")
  | otherwise                  = reverse (nnBackwardStep layers orderedSteps dAL)
    where
      layers = reverse (nnLayers nn)
      orderedSteps = reverse steps
      numLayers = length layers
      stepsSize = length steps

      y = exampleSetY examples
      al = forwardPropA (head orderedSteps)
      unitRowVector = row (replicate (cols y) 1)
      dAL1 = y / al
      dAL2 = (unitRowVector - y) / (unitRowVector - al)
      dAL = -(dAL1 - dAL2)

nnBackwardStep :: [NeuronLayer] -> [ForwardPropStep] -> Matrix Double -> [(Matrix Double, Matrix Double)]
nnBackwardStep [] _ _ = []
nnBackwardStep _ [] _ = error "Shouldn't get here"
nnBackwardStep _ [_] _ = error "Shouldn't get here"
nnBackwardStep (l:ls) (s:ss@(ps:_)) prevDA = (dW, db) : nnBackwardStep ls ss dA
  where
    (dA, dW, db) = nnBackwardLayer l prevDA (forwardPropA ps) (forwardPropZ s)

nnBackwardLayer :: NeuronLayer -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
nnBackwardLayer l dA prevA z = (dANext, dW, db)
  where
    m = fromIntegral (cols prevA)
    dZ = backward (layerActivation l) dA z
    dW = scale (1/m) (dZ <> tr prevA)
    db = scale (1/m) (sumRows dZ)
    dANext = tr (layerW l) <> dZ

updateNNParams :: NeuralNet -> [(Matrix Double, Matrix Double)] -> Double -> NeuralNet
updateNNParams nn grads learningRate
  | length layers /= length grads = error "Must provide one grad per layer"
  | otherwise                     = updateNNLayers nn newLayers
    where
      layers = nnLayers nn
      layerSets = zip layers grads
      newLayers = map (\(l, (dW, dB)) -> updateLayerParamsFromGrads l dW dB learningRate) layerSets

updateLayerParamsFromGrads :: NeuronLayer -> Matrix Double -> Matrix Double -> Double -> NeuronLayer
updateLayerParamsFromGrads l dW db learningRate = updateLayerParams l newDw newDb
  where
    newDw = layerW l - scale learningRate dW
    newDb = layerB l - scale learningRate db


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

module Math.Probably.EMPCA where

import Numeric.LinearAlgebra hiding (orth)
import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix ((#>), app, orth)
import Math.Probably.FoldingStats
import qualified Data.Vector.Storable as VS
import Foreign.Storable.Tuple ()
import Debug.Trace

import Data.Ord (comparing)
import Data.List (sortBy)

import Math.Probably.PCA
import Data.Random
import Unsafe.Coerce
import qualified Data.Binary as B
import Data.Word
import Control.Applicative
import Control.DeepSeq

data EmPcaBasis = EmPcaBasis {
   centering :: VS.Vector (Double,Double),
   emEvecs :: Mat
   } deriving (Show, Read)

instance B.Binary EmPcaBasis where
  put (EmPcaBasis c m) = do
      B.put (VS.length c)
      VS.mapM_ (\(c1,c2) -> B.put (toWord64 c1) >> B.put (toWord64 c2)) c
      B.put $ rows m
      B.put $ cols m
      VS.mapM_ (B.put . toWord64) $ flatten m
  get = do nc <- B.get
           c <- VS.replicateM nc $ (,) <$> fmap fromWord64 B.get <*> fmap fromWord64 B.get
           nrows <- B.get
           ncols <- B.get
           mflat <- VS.replicateM (nrows*ncols) $ fmap fromWord64 B.get
           return $ EmPcaBasis c $ reshape ncols mflat

instance NFData EmPcaBasis where
  rnf (EmPcaBasis x y) = rnf x `seq` rnf y

toWord64 :: a -> Word64
toWord64 = unsafeCoerce

fromWord64 :: Word64 -> a
fromWord64 = unsafeCoerce

emPca :: Int -> Int -> Maybe (Matrix Double) -> [Vec] -> RVar EmPcaBasis
emPca k iters mcinit vecs = do
  let meansds = findCentre vecs

      n = length vecs
      p = VS.length $ head vecs

      dat =  fromColumns $ map (centre meansds) vecs

      go c 0 = c
      go c iter =
        let x = inv (tr c <> c)<> tr c <> dat :: Matrix Double
            c1 = dat <> tr x <> inv (x<> tr x) :: Matrix Double
        in go (trace ("EM "++show iter ++ " C00="++show (c1!0!0)) c1) (iter-1)

  cinit <- case mcinit of
             Nothing -> fmap (p><k) $ sequence $ replicate (p*k) $ uniform 0 1
             Just c -> return c
  let cfinal = trdimit "cfinal" $ go cinit iters
      cOrth :: Matrix Double
      cOrth =  trdimit "cOrth" $ orth cfinal
      xevec = truePca $ trdimit "pcainput" $ tr cOrth <> dat
      evec = cOrth <>  xevec
--      truepca m = stat m

  --    (xevec, eval) = truepca (tr cOrth <> dat)
  return $ EmPcaBasis meansds evec

showC (m, iter) = "iter "++show iter++" C ="++dispf 2 m

applyEmPca :: EmPcaBasis -> Vec -> Vec
applyEmPca (EmPcaBasis cent vp) v =  tr vp <> centre cent v

findCentre :: [Vec] -> VS.Vector (Double,Double)
findCentre  = uncurry (VS.zipWith (,)) . runStat meanSDF

centre :: VS.Vector (Double,Double) -> Vec -> Vec
centre meansds v = VS.zipWith (\x (mn,_) -> x-mn) v meansds

trdims s m = trace (s++": "++show (rows m, cols m))

trdimit s m = trace (s++": "++show (rows m, cols m)) m

trdispit s m = trace (s++": "++dispf 3 m) m

trdisp s m = trace (s++": "++dispf 3 m)

truePca dataset =
  let --mn = trdispit "mn" $ asColumn (mean $ tr dataset_mn)
      --dataset = trdimit "dataset" $dataset_mn - mn
      cc = trdimit "cc" $ covN (tr dataset)
      (cdd,cvv) = eigSH cc
      ii = reverse $ map fst $ sortBy (comparing snd) $ zip [0..] $ VS.toList cdd
      evects = cvv ? ii
  in evects


{-# LANGUAGE ViewPatterns #-}

module Statistics.Covariance.Shrinkage.MMSE (
    mmseCov
  , shrinkCov
  , sampleCov
  , naiveCov
  , pinCoef
  , lwShrinkageCoef
  , rblwShrinkCovCoef
  , rblwShrinkageCoef
  , oracleShrinkCovCoef
  , oracleShrinkageCoef
  ) where

import Numeric.LinearAlgebra hiding ((<>))
import Numeric.Sum (kbn, sumVector)
import Statistics.Sample (mean)
import qualified Data.Vector.Unboxed as VU

mmseCov :: (Matrix Double -> Double) -> Matrix Double -> Herm Double
mmseCov shrink x = shrinkCov s p f
  where
    s = sampleCov x
    f = naiveCov x
    p = shrink x

shrinkCov :: Herm Double -> Double -> Herm Double -> Herm Double
shrinkCov sample coef target = scale (1 - coef) sample `add` scale coef target

sampleCov :: Matrix Double -> Herm Double
sampleCov = snd . meanCov

avgVarCov :: Herm Double -> Herm Double
avgVarCov (unSym -> s) = trustSym $ scale s' (ident p)
  where
    p = cols s
    s' = mean $ takeDiag s

naiveCov :: Matrix Double -> Herm Double
naiveCov = avgVarCov . sampleCov

pinCoef :: Double -> Double
pinCoef = max 0 . min 1

lwShrinkageCoef :: Matrix Double -> Double
lwShrinkageCoef x = numer / denom
  where
    n = rows x
    p = cols x
    s = unSym $ sampleCov x
    trs2 = trace $ s <> s
    tr2s = sq $ trace s
    frob ix = let c = x ? [ix] in sq . norm_Frob $ unSym (mTm c) - s
    numer = sumVector kbn (VU.generate n frob) / fromIntegral (sq n)
    denom = trs2 - tr2s / fromIntegral p

rblwShrinkCovCoef :: Int -> Herm Double -> Double
rblwShrinkCovCoef (fromIntegral -> n) (unSym -> s) = numer / denom
  where
    p = cols s
    trs2 = trace $ s <> s
    tr2s = sq $ trace s
    numer = (n - 2) / n * trs2 + tr2s
    denom = (n + 2) * (trs2 - tr2s / fromIntegral p)

rblwShrinkageCoef :: Matrix Double -> Double
rblwShrinkageCoef x = rblwShrinkCovCoef n (sampleCov x)
  where n = rows x

data OASIter = OASIter Double (Herm Double)

data OAS = OAS Int Int (Herm Double) (Herm Double)

oracleStep :: OAS -> Herm Double -> OASIter
oracleStep (OAS n p s f) sj = OASIter pj' sj'
  where
    n' = fromIntegral n
    p' = fromIntegral p
    trss = trace $ unSym sj <> unSym s
    tr2s = sq . trace $ unSym sj
    numer = (1 - 2 / p') * trss  + tr2s
    denom = (n' + 1 - 2 / p') * trss + (1 - n' / p') * tr2s
    pj' = numer / denom
    sj' = scale (1 - pj') s `add` scale pj' f

oracleIter :: Double -> OAS -> Herm Double -> Double
oracleIter tol oas = go 1 . oracleStep oas
  where
    go p (OASIter p' sj')
      | abs (p - p') < tol = p'
      | otherwise = go p' $ oracleStep oas sj'

oracleShrinkCovCoef :: Int -> Herm Double -> Double
oracleShrinkCovCoef n s = oracleIter 1e-6 (OAS n p s f) s
  where
    p = cols $ unSym s
    f = avgVarCov s

oracleShrinkageCoef :: Matrix Double -> Double
oracleShrinkageCoef x = oracleShrinkCovCoef n (sampleCov x)
  where n = rows x

sq :: Num a => a -> a
sq a = a * a

trace :: Matrix Double -> Double
trace = sumVector kbn . takeDiag


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

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Container
-- Copyright   :  (c) Alberto Ruiz 2010
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable
--
-- Basic numeric operations on 'Vector' and 'Matrix', including conversion routines.
--
-- The 'Container' class is used to define optimized generic functions which work
-- on 'Vector' and 'Matrix' with real or complex elements.
--
-- Some of these functions are also available in the instances of the standard
-- numeric Haskell classes provided by "Numeric.LinearAlgebra".
--
-----------------------------------------------------------------------------

module Numeric.Container (
    -- * Basic functions
    module Data.Packed,
    constant, linspace,
    diag, ident,
    ctrans,
    -- * Generic operations
    Container(..),
    -- * Matrix product
    Product(..),
    optimiseMult,
    mXm,mXv,vXm,(<.>),Mul(..),LSDiv(..),
    outer, kronecker,
    -- * Random numbers
    RandDist(..),
    randomVector,
    gaussianSample,
    uniformSample,
    meanCov,
    -- * Element conversion
    Convert(..),
    Complexable(),
    RealElement(),

    RealOf, ComplexOf, SingleOf, DoubleOf,

    IndexOf,
    module Data.Complex,
    -- * Input / Output
    dispf, disps, dispcf, vecdisp, latexFormat, format,
    loadMatrix, saveMatrix, fromFile, fileDimensions,
    readMatrix,
    fscanfVector, fprintfVector, freadVector, fwriteVector,
    -- * Experimental
    build', konst'
) where

import Data.Packed
import Data.Packed.Internal(constantD)
import Numeric.ContainerBoot
import Numeric.Chain
import Numeric.IO
import Data.Complex
import Numeric.LinearAlgebra.Algorithms(Field,linearSolveSVD)
import Data.Packed.Random

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

{- | creates a vector with a given number of equal components:

@> constant 2 7
7 |> [2.0,2.0,2.0,2.0,2.0,2.0,2.0]@
-}
constant :: Element a => a -> Int -> Vector a
-- constant x n = runSTVector (newVector x n)
constant = constantD-- about 2x faster

{- | Creates a real vector containing a range of values:

@\> linspace 5 (-3,7)
5 |> [-3.0,-0.5,2.0,4.5,7.0]@

Logarithmic spacing can be defined as follows:

@logspace n (a,b) = 10 ** linspace n (a,b)@
-}
linspace :: (Enum e, Container Vector e) => Int -> (e, e) -> Vector e
linspace n (a,b) = addConstant a $ scale s $ fromList [0 .. fromIntegral n-1]
    where s = (b-a)/fromIntegral (n-1)

-- | Dot product: @u \<.\> v = dot u v@
(<.>) :: Product t => Vector t -> Vector t -> t
infixl 7 <.>
(<.>) = dot

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

class Mul a b c | a b -> c where
 infixl 7 <>
 -- | Matrix-matrix, matrix-vector, and vector-matrix products.
 (<>)  :: Product t => a t -> b t -> c t

instance Mul Matrix Matrix Matrix where
    (<>) = mXm

instance Mul Matrix Vector Vector where
    (<>) m v = flatten $ m <> asColumn v

instance Mul Vector Matrix Vector where
    (<>) v m = flatten $ asRow v <> m

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

class LSDiv b c | b -> c, c->b where
 infixl 7 <\>
 -- | least squares solution of a linear system, similar to the \\ operator of Matlab\/Octave (based on linearSolveSVD)
 (<\>)  :: Field t => Matrix t -> b t -> c t

instance LSDiv Vector Vector where
    m <\> v = flatten (linearSolveSVD m (reshape 1 v))

instance LSDiv Matrix Matrix where
    (<\>) = linearSolveSVD

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

-- | Compute mean vector and covariance matrix of the rows of a matrix.
meanCov :: Matrix Double -> (Vector Double, Matrix Double)
meanCov x = (med,cov) where
    r    = rows x
    k    = 1 / fromIntegral r
    med  = konst k r `vXm` x
    meds = konst 1 r `outer` med
    xc   = x `sub` meds
    cov  = scale (recip (fromIntegral (r-1))) (trans xc `mXm` xc)


module Math.IRT.MLE.Fenced
  ( DF (..)
  , MLEResult (..)
  , mleEst
  ) where

import Data.Default.Class

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.MLE.Internal.Generic
import Math.IRT.Model.Generic

data DF = DF { steps                :: !Int
             , thetaEstimate        :: !Double
             , lower_fence          :: !Double
             , upper_fence          :: !Double
             , fence_discrimination :: !Double }

instance Default DF where
  def = DF 10 0.0 (-3.5) 3.5 3.0

mleEst :: (ContDistr d, DensityDeriv d, LogLikelihood d, GenericModel d) => DF -> [Bool] -> [d] -> MLEResult
mleEst (DF n vTheta lf uf fd) rs params =
    let resp = True : False : rs
        pars = fromThreePLM fd lf 0 : fromThreePLM fd uf 0 : params
    in generic_mleEst resp pars n vTheta


module Haskov where

import System.Random
import Data.Tuple (swap)
import Data.List (foldl')
import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seqq
import Data.Map.Strict (Map, keys, foldlWithKey', foldrWithKey', mapWithKey)
import qualified Data.Map.Strict as Map
import Numeric.LinearAlgebra (scale, luSolve, luPacked)
--import qualified Numeric.LinearAlgebra as Lin
import Numeric.LinearAlgebra.Data (Matrix, matrix, col, cols, ident, tr)
import qualified Numeric.LinearAlgebra.Data as Dat
import Numeric.LinearAlgebra.HMatrix (sumElements, cmap)
import qualified Numeric.LinearAlgebra.HMatrix as Hma

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

-- Markov Type --
type IndexMap a = Map a Int
type ReverseIndexMap a = Map Int a
type HMatrixMap = Map (Int, Int) Double

data Markov a = Markov {  imap :: IndexMap a
                        , hmap :: HMatrixMap }

instance (Show a, Ord a) => Show (Markov a) where
    show haskov = show $ toList haskov

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

-- Query --
-- Safe query
lookup :: (Ord a) => a -> a -> Markov a -> Maybe Double
lookup i j (Markov imap hmap)
    | (Map.member i imap) && (Map.member j imap) = Just (hmap Map.! ((imap Map.!i), (imap Map.! j)))
    | otherwise                                  = Nothing

-- Returns number of states
size :: (Ord a) => Markov a -> Int
size (Markov imap hmap) = Map.size imap

-- Returns true if the chain is empty
null :: (Ord a) => Markov a -> Bool
null (Markov imap hmap) = Map.null imap

-- Returns true if input is a state of the chain
member :: (Ord a) => a -> Markov a -> Bool
member i (Markov imap hmap)
    | (Map.member i imap) = True
    | otherwise = False

-- Returns true if input is not a state of the chain
notMember :: (Ord a) => a -> Markov a -> Bool
notMember i haskov = not (member i haskov)

-- States of the Markov chain
states :: (Ord a) => Markov a -> [a]
states (Markov imap hmap) = keys imap

statesI :: (Ord a) => Markov a -> IndexMap a
statesI (Markov imap hmap) = imap

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

-- Construction --
empty :: Markov a
empty = Markov (Map.empty) (Map.empty)

insert :: (Ord a) => a -> a -> Double -> Markov a -> Markov a
insert i j d (Markov imap hmap)
    | hasI && (not hasJ) = Markov (Map.insert j s imap) (Map.insert (imap Map.! i, s) d hmap)
    | (not hasI) && hasJ = Markov (Map.insert i s imap) (Map.insert (s, imap Map.! j) d hmap)
    | (not hasI) && (not hasJ) && (i /= j) = Markov (Map.insert j (s+1) (Map.insert i s imap)) (Map.insert (s, (s+1)) d hmap)
    | (not hasI) && (not hasJ) && (i == j) = Markov (Map.insert i s imap) (Map.insert (s, s) d hmap)
    | otherwise          = Markov imap (Map.insert (imap Map.! i, imap Map.! j) d hmap)
    where
        hasI = Map.member i imap
        hasJ = Map.member j imap
        s = Map.size imap

insertWith :: (Ord a) => (Double -> Double -> Double) -> a -> a -> Double -> Markov a -> Markov a
insertWith f i j d (Markov imap hmap)
    | hasI && (not hasJ) = Markov (Map.insert j s imap) (Map.insert (imap Map.! i, s) d hmap)
    | (not hasI) && hasJ = Markov (Map.insert i s imap) (Map.insert (s, imap Map.! j) d hmap)
    | (not hasI) && (not hasJ) && (i /= j) = Markov (Map.insert j (s+1) (Map.insert i s imap)) (Map.insert (s, (s+1)) d hmap)
    | (not hasI) && (not hasJ) && (i == j) = Markov (Map.insert i s imap) (Map.insert (s, s) d hmap)
    | otherwise          = Markov imap (Map.insertWith f (imap Map.! i, imap Map.! j) d hmap)
    where
        hasI = Map.member i imap
        hasJ = Map.member j imap
        s = Map.size imap

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

-- Lists --
fromList :: (Ord a) => [((a, a), Double)] -> Markov a
fromList list =
    let imap = foldl' toimap Map.empty list
        hmap = foldl' (tohmap imap) Map.empty list
    in Markov imap hmap

toList :: (Ord a) => Markov a -> [((a, a), Double)]
toList (Markov imap hmap) =
    let rIMap = reverseIMap imap
    in foldrWithKey' (combineImapHmap rIMap) [] hmap

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

-- Maps --
fromMap :: (Ord a) => Map (a, a) Double -> Markov a
fromMap mapp =
    let imap = foldlWithKey' toimap' Map.empty mapp
        hmap = foldlWithKey' (tohmap' imap) Map.empty mapp
    in Markov imap hmap

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

-- Matrices --
hmatrix :: (Ord a) => Markov a -> Matrix Double
hmatrix (Markov imap hmap) =
    let m = Map.size imap
        zeros = matrix m (replicate (m*m) 0)
        hlist = Map.toList hmap
    in  Dat.accum zeros (+) hlist

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

-- Chains --

walkFrom :: (Ord a, Show a) => a -> Int -> Markov a -> StdGen -> IO [a]
walkFrom start n haskov gen = do
    let mat = hmatrix haskov
        imap = statesI haskov
        rimap = transMap imap
        nextRow = mat Dat.? [(imap Map.! start)]
    return (start : (steps rimap mat nextRow n gen))

walk :: (Ord a, Show a) => Int -> Markov a -> StdGen -> IO [a]
walk n haskov gen = do
    let mat = hmatrix haskov
        ss = tr . steady $ mat
        imap = statesI haskov
        rimap = transMap imap
    return (steps rimap mat ss n gen)

steps :: (Ord a, Show a) => ReverseIndexMap a -> Matrix Double -> Matrix Double -> Int -> StdGen -> [a]
steps _ _ _ 0 _ = []
steps rimap mat row n gen
    | (sumElements row) == 0 = []
    | otherwise = choice : steps rimap mat newRow (n-1) (snd rand)
    where
        rand = random gen
        index = randomStep row (fst rand) 0.0 0
        choice = rimap Map.! index
        newRow = mat Dat.? [index]

randomStep :: Matrix Double -> Double -> Double -> Int -> Int
randomStep row rand total j
    | newTotal < rand = randomStep row rand newTotal newJ
    | otherwise       = j
    where
        newTotal = total + (row Dat.! 0 Dat.! j)
        newJ = (j + 1) `mod` (cols row)

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

-- Steady State --
steadyState :: (Ord a) =>  Markov a -> [(a, Double)]
steadyState haskov = zip (states haskov) (concat . Dat.toLists $ steady (hmatrix haskov))

steady :: Matrix Double -> Matrix Double
steady mat =
    let m = Dat.rows mat -- Number of states
        q = tr $ ident m - mat -- Transpose (I-matrix size m - hmatrix)
        e = col ((replicate (m-1) 0) ++ [1.0]) -- col of zeros size m, last element is 1.0
        x = luSolve (luPacked q) e -- Solve linear system of Q and e (Qx = e)
    in  scale (1 / (sumElements x)) x  -- Divide x by sum of elements

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

-- Normalize --
normalize :: (Ord a) => Markov a -> Markov a
normalize (Markov imap hmap) =
    let ss = foldlWithKey' sumSeq Seqq.empty hmap
    in Markov imap (mapWithKey (calcNorm ss) hmap)

calcNorm :: Seq Double -> (Int, Int) -> Double -> Double
calcNorm ss (i, j) n = n / (Seqq.index ss i)

sumSeq :: Seq Double -> (Int, Int) -> Double -> Seq Double
sumSeq s (i, j) n =
    if Seqq.length s < (i+1)
    then s |> n
    else Seqq.adjust (+n) i s

norm :: Matrix Double -> Matrix Double
norm row = cmap (/ (sumElements row)) row

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

---- Mapping --
--map :: (Ord a) => (Double -> Double) -> Markov a -> Markov a
--map f (Markov imap hmap) = Markov imap (fmap f hmap)

--mapWithState :: (Ord a) => (a -> Double -> Double) -> Markov a -> Markov a
--mapWithState f (Markov imap hmap) =
    --let rIMap = reverseIMap imap

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

-- Delete/Update --
--delete :: (Ord a) => a -> a -> Markov a -> Markov a
--delete i j (Markov imap hmap) =

adjust :: (Ord a) => (Double -> Double) -> a -> a -> Markov a -> Markov a
adjust f i j (Markov imap hmap) =
    if Map.member i imap && Map.member j imap
    then Markov imap (Map.adjust f (imap Map.! i, imap Map.! j) hmap)
    else Markov imap hmap

-- Helper Functions --
-- for fromList
toimap :: (Ord a) => IndexMap a -> ((a, a), Double) -> IndexMap a
toimap acc ((i, j), d)
    | i /= j && Map.notMember i acc && Map.notMember j acc = Map.insert j (Map.size acc+1) (Map.insert i (Map.size acc) acc)
    | Map.notMember i acc = Map.insert i (Map.size acc) acc
    | Map.notMember j acc = Map.insert j (Map.size acc) acc
    | otherwise = acc

-- For fromMap
toimap' :: (Ord a) => IndexMap a -> (a, a) -> Double -> IndexMap a
toimap' acc (i, j) d
    | i /= j && Map.notMember i acc && Map.notMember j acc = Map.insert j (Map.size acc+1) (Map.insert i (Map.size acc) acc)
    | Map.notMember i acc = Map.insert i (Map.size acc) acc
    | Map.notMember j acc = Map.insert j (Map.size acc) acc
    | otherwise = acc

-- For fromList
tohmap :: (Ord a) => IndexMap a -> HMatrixMap -> ((a, a), Double) -> HMatrixMap
tohmap imap acc ((i, j), n) = Map.insert (imap Map.! i, imap Map.! j) n acc

-- For fromMap
tohmap' :: (Ord a) => IndexMap a -> HMatrixMap -> (a, a) -> Double -> HMatrixMap
tohmap' imap acc (i, j) n = Map.insert (imap Map.! i, imap Map.! j) n acc

combineImapHmap :: (Ord a) => Map Int a -> (Int, Int) -> Double -> [((a, a), Double)] -> [((a, a), Double)]
combineImapHmap rIMap (i, j) d acc = ((rIMap Map.! i, rIMap Map.! j), d) : acc

reverseIMap :: (Ord a) => Map a Int -> Map Int a
reverseIMap imap = foldrWithKey' reverseIMap' Map.empty imap
    where reverseIMap' state i m = Map.insert i state m

-- swap map key and index
transMap :: (Ord k, Ord a) => Map k a -> Map a k
transMap = Map.fromList . map swap . Map.toList

-- Machine Epsilon --

machineE :: Double
machineE = until ((== 1) . (+1)) (/2) 1


{-# LANGUAGE RecordWildCards #-}

module View where

import Constants
import Model

import Control.Arrow ((&&&))
import Data.Array
import Data.Bifunctor
import Data.Complex
import Graphics.Gloss
import Graphics.Gloss.Interface.IO.Interact

(<>) = mappend

draw :: World -> IO Picture
draw (World {..}) = pure $ (uncurry scale scaling . uncurry translate coords $ tiling board)

shearTile :: (Float, Float) -> (Float, Float)
shearTile (x, y) = (a*x + b*y, c*y)
  where
    a = rad * sqrt 3 + 1
    b = a / 2
    c = 15.866

-- | draw a board at each point in the lattice
tiling :: Board -> Picture
tiling board = pictures
             . map (\(x,y) -> translate x y (boardWithTiles rad spc board))
             . map shearBoard
             $ lattice

boardWithTiles :: Float -> Float -> Board -> Picture
boardWithTiles r s board = borderBoard r s <> tiles
  where
      tiles = pictures
            . map (\(i,e) -> drawTile (shearTile (bimap fromIntegral fromIntegral i)) e)
            . assocs
            $ board

drawTile :: (Float, Float) -> Tile -> Picture
drawTile _ Empty = blank
drawTile coords Black = color black (uncurry translate coords (circle (rad / 2)))
drawTile coords White = color white (uncurry translate coords (circle (rad / 2)))

shearBoard :: (Float, Float) -> (Float, Float)
shearBoard (x,y) = (a*x + b*y, c*y)
  where
      a = 11 * (2 * rad * cos (pi / 6) + spc)
      c = 11 * 15.866 -- TODO
      --c = 11 * (1.5 * rad + spc)
      b = 11 * (rad * cos (pi / 6) + spc / 2)

-- | square unit lattice corresponding to the boards we need to show
lattice :: [(Float, Float)]
--lattice = (,) <$> [-8..8] <*> [-4..3]
lattice = (,) <$> [-6..5] <*> [-4..3]

-- | board with the border
borderBoard :: Float -> Float -> Picture
borderBoard r s =
    color black
      (translate (-r * cos (pi / 6)) (-r * sin (pi / 6))
        (background r s)) <>
    color orange
      (boardWithSpace r s)

-- | a rhombus that goes behind the board to serve as the background
background :: Float -> Float -> Picture
background r s = polygon
               [ (0,0)
               , (10 * w + d, 0)
               , (10 * 1.5 * w  + d + r / sqrt 3, 10 * 15.866 + r) -- TODO: figure out where i got this number
               , (10 * 0.5 * w + r / sqrt 3, 10 * 15.866 + r)
               ]
  where w = 2 * r * cos (pi / 6) + s
        d = 5 / sqrt 3 + 5 * sqrt 3

-- | board with space in between
boardWithSpace :: Float -> Float -> Picture
boardWithSpace r s = pictures
                   . map (\(p,n) -> translate (n * a * cos (pi / 3)) (n * a * sin (pi / 3)) p)
                   . zip (repeat (rowWithSpace r s))
                   $ [0..10]
  where a = 2 * r * cos (pi / 6) + s

-- | several hexagons of given radius in a row with space in between
rowWithSpace :: Float -> Float -> Picture
rowWithSpace r s = pictures
                 . map (\(p,n) -> translate (n * (2 * r * cos (pi / 6) + s)) 0 p)
                 . zip (repeat (hexagon (r :+ 0)))
                 $ [0..10]

-- | draw a regular hexagon with the given length from the center to a vertex
hexagon :: Complex Float -> Picture
hexagon r = rotate 90
          . polygon
          . take 6
          . map ((realPart &&& imagPart) . (* r) . (\k -> exp (0 :+ 2 * pi * k / 6)))
          $ [0..]

{-
row :: Picture
row = pictures $ map (\(p,n) -> translate (n*18.3205) 0 p) $ zip (repeat (hexagon 10)) [0..10]

board :: Picture
board = pictures $ map (\(p,n) -> translate (n * 9.1603) (n * 15.866) p) $ zip (repeat row) [0..10]
-}


module Lib
    ( forwardPropogate,
      weightsWithInitializer,
      biasesWithInitializer,
      randomWeights,
      sigmoid,
      runBatch,
      updateParameters,
      getLayer,
      runEpochs,
      ActivationFunction (..),
      Activation (..),
      Parameters,
      Gradients,
      Layer (..)
    ) where

import Numeric.LinearAlgebra

-- |The Layer type represents a layer in an artificial neural network
data Layer = Layer { layerParameters  :: Parameters
                   , activationFunction :: ActivationFunction
                   }

-- |Wraps activation function and derivative
data ActivationFunction = ActivationFunction { f  :: Double -> Double
                                             , f' :: Double -> Double
                                             }
-- |Wraps values from layer activations
data Activation = Activation { -- | Values after matrix multiplication but before activation function is applied
                               preActivation     :: Matrix Double
                               -- | Values after activation function is applied
                             , postActivation    :: Matrix Double
                             }

type Parameters = (Matrix Double, Vector Double)
type Gradients = Parameters

-- |Computes values for each layer
forwardPropogate :: Matrix Double -> [Layer] -> [Activation]
forwardPropogate inputs = scanl activate firstActivation
  where
    firstActivation = Activation {preActivation=(numInputRows><numInputCols) [0..], postActivation=inputs}
    numInputRows = rows inputs
    numInputCols = cols inputs

activate :: Activation -> Layer -> Activation
activate previousActivation layer = Activation {preActivation=layerPreActivation, postActivation=layerPostActivation}
  where
    layerPostActivation = cmap layerActivationFunction layerPreActivation
    layerPreActivation = (previousPostActivation <> layerWeights) `add` fromRows (replicate (rows previousPostActivation) layerBiases)
    previousPostActivation = postActivation previousActivation
    (layerWeights,layerBiases) = layerParameters layer
    layerActivationFunction = (f . activationFunction) layer

-- |Get random weights for a given network structure
randomWeights :: [Int] -> Maybe (IO [Matrix Double])
randomWeights []  = Nothing
randomWeights [x] = Nothing
randomWeights xs  = Just $ traverse (uncurry randn) $ steps xs

-- |Get weights all with initializer value for given network structure
weightsWithInitializer :: Double -> [Int] -> [Matrix Double]
weightsWithInitializer initializer layerSizes = map (\(from,to) -> (from><to) $ repeat initializer) $ steps layerSizes

biasesWithInitializer :: Double -> [Int] -> [Vector Double]
biasesWithInitializer initializer layerSizes = map ((fromList . ($ initializer)) . replicate) $ tail layerSizes

steps :: [a] -> [(a,a)]
steps [] = []
steps [x] = []
steps (x:y:xs) = (x,y) : steps (y:xs)

-- |Sigmoid activation function
sigmoid :: ActivationFunction
sigmoid = ActivationFunction {f=s,f'=s'}
  where
    s x = 1 / (1 + exp (-x))
    s' x = s x * (1 - s x)

-- |Run a batch for training neural network model
runBatch :: Matrix Double -> Matrix Double -> [Layer] -> [Gradients]
runBatch inputs expected initialParameters = runBatch' (outputPreviousActivations:remainingActivations) (reverse initialWeights) (tail reversedActivationDerivatives) outputDelta [(outputWeightGradients, outputBiasGradients)]
  where
    outputBiasGradients = getBiasGradient outputDelta
    outputWeightGradients = getWeightGradients outputDelta outputPreviousActivations
    outputDelta = getOutputDelta expected output outputActivationDerivative
    (output:outputPreviousActivations:remainingActivations) = reverse activations
    outputActivationDerivative = head reversedActivationDerivatives
    activations = forwardPropogate inputs initialParameters
    initialWeights = map (fst . layerParameters) initialParameters
    initialBiases = map (snd . layerParameters) initialParameters
    reversedActivationDerivatives = reverse $ map (f' . activationFunction) initialParameters

runBatch' :: [Activation] -> [Matrix Double] -> [Double -> Double] -> Matrix Double -> [Gradients] -> [Gradients]
runBatch' [input] _ _ _ gradients = gradients
runBatch' reversedActivations reversedWeights activationDerivatives nextDelta gradients = runBatch' (previousActivations:remaining) currentWeights remainingActivationDerivatives currentDelta (gradient:gradients)
  where
    gradient = (getWeightGradients currentDelta previousActivations, getBiasGradient currentDelta)
    currentDelta = getDelta currentActivations nextDelta nextWeights currentActivationDerivative
    (currentActivations:previousActivations:remaining) = reversedActivations
    (nextWeights:currentWeights) = reversedWeights
    (currentActivationDerivative:remainingActivationDerivatives) = activationDerivatives

getDelta :: Activation -> Matrix Double -> Matrix Double -> (Double -> Double) -> Matrix Double
getDelta activation nextDelta nextWeights activationDerivative = (nextDelta <> tr nextWeights) * cmap activationDerivative (preActivation activation)

getBiasGradient :: Matrix Double -> Vector Double
getBiasGradient deltaMatrix = fromList $ map (\column -> sum column / fromIntegral (length column)) columns
  where
    columns = map toList $ toColumns deltaMatrix

getOutputDelta :: Matrix Double -> Activation -> (Double -> Double) -> Matrix Double
getOutputDelta expected output activationDerivative = (postActivation output - expected) * cmap activationDerivative (preActivation output)

getWeightGradients ::  Matrix Double -> Activation -> Matrix Double
getWeightGradients delta previousActivations = tr previousPostActivation <> delta
  where
    previousPostActivation = postActivation previousActivations

updateParameters :: Double -> Double -> [Gradients] -> [Parameters] -> [Parameters]
updateParameters learningRate regularizationTerm = zipWith updateParameters
  where
    updateParameters :: Gradients -> Parameters -> Parameters
    updateParameters (weightGradients,biasGradients) (weights,biases) =
      (
        weights - cmap (*learningRate) (applyRegularization regularizationTerm weightGradients),
        biases - cmap (*learningRate) biasGradients
      )

applyRegularization :: Double -> Matrix Double -> Matrix Double
applyRegularization regularizationTerm weightGradients = weightGradients + cmap (*regularizationTerm) weightGradients

runEpochs :: Int -> Double -> Double -> Matrix Double -> Matrix Double -> [Layer] -> [Parameters]
runEpochs numEpochs learningRate regularizationTerm inputs expected = runEpochs' numEpochs
  where
    runEpochs' 0 layers = map layerParameters layers
    runEpochs' remainingEpochs layers = runEpochs' (remainingEpochs - 1) (runEpoch learningRate regularizationTerm layers inputs expected)

runEpoch :: Double -> Double -> [Layer] -> Matrix Double -> Matrix Double -> [Layer]
runEpoch learningRate regularizationTerm layers inputs expected = updatedLayers
  where
    gradientMatrices = runBatch inputs expected layers
    parameters = map layerParameters layers
    functions = map activationFunction layers
    updatedParameters = updateParameters learningRate regularizationTerm gradientMatrices parameters
    updatedLayers = zipWith getLayer updatedParameters functions

getLayer :: Parameters -> ActivationFunction -> Layer
getLayer parameters activationFunction = Layer{layerParameters=parameters, activationFunction=activationFunction}


{-# LANGUAGE GADTs #-}

module HVX.Internal.SymbolicSubgrad where

import Numeric.LinearAlgebra hiding (i)
import Numeric.LinearAlgebra.Util

import HVX.Internal.Matrix
import HVX.Internal.Primitives

type Vars = [(Var, Mat)]

evaluate :: Expr vex mon -> Vars -> Mat
evaluate (EConst c) _ = c
evaluate (EVar var) vars
    | Just val <- lookup var vars = val
    | otherwise = error $ "Variable " ++ show var ++ " not set."
evaluate (EFun f e) vars = getFun f (evaluate e vars)
evaluate (EAdd l r) vars = evaluate l vars + evaluate r vars

jacobianWrtVar :: Expr vex mon  -- ^ Expression whose jacobian to take.
  -> Vars                       -- ^ Variables appearing in the expression.
  -> Var                        -- ^ Variable to take jacobian with respect to.
  -> Mat                        -- ^ The jacobian.
jacobianWrtVar (EConst c) vars wrtVar
    | Just wrtVal <- lookup wrtVar vars = zeros (rows c) (rows wrtVal)
    | otherwise = error $ "Variable " ++ show wrtVar ++ " not set."
jacobianWrtVar (EVar var) vars wrtVar
    | var == wrtVar, Just val <- lookup var vars = ident $ rows val
    | Just val <- lookup var vars, Just wrtVal <- lookup wrtVar vars = zeros (rows val) (rows wrtVal)
    | otherwise = error $ "Variable " ++ show var ++ " or " ++ show wrtVar ++ " not set."
-- Chain rule: ddx f(e) = dde f * ddx e
jacobianWrtVar (EFun f e) vars var = dde_f <> ddx_e
  where
    dde_f = getJacobian f val
    ddx_e = jacobianWrtVar e vars var
    val   = evaluate e vars
-- Special case for addition.
jacobianWrtVar (EAdd l r) vars var = ddx_l + ddx_r
  where
    ddx_l = jacobianWrtVar l vars var
    ddx_r = jacobianWrtVar r vars var


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

module NQS.CG
  ( Operator
  , NQS.CG.cg
  , test
  ) where

import Data.Complex
import Control.Monad.Primitive

import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV

import Foreign.Storable

import Lens.Micro

import NQS.Internal.Types
import NQS.Internal.BLAS

type Operator m el = -- forall m. PrimMonad m
     MDenseVector 'Direct (PrimState m) el -- ^ Input
  -> MDenseVector 'Direct (PrimState m) el -- ^ Output
  -> m ()

class ToComplex a where
  promoteToComplex :: RealOf a -> a

instance {-# OVERLAPPABLE #-} (a ~ RealOf a) => ToComplex a where
  promoteToComplex = id

instance {-# OVERLAPS #-} (Num a, a ~ RealOf (Complex a)) => ToComplex (Complex a) where
  promoteToComplex = (:+ 0)

cg :: forall v el m.
     ( PrimMonad m, v ~ MDenseVector 'Direct
     , DOTC v el el, AXPY v el, COPY v el
     , NRM2 v el (RealOf el), SCAL (RealOf el) v el
     , Storable el, Fractional el, ToComplex el, RealFloat (RealOf el)
     )
  => Int -- ^ Max number of iterations
  -> RealOf el -- ^ Tolerance
  -> Operator m el -- ^ A
  -> MDenseVector 'Direct (PrimState m) el -- ^ b
  -> MDenseVector 'Direct (PrimState m) el -- ^ Initial guess x
  -> m (Int, RealOf el)
cg !maxIter !tol operator !b !x0 = do
  r <- newTempVector n
  p <- newTempVector n
  q <- newTempVector n
  bNorm <- norm b
  -- This will probably never ever happen, but we should check anyway
  if bNorm == 0
    then fill x0 0 >> return (0, 0)
    else do
      -- r := b - max * x0
      copy b r
      operator x0 q
      axpy (-1) q r
      -- Are we done already?
      norm r >>= \rNorm -> if rNorm < tol * tol * bNorm
        then return $! (0, sqrt (rNorm / bNorm))
        else do
          -- Start the loop
          copy r p
          go bNorm x0 r p q rNorm 0
  where
    n = b ^. dim
    newTempVector size = MDenseVector size 1 <$> newVectorAligned size 64
    norm x = (\y -> y * y) <$> nrm2 x
    go !bNorm !x !r !p !q !ρ !i
      | i >= maxIter = nrm2 r >>= \rNorm -> return (i, rNorm / bNorm)
      | otherwise = do
        operator p q
        α <- (promoteToComplex ρ /) <$> dotc p q
        axpy   α  p x
        axpy (-α) q r
        norm r >>= \rNorm -> if rNorm < threshold
          then return (i, sqrt (rNorm / bNorm))
          else do copy r q
                  let ρ' = rNorm
                      β  = ρ' / ρ
                  scal β p
                  axpy 1 q p
                  go bNorm x r p q ρ' (i + 1)
      where threshold :: RealOf el
            threshold = tol * tol * bNorm

test :: IO ()
test = do
  let -- operator :: Operator Float
      operator x y = do
        aBuff <- newVectorAligned 9 64 :: IO (MV.MVector (PrimState IO) Float)
        V.copy aBuff $ V.fromList [1, 2, 3, 4, 5, 6, 7, 8, 9]
        let a = MDenseMatrix 3 3 3 aBuff :: MDenseMatrix 'Row (PrimState IO) Float
        gemv NoTranspose 1.0 a x 0.0 y
  bBuff <- V.unsafeThaw $ V.fromList [2, 2, 2 :: Float]
  xBuff <- V.unsafeThaw $ V.fromList [0, 0, 0 :: Float]
  let x = MDenseVector 3 1 xBuff :: MDenseVector 'Direct (PrimState IO) Float
  let b = MDenseVector 3 1 bBuff :: MDenseVector 'Direct (PrimState IO) Float
  print =<< (NQS.CG.cg 30 (1.0E-8 :: Float) operator b x)


{-# OPTIONS_GHC -fplugin Numeric.Units.Dimensional.DK.Solver #-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeOperators #-}

module Controls
where

import Numeric.Units.Dimensional.DK.Prelude
import Numeric.LinearAlgebra.Dimensional.DK.Shapes
import qualified Prelude as P

data DimMat (s :: Shape) (e :: *) = DimMat

(<>) :: DimMat s1 e -> DimMat s2 e -> DimMat (ShapeProduct s1 s2) e
_ <> _ = DimMat

add :: DimMat s e -> DimMat s e -> DimMat s e
add _ _ = DimMat

data ContinuousLiSystem (iv :: Dimension) (xs :: [Dimension]) (ys :: [Dimension]) (us :: [Dimension]) e = ContinuousLiSystem
                                                                                {
                                                                                a'' :: DimMat (DivideVectorLists (MapDiv iv xs) xs) e,
                                                                                b'' :: DimMat (DivideVectorLists (MapDiv iv xs) us) e,
                                                                                c'' :: DimMat (DivideVectorLists ys xs) e,
                                                                                d'' :: DimMat (DivideVectorLists ys us) e
                                                                                }

type ContinuousLtiSystem = ContinuousLiSystem DTime

type ExampleSystem = ContinuousLtiSystem
    '[DLength, DVelocity, DPlaneAngle, DAngularVelocity]
    '[DLength, DPlaneAngle]
    '[DForce]
    Double

evaluate :: ContinuousLiSystem iv (x ': xs) (y ': ys) (u ': us) e -> DimMat (VectorShape x xs) e -> DimMat (VectorShape u us) e -> (DimMat (VectorShape (x / iv) (MapDiv iv xs)) e,  DimMat (VectorShape y ys) e)
evaluate sys x u = let
                      a = a'' sys
                      b = b'' sys
                      c = c'' sys
                      d = d'' sys
                      xDot = (a <> x) `add` (b <> u)
                      y = (c <> x) `add` (d <> u)
                    in (xDot, y)

{- Example from http://ctms.engin.umich.edu/CTMS/index.php?example=InvertedPendulum&section=ControlStateSpace -}

-- not really stated on that page, but if you go back a couple of pages in their derivation
-- you can see that the type of u is a 1x1 matrix whose sole element is a force

massOfCart = (0.5 :: Double) *~ (kilo gram)
massOfPendulum = (0.2 :: Double) *~ (kilo gram)
coefficientOfFrictionForCart = (0.1 :: Double) *~ (newton / (meter / second))
lengthToPendulumCenterOfMass = (0.3 :: Double) *~ meter
massMomentOfInertiaOfPendulum = (0.006 :: Double) *~ (kilo gram * meter^pos2)
g = (9.8 :: Double) *~ (meter / second^pos2)

p = massMomentOfInertiaOfPendulum*(massOfCart+massOfPendulum)+(massOfCart*massOfPendulum*lengthToPendulumCenterOfMass*lengthToPendulumCenterOfMass)

a22 = negate (massMomentOfInertiaOfPendulum+massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass) * coefficientOfFrictionForCart / p
a23 = (massOfPendulum * massOfPendulum * g * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass) / p
a42 = negate (massOfPendulum * lengthToPendulumCenterOfMass * coefficientOfFrictionForCart) / p
a43 = massOfPendulum * g * lengthToPendulumCenterOfMass*(massOfCart + massOfPendulum)/p

b21 = (massMomentOfInertiaOfPendulum + (massOfPendulum * lengthToPendulumCenterOfMass * lengthToPendulumCenterOfMass)) / p
b41 = massOfPendulum * lengthToPendulumCenterOfMass / p

{-
pendulum = ContinuousLiSystem {
           a'' = [matD| _0, _1, _0, _0;
                        _0, a22, a23, _0;
                        _0, _0, _0, _1;
                        _0, a42, a43, _0 |],
           b'' = [matD| _0;
                        b21;
                        _0;
                        b41 |],
           c'' = [matD| _1, _0, _0, _0;
                        _0, _0, _1, _0 |],
           d'' = zeroes
         } :: ExampleSystem
-}

pendulum = ContinuousLiSystem {
           a'' = DimMat,
           b'' = DimMat,
           c'' = DimMat,
           d'' = DimMat
         } :: ExampleSystem


import System.Environment (getArgs)
import Control.Arrow (first)
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Diagrams
import Data.Time
import qualified Statistics.Sample as S
import qualified Data.Vector as V
import System.Process

import Data.Text.Time (parseISODateTime)
import qualified Data.TimeSeries as TS

signal :: Num a => TS.Series a -> [(LocalTime, a)]
signal ts = map (first (utcToLocalTime utc)) (TS.toList ts)

main = do
    args <- getArgs
    let fn = head args
    ts <- TS.loadCSV TS.HasHeader parseISODateTime fn
    let y3 = TS.years 3
    let xs = TS.rolling y3 (S.mean . V.fromList) ts
    toFile def "dist/smooth.svg" $ do
        layout_title .= fn
        setColors [opaque blue, opaque red]
        plot (line "Orignal" [signal ts])
        plot (line "Smoothed" [signal xs])
    putStrLn "Plot saved to: dist/smooth.svg"
    createProcess (shell "firefox dist/smooth.svg")


module Cas.Instances.Other where

import PreludeCustom
import Cas.Internal.Interface
import Cas.Internal.Normalization
import Data.Complex.Generic

instance Enum X where
    fromEnum = xI
    toEnum = X

instance Fractional T where
    recip = normalize . mInv
    fromRational = TF . fromRational

instance Num T where
    negate      = normalize . aInv
    (+)         = (normalize .) .  addBin
    (*)         = (normalize .) .  mulBin
    fromInteger = TF . fromInteger

instance Floating T where
    (**) = (normalize .) . pow

-- instance Fractional T where
--     recip = mInv
--     fromRational = TF . fromRational

-- instance Num T where
--     negate      =  aInv
--     (+)         =  addBin
--     (*)         =  mulBin
--     fromInteger = TF . fromInteger

-- instance Floating T where
--     (**) = pow


module Statistics.Information.Continuous.TotalCorrelation where

import Data.Matrix
import Statistics.Information.Continuous.Entropy
import Statistics.Information.Continuous.MutualInfo
import Statistics.Information.Utils.Matrix
import Statistics.Information.Utils.Random
import System.Random

tc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
tc g k base xs = sum hxs - hx where
  hxs = [entropy gcol k base col | (col, gcol) <- cols]
  hx = entropy g2 k base xs
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

ctc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
       Double
ctc g k base xs ys = sum hxs - hx where
  hxs = [centropy gcol k base col ys | (col, gcol) <- cols]
  hx = centropy g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

corex_tcs :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_tcs g k base xs ys = tc g1 k base xs - ctc g2 k base xs ys where
  (g1, g2) = split g

corex_mis :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
             Double
corex_mis g k base xs ys = sum mixs - mi_all where
  mixs = [mi gcol k base col ys | (col, gcol) <- cols]
  mi_all = mi g2 k base xs ys
  (g1, g2) = split g
  cols = (map colVector (columns xs)) `zip` (splitN (ncols xs) g1)

residual :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
residual g k base xs = sum $ [hlessxi i gi | (i, gi) <- igs] where
  hlessxi i gi = let xi = (colVector $ getCol i xs) in
    centropy gi k base xi $ residualMatrix xs i
  igs = [1..(ncols xs)] `zip` (splitN (ncols xs) g)

dtc :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
dtc g k base xs = entropy g1 k base xs - residual g2 k base xs where
  (g1, g2) = split g


-- conversion to/from Data.Vector.Storable
-- from Roman Leshchinskiy "vector" package
--
-- In the future Data.Packed.Vector will be replaced by Data.Vector.Storable

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

import Numeric.LinearAlgebra as H
import Data.Packed.Development(unsafeFromForeignPtr, unsafeToForeignPtr)
import Foreign.Storable
import qualified Data.Vector.Storable as V

fromVector :: Storable t => V.Vector t -> H.Vector t
fromVector v = unsafeFromForeignPtr p i n where
    (p,i,n) = V.unsafeToForeignPtr v

toVector :: Storable t => H.Vector t -> V.Vector t
toVector v = V.unsafeFromForeignPtr p i n where
    (p,i,n) = unsafeToForeignPtr v

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

v = V.slice 5 10 (V.fromList [1 .. 10::Double] V.++ V.replicate 10 7)

w = subVector 2 3 (linspace 5 (0,1)) :: Vector Double

main = do
    print v
    print $ fromVector v
    print w
    print $ toVector w