module Neural
  ( Layer(Layer)
  , Network
  , evalLayer
  , evalLayerM
  , evalNetwork
  , evalNetworkM
  , evalOutputError
  , evalOutputErrorM
  , evalLayerError
  , evalLayerErrorM
  , evalNetworkError
  , evalNetworkErrorM
  , gradientDescent
  , gradientDescentM
  , sumColumns
  , vecToMat
  )
where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data

-- Dependent types would be great to represent some of these constraints.

-- A layer in the neural network. The weights are represented by a matrix with
-- 1 row for each neuron, and 1 column for each weight. The biases are represented
-- by a vector. The number of biases must match the number of weights.
data Layer = Layer (Matrix R) (Vector R) deriving Show

-- A network is a list of layers, where each layer has the same number of weights as
-- the previous one has neurons. Otherwise, evalNetwork will fail.
type Network = [Layer]

-- A utility function that converts a vector to a row matrix
vecToMat :: Vector R -> Matrix R
vecToMat = fromRows . (:[])

-- A utility function that sums the colums in a matrix to produce a row vector
sumColumns :: Matrix R -> Vector R
sumColumns m = (vector . replicate (rows m) $ 1) <# m

-- A sigmoid function which is at 0.5 for the input 0, and tends to 1 as the input
-- increases, and to 0 as the input decreases.
sigmoid :: Floating f => f -> f
sigmoid z = 1 / (1 + exp (-z))

-- Differential of the sigmoid fuunction
sigmoid' :: Floating f => f -> f
sigmoid' z = sigmoid z * (1 - sigmoid z)

-- Evaluates a layer of the neural network given a set of inputs. The number of inputs
-- must match the number of weights in the input network or this function will fail
-- with an exception. Returns the outputs as well as the weighted inputs.
evalLayer :: Layer -> Vector R -> (Vector R, Vector R)
evalLayer layer inputs = (flatten output, flatten weightedInput)
  where
    (output, weightedInput) = evalLayerM layer (fromRows [inputs])

-- Evaluates a layer of the neural network given a set of inputs. Contrary to evalLayer,
-- the set of inputs is provided as a matrix where each row is a vector of inputs as
-- given in evalLayer.
evalLayerM :: Layer -> Matrix R -> (Matrix R, Matrix R)
evalLayerM (Layer weights biases) inputs = (sigmoid weightedInputs, weightedInputs)
  where
    inputRows = rows inputs
    biasMatrix = fromColumns . replicate inputRows $ biases
    weightedInputs = tr' $ (weights <> tr' inputs) + biasMatrix

-- Feed forward the given function (e.g. evalLayer) over the network
feedForward :: (Layer -> a -> (a, a)) -> a -> Network -> [(a, a)]
feedForward f inputs network = tail . reverse $ feedForward' f inputs network
  where
    feedForward' f inputs network = foldl (\acc@((a,_):_) l -> (f l a) : acc) [(inputs, undefined)] network

-- Evaluates a neural network by folding a set of inputs through each layer.
-- The number of inputs must match the number of weights in the first layer,
-- and the number of elements in the output vector depends on the number of neurons
-- in the final layer of the network. Returns the output and weighted input of each
-- layer in the network.
evalNetwork :: Network -> Vector R -> [(Vector R, Vector R)]
evalNetwork network inputs = feedForward evalLayer inputs network

-- Evaluates a neural network. Contrary to evalNetwork, the input is taken in the same form
-- as evalLayerM, i.e. a matrix where the rows are a single input vector.
evalNetworkM :: Network -> Matrix R -> [(Matrix R, Matrix R)]
evalNetworkM network inputs = feedForward evalLayerM inputs network

-- Evaluates the error in the output layer of a network
evalOutputError :: Vector R -> Vector R -> Vector R -> Vector R -> Vector R
evalOutputError output weightedInput inputs training = (output - training) * sigmoid' weightedInput

-- Evaluates the error in the output layer of a network. Input is a matrix where rows are inputs.
evalOutputErrorM :: Matrix R -> Matrix R -> Matrix R -> Matrix R -> Matrix R
evalOutputErrorM output weightedInput inputs training = (output - training) * sigmoid' weightedInput

-- Evaluates the error in a layer of a network given the weighted cost, and the next layer's weight
-- and error.
evalLayerError :: Vector R -> Matrix R -> Vector R -> Vector R
evalLayerError weightedCost nextLayerWeights nextLayerError =
  ((tr' nextLayerWeights) #> nextLayerError) * (sigmoid' weightedCost)

-- Evaluates the error in a layer like evalLayerError, but takes a matrix of inputs as rows.
evalLayerErrorM :: Matrix R -> Matrix R -> Matrix R -> Matrix R
evalLayerErrorM weightedCost nextLayerWeights nextLayerError =
  (nextLayerError <> nextLayerWeights) * (sigmoid' weightedCost)

-- Evaluates the error for each layer in the network by means of backpropogation.
evalNetworkError :: Network -> [(Vector R, Vector R)] -> Vector R -> Vector R -> [Vector R]
evalNetworkError network networkValues inputs training =
  map flatten $ evalNetworkErrorM network networkValuesM (vecToMat inputs) (vecToMat training)
    where
      networkValuesM = map (\(a,b) -> (vecToMat a, vecToMat b)) networkValues

-- Evaluates the error for each layer in the network by means of backpropogation.
evalNetworkErrorM :: Network -> [(Matrix R, Matrix R)] -> Matrix R -> Matrix R -> [Matrix R]
evalNetworkErrorM network networkValues inputs training = map fst backPropogation
  where
    -- The layers of the network in reverse order
    reverseNetwork = reverse network
    -- The activations of the layers of the network in reverse order
    reverseLayerVals = reverse networkValues
    -- (reverse layers, reverse activations)
    reverseLayerData = zip reverseNetwork reverseLayerVals
    -- The activations of the output layer of the network
    outputLayer = tail reverseLayerVals
    -- The weights of the output layer
    Layer outputWeights _ = head reverseNetwork
    -- The activations and weighted inputs of the output layer
    (outputValues, outputLayerWeightedInput) = head reverseLayerVals
    -- The error of the output layer
    outputError = evalOutputErrorM outputValues outputLayerWeightedInput inputs training
    -- An accumulator for backpropogation, started with the output layer's error
    acc = [(outputError, outputWeights)]
    -- The backpropogation algorithhm (apply evalOutputError to each layer in reverse order)
    backPropogate = \acc@((loe,lw):_) (Layer w _, (_,z)) -> (evalLayerErrorM z lw loe, w):acc
    backPropogation = foldl backPropogate acc (tail reverseLayerData)

-- Gradient descent function. Given a network, its inputs, its activations, and its error,
-- produces a new network corrected through gradient descent.
gradientDescent :: Network -> Vector R -> [Vector R] -> [Vector R] -> R -> Network
gradientDescent network input activations error learnRate =
  gradientDescentM network (vecToMat input) (map vecToMat activations) (map vecToMat error) learnRate

-- Gradient descent function. Given a network, its inputs, its activations, and its error,
-- produces a new network corrected through gradient descent.
gradientDescentM :: Network -> Matrix R -> [Matrix R] -> [Matrix R] -> R -> Network
gradientDescentM network input activations error learnRate =
  map gradientDescent' layerData
    where
      -- Get the activations from the previous layer of neurons
      -- by offsetting the activations with the input at the start
      -- and discarding the final element
      previousActivations = input : (init activations)
      -- Layers zipped with their errors and the previous layer's activation
      layerData = zip3 network error previousActivations
      -- Gradient descent for one layer
      gradientDescent' (Layer weight bias, error, previousActivation) =
        -- New layer is evaluated by calculating the gradient for the weight and bias and
        -- subtracting gradient * learnRate from the current value
        Layer
          (weight - (weightGradients * realToFrac learnRate))
          (bias - (biasGradients * realToFrac learnRate))
            where
              -- Calculate gradient for each rate
              -- Calculated using the kronecker product, a generalisation of tensor product to matrices
              --weightGradients = kronecker previousActivation (tr' error)
              --weightGradients = (previousActivation <> (tr' error)) -- / (fromIntegral $ cols previousActivation)
              weightGradients = (tr' error <> previousActivation) / (fromIntegral . cols $ previousActivation)
              -- Sum the bias gradient components for each neuron (the columns)
              biasGradients = sumColumns error / (fromIntegral . size $ bias)


{-# LANGUAGE FlexibleContexts #-}

module Spectral(layout,default_layout) where

import Numeric.LinearAlgebra.HMatrix
import Numeric.LinearAlgebra.Data

degree_matrix m =
  let a = foldl1 (+) (toRows m)
      b = takeDiag m
      in
    a - b

compute_laplacian adj_matrix node_weights =
  let (n,m) = size adj_matrix
      weighted_adjency = adj_matrix
      mat_D = diag $ degree_matrix weighted_adjency
      mat_L = mat_D - weighted_adjency
      in
      (mat_L,mat_D)

symmetrize m = m + tr m

layout adjacent_matrix node_weights =
  let sym_adjacent_matrix = symmetrize adjacent_matrix
      (mat_L,mat_D) = compute_laplacian sym_adjacent_matrix node_weights
      eigenVectors = toColumns $ snd $ geigSH' mat_L mat_D
      v2 = eigenVectors !! 1
      v3 = eigenVectors !! 2
  in
    [(v2  ! i, v3 ! i) | i <- [0..size v2-1] ]

default_layout adjacent_matrix =
  let (n,_) = size adjacent_matrix
      ones = matrix n [1 | i <- [1..n]]
  in
    layout adjacent_matrix ones


module Numeric.FFT.Utils
       ( omega, slicevecs, slicemvecs, primes, isPrime
       , allFactors, factors
       , primitiveRoot, invModN, log2, isPow2, dupperm, (%.%)
       , compositions, makeComp, multisetPerms
       , backpermuteM
       ) where

import qualified Control.Monad               as CM
import           Control.Monad.ST
import           Data.Bits
import           Data.Complex
import           Data.List                   (nub)
import qualified Data.List                   as L
import qualified Data.Vector                 as V
import           Data.Vector.Unboxed
import qualified Data.Vector.Unboxed.Mutable as MV
import           Prelude                     hiding (all, concatMap, dropWhile,
                                              enumFromTo, filter, head, length,
                                              map, maximum, null, reverse)
import qualified Prelude                     as P

import           Numeric.FFT.Types

-- | Roots of unity.
omega :: Int -> Complex Double
omega n = cis (2 * pi / fromIntegral n)

-- | Slice a vector @v@ into equally sized parts, each of length @m@.
slicevecs :: Int -> VCD -> VVCD
slicevecs m v = V.map (\i -> slice (i * m) m v) $
                V.enumFromN 0 (length v `div` m)

-- | Slice a mutable vector @v@ into equally sized parts, each of
-- length @m@.
slicemvecs :: Int -> MVCD a -> VMVCD a
slicemvecs m v = V.map (\i -> MV.slice (i * m) m v) $
                 V.enumFromN 0 (MV.length v `div` m)

-- | Determine primitive roots modulo n.
--
-- From Wikipedia (https://en.wikipedia.org/wiki/Primitive_root_modulo_n):
--
-- No simple general formula to compute primitive roots modulo n is
-- known.  There are however methods to locate a primitive root that
-- are faster than simply trying out all candidates.  If the
-- multiplicative order of a number m modulo n is equal to phi(n) (the
-- order of Z_n^x), then it is a primitive root.  In fact the converse
-- is true: if m is a primitive root modulo n, then the multiplicative
-- order of m is phi(n).  We can use this to test for primitive roots.
-- [Here, phi(n) is Euler's totient function, and Z_n^x is the
-- multiplicative group of integers modulo n.]
--
-- First, compute phi(n).  Then determine the different prime factors
-- of phi(n), say p1, ..., pk.  Now, for every element m of Z_n^x,
-- compute
--
--    m^(phi(n) / pi) mod n   for i = 1, ..., k
--
-- using a fast algorithm for modular exponentiation such as
-- exponentiation by squaring.  A number m for which these k results
-- are all different from 1 is a primitive root.
--
-- [In our case, n is restricted to being prime, and phi(p) = p - 1
-- for prime p.]
--
primitiveRoot :: Int -> Int
primitiveRoot p
  | isPrime p =
    let tot = p - 1
        -- ^ Euler's totient function for prime values.
        totpows = map (tot `div`) $ fromList $ nub $ toList $ allFactors tot
        -- ^ Powers to check.
        check n = all (/=1) $ map (expt p n) totpows
        -- ^ All powers are different from 1 => primitive root.
    in fromIntegral $ head $ dropWhile (not . check) $ fromList [1..p-1]
  | otherwise = error "Attempt to take primitive root of non-prime value"

-- | Fast exponentation modulo n by squaring.
expt :: Int -> Int -> Int -> Int
expt n b pow = fromIntegral $ go pow
  where bb = fromIntegral b
        nb = fromIntegral n
        go :: Int -> Integer
        go p
          | p == 0 = 1
          | p `mod` 2 == 1 = (bb * go (p - 1)) `mod` nb
          | otherwise = let h = go (p `div` 2) in (h * h) `mod` nb

-- | Find inverse element in multiplicative integer group modulo n.
invModN :: Int -> Int -> Int
invModN n g = head $ filter (\iv -> (g * iv) `mod` n == 1) $ enumFromTo 1 (n-1)

-- | Prime sieve from Haskell wiki.
primes :: Integral a => [a]
primes = 2 : primes'
  where primes' = sieve [3, 5 ..] 9 primes'
        sieve (x:xs) q ps@ ~(p:t)
          | x < q     = x : sieve xs q ps
          | otherwise =     sieve [n | n <- xs, rem n p /= 0] (P.head t^(2::Int)) t

-- | Naive primality testing.
isPrime :: Integral a => a -> Bool
isPrime n = n `P.elem` P.takeWhile (<= n) primes

-- | Simple prime factorisation.
allFactors :: (Integral a, Unbox a) => a -> Vector a
allFactors n = fromList $ go n primes
  where go cur pss@(p:ps)
          | cur == p         = [p]
          | cur `mod` p == 0 = p : go (cur `div` p) pss
          | otherwise        = go cur ps

-- | Simple prime factorisation: small factors only; largest/last
-- factor picked out as "special".
factors :: (Integral a, Unbox a) => a -> (a, Vector a)
factors n = let (lst, rest) = go n primes in (lst, fromList rest)
  where go cur pss@(p:ps)
          | cur == p         = (p, [])
          | cur `mod` p == 0 = let (lst, rest) = go (cur `div` p) pss
                               in (lst, p : rest)
          | otherwise        = go cur ps

-- | Base-2 logarithm.
log2 :: Int -> Int
log2 1 = 0
log2 n = 1 + log2 (n `div` 2)

-- | Check for powers of two.
isPow2 :: Int -> Bool
isPow2 1 = True
isPow2 n
  | n `mod` 2 == 0 = isPow2 $ n `div` 2
  | otherwise      = False

-- | Duplicate a sub-permutation to fill a given vector length.
dupperm :: Int -> VI -> VI
dupperm n p =
  let sublen = length p
      shift di = map (+(sublen * di)) p
  in concatMap shift $ enumFromN 0 (n `div` sublen)

-- | Composition of permutations.
(%.%) :: VI -> VI -> VI
p1 %.% p2 = backpermute p2 p1

-- | Generate all compositions of a given integer.
compositions :: Int -> V.Vector (Vector Int)
compositions 0 = V.empty
compositions n = let fs = allFactors n
                 in V.reverse $ V.map (makeComp fs) $ V.enumFromN 0 (2^(n-1))

-- | Generate a single composition of a given integer.
makeComp :: Vector Int -> Int -> Vector Int
makeComp fs i = fromList $ foldOps (toList fs) $ makeOps (length fs) i
  where foldOps :: [Int] -> [Bool] -> [Int]
        foldOps (f:fs) ops = go f fs ops
          where go acc [] [] = [acc]
                go acc (f:fs) (op:ops) = if op then go (acc * f) fs ops
                                         else acc : go f fs ops
        makeOps :: Int -> Int -> [Bool]
        makeOps n i = P.replicate (n - 1 - P.length bs) False P.++ bs
          where bs = P.dropWhile not $ P.reverse $
                     P.map (testBit i) [0..finiteBitSize i-1]

-- | Generate all distinct permutations of a multiset in lexicographic
-- order.
multisetPerms :: Vector Int -> [Vector Int]
multisetPerms idp = sidp : L.unfoldr step sidp
  where sidp = fromList $ L.sort $ toList idp
        step v = case permStep v of
          Nothing -> Nothing
          Just p  -> Just (p, p)
        permStep :: Vector Int -> Maybe (Vector Int)
        permStep v =
          if null ks
          then Nothing
          else let k = maximum ks
                   ls = filter (\i -> v ! k < v ! i) $ enumFromN 0 n
                   l = maximum ls
               in Just $ revEnd k (swap k l)
          where n = length v
                ks = filter (\i -> v ! i < v ! (i+1)) $ enumFromN 0 (n-1)
                swap a b = generate n $ \i ->
                  if i == a then v ! b
                  else if i == b then v ! a
                       else v ! i
                revEnd f vv = generate n $ \i ->
                  if i <= f
                  then vv ! i
                  else vv ! (n - i + f)

-- | ST monad version of vector permutation.
backpermuteM :: Int -> VI -> MVCD s -> MVCD s -> ST s ()
backpermuteM n perm vin vout =
  CM.forM_ [0..n-1] $ \i -> do
    idx <- indexM perm i
    x <- MV.unsafeRead vin idx
    MV.unsafeWrite vout i x


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict #-}
module Pinwheel.Base where

import           Data.Complex

-- {-# INLINE fourierMellin #-}
-- fourierMellin :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
-- fourierMellin sigma angularFreq radialFreq (!x, !y) =
--   if x == 0 && y == 0
--     then 0
--     else (x :+ y) ** (fromIntegral (-angularFreq) :+ 0) *
--          ((x ^ 2 + y ^ 2) :+ 0) **
--          (((fromIntegral angularFreq - 2 + sigma) :+ (-fromIntegral radialFreq)) /
--           2)

{-# INLINE fourierMellin #-}
fourierMellin :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
fourierMellin sigma angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else let r = sqrt $ x ^ 2 + y ^ 2
             theta = atan2 y x
         in ((r :+ 0) ** ((sigma - 1) :+ fromIntegral (-radialFreq))) *
            (cis $ fromIntegral (-angularFreq) * theta)

{-# INLINE fourierMellinInv #-}
fourierMellinInv :: (Eq e, RealFloat e) => e -> Int -> Int -> (e, e) -> Complex e
fourierMellinInv sigma angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else (x :+ y) ** (fromIntegral angularFreq :+ 0) *
         ((x ^ 2 + y ^ 2) :+ 0) **
         (((-(fromIntegral angularFreq) - sigma) :+ fromIntegral radialFreq) /
          2)

{-# INLINE fourierMellinInvPeriod #-}
fourierMellinInvPeriod ::
     (Eq e, RealFloat e) => e -> e -> Int -> Int -> (e, e) -> Complex e
fourierMellinInvPeriod sigma periodEnv angularFreq radialFreq (!x, !y) =
  if x == 0 && y == 0
    then 0
    else (x :+ y) ** (fromIntegral angularFreq :+ 0) *
         ((x ^ 2 + y ^ 2) :+ 0) **
         (((-(fromIntegral angularFreq) - sigma) :+ 2 * pi * fromIntegral radialFreq / (log periodEnv)) /
          2)


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

module Main where

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

{-
 Weights, Network & neural net implementation by Justin Le (not mine!)

 Tutorial by Justin Le:

 https://blog.jle.im/entry/practical-dependent-types-in-haskell-1.html
-}
data Weights = W { wBiases :: !(Vector Double)  -- n
                 , wNodes  :: !(Matrix Double)  -- n x m
                 }                              -- "m to n" layer
 deriving Show

data Network :: * where
    O     :: !Weights
          -> Network
    (:&~) :: !Weights
          -> !Network
          -> Network
  deriving Show -- ?

infixr 5 :&~

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

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

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

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

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

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

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

{-
 Inserting my own training exercise here.
 This is inspired by the language-model task of predicting the following
 element in a 4-gram, based upon the last three elements available.
 The task: train a randomly generated system of weights to recognize
 the following pattern & predict the following item according to it:

 0 1 2 0 1 2 0 1 2 ..
-}

inf_pattern :: [Int]
inf_pattern = cycle [0..2]

fin_pttrn_len :: Int
fin_pttrn_len = 99999

{-
 Forming training data from the inf_pattern.
 The training data will be a list of pairs, each pair containing
 * a training input, 3-element vector;
 * a target output, 1-element vector.
-}
training_data :: [(Vector Double,  -- 3 * 1
                   Vector Double)] -- 1 * 1
training_data =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 3 sub_pattern, [sub_pattern !! 3])
            | j <- [0..(length fin_pattern - 4)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (vector $ fromIntegral <$> inp,
                          vector $ fromIntegral <$> out))
         list_pairs

-- Actually, if we were really being realistic about this then we would
-- make our inputs and outputs work according to some sort of one-hot encoding.
-- But never mind.

-- 1. Get an initial random network. We'll pick a basic enough size:
myRandNet :: MonadRandom m => m Network
myRandNet = randomNet 3 [4, 5] 1

-- 2. How can I get a pre-training prediction, then iteratively
--    train on the next available sample, followed by getting the
--    revised prediction?

-- An auxiliary function for "sequencing together"
-- training a network & printing info about it.
foldTraining :: Double              -- training rate
             -> Network             -- untrained network
             -> [(Vector Double,    -- pairs of training input
                  Vector Double)]   --          and training target output
             -> IO Network          -- train steps & print steps alternate in IO
foldTraining rate net pairs =
  foldl (\io_net sample@(inp, tgt) -> do
            n <- io_net
            putStrLn $ "Will train on: " ++ show inp
            let prediction = runNet n inp
            putStrLn $ "Predicted following item: " ++ show prediction
            -- now trying 10 rounds each per data sample
            let n' = last $
                     take 10 $
                     iterate (train rate inp tgt) n
                          -- ^ :: Network -> Network
            return n')
        (return net)
        pairs

-- We can come back later & modify this to give back the resulting
-- trained network, for use elsewhere (i.e. recurrent NN's!)
displayTraining :: Double              -- training rate
                -> [(Vector Double,    -- pairs of training input
                     Vector Double)]   --          and training target output
                -> IO ()
displayTraining rate pairs = do
  -- pluck out a random net
  randNet <- evalRandIO myRandNet
  -- show what it predicts for 1st sample *before* any training is done
  let inp_0 = fst $ head training_data
  let tgt_0 = snd $ head training_data
  let predict_0 = runNet randNet inp_0
  putStrLn $ "The untrained network predicts that following: "
             ++ show inp_0 ++ ", we will get: " ++ show predict_0
  -- now, alternate showing predictions & performing training on some input
  foldTraining rate randNet pairs
  return ()

{-
 Remark.
 The network is obviously learning *something*, because early predictions
 (which are basically noise) eventually give way to predicting minuscule
 numbers to follow every instance of "2", and numbers infinitesimally
 less than "1" following every instance of "0". But learning that every
 instance of "1" is followed by a "2" seems to take much longer! I haven't
 witnessed it yet despite using up to 10K training sample, with 1 learning
 round per sample.

 What if I try more rounds? It doesn't seem to help as far as training the
  network to detect the "2 is followed by 0" part of the pattern.
-}

{-
 RECURRENT NEURAL NETWORKS.
 We study a simple case of these, Elman network a.k.a. simple RNN.
 * for ea. training sample, we simply take *one element* of the training
   data at a particular point in the sequence (rather than the window-base
   n-gram approach of before.)
 * for the training step at current time t, we rely upon:
    a) the training input at current time t, AND
    b) the activation from previous time (t-1)
        (so we have to initialize a "blank activation" for use
         at the very first step, t := 1; this is h := the zero vector.)
    c) therefore, in a simple model with just *two* sets of weights
       (besides the one which handles the previous activation), the
       back-propagation due to comparing current output vs. current target
       will back-propagate a bundle of "nudges" dh_t given what W2_t
       is at the current time;
       this in turn must back-prop *weight deltas* to
       BOTH U_t (which applies to the previous activation)
       AND W1_t (which applies to the current training input.)
-}

-- RNN ~ Weights x Network.
-- We will model this using a separate collection of weights,
-- U, that we simply carry along while remembering to train.
data RNN = MkRNN {
  wts :: Weights,   -- <- U: h-to-h
  net :: Network }  -- <- (W: i-to-h) :&~ (V: h-to-o)

runRNN :: RNN
       -> Vector Double -- input vector
       -> Vector Double -- previous-step activation
       -> Vector Double -- prediction
runRNN rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) x_t h_t_1 =
    let z_u   = (uB + uN #> h_t_1) + -- weighted previous activation
                (wB + wN #> x_t)     -- weighted current training input
        h_t   = logistic z_u         -- current activation.
                                     -- *Need to include somewhere in output!*
        z_h   = (vB + vN #> h_t)     -- next layer of weights
        y_t   = logistic z_h         -- current final activation i.e. output
    in  y_t  -- ^ softmax usually, for tagging problems & similar
             --   (multiple way classification)

{-
 The below is *not* actually a full implementation of the simple RNN
 described at the beginning of Jurafsky Ch. 9, because it does not
 perform the overall training in the two-pass algorithm which is
 explained there.
 In particular, the first pass is more or less all here, but the
 second pass (Backpropgataion Through Time) is *not* taking place right now.
 The idea is actually to *wait* (I think?) until all of y_1, ..., y_l have been
 attempted, *then* (in phase #2) work backwards & successive apply gradient
 terms to V, U, W, U, W, U, W ... (? sequence of matrices being adjusted),
 propagating error back through h_l, h_(l-1), h_(l-2) ... , h_1 *all in one phase*!

 The idea is that, instead of just nudging the matrix entries of U (the recurrent
 part) once per training sample, you would actually wait until the end, and *then*
 have steps iteratively unfolding along the lines of:

 dh_t       = tr vN #> dEdz
 dh_(t-1)   = tr uN #> dEdh_t
 dh_(t-2)   = tr uN #> dEdh_(t-1)
 dh_(t-3)   = tr uN #> dEdh_(t-2)
 ...

-}

trainRNN :: RNN -- ~ U x (W :&~ O V) :
                -- the  graph for this neural network is
                -- as follows:
                {-
                  h_t-1 -> U -> h_t -> V -> y_t
                              ^
                  x     -> W /
                -}
          -> Double        -- training rate
          -> Vector Double -- activation, h_t-1
          -> Vector Double -- training input, x_t
          -> Vector Double -- training target output, tgt_t
          -> (Vector Double, RNN) -- the new activation, h_t, and
                                  -- all three weight-sets of the RNN trained
-- -- don't write recursions, just do everything by hand
trainRNN rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) rate h_t_1 x_t tgt_t =
  -- compute the model output
  let z_u   = (uB + uN #> h_t_1) + -- weighted previous activation
              (wB + wN #> x_t)     -- weighted current training input
      h_t   = logistic z_u         -- current activation.
                                   -- *Need to include somewhere in output!*
      z_h   = (vB + vN #> h_t)     -- next layer of weights
      y_t   = logistic z_h         -- current final activation i.e. output
      -- Here, E = (y_t - tgt_t) ^ 2 is implied.
      dEdz = logistic' z_h * (y_t - tgt_t) -- this vector reports, for each
                                            -- component of the final activation (y_t)_j,
                                            -- how much a unit "bump" to it will impact the
                                            -- matching component (E_t)_j of the error,
                                            -- which is not explicitly calculated here.
      -- Bump the preceding-weights bias-vector vB in the opposite direction
      -- of dEdy, & (per component) in proportion with the impact that its change
      -- has on the error E. This is (- scale rate dEdy) (we scale by learning rate)
      vB'   = vB - scale rate dEdz
      -- Bump every matrix entry in vN by the appropriate amount.
      -- Row #1's matrix-entry bumps are given by transposing the layer input (h_t)'s
      -- entries by (- scale rate dEdy)_1, then transposing (so that connections from the
      -- 1st component of h_t will be reinforced appropriately.)
      -- Do this for all entries of "d_vB / d_y = -scale rate dEdy", and you will get
      vN'   = vN + ((- scale rate dEdz) `outer` h_t)
      -- The above two terms vB', vN' specified how to "bump" the entries of the V-layer
      -- bias & weights matrix, given what the layer input h_t currently is.
      -- Now, figure out how to "bump" the entries of the V-layer input h_t,
      -- given what the current V-layer weights vB / vN are.
      -- In fact, only vN will play any role.
      dh_t  = tr vN #> dEdz -- should it be (- scale rate dEdy)?
      -- I think that the learning rate factor is omitted when back-propagating
      -- a new error to the preceding layer, because otherwise the n-th layer back
      -- would get a factor of (rate ^ n) -> 0 as n increases, causing only the
      -- very few final layers to really be modified.

      -- Now, backpropagate this desired change in h_t to train the wts in U, W.
      -- Just as we previously had      "dy := - scale rate dEdy"
      -- so also will we now have       "dEdh := logistic' h_t * dh_t"
      --                                "dh := - scale rate dEdh"
      -- Train both U (whose input was h_t_1) and W (whose input was x_t).
      -- The activation, h_t, on these linear maps' composite result is now
      -- the vector being nudged, because it was the input of the final layer.
      dEdh  = logistic' h_t * dh_t

      uB'   = uB - scale rate dEdh
      uN'   = uN + ((- scale rate dEdh) `outer` h_t_1)
      wB'   = wB - scale rate dEdh
      wN'   = wN + ((- scale rate dEdh) `outer` x_t)  -- x_t *not* h_t / h_t_1
  in  ( h_t,                  -- the new activation
        MkRNN                 -- the new RNN
          (W uB' uN')
          ((W wB' wN') :&~
            O (W vB' vN')))

-- Sample function for using trainRNN one time.
f = do
    rnn_0 <- evalRandIO (randRNN 1 2 5)
    let h_rnn' =
          trainRNN rnn_0 0.25
                   (vector $ fromIntegral <$> [0,0,0,0,0])
                   (vector $ fromIntegral <$> [0])
                   (vector $ fromIntegral <$> [1,4])
    return h_rnn'

-- Need function to product a random RNN.
-- The relevant dimensions need to be:

{-
h_t-1 -> U -> h_t -> V -> y_t
            ^
x     -> W /

For x (input vector)  - dim(x): m
    h (activation)    - dim(h): h
    y (output vector) - dim(y): n
we have that the matrix sizes should be
  W:  h * m
  U:  h * h
  V:  n * h

rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN)))
-}
randRNN :: MonadRandom m => Int   -- input vector dim
                         -> Int   -- output vector dim
                         -> Int   -- hidden vector dim
                         -> m RNN
randRNN i o h = do
  u <- randomWeights h h
  w <- randomWeights i h
  v <- randomWeights h o
  return $ MkRNN u (w :&~ O v)

{-
 A typical form of training data to take will be
 a list [(inp, outp)] with inp  :: Vector Double
                           outp :: Vector Double
-}

foldRNNTraining :: Double
                -> (Vector Double, RNN)    -- initial hidden state & RNN
                -> [(Vector Double,        -- training data
                     Vector Double)]
                -> IO (Vector Double, RNN) -- final hidden state & trained RNN
foldRNNTraining rate (h_0, rnn_0) pairs =
  foldl (\io p@(x, tgt) -> do
              (h, rnn) <- io
              putStrLn $ "Current prediction given: " ++ show x
                         ++ " is: " ++ show (runRNN rnn x h)
              let h_rnn' = trainRNN rnn rate h x tgt
              return h_rnn')
          (return (h_0, rnn_0))
          pairs

displayRNNTraining :: Double            -- training rate
                   -- -> Vector Double     -- initial activation vector h_0
                   -> [(Vector Double,  -- lists of training inputs
                        Vector Double)] -- paired with target outputs
                   -> IO ()
displayRNNTraining rate pairs = do
  -- get a random RNN
  randRNN <- evalRandIO (randRNN 1 1 5)
  -- we will always use an initial activation h_0 of zeroes
  let h_0 = vector $ fromIntegral <$> (take 5 $ repeat 0)
  foldRNNTraining rate (h_0, randRNN) pairs
  return ()

-- no n-gram window this time, just inputs & outputs one at a time each.
rnn_training_data :: [(Vector Double,  -- 1 * 1
                       Vector Double)] -- 1 * 1
rnn_training_data =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 1 sub_pattern, [sub_pattern !! 1])
            | j <- [0..(length fin_pattern - 2)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (vector $ fromIntegral <$> inp,
                          vector $ fromIntegral <$> out))
         list_pairs

-- Retry the simple RNN training using softmax.
-- But at the same time, the 1-vectors in the training data
-- [(inp, out)] pairs should become one-hot encoding 3-vectors.
-- Likewise, we will want 'randRNN 3 3 5' (5 the size of the hidden layer.)
softmax :: Vector Double -> Vector Double
softmax v =
  let vexp = exp v :: Vector Double
  in  vexp / (scalar $ sumElements vexp)

-- Note that the derivative of (vector) softmax is matrix-valued
softmax' :: Vector Double -> Matrix Double
softmax' v =
  let s = softmax v
  in  diag s - s `outer` s

runRNN_softmax :: RNN
               -> Vector Double -- input vector
               -> Vector Double -- previous-step activation
               -> Vector Double -- prediction
runRNN_softmax rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) x_t h_t_1 =
    let z_u   = (uB + uN #> h_t_1) +
                (wB + wN #> x_t)
        h_t   = logistic z_u
        z_h   = (vB + vN #> h_t)
        y_t   = softmax z_h
    in  y_t  -- ^ softmax usually, for tagging problems & similar
             --   (multiple way classification)

trainRNN_softmax  :: RNN
                  -> Double
                  -> Vector Double
                  -> Vector Double
                  -> Vector Double
                  -> (Vector Double, RNN)
trainRNN_softmax rnn@(MkRNN
  u@(W uB uN)
  n@(
    w@(W wB wN) :&~
    O v@(W vB vN))) rate h_t_1 x_t tgt_t =
  let z_u   = (uB + uN #> h_t_1) +
              (wB + wN #> x_t)
      h_t   = logistic z_u
      z_h   = (vB + vN #> h_t)
      y_t   = softmax z_h                   -- <-
      dEdz  = softmax' z_h #> (y_t - tgt_t)  -- <-
      vB'   = vB - scale rate dEdz
      vN'   = vN + ((- scale rate dEdz) `outer` h_t)
      dh_t  = tr vN #> dEdz
      dEdh  = logistic' h_t * dh_t

      uB'   = uB - scale rate dEdh
      uN'   = uN + ((- scale rate dEdh) `outer` h_t_1)
      wB'   = wB - scale rate dEdh
      wN'   = wN + ((- scale rate dEdh) `outer` x_t)
  in  ( h_t,
        MkRNN
          (W uB' uN')
          ((W wB' wN') :&~
            O (W vB' vN')))

foldRNNTraining_softmax :: Double
                        -> (Vector Double, RNN)    -- initial hidden state & RNN
                        -> [(Vector Double,        -- training data
                             Vector Double)]
                        -> IO (Vector Double, RNN) -- final hidden state & trained RNN
foldRNNTraining_softmax rate (h_0, rnn_0) pairs =
  foldl (\io p@(x, tgt) -> do
              (h, rnn) <- io
              putStrLn $ "Current prediction given: " ++ show x
                         ++ " is: " ++ show (runRNN_softmax rnn x h)
              let h_rnn' = trainRNN_softmax rnn rate h x tgt
              return h_rnn')
          (return (h_0, rnn_0))
          pairs

displayRNNTraining_softmax :: Double
                           -> [(Vector Double,  -- lists of training inputs
                                Vector Double)] -- paired with target outputs
                           -> IO ()
displayRNNTraining_softmax rate pairs = do
  -- get a random RNN
  randRNN <- evalRandIO (randRNN 3 3 7)
  -- we will always use an initial activation h_0 of zeroes
  let h_0 = vector $ fromIntegral <$> (take 7 $ repeat 0)
  foldRNNTraining_softmax rate (h_0, randRNN) pairs
  return ()

--
to_onehot :: [Int] -> Vector Double
to_onehot m = case m of
  [0] -> vector $ fromIntegral <$> [1, 0, 0]
  [1] -> vector $ fromIntegral <$> [0, 1, 0]
  [2] -> vector $ fromIntegral <$> [0, 0, 1]
  -- so that to_onehot won't be partial. Not using this branch
  _   -> vector $ fromIntegral <$> [0, 0, 0]

rnn_training_data_onehot :: [(Vector Double,  -- 3 * 1
                              Vector Double)] -- 3 * 1
rnn_training_data_onehot =
  let fin_pattern = take fin_pttrn_len inf_pattern
      list_pairs =
        [ (take 1 sub_pattern, [sub_pattern !! 1])
            | j <- [0..(length fin_pattern - 2)],
              let sub_pattern = drop j fin_pattern ]
  in map (\(inp, out) -> (to_onehot inp, to_onehot out))
         list_pairs

{-
 To see these examples in action:

 > displayTraining 0.25 training_data
 > displayRNNTraining 0.25 rnn_training_data
 > displayRNNTraining_softmax 0.25 rnn_training_data_onehot

-}

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

main :: IO ()
main = return ()


{-# LANGUAGE CPP #-}
#if __GLASGOW_HASKELL__ > 707
{-# LANGUAGE AllowAmbiguousTypes #-}
#endif
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
module T1 where
import GHC.Exts (Constraint)
import DimMat
import Numeric.Units.Dimensional.Prelude
import Numeric.Units.Dimensional

import Text.PrettyPrint.ANSI.Leijen hiding (dot)

import GHC.TypeLits

import qualified Numeric.LinearAlgebra as H
import qualified Prelude as P

{- 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 *~ kilo gram
coefficientOfFrictionForCart = 0.1 *~ (newton / meter / second)
lengthToPendulumCenterOfMass = 0.3 *~ meter
massMomentOfInertiaOfPendulum = 0.006 *~ (kilo gram * meter * meter)
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

{-

>>> inv aSmall `sub` aInvSmall
2><2     1                      m
s        -8.881784197001252e-16 -2.220446049250313e-16
m^-1 s^2 0.0                    0.0

>>> inv aSmall `multiply` aSmall
2><2   1                     m s^-1
1      1.0                   -6.732461810265988e-15
m^-1 s 5.014435047745458e-19 0.9999999999999999

-}
aSmall = [matD| a22, a23; a42, a43 |]

aInvSmall = scale (_1 / det aSmall)
    [matD| a43, negate a23; negate a42, a22 |]

-- id2 = Proxy :: Proxy (0 * a)

colMat :: D '(a, '[ra, '[]]) e -> D '(a, '[ra, '[]]) e
colMat x = x

{-
charEq lam1 lam2 x1 x2 =
        (det ( aSmall `sub` [matD| lam1, _0; _0, lam2 |]),
         colMat $ scale lam1 x1 `sub` dot aSmall x1,
         colMat $ scale lam2 x2 `sub` dot aSmall x2,
         hconcat x1 x2)
         --
         -- scale lam1 x1 `sub` dot aSmall x1, -- `asTypeOf` (undefined :: D '(s, ['[t],'[ ]]) Double),
         -- scale lam2 x2 `sub` dot aSmall x2) --  `asTypeOf` (undefined :: D '(s, ['[t],'[ ]]) Double),
         -- hconcat x1 x2 `asTypeOf` (undefined :: D '(t, ['[v],'[u]]) Double))
-- still has ambiguity. Do the SVD instead?
-}

{-
>>> cmap (Scale' ((2::Double) *~ second)) aSmall
2><2 1                    m s^-1
1    -0.36363636363636365 5.3454545454545475
m^-1 -0.9090909090909091  62.36363636363637

alternative to scale. Also more fancy examples (say swap metres and seconds
powers)...

-}
newtype Scale' f a = Scale' (Quantity f a)
instance (Num a, a ~ Double,
          x ~ Quantity d a,
          y ~ Quantity d' a,
          Mul f d d')
        => ApplyAB (Scale' f a) x y where
  applyAB (Scale' n) x = n * x

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

-- example state value
x = [matD| 1.0 *~ meter;
           0.2 *~ (meter / second);
           _0 :: Dimensionless Double;
           0.1 *~ (second^neg1) |]

dx = scale (_1 / (1 *~ second)) x

-- example control input
u = [matD| (0 :: Double) *~ newton |]

--xDot = (a `multiply` x) `add` (b `multiply` u)

y = [matD| 1 *~ meter; _0 :: Dimensionless Double  |]

-- :t isLTI (1 *~ second) x u y, given that x u and y have monomorphic
-- types properly infers constraints on the a,b,c,d arguments!
isLTI time x u y a b c d =
    (scale (_1 /time) (colMat x) `add` dot a x `add` dot b (colMat u),
     colMat $ y `add` dot c x `add` dot d u)

testIsLTI =
  (\ a b c d -> case isLTI (1 *~ second) x u y a (colMat b) c (colMat d) of
   _ -> do
    print $ vsep
        [text "A = " </> indent 0 (pretty (zeroes `asTypeOf` a)),
         text "B = " </> indent 0 (pretty (zeroes `asTypeOf` b)),
         text "C = " </> indent 0 (pretty (zeroes `asTypeOf` c)),
         text "D = " </> indent 0 (pretty (zeroes `asTypeOf` d))]
    ) undefined undefined undefined undefined

{- | data type encoding units required by
http://en.wikibooks.org/wiki/Control_Systems/State-Space_Equations#State-Space_Equations

To refresh:

> dxs/dt = A xs + B us
> ys = C xs + D us

the units of d/dt are 1/iv
-}
class LiSystem (iv :: *) (xs :: (*, [[*]])) (ys :: (*, [[*]])) (us :: (*, [[*]]))
            (a :: (*, [[*]])) (b :: (*, [[*]])) (c :: (*, [[*]])) (d :: (*, [[*]]))
instance LiSystemCxt dxs iv xs ys us a b c d
    => LiSystem iv xs ys us a b c d

type LiSystemCxt dxs iv xs ys us a b c d =
   (Dot a xs dxs,
    Dot b us dxs,
    Dot c xs ys,
    Dot d us ys,
    DotS iv dxs xs
    )

-- | identity function but constrains types
isLiTuple ::
    (LiSystem iv xs ys us a b c d,
    t ~ (D a e, D b e, D c e, D d e)) =>
    t -> t
isLiTuple x = x

isExample ::
     (LiSystem DTime
        '(DLength , '[[DVelocity, DPlaneAngle, DAngularVelocity]])
        '(DLength, '[ '[DPlaneAngle]])
        '(DForce, '[ '[] ] )
        a b c d, t ~ (D a e, D b e, D c e, D d e)) =>
    t -> t
isExample x = x

pendulum = case (a',b',c',d') of
             sys -> isExample sys -- `const` evaluate sys
         where
           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

{-
-- poles (a,_,_,_) = eigenvalues a

-}
-- causes no problems for AV
-- evaluatePendulum = evaluate pendulum

{-
evaluate ::
    ( -- require all matrices to be at least 1x1
      -- really ought to be part of the MultiplyCxt
     a ~ [_1 ': _2, '[] ': _3],
     b ~ [_4 ': _5, '[] ': _6],
     c ~ [_7 ': _8, '[] ': _9],
     d ~ [_a ': _b, '[] ': _c],
     H.Field e,
     LiSystemCxt dxs iv xs ys us a b c d) =>
    (DimMat a e, DimMat b e, DimMat c e, DimMat d e)
    -> DimMat [xs, '[ '[] ]] e
    -> DimMat [us, '[ '[] ]] e
    -> (DimMat [dxs, '[ '[] ]] e, DimMat [ys, '[ '[] ]] e)
    -}
evaluate (a,b,c,d) x u = case (a `dot` x) `add` (b `dot` u) of
    xDot -> case (c `dot` x) `add` (d `dot` u) of
     y -> (xDot, y)

{-

ctrb = let (a,b,_,_) = pendulum
        in [blockD| b, a `multiply` b, a `multiply` a `multiply` b, a `multiply` a `multiply` a `multiply` b |]
        -}


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

module Section_4_5.TwoLayer where

import Data.List

import Numeric.LinearAlgebra
import System.Random.MWC
import System.Random.MWC.Distributions (normal)

crossEntropyError :: Vector Double -> Vector Double -> Double
crossEntropyError y t =
  let delta = 1.0e-7
   in -1 * (sumElements $ t * cmap (\x -> log (x + delta)) y)

class Indexes c t where
  indexes :: c t -> [IndexOf c]

instance Container Vector t => Indexes Vector t where
  indexes v = [0 .. size v - 1]

instance (Num t, Container Vector t) => Indexes Matrix t where
  indexes m = let (x, y) = size m in (,) <$> [0..x-1] <*> [0..y-1]

numericalGradient :: (Container c Double, Additive (c Double), Indexes c Double) => (c Double -> Double) -> c Double -> c Double
numericalGradient f x =
  foldr add (assoc (size x) 0.0 []) $
    flip map (indexes x) $ \i ->
      let df1 = f (x `add` assoc (size x) 0.0 [(i,  h)])
          df2 = f (x `add` assoc (size x) 0.0 [(i, -h)])
          dfdx = (df1 - df2) / (2*h)
       in assoc (size x) 0.0 [(i, dfdx)]
    where h = 1.0e-4 -- 0.0001

type ActivationFunction = Vector Double -> Vector Double

-- | シグモイド関数
sigmoid :: ActivationFunction
sigmoid = cmap (\x -> 1 / (1 + exp (-x)))

-- | ソフトマックス関数
softmax :: ActivationFunction
softmax a =
  let c = maxElement a
      exp_a = cmap (\x -> exp (x - c)) a
      sum_exp_a = sumElements exp_a
   in cmap (/sum_exp_a) exp_a

type Layer = ( Matrix Double      -- 重み行列
             , Vector Double      -- バイアス
             , ActivationFunction -- 活性化関数
             )
type Network = [Layer]

-- | 学習前のニューラルネットワークを生成する
initNetwork :: [Int]                -- 各層のユニット数
            -> [ActivationFunction] -- 各層の活性化関数
            -> IO Network
initNetwork inputs as =
  withSystemRandom $ \gen ->
    sequence (zipWith (build gen) inouts as)
  where
    inouts :: [(Int, Int)]
    inouts = zip inputs (tail inputs)
    build :: GenIO -> (Int, Int) -> ActivationFunction -> IO Layer
    build gen (x, y) a = do
      w <- (x >< y) <$> (sequence $ replicate (x * y) (normal 0 0.01 gen))
      b <- vector <$> (sequence $ replicate y (normal 0 0.01 gen))
      pure (w, b, a)

-- | 認識（推論）を行う
predict :: Network -> Vector Double -> Vector Double
predict nw x = foldl (\x (w, b, h) -> h (x <# w + b)) x nw

-- | 損失関数の値を求める
loss :: [Vector Double] -> [Vector Double] -> Network -> Double
loss xs ts nw =
  let batchSize = genericLength xs
      ys = map (predict nw) xs
   in sum (zipWith crossEntropyError ys ts) / batchSize

-- | 認識精度を求める
accuracy :: [Vector Double] -> [Vector Double] -> Network -> Double
accuracy xs ts nw =
  let batchSize = genericLength xs
      ys = map (predict nw) xs
      yis = map maxIndex ys
      tis = map maxIndex ts
      correctCount = genericLength (filter id (zipWith (==) yis tis))
   in correctCount / batchSize

dismantle :: [a] -> [(a, a -> [a])]
dismantle xs = map (\i -> (xs!!i, (\x -> (take i xs) ++ [x] ++ (drop (i+1) xs)))) [0..length xs - 1] -- TODO more efficient implementation

updateNetwork :: ([Vector Double] -> [Vector Double] -> Network -> Double) -> [Vector Double] -> [Vector Double] -> Network -> Network
updateNetwork loss xs ts network = map (uncurry $ updateLayer xs ts) $ dismantle network
  where
    rate :: Double
    rate = 0.1
    updateLayer :: [Vector Double] -> [Vector Double] -> Layer -> (Layer -> Network) -> Layer
    updateLayer xs ts (w, b, a) mkNw =
      let dw = numericalGradient (\w -> loss xs ts (mkNw (w, b, a))) w
          db = numericalGradient (\b -> loss xs ts (mkNw (w, b, a))) b
       in (w - rate `scale` dw, b - rate `scale` db, a)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE TupleSections #-}

module Main where

import Data.Random
import Data.Random.Distribution.MultiNormal
import Numeric.Kalman

import Control.Monad (replicateM)
--import GHC.TypeLits

import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo

import Numeric.LinearAlgebra.Static hiding (mean)

--import Control.Monad (replicateM)
--import Data.Maybe (fromJust)

import qualified Data.List as L
import qualified Numeric.LinearAlgebra as LA

main :: IO ()
--main = putStrLn "Test suite not yet implemented"

--main = runUKFTest1
--main = runNoisySine1
--main = runGSTest1
main = runGSTest2

---------------------------
--- Gaussian sum filter ---
---------------------------
runGSTest2 :: IO ()
runGSTest2 = do
  putStrLn "Running Gaussian sum test 2"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> R 3
      sysModel t cur =
        let [b,m,y] = LA.toList (extract cur)
        in vector [b, m, m*t + (2*b^2+b-1)]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runUKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) like1) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runUKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) like2) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs2-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs2-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs2-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs2-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

runGSTest1 :: IO ()
runGSTest1 = do
  putStrLn "Running Gaussian sum test 1"
  let numSamples = 1000
      dt = 0.01
      tVals = take numSamples [0, dt ..]
      truey = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truem = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      trueb = [(head tVals, 0), (last tVals, 0)] :: [(Double, Double)]
      truew = [(head tVals, 1), (last tVals, 1)] :: [(Double, Double)]

  let intErr = 0.04
      slpErr = 0.01
      yErr   = 0.03

  let initEst = [(1/2, MultiNormal (vector [2,-0.5,0.1])  (sym . diag . vector $ [intErr, slpErr, yErr])),
                 (1/2, MultiNormal (vector [-2,0.5,0.1]) (sym . diag . vector $ [intErr, slpErr, yErr]))]
                :: [(Double, MultiNormal (R 3))]

      measMat = row . vector $ [0,0,1]
      measModel = const (const measMat)
      measErr = diag . vector $ [0.0001] :: Sq 1

      sysModel :: Double -> R 3 -> Sq 3
      sysModel t _cur = matrix [1,0,0, 0,1,0, 1,t,0]
      sysErr = sym . diag . vector $ [0,0,0]

  measurements <- replicateM numSamples (sample $ MultiNormal (vector [0]) (sym measErr))

  let samples = [(t, head . LA.toList . extract $ s) | t <- tVals | s <- measurements]
      gsfilter = scanl
                 (\[(_w1, est1), (_w2, est2)] (t, meas) ->
                    let pred1 = runKFPrediction sysModel (const sysErr) t est1
                        pmu1  = mu pred1
                        pcov1 = unSym . cov $ pred1
                        like1 = sym $ measMat <> pcov1 <> tr measMat + measErr
                        neww1 = pdf (MultiNormal (measModel t pmu1 #> pmu1) (like1)) meas
                        updt1 = runKFUpdate measModel (const (sym measErr)) t pred1 meas

                        pred2 = runKFPrediction sysModel (const sysErr) t est2
                        pmu2  = mu pred2
                        pcov2 = unSym . cov $ pred2
                        like2 = sym $ measMat <> pcov2 <> tr measMat + measErr
                        neww2 = pdf (MultiNormal (measModel t pmu2 #> pmu2) (like2)) meas
                        updt2 = runKFUpdate measModel (const (sym measErr)) t pred2 meas

                        scale = neww1 + neww2

                    in [(neww1 / scale, updt1), (neww2 / scale, updt2)]
                 ) initEst $ zip tVals measurements

  let gsy1 = [(t, last . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsy2 = [(t, last . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsy  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = last . LA.toList . extract . mu $ est1
                            p2 = last . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsm1 = [(t, (!!1) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsm2 = [(t, (!!1) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsm  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!1) . LA.toList . extract . mu $ est1
                            p2 = (!!1) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsb1 = [(t, (!!0) . LA.toList . extract . mu . snd . head $ y) | t <- tVals | y <- gsfilter]
      gsb2 = [(t, (!!0) . LA.toList . extract . mu . snd . last $ y) | t <- tVals | y <- gsfilter]
      gsb  = zipWith (\t [(w1, est1), (w2, est2)] ->
                        let p1 = (!!0) . LA.toList . extract . mu $ est1
                            p2 = (!!0) . LA.toList . extract . mu $ est2
                        in (t, w1 * p1 + w2 * p2))
             tVals gsfilter

  let gsw1 = zipWith (\t [(w1,_),_] -> (t,w1)) tVals gsfilter
      gsw2 = zipWith (\t [_,(w2,_)] -> (t,w2)) tVals gsfilter
      gsw  = zipWith (\t [(w1, _), (w2, _)] -> (t,w1 + w2))
             tVals gsfilter

  let [(w1,cur1), (w2, cur2)] = last gsfilter
      final1 = (w1, LA.toList . extract . mu $ cur1)
      final2 = (w2, LA.toList . extract . mu $ cur2)

  print final1
  print final2
  print $ w1 + w2

  toFile def "gs1-y.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truey])
    plot (points "Measurements" samples)
    plot (points "G1 estimates" gsy1)
    plot (points "G2 estimates" gsy2)
    plot (points "GS estimates" gsy)

  toFile def "gs1-m.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truem])
    plot (points "G1 estimates" gsm1)
    plot (points "G2 estimates" gsm2)
    plot (points "GS estimates" gsm)

  toFile def "gs1-b.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [trueb])
    plot (points "G1 estimates" gsb1)
    plot (points "G2 estimates" gsb2)
    plot (points "GS estimates" gsb)

  toFile def "gs1-w.png" $ do
    layout_title .= "Gaussian sums"
    setColors [opaque black --, opaque blue
              ,opaque red, opaque green, opaque orange
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truew])
    plot (points "G1 estimates" gsw1)
    plot (points "G2 estimates" gsw2)
    plot (points "GS estimates" gsw)

---------------------------------------
--- Unscented Kalman filter testing ---
---------------------------------------
runUKFTest1 :: IO ()
runUKFTest1 = do
  -- polar coordinate transform test
  let numSamples = 500
      dt = 0.01
      tVals = take (numSamples + 1) [0,dt..] :: [Double] -- +1 because t = 0 is our initial estimate

      initSys = MultiNormal (vector [0, 1]) (sym . diag . vector $ [4e-4, 2.7e-5] :: Sym 2)
      initEst = MultiNormal (vector [1, pi/2]) (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2)

  let coordTrans :: R 2 -> R 2
      coordTrans polar =
        let [r,t] = LA.toList $ extract polar
        in vector [r * cos t, r * sin t]

  measurements <- map coordTrans <$> replicateM numSamples (sample initEst)

  let samples = map (\meas ->
                       let [x,y] = LA.toList $ extract meas
                       in (x,y)
                    ) measurements

      ukf   = scanl (runUKF (const coordTrans)
                            (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                            (const id)
                            (const (sym . diag . vector $ [0,0])) 1)
                    initSys measurements

      ukpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                  ) ukf

  let linMeas :: R 2 -> L 2 2
      linMeas est =
        let [x, _] = LA.toList $ extract est
        in matrix [1,0, 0,x]

  let ekf = scanl (runEKF (const coordTrans)
                          (const linMeas)
                          (const (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2))
                          (const id)
                          (const (const eye))
                          (const (sym . diag . vector $ [0,0])) 1)
                  initSys measurements

      ekpts = map (\kf ->
                       let [x,y] = LA.toList . extract $ mu kf
                       in (x,y)
                    ) ekf

      -- predTest = runUKFPrediction id (sym . diag . vector $ [0,0]) initSys
      -- updtTest = runUKFUpdate coordTrans (sym . diag . vector $ [0.02^2, (15*pi/180)^2] :: Sym 2) initSys (head measurements)

  -- print $ extract . mu $ predTest
  -- print $ extract . unSym . cov $ predTest
  -- putStrLn "\n"
  -- print $ extract . mu $ updtTest
  -- print $ extract . unSym . cov $ updtTest
  -- putStrLn "\n"
  -- print $ last ukpts
  -- print $ last ekpts

  toFile def "noisy-coords.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque blue
              ,opaque red, opaque green
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (points "True signal"  [(0,1)])
    plot (points "Measurements" samples)
    plot (points "UKF estimates" ukpts)
    plot (points "EKF estimates" ekpts)

  toFile def "noisy-params.png" $ do
    layout_title .= "Noisy coord transform"
    setColors [opaque black, opaque black
              ,opaque green, opaque green
              ,opaque red, opaque red
              --,opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line "True X signal"  [[(head tVals, 0), (last tVals, 0)]])
    plot (line "True Y signal"  [[(head tVals, 1), (last tVals, 1)]])
    plot (points "UKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ukpts)
    plot (points "UKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ukpts)
    plot (points "EKF X estimates" $ zipWith (\t (x,_) -> (t, x)) tVals ekpts)
    plot (points "EKF Y estimates" $ zipWith (\t (_,y) -> (t, y)) tVals ekpts)

---------------------
--- Noisy sine # 1---
---------------------
-- In this example, we estimate the parameters (amplitude, period, vertical shift, horizontal shift)
-- defining sinusoidal evolution.
-- Figure 3 in my notes

runNoisySine1 :: IO ()
runNoisySine1 = do

  let (a, p, v) = (0.7, 2.2*pi, -0.5) -- (hidden) internal system parameters
      (a0, p0, v0) = (0.8, 2*pi, -0.6)    -- initial estimates
      --(aSig, pSig, vSig) = (0.2, 0.2*pi, 0.001) -- initial uncertainties
      measCov = 0.01

      --- Figure 1 in my notes
      --numSamples = 499
      --dt = 0.002
      --- Figure 2 in my notes
      numSamples = 199
      dt = 0.02
      tVals = take (numSamples + 1) [0,dt..] -- +1 because t = 0 is our initial estimate

      --initSys = vector [a, p, v] -- initial (true) system: (parameter1, parameter2, state)
      --initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0]) (sym . diag . vector $ [0.04, 0.4, 0.04, 0.001, 0.001] :: Sym 5)
      initEst = MultiNormal (vector [a0, p0, v0, a0*p0, v0])
                (sym . matrix $ [0.04,  0,    0,    1e-5, -2e-3,
                                 0,     0.4,  0,    1e-2, 12e-3,
                                 0,     0,    0.04, 1e-4, 1e-3,
                                 1e-5,  1e-2, 1e-4, 1e-3, 1e-4,
                                 -2e-3, 12e-3, 1e-3, 1e-4, 1e-3] :: Sym 5)

      evolMat :: R 5 -> Sq 5
      evolMat cur =
        let pCur  = (extract cur) LA.! 1
            p2dt  = pCur * pCur * dt
        in matrix [1,0,0,0,0, 0,1,0,0,0, 0,0,1,0,0, 0,0,p2dt,1,-p2dt, 0,0,0,dt,1]

      evolSys :: Double -> R 5 -> R 5
      evolSys t cur =
        let [a,p,v,ap,b] = LA.toList $ extract cur
        in vector [a, p, v, a * p * cos (p * t), a * sin (p * t) + v]

      sysCovariance = sym . diag . vector $ [0,0,0,5e-4,7e-5] :: Sym 5

      measMat = row (vector [0,0,0,0,1]) :: L 1 5
      measTrans = (measMat #>)
      measCovariance = sym . diag . vector $ [measCov] :: Sym 1
      measurementNoise = MultiNormal (vector [0]) measCovariance

      trueSystem = [vector [a, p, v, a * p * cos(p*t), a * sin(p*t) + v] | t <- tVals] :: [R 5]

  --trueSystem   <- simulateSystem numSamples sysModel initSys
  measurements <- mapM (\sys -> do
                            noise <- sample measurementNoise
                            return $ measTrans sys + noise
                       ) $ tail trueSystem

  --mapM_ print $ map (LA.toList . extract) trueSystem
  -- print $ map ((LA.! 0) . extract) measurements

  let truesine = zip tVals $ map ((LA.! 4) . extract) trueSystem
      samples  = zip (tail tVals) $ map ((LA.! 0) . extract) measurements

      kfiltered = scanl (runKF (const (const measMat)) (const measCovariance) (const evolMat) (const sysCovariance) 1) initEst measurements
      ekf = scanl (\sys (t,meas) ->
                     runEKF (\_t sys -> measMat #> sys)
                     (\_t _sys -> measMat)
                     (const measCovariance)
                     evolSys
                     (const evolMat)
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      ukf = scanl (\sys (t,meas) ->
                     runUKF (\_t sys -> measMat #> sys)
                     (const measCovariance)
                     evolSys
                     (const sysCovariance)
                     t sys meas
                  ) initEst $ zip (tail tVals) measurements

      kfstates  = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals kfiltered
      ekfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ekf
      ukfstates = zipWith (\t kf -> (t, (extract . mu $ kf) LA.! 4)) tVals ukf

      [kfa,kfp,kfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals kfiltered
      [ekfa,ekfp,ekfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ekf
      [ukfa,ukfp,ukfv] = L.transpose $ zipWith (\t kf -> map (t,) . take 3 . LA.toList . extract . mu $ kf) tVals ukf

      [kfaCovs, kfpCovs, kfvCovs, kfsCovs] = L.transpose $
                                             map (\kf -> let kfcov = extract . unSym $ cov kf
                                                         in [kfcov LA.! 0 LA.! 0, kfcov LA.! 1 LA.! 1,
                                                             kfcov LA.! 2 LA.! 2, kfcov LA.! 4 LA.! 4])
                                             kfiltered

      kfsUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfstates kfsCovs
      kfsLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfstates kfsCovs

      kfsUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfstates kfsCovs
      kfsLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfstates kfsCovs

      kfaUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfa kfaCovs
      kfaLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfa kfaCovs

      kfaUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfa kfaCovs
      kfaLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfa kfaCovs

      kfpUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfp kfpCovs
      kfpLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfp kfpCovs

      kfpUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfp kfpCovs
      kfpLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfp kfpCovs

      kfvUpper1 = zipWith (\(t, kfs) kfsCov -> (t, kfs + sqrt kfsCov)) kfv kfvCovs
      kfvLower1 = zipWith (\(t, kfs) kfsCov -> (t, kfs - sqrt kfsCov)) kfv kfvCovs

      kfvUpper2 = zipWith (\(t, kfs) kfsCov -> (t, kfs + 2 * sqrt kfsCov)) kfv kfvCovs
      kfvLower2 = zipWith (\(t, kfs) kfsCov -> (t, kfs - 2 * sqrt kfsCov)) kfv kfvCovs

      ksmoothed = scanr1 (\pres fut -> runKS (const evolMat) (const sysCovariance) 1 fut pres) kfiltered
      ksstates  = zipWith (\t ks -> (t, (extract . mu $ ks) LA.! 4)) tVals ksmoothed

  --mapM_ print $ map (LA.toList . extract . mu) kfiltered
  --mapM_ print kfa

  toFile def "noisy-sine-points.png" $ do
    layout_title .= "Noisy sine (state estimate)"
    setColors [opaque black, opaque blue
              ,opaque red, opaque purple, opaque hotpink, opaque green
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey]
    plot (line   "True signal"  [truesine])
    plot (points "Measurements" samples)
    plot (line   "KF estimated signal" [kfstates])
    plot (line   "EKF estimated signal" [ekfstates])
    plot (line   "UKF estimated signal" [ukfstates])
    plot (line   "KS smoothed signal"  [ksstates])
    plot (line   "Filtered 1 standard deviation" [kfsUpper1])
    plot (line   "Filtered 1 standard deviation" [kfsLower1])
    plot (line   "Filtered 2 standard deviations" [kfsUpper2])
    plot (line   "Filtered 2 standard deviations" [kfsLower2])

  let truea = [(head tVals, a), (last tVals, a)]
      truep = [(head tVals, p), (last tVals, p)]
      truev = [(head tVals, v), (last tVals, v)]

  toFile def "noisy-sine-amplitude.png" $ do
    layout_title .= "Noisy sine (amplitude estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truea])
    plot (line   "KF estimated signal" [kfa])
    plot (line   "EKF estimated signal" [ekfa])
    plot (line   "UKF estimated signal" [ukfa])
    plot (line   "Filtered 1 standard deviation" [kfaUpper1])
    plot (line   "Filtered 1 standard deviation" [kfaLower1])
    plot (line   "Filtered 2 standard deviations" [kfaUpper2])
    plot (line   "Filtered 2 standard deviations" [kfaLower2])

  toFile def "noisy-sine-period.png" $ do
    layout_title .= "Noisy sine (period estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truep])
    plot (line   "KF estimated signal" [kfp])
    plot (line   "EKF estimated signal" [ekfp])
    plot (line   "UKF estimated signal" [ukfp])
    plot (line   "Filtered 1 standard deviation" [kfpUpper1])
    plot (line   "Filtered 1 standard deviation" [kfpLower1])
    plot (line   "Filtered 2 standard deviations" [kfpUpper2])
    plot (line   "Filtered 2 standard deviations" [kfpLower2])

  toFile def "noisy-sine-vertical.png" $ do
    layout_title .= "Noisy sine (vertical estimate)"
    setColors [opaque black, opaque red, opaque purple, opaque hotpink
              ,opaque lightgrey, opaque lightgrey
              ,opaque lightslategrey, opaque lightslategrey
              ]
    plot (line   "True signal"  [truev])
    plot (line   "KF estimated signal" [kfv])
    plot (line   "EKF estimated signal" [ekfv])
    plot (line   "UKF estimated signal" [ukfv])
    plot (line   "Filtered 1 standard deviation" [kfvUpper1])
    plot (line   "Filtered 1 standard deviation" [kfvLower1])
    plot (line   "Filtered 2 standard deviations" [kfvUpper2])
    plot (line   "Filtered 2 standard deviations" [kfvLower2])

  putStrLn $ "Kalman processing complete."


--{-# LANGUAGE UndecidableInstances,
--             FlexibleInstances,
--             FlexibleContexts,
--             TypeFamilies,
--             ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal.Noise
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  uses Concurrency
--
-- Noise generation functions
--
-----------------------------------------------------------------------------

--             IncoherentInstances,

module Numeric.Signal.Noise (
                             pinkNoise
                            , spatialNoise
                            , powerNoise
                            ) where

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

--import qualified Numeric.Signal as S

--import Complex

--import qualified Data.Array.IArray as I
--import Data.Ix

--import Data.Word

--import System.IO.Unsafe(unsafePerformIO)

--import qualified Data.List as L

--import Data.Binary

--import Foreign.Storable

import Numeric.LinearAlgebra

import qualified Data.Vector.Generic as GV

import qualified Numeric.GSL.Fourier as F

--import qualified Numeric.GSL.Histogram as H
--import qualified Numeric.GSL.Histogram2D as H2

--import qualified Numeric.Statistics.Information as SI

import Prelude hiding(filter)

--import Control.Monad(replicateM)

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

-- | The method is briefly descirbed in Lennon, J.L. "Red-shifts and red
-- herrings in geographical ecology", Ecography, Vol. 23, p101-113 (2000)
--
-- Matlab version Written by Jon Yearsley  1 May 2004
-- j.yearsley@macaulay.ac.uk
--
-- Creates 1/f scale invariant spatial noise
spatialNoise :: Double    -- ^  β: spectral distribution
                         --    0: White noise
                         --   -1: Pink noise
                         --   -2: Brownian noise
             -> Int -> Int -- ^ matrix dimensions
             -> Int       -- ^ random seed
             -> Matrix Double
spatialNoise b r' c' s = let c = fromIntegral c'
                             r = fromIntegral r'
                             pre_x = linspace c' (0::Double,c-1)
                             post_x = linspace c' (c,1)
                             freq_x = GV.map (/c) $ vjoin [pre_x,post_x]
                             u = fromRows (replicate (2*r') freq_x)
                             pre_y = linspace r' (0::Double,r-1)
                             post_y = linspace r' (r,1)
                             freq_y = GV.map (/c) $ vjoin [pre_y,post_y]
                             v = fromColumns (replicate (2*c') freq_y)
                             s_f = cmap (**(b/2)) ((u**2) + (v**2))
                             s_f' = cmap (\x -> if isInfinite x then 0 else x) s_f
                             phi = reshape (2*c') (randomVector s Uniform (4*r'*c'))
                         in subMatrix (1,1) (r',c') $ fst $ fromComplex $ fromRows $ map F.ifft $ toRows $ ((complex $ s_f'**0.5) * (toComplex (cos(2*pi*phi),sin(2*pi*phi))))

-- | 1/f scale invariant noise
pinkNoise ::
    Double    -- ^  β: spectral distribution
              --    0: White noise
              --   -1: Pink noise
              --   -2: Brownian (red) noise
    -> Int       -- ^ samples
    -> Int       -- ^ random seed
    -> Vector Double
pinkNoise b s r = let pre = linspace s (0::Double,fromIntegral (s-1))
                      post = linspace s (fromIntegral s,1)
                      freq = vjoin [pre/(fromIntegral s),post/(fromIntegral s)]
                      s_f = GV.map (**(b/2)) (freq**2)
                      s_f' = GV.map (\x -> if isInfinite x then 0 else x) s_f
                      phi = randomVector r Uniform (2*s)
                  in subVector 0 s $ fst $ fromComplex $ F.ifft ((complex $ s_f'**0.5) * (toComplex (cos(2*pi*phi),sin(2*pi*phi))))

-- | generate noise from a power spectrum
powerNoise :: Vector Double   -- ^ the power spectrum
           -> Int             -- ^ random seed
           -> Vector Double
powerNoise psd r = let ln = GV.length psd
                       freq = vjoin [fromList [0],psd, (fromList . reverse . tail . toList) psd]
                       phi = randomVector r Uniform (2*ln)
                   in (fromIntegral ln) * (subVector 0 (ln-1) $ fst $ fromComplex $ F.ifft ((complex $ freq) * (toComplex (cos(2*pi*phi),sin(2*pi*phi)))))


module Math.IRT.Fisher
    ( FisherInfo(..)
    , fisherInfoObserved
    , fisherInfoExpected
    , l'
    , l''
    ) where

import Statistics.Distribution
import Math.IRT.Internal.Distribution

data FisherInfo = FisherInfo { items :: [Double]
                             , test  :: !Double
                             , sem   :: !Double
                             } deriving (Show)

-- |The observed Fisher Information
--  Applies a min/max prior to ensure a real-valued SEM
fisherInfoObserved :: (ContDistr d, DensityDeriv d) => Double -> [Bool] -> [d] -> FisherInfo
fisherInfoObserved theta resps params =
    let is = zipWith go resps params
        t  = sum is
        s  = sqrt (1 / t)
    in FisherInfo is t s
    where go u x = negate $ l'' u x theta

-- |The expected Fisher Information
fisherInfoExpected :: ContDistr d => Double -> [d] -> FisherInfo
fisherInfoExpected theta params =
    let is = map go params
        t  = sum is
        s  = sqrt (1 / t)
    in FisherInfo is t s
    where go x = let (cdf, ccdf, pdf) = pqDers x theta
                 in (square pdf) / (cdf * ccdf)

-- |The first derivative of `l`, whatever `l` is... Once again, the catIrt documentation disappoints.
-- According to catIrt, "u is the response, and x are the parameters."
-- We only implement the MLE route
l' :: ContDistr d => Bool -> d -> Double -> Double
l' u x theta =
    let (cdf, ccdf, pdf) = pqDers x theta
        denom = cdf * ccdf
        f k   = k * pdf / denom
    in case u of
         True  -> f ccdf
         False -> f cdf

-- |The second derivative of `l` (same one as in `l'`)
l'' :: (ContDistr d, DensityDeriv d) => Bool -> d -> Double -> Double
l'' u x theta =
    let (cdf, ccdf, pdf) = pqDers x theta
        pdf'             = densityDeriv x theta
    in case u of
         True  -> (((-1) / square cdf) * (square pdf))
                  + (1 / cdf * pdf')
         False -> negate $ ((1 / square ccdf) * square pdf)
                           + (1 / ccdf * pdf')

pqDers :: ContDistr d => d -> Double -> (Double, Double, Double)
pqDers x theta = let pComp = cumulative x theta
                     p'    = density x theta
                 in ( pComp
                    , (1 - pComp)
                    , p')

square :: Double -> Double
square = (^^ (2 :: Int))


module Mnist (Image(..), normalisedData, readImages, readLabels, toMatrix)
  where

import qualified Data.ByteString.Lazy as BL
import Data.Binary.Get
import Data.Word
import qualified Data.List.Split as S
import Numeric.LinearAlgebra

data Image = Image {
      iRows :: Int
    , iColumns :: Int
    , iPixels :: [Word8]
    } deriving (Eq, Show)

toMatrix :: Image -> Matrix Double
toMatrix image = (r><c) p :: Matrix Double
  where r = iRows image
        c = iColumns image
        p = map fromIntegral (iPixels image)

{-
toColumnVector :: Image -> Matrix Double
toColumnVector i = (r><1) q :: Matrix Double
  where r = Mnist.rows i * Mnist.columns i
        p = map fromIntegral (pixels i)
        q = map normalise p
-}

normalisedData :: Image -> [Double]
normalisedData image = map normalisePixel (iPixels image)

--normalisedData :: Image -> [Double]
--normalisedData i = map (/m) x
--    where x = map normalisePixel (pixels i)
--          m = sqrt( sum (zipWith (*) x x))

normalisePixel :: Word8 -> Double
normalisePixel p = (fromIntegral p) / 255.0

-- MNIST label file format
--
-- [offset] [type]          [value]          [description]
-- 0000     32 bit integer  0x00000801(2049) magic number (MSB first)
-- 0004     32 bit integer  10000            number of items
-- 0008     unsigned byte   ??               label
-- 0009     unsigned byte   ??               label
-- ........
-- xxxx     unsigned byte   ??               label
--
-- The labels values are 0 to 9.

deserialiseLabels :: Get (Word32, Word32, [Word8])
deserialiseLabels = do
  magicNumber <- getWord32be
  count <- getWord32be
  labelData <- getRemainingLazyByteString
  let labels = BL.unpack labelData
  return (magicNumber, count, labels)

readLabels :: FilePath -> IO [Int]
readLabels filename = do
  content <- BL.readFile filename
  let (_, _, labels) = runGet deserialiseLabels content
  return (map fromIntegral labels)

-- MNIST Image file format
--
-- [offset] [type]          [value]          [description]
-- 0000     32 bit integer  0x00000803(2051) magic number
-- 0004     32 bit integer  ??               number of images
-- 0008     32 bit integer  28               number of rows
-- 0012     32 bit integer  28               number of columns
-- 0016     unsigned byte   ??               pixel
-- 0017     unsigned byte   ??               pixel
-- ........
-- xxxx     unsigned byte   ??               pixel
--
-- Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255
-- means foreground (black).

deserialiseHeader :: Get (Word32, Word32, Word32, Word32, [[Word8]])
deserialiseHeader = do
  magicNumber <- getWord32be
  imageCount <- getWord32be
  r <- getWord32be
  c <- getWord32be
  packedData <- getRemainingLazyByteString
  let len = fromIntegral (r * c)
  let unpackedData = S.chunksOf len (BL.unpack packedData)
  return (magicNumber, imageCount, r, c, unpackedData)

readImages :: FilePath -> IO [Image]
readImages filename = do
  content <- BL.readFile filename
  let (_, _, r, c, unpackedData) = runGet deserialiseHeader content
  return (map (Image (fromIntegral r) (fromIntegral c)) unpackedData)


module Isumi.Math.FFT.Image
  ( fftWithImage
  ) where

import           Codec.Picture         (Image, Pixel8)
import           Data.Complex
import           Isumi.Math.FFT.Bruun  (fft2dBruun)
import           Isumi.Math.MatrixUtil (imageToMatrix)
import           Numeric.LinearAlgebra (Matrix)

fftWithImage :: Image Pixel8 -> Maybe (Matrix (Complex Double))
fftWithImage = fft2dBruun . imageToMatrix


module Scenes.Resch (getScene, vectorToFloats) where

import Numeric.LinearAlgebra (vector)

import qualified Group as G
import Group (pattern Group, optimizeGroup)
import Helpers (vectorToFloats)
import Layer (mkLayer)
import Line (pattern Line)
import Point (pattern Point, cartesianProduct, fromFloat)
import Rectangle (mkRectangle)
import Scene (Scene, mkSceneWithStyle)
import Shape (mirror, rotate)
import Style

getScene :: IO Scene
getScene = do
    let width      = 3 * (m + 0.5) * l
    let height     = sqrt 3 * (n + 1) * l
    let layerStyle = StyleAttrs { strokeColor=Just "#03161d"
                                , strokeWidth=Just 0.01
                                , fillColor=Nothing }
    pure $ mkSceneWithStyle width height layerStyle [triangles, frame]
  where
    l = 0.5  -- triangle circumradius
    n = 6    -- twice number in height
    m = 6    -- twice number in width

    lp = fromFloat l
    p0 = Point (-1) 0 * lp
    p1 = Point 0 (sqrt 3  / 3) * lp
    p2 = Point ((-1) / 2) (-sqrt 3 / 6) * lp

    -- Create a triangle from our three points.
    baseTriangle = Group [ Line p0 p1, Line p1 p2, Line p0 p2]
    -- Create a group from the base triangle rotated three times.
    rotationalVerticies = map (\r -> 2 * pi * r / 3 ) [0..2]
    t0 = mconcat $ map (\t -> rotate (Point 0 0) t baseTriangle) rotationalVerticies
    -- Mirror this group twice.
    t1 = t0 <> mirror p0 (p1 - p0) t0
    t2 = t1 <> mirror (Point (0.5 * l) 0) (Point 0 1) t1

    -- Clone the above group across a series of positions on our convas.
    xs = vectorToFloats $ (vector [0..5] - 6 / 2 + 0.25) * 3 * 0.5
    ys = vectorToFloats $ (vector [0..5] - 6 / 2 + 0.25) * sqrt 3 * 0.5
    t3 = optimizeGroup (G.translateGroupOverPoints (cartesianProduct xs ys) t2) 0.01
    triangles = G.toLayer "triangles" t3

  -- Create a framing rectangle around the group.
    topLeft = Point ((-1.5 * m - 0.25) * l) (0.5 * sqrt 3 * (n + 0.5) * l)
    bottomRight = Point ((1.5 * m - 0.25) * l) (-0.5 * sqrt 3 * (n + 0.5) * l)
    frame = mkLayer "frame" [mkRectangle topLeft bottomRight]


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.BLAS.Level3
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Matrix-Matrix operations.
--

module Foreign.BLAS.Level3 (
    BLAS3(..),
    ) where

import Data.Complex
import Foreign( Ptr, Storable, with )

import Foreign.BLAS.Types
import Foreign.BLAS.Level2
import Foreign.BLAS.Double
import Foreign.BLAS.Zomplex

-- | Types with matrix-matrix operations.
class (BLAS2 a) => BLAS3 a where
    gemm  :: Trans -> Trans -> Int -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    symm  :: Side -> Uplo -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    hemm  :: Side -> Uplo -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    trmm  :: Side -> Uplo -> Trans -> Diag -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> IO ()
    trsm  :: Side -> Uplo -> Trans -> Diag -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> IO ()
    syrk  :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    syr2k :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    herk  :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()
    her2k :: Uplo -> Trans -> Int -> Int -> a -> Ptr a -> Int -> Ptr a -> Int -> a -> Ptr a -> Int -> IO ()

withEnum :: (Enum a, Storable a) => Int -> (Ptr a -> IO b) -> IO b
withEnum = with . toEnum
{-# INLINE withEnum #-}

instance BLAS3 Double where
    gemm transa transb m n k alpha pa lda pb ldb beta pc ldc =
        withTrans transa $ \ptransa ->
        withTrans transb $ \ptransb ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dgemm ptransa ptransb pm pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE gemm #-}

    symm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsymm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE symm #-}

    hemm = symm
    {-# INLINE hemm #-}

    trmm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            dtrmm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trmm #-}

    trsm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            dtrsm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trsm #-}

    syrk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsyrk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE syrk #-}

    syr2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            dsyr2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE syr2k #-}

    herk = syrk
    {-# INLINE herk #-}

    her2k = syr2k
    {-# INLINE her2k #-}

instance BLAS3 (Complex Double) where
    gemm transa transb m n k alpha pa lda pb ldb beta pc ldc =
        withTrans transa $ \ptransa ->
        withTrans transb $ \ptransb ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zgemm ptransa ptransb pm pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE gemm #-}

    symm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsymm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE symm #-}

    hemm side uplo m n alpha pa lda pb ldb beta pc ldc =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zhemm pside puplo pm pn palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE hemm #-}

    trmm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            ztrmm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trmm #-}

    trsm side uplo transa diag m n alpha pa lda pb ldb =
        withSide side $ \pside ->
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withDiag diag $ \pdiag ->
        withEnum m $ \pm ->
        withEnum n $ \pn ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
            ztrsm pside puplo ptransa pdiag pm pn palpha pa plda pb pldb
    {-# INLINE trsm #-}

    syrk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsyrk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE syrk #-}

    syr2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zsyr2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE syr2k #-}

    herk uplo transa n k alpha pa lda beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zherk puplo ptransa pn pk palpha pa plda pbeta pc pldc
    {-# INLINE herk #-}

    her2k uplo transa n k alpha pa lda pb ldb beta pc ldc =
        withUplo uplo $ \puplo ->
        withTrans transa $ \ptransa ->
        withEnum n $ \pn ->
        withEnum k $ \pk ->
        with alpha $ \palpha ->
        withEnum lda $ \plda ->
        withEnum ldb $ \pldb ->
        with beta $ \pbeta ->
        withEnum ldc $ \pldc ->
            zher2k puplo ptransa pn pk palpha pa plda pb pldb pbeta pc pldc
    {-# INLINE her2k #-}