{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
module Test.Utils where
import Numeric.LinearAlgebra
import Control.Exception
import Control.Monad
import qualified Data.NeuralNetwork.Backend.BLASHS.Utils as U
import qualified Numeric.LinearAlgebra as L
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector as BV
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import System.IO.Unsafe

asHM (U.DenseMatrix r c v) = L.reshape c $ unsafePerformIO $ V.freeze v
asDM m = let (r,c) = size m in U.DenseMatrix r c (unsafePerformIO $ V.thaw $ L.flatten m)

good_corr2 :: Int -> L.Matrix Float -> L.Matrix Float -> L.Matrix Float

good_corr2 p k m | w > s     = good_corr2 p m k
                 | otherwise = corr2 k padded
  where
    (w,h) = L.size k
    (s,t) = L.size m
    padded = fromBlocks [[z,0,0]
                        ,[0,m,0]
                        ,[0,0,z]]
    z = konst 0 (p, p)

test_corr2 :: Int -> L.Matrix Float -> L.Matrix Float -> IO (L.Matrix Float)
test_corr2 p k m | w > s     = test_corr2 p m k
                 | otherwise = do x@(U.DenseMatrixArray _ _ _ vx) <- U.newDenseMatrixArray 1 r c
                                  k' <- U.DenseMatrix w h <$> V.thaw (flatten k)
                                  m' <- U.DenseMatrix s t <$> V.thaw (flatten m)
                                  U.corr2 p (BV.singleton k') m' (x U.<<=)
                                  reshape c <$> V.freeze vx
  where
    (w,h) = L.size k
    (s,t) = L.size m
    (r,c) = (s-w+2*p+1, t-h+2*p+1)

test_corr2_arr :: Int -> [L.Matrix Float] -> L.Matrix Float -> IO [L.Matrix Float]
test_corr2_arr p ks m = do x@(U.DenseMatrixArray _ _ _ vx) <- U.newDenseMatrixArray n r c
                           print ("test", n, r, c, MV.length vx)
                           ks' <- mapM (\k -> U.DenseMatrix w h <$> V.thaw (flatten k)) ks
                           m'  <- U.DenseMatrix s t <$> V.thaw (flatten m)
                           U.corr2 p (BV.fromList ks') m' (x U.<<=)
                           let vm = U.denseMatrixArrayToVector x
                           vhm <- BV.mapM (\(U.DenseMatrix _ _ vx) -> reshape c <$> V.freeze vx) vm
                           return $ BV.toList vhm

  where
    n     = length ks
    (w,h) = L.size (head ks)
    (s,t) = L.size m
    (r,c) = (s-w+2*p+1, t-h+2*p+1)

eq :: L.Matrix Float -> L.Matrix Float -> Bool
eq a b = V.all id $ ratio a b

ratio a b =
  let va = flatten a
      vb = flatten b
      ae :: V.Vector Float
      ae = V.zipWith (\a b -> abs (a - b)) va vb
      aa = V.sum ae / fromIntegral (V.length ae)
      err x 0 = x < 0.1
      err x y = let e = x-y
                in (abs (e / y) < 0.02)
  in V.zipWith err va vb


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables   #-}
module Bio.ChromVAR.Utils
    ( weight
    , Whitening(..)
    , covariance
    , whiten
    ) where

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as S
import Statistics.Sample (mean)
import Data.Matrix.Static.LinearAlgebra
import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Generic as D
import Data.Singletons (SingI)

weight :: (Double, Double) -> (Double, Double) -> Double
weight (x1, y1) (x2, y2) = gaussian $ sqrt $ (x1 - x2)**2 + (y1 - y2)**2
  where
    gaussian d = let sigma = 0.1 in exp $ negate $ d**2 / (2 * sigma**2)
    --gaussian d = d
{-# INLINE weight #-}

data Whitening = ZCA
               | Cholesky

-- | Rows are samples, columns are features.
whiten :: (SingI n, SingI m) => Whitening -> Matrix n m Double -> Matrix n m Double
whiten method mat = case method of
    Cholesky -> D.transpose $ inverse (cholesky cov) @@ D.transpose mat
  where
    cov = covariance mat

covariance :: forall n m. (SingI n, SingI m) => Matrix n m Double -> Matrix m m Double
covariance mat = D.map (/n) $ D.transpose cs @@ cs
  where
    n = fromIntegral $ D.rows mat - 1
    cs = D.fromColumns $ map f $ D.toColumns mat :: Matrix n m Double
    f x = let m = mean x in S.map (subtract m) x
{-# INLINE covariance #-}

-- | Mahalanobis or ZCA whitening
zca :: U.Vector (Double, Double) -> U.Vector (Double, Double)
zca points = U.map (`dot'` w) $ points -- U.zip xs' ys'
  where
    w = t $ evec `dot` d `dot` t evec
    d = (1 / sqrt (eval1 + eps), 0, 0, 1 / sqrt (eval2 + eps))
    (evec, (eval1, eval2)) = eigendecomposition (sx, sxy, sxy, sy)
    eps = 10e-6
    sx = U.sum (U.zipWith (*) xs' xs') / n
    sy = U.sum (U.zipWith (*) ys' ys') / n
    sxy = U.sum (U.zipWith (*) xs' ys') / n
    xs' = U.map (\x -> x - mx) xs
    ys' = U.map (\x -> x - my) ys
    mx = mean xs
    my = mean ys
    n = fromIntegral $ U.length points - 1
    (xs, ys) = U.unzip points
{-# INLINE zca #-}

type Mat2x2 = (Double, Double, Double, Double)

-- | Eigendecomposition of 2 x 2 symetric matrix
eigendecomposition :: Mat2x2
                   -> ( Mat2x2   -- ^ Eigenvectors in the columns
                      , (Double, Double) ) -- ^ Eigenvalues
eigendecomposition (a,b,_,d) =
    let (x1, y1) = evector eval1
        (x2, y2) = evector eval2
    in ((x1, x2, y1, y2), (eval1, eval2))
  where
    eval1 = (a + d + s) / 2
    eval2 = (a + d - s) / 2
    s = sqrt $ (a - d)**2 + 4 * b**2
    evector lambda =
        let denominator = sqrt $ b**2 + (lambda - d)**2
        in ((lambda - d) / denominator, b / denominator)
{-# INLINE eigendecomposition #-}

dot :: Mat2x2 -> Mat2x2 -> Mat2x2
dot (a,b,c,d) (a',b',c',d') =
    ( a * a' + b * c'
    , a * b' + b * d'
    , c * a' + d * c'
    , c * b' + d * d' )
{-# INLINE dot #-}

dot' :: (Double, Double) -> Mat2x2 -> (Double, Double)
dot' (x,y) (a,b,c,d) = (x * a + y * c, x * b + y * d)
{-# INLINE dot' #-}

t :: Mat2x2 -> Mat2x2
t (a,b,c,d) = (a,c,b,d)
{-# INLINE t #-}


module Utils where

import Data.Function (on)
import Data.List (sortBy)
import System.Random.MWC (GenIO,uniform)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Numeric.LinearAlgebra.HMatrix (Vector)
import qualified Data.Vector.Storable as V

multinomial :: Vector Double  -- ^ Probabilities of each of the p different outcomes. These should sum to 1.
            -> GenIO
            -> IO (Vector Int)
multinomial pvals gen
    | V.sum pvals > 1.01 =
        error "Utils.multinomial: sum of probabilities should sum to 1"
    | otherwise = do
        y <- uniform gen
        let (_,sample,_) = prepvals y
        return sample
    where prepvals :: Double -> (Double, Vector Int, Bool)
          prepvals x = foldr
                  (\(j,p) (s,v,z) ->
                        if z
                            then (s, v, z)
                            else if s + p >= x
                                     then ( s
                                          , v V.//
                                            [(j, 1)]
                                          , True)
                                     else (s + p, v, z))
                  (0.0, V.replicate k 0, False)
                  (sortBy
                       (compare `on` snd)
                       (zip
                            [0 ..]
                            (V.toList pvals)))
          k = V.length pvals


{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
module Neural.Layer
  ( Layer(..)
  , newLayer
  , randomLayer
  , feed
  )
where

import           System.Random
import           Data.Random.Normal
import           Control.Monad
import           Control.Monad.State
import           Numeric.LinearAlgebra
import           GHC.Generics                   ( Generic )
import           Control.DeepSeq

import           Neural.Activation              ( ActivationFunction(..) )

-- | A layer in a neural network.
data Layer a = Layer
    { weights :: Matrix a    -- Weight matrix
    , biases  :: Vector a    -- Bias vector
    } deriving (Show, Generic, NFData)

-- | Construct a new layer from a weigth matrix and a bias vector.
newLayer :: Matrix a -> Vector a -> Layer a
newLayer ws bs = Layer { weights = ws, biases = bs }

-- | Construct a new random layer given its,
-- size and the size of the previous layer.
-- Each bias and weight is given a random value
-- normally distributed with mean 0 and standard deviation 1.
randomLayer
  :: (RandomGen g, Random a, Element a, Floating a)
  => Int         -- ^ The size of the previous layer
  -> Int         -- ^ The size of the layer
  -> State g (Layer a)
randomLayer m n = do
  _biases  <- replicateM n $ state normal
  _weights <- replicateM n $ replicateM m $ state normal
  return $ Layer (fromLists _weights) (fromList _biases)

-- | Calculate the activation and weighted input of a layer
feed
  :: (Numeric a, Num (Vector a))
  => ActivationFunction a  -- ^ The activation function
  -> Vector a              -- ^ The layer input
  -> Layer a               -- ^ The layer
  -> (Vector a, Vector a)  -- ^ The layer output
feed _σ x (Layer w b) = (σ y, y)
 where
  y = w #> x + b
  σ = cmap _σ


{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}

{- |
Module      :  Internal.Modular
Copyright   :  (c) Alberto Ruiz 2015
License     :  BSD3
Stability   :  experimental

Proof of concept of statically checked modular arithmetic.

-}

module Internal.Modular(
    Mod, type (./.)
) where

import Internal.Vector
import Internal.Matrix hiding (size)
import Internal.Numeric
import Internal.Element
import Internal.Container
import Internal.Vectorized (prodI,sumI,prodL,sumL)
import Internal.LAPACK (multiplyI, multiplyL)
import Internal.Algorithms(luFact,LU(..))
import Internal.Util(Normed(..),Indexable(..),
                     gaussElim, gaussElim_1, gaussElim_2,
                     luST, luSolve', luPacked', magnit, invershur)
import Internal.ST(mutable)
#if MIN_VERSION_base(4,11,0)
import GHC.TypeLits hiding (Mod)
#else
import GHC.TypeLits
#endif
import Data.Proxy(Proxy)
import Foreign.ForeignPtr(castForeignPtr)
import Foreign.Storable
import Data.Ratio
import Data.Complex
import Control.DeepSeq ( NFData(..) )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

-- | Wrapper with a phantom integer for statically checked modular arithmetic.
newtype Mod (n :: Nat) t = Mod {unMod:: t}
  deriving (Storable)

instance (NFData t) => NFData (Mod n t)
  where
    rnf (Mod x) = rnf x

infixr 5 ./.
type (./.) x n = Mod n x

instance (Integral t, Enum t, KnownNat m) => Enum (Mod m t)
  where
    toEnum = l0 (\m x -> fromIntegral $ x `mod` (fromIntegral m))
    fromEnum = fromIntegral . unMod

instance (Eq t, KnownNat m) => Eq (Mod m t)
  where
    a == b = (unMod a) == (unMod b)

instance (Ord t, KnownNat m) => Ord (Mod m t)
  where
    compare a b = compare (unMod a) (unMod b)

instance (Integral t, Real t, KnownNat m, Integral (Mod m t)) => Real (Mod m t)
  where
    toRational x = toInteger x % 1

instance (Integral t, KnownNat m, Num (Mod m t)) => Integral (Mod m t)
  where
    toInteger = toInteger . unMod
    quotRem a b = (Mod q, Mod r)
      where
         (q,r) = quotRem (unMod a) (unMod b)

-- | this instance is only valid for prime m
instance (Show (Mod m t), Num (Mod m t), Eq t, KnownNat m) => Fractional (Mod m t)
  where
    recip x
        | x*r == 1  = r
        | otherwise = error $ show x ++" does not have a multiplicative inverse mod "++show m'
      where
        r = x^(m'-2 :: Integer)
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    fromRational x = fromInteger (numerator x) / fromInteger (denominator x)

l2 :: forall m a b c. (Num c, KnownNat m) => (c -> a -> b -> c) -> Mod m a -> Mod m b -> Mod m c
l2 f (Mod u) (Mod v) = Mod (f m' u v)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

l1 :: forall m a b . (Num b, KnownNat m) => (b -> a -> b) -> Mod m a -> Mod m b
l1 f (Mod u) = Mod (f m' u)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

l0 :: forall m a b . (Num b, KnownNat m) => (b -> a -> b) -> a -> Mod m b
l0 f u = Mod (f m' u)
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance Show t => Show (Mod n t)
  where
    show = show . unMod

instance (Integral t, KnownNat n) => Num (Mod n t)
  where
    (+) = l2 (\m a b -> (a + b) `mod` (fromIntegral m))
    (*) = l2 (\m a b -> (a * b) `mod` (fromIntegral m))
    (-) = l2 (\m a b -> (a - b) `mod` (fromIntegral m))
    abs = l1 (const abs)
    signum = l1 (const signum)
    fromInteger = l0 (\m x -> fromInteger x `mod` (fromIntegral m))

instance KnownNat m => Element (Mod m I)
  where
    constantD x n = i2f (constantD (unMod x) n)
    extractR ord m mi is mj js = i2fM <$> extractR ord (f2iM m) mi is mj js
    setRect i j m x = setRect i j (f2iM m) (f2iM x)
    sortI = sortI . f2i
    sortV = i2f . sortV . f2i
    compareV u v = compareV (f2i u) (f2i v)
    selectV c l e g = i2f (selectV c (f2i l) (f2i e) (f2i g))
    remapM i j m = i2fM (remap i j (f2iM m))
    rowOp c a i1 i2 j1 j2 x = rowOpAux (c_rowOpMI m') c (unMod a) i1 i2 j1 j2 (f2iM x)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    gemm u a b c = gemmg (c_gemmMI m') (f2i u) (f2iM a) (f2iM b) (f2iM c)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Element (Mod m Z)
  where
    constantD x n = i2f (constantD (unMod x) n)
    extractR ord m mi is mj js = i2fM <$> extractR ord (f2iM m) mi is mj js
    setRect i j m x = setRect i j (f2iM m) (f2iM x)
    sortI = sortI . f2i
    sortV = i2f . sortV . f2i
    compareV u v = compareV (f2i u) (f2i v)
    selectV c l e g = i2f (selectV c (f2i l) (f2i e) (f2i g))
    remapM i j m = i2fM (remap i j (f2iM m))
    rowOp c a i1 i2 j1 j2 x = rowOpAux (c_rowOpML m') c (unMod a) i1 i2 j1 j2 (f2iM x)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    gemm u a b c = gemmg (c_gemmML m') (f2i u) (f2iM a) (f2iM b) (f2iM c)
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => CTrans (Mod m I)
instance KnownNat m => CTrans (Mod m Z)

instance KnownNat m => Container Vector (Mod m I)
  where
    conj' = id
    size' = dim
    scale' s x = vmod (scale (unMod s) (f2i x))
    addConstant c x = vmod (addConstant (unMod c) (f2i x))
    add' a b = vmod (add' (f2i a) (f2i b))
    sub a b = vmod (sub (f2i a) (f2i b))
    mul a b = vmod (mul (f2i a) (f2i b))
    equal u v = equal (f2i u) (f2i v)
    scalar' x = fromList [x]
    konst' x = i2f . konst (unMod x)
    build' n f = build n (fromIntegral . f)
    cmap' = mapVector
    atIndex' x k = fromIntegral (atIndex (f2i x) k)
    minIndex'     = minIndex . f2i
    maxIndex'     = maxIndex . f2i
    minElement'   = Mod . minElement . f2i
    maxElement'   = Mod . maxElement . f2i
    sumElements'  = fromIntegral . sumI m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    prodElements' = fromIntegral . prodI m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    step'         = i2f . step . f2i
    find' = findV
    assoc' = assocV
    accum' = accumV
    ccompare' a b = ccompare (f2i a) (f2i b)
    cselect' c l e g = i2f $ cselect c (f2i l) (f2i e) (f2i g)
    scaleRecip s x = scale' s (cmap recip x)
    divide x y = mul x (cmap recip y)
    arctan2' = undefined
    cmod' m = vmod . cmod' (unMod m) . f2i
    fromInt' = vmod
    toInt'   = f2i
    fromZ'   = vmod . fromZ'
    toZ'     = toZ' . f2i

instance KnownNat m => Container Vector (Mod m Z)
  where
    conj' = id
    size' = dim
    scale' s x = vmod (scale (unMod s) (f2i x))
    addConstant c x = vmod (addConstant (unMod c) (f2i x))
    add' a b = vmod (add' (f2i a) (f2i b))
    sub a b = vmod (sub (f2i a) (f2i b))
    mul a b = vmod (mul (f2i a) (f2i b))
    equal u v = equal (f2i u) (f2i v)
    scalar' x = fromList [x]
    konst' x = i2f . konst (unMod x)
    build' n f = build n (fromIntegral . f)
    cmap' = mapVector
    atIndex' x k = fromIntegral (atIndex (f2i x) k)
    minIndex'     = minIndex . f2i
    maxIndex'     = maxIndex . f2i
    minElement'   = Mod . minElement . f2i
    maxElement'   = Mod . maxElement . f2i
    sumElements'  = fromIntegral . sumL m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    prodElements' = fromIntegral . prodL m' . f2i
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)
    step'         = i2f . step . f2i
    find' = findV
    assoc' = assocV
    accum' = accumV
    ccompare' a b = ccompare (f2i a) (f2i b)
    cselect' c l e g = i2f $ cselect c (f2i l) (f2i e) (f2i g)
    scaleRecip s x = scale' s (cmap recip x)
    divide x y = mul x (cmap recip y)
    arctan2' = undefined
    cmod' m = vmod . cmod' (unMod m) . f2i
    fromInt' = vmod . fromInt'
    toInt'   = toInt . f2i
    fromZ'   = vmod
    toZ'     = f2i

instance (Storable t, Indexable (Vector t) t) => Indexable (Vector (Mod m t)) (Mod m t)
  where
    (!) = (@>)

type instance RealOf (Mod n I) = I
type instance RealOf (Mod n Z) = Z

instance KnownNat m => Product (Mod m I) where
    norm2      = undefined
    absSum     = undefined
    norm1      = undefined
    normInf    = undefined
    multiply   = lift2m (multiplyI m')
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Product (Mod m Z) where
    norm2      = undefined
    absSum     = undefined
    norm1      = undefined
    normInf    = undefined
    multiply   = lift2m (multiplyL m')
      where
        m' = fromIntegral . natVal $ (undefined :: Proxy m)

instance KnownNat m => Normed (Vector (Mod m I))
  where
    norm_0 = norm_0 . toInt
    norm_1 = norm_1 . toInt
    norm_2 = norm_2 . toInt
    norm_Inf = norm_Inf . toInt

instance KnownNat m => Normed (Vector (Mod m Z))
  where
    norm_0 = norm_0 . toZ
    norm_1 = norm_1 . toZ
    norm_2 = norm_2 . toZ
    norm_Inf = norm_Inf . toZ

instance KnownNat m => Numeric (Mod m I)
instance KnownNat m => Numeric (Mod m Z)

i2f :: Storable t => Vector t -> Vector (Mod n t)
i2f v = unsafeFromForeignPtr (castForeignPtr fp) (i) (n)
    where (fp,i,n) = unsafeToForeignPtr v

f2i :: Storable t => Vector (Mod n t) -> Vector t
f2i v = unsafeFromForeignPtr (castForeignPtr fp) (i) (n)
    where (fp,i,n) = unsafeToForeignPtr v

f2iM :: (Element t, Element (Mod n t)) => Matrix (Mod n t) -> Matrix t
f2iM m = m { xdat = f2i (xdat m) }

i2fM :: (Element t, Element (Mod n t)) => Matrix t -> Matrix (Mod n t)
i2fM m = m { xdat = i2f (xdat m) }

vmod :: forall m t. (KnownNat m, Storable t, Integral t, Numeric t) => Vector t -> Vector (Mod m t)
vmod = i2f . cmod' m'
  where
    m' = fromIntegral . natVal $ (undefined :: Proxy m)

lift1 f a   = vmod (f (f2i a))
lift2 f a b = vmod (f (f2i a) (f2i b))

lift2m f a b = liftMatrix vmod (f (f2iM a) (f2iM b))

instance KnownNat m => Num (Vector (Mod m I))
  where
    (+) = lift2 (+)
    (*) = lift2 (*)
    (-) = lift2 (-)
    abs = lift1 abs
    signum = lift1 signum
    negate = lift1 negate
    fromInteger x = fromInt (fromInteger x)

instance KnownNat m => Num (Vector (Mod m Z))
  where
    (+) = lift2 (+)
    (*) = lift2 (*)
    (-) = lift2 (-)
    abs = lift1 abs
    signum = lift1 signum
    negate = lift1 negate
    fromInteger x = fromZ (fromInteger x)

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

instance (KnownNat m) => Testable (Matrix (Mod m I))
  where
    checkT _ = test

test = (ok, info)
  where
    v = fromList [3,-5,75] :: Vector (Mod 11 I)
    m = (3><3) [1..]   :: Matrix (Mod 11 I)

    a = (3><3) [1,2 , 3
               ,4,5 , 6
               ,0,10,-3] :: Matrix I

    b = (3><2) [0..] :: Matrix I

    am = fromInt a :: Matrix (Mod 13 I)
    bm = fromInt b :: Matrix (Mod 13 I)
    ad = fromInt a :: Matrix Float
    bd = fromInt b :: Matrix Float

    g = (3><3) (repeat (40000)) :: Matrix I
    gm = fromInt g :: Matrix (Mod 100000 I)

    lg = (3><3) (repeat (3*10^(9::Int))) :: Matrix Z
    lgm = fromZ lg :: Matrix (Mod 10000000000 Z)

    gen  n = diagRect 1 (konst 5 n) n n :: Numeric t => Matrix t

    rgen n = gen n :: Matrix R
    cgen n = complex (rgen n) + fliprl (complex (rgen n)) * scalar (0:+1) :: Matrix C
    sgen n = single (cgen n)

    checkGen x = norm_Inf $ flatten $ invg x <> x - ident (rows x)

    invg t = gaussElim t (ident (rows t))

    checkLU okf t = norm_Inf $ flatten (l <> u <> p - t)
      where
        (l,u,p,_) = luFact (LU x' p')
          where
            (x',p') = mutable (luST okf) t

    checkSolve aa = norm_Inf $ flatten (aa <> x - bb)
       where
         bb = flipud aa
         x = luSolve' (luPacked' aa) bb

    tmm = diagRect 1 (fromList [2..6]) 5 5 :: Matrix (Mod 19 I)

    info = do
        print v
        print m
        print (tr m)
        print $ v+v
        print $ m+m
        print $ m <> m
        print $ m #> v

        print $ am <> gaussElim am bm - bm
        print $ ad <> gaussElim ad bd - bd

        print g
        print $ g <> g
        print gm
        print $ gm <> gm

        print lg
        print $ lg <> lg
        print lgm
        print $ lgm <> lgm

        putStrLn "checkGen"
        print (checkGen (gen 5 :: Matrix R))
        print (checkGen (gen 5 :: Matrix Float))
        print (checkGen (cgen 5 :: Matrix C))
        print (checkGen (sgen 5 :: Matrix (Complex Float)))
        print (invg (gen 5) :: Matrix (Mod 7 I))
        print (invg (gen 5) :: Matrix (Mod 7 Z))

        print $ mutable (luST (const True)) (gen 5 :: Matrix R)
        print $ mutable (luST (const True)) (gen 5 :: Matrix (Mod 11 Z))

        putStrLn "checkLU"
        print $ checkLU (magnit 0) (gen 5 :: Matrix R)
        print $ checkLU (magnit 0) (gen 5 :: Matrix Float)
        print $ checkLU (magnit 0) (cgen 5 :: Matrix C)
        print $ checkLU (magnit 0) (sgen 5 :: Matrix (Complex Float))
        print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))
        print $ checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))

        putStrLn "checkSolve"
        print $ checkSolve (gen 5 :: Matrix R)
        print $ checkSolve (gen 5 :: Matrix Float)
        print $ checkSolve (cgen 5 :: Matrix C)
        print $ checkSolve (sgen 5 :: Matrix (Complex Float))
        print $ checkSolve (gen 5 :: Matrix (Mod 7 I))
        print $ checkSolve (gen 5 :: Matrix (Mod 7 Z))

        putStrLn "luSolve'"
        print $ luSolve' (luPacked' tmm) (ident (rows tmm))
        print $ invershur tmm

    ok = and
      [ toInt (m #> v) == cmod 11 (toInt m #> toInt v )
      , am <> gaussElim_1 am bm == bm
      , am <> gaussElim_2 am bm == bm
      , am <> gaussElim   am bm == bm
      , (checkGen (gen 5 :: Matrix R)) < 1E-15
      , (checkGen (gen 5 :: Matrix Float)) < 2E-7
      , (checkGen (cgen 5 :: Matrix C)) < 1E-15
      , (checkGen (sgen 5 :: Matrix (Complex Float))) < 3E-7
      , (checkGen (gen 5 :: Matrix (Mod 7 I))) == 0
      , (checkGen (gen 5 :: Matrix (Mod 7 Z))) == 0
      , (checkLU (magnit 1E-10) (gen 5 :: Matrix R)) < 2E-15
      , (checkLU (magnit 1E-5) (gen 5 :: Matrix Float)) < 1E-6
      , (checkLU (magnit 1E-10) (cgen 5 :: Matrix C)) < 5E-15
      , (checkLU (magnit 1E-5) (sgen 5 :: Matrix (Complex Float))) < 1E-6
      , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 I))) == 0
      , (checkLU (magnit 0) (gen 5 :: Matrix (Mod 7 Z))) == 0
      , checkSolve (gen 5 :: Matrix R) < 2E-15
      , checkSolve (gen 5 :: Matrix Float) < 1E-6
      , checkSolve (cgen 5 :: Matrix C) < 4E-15
      , checkSolve (sgen 5 :: Matrix (Complex Float)) < 1E-6
      , checkSolve (gen 5 :: Matrix (Mod 7 I)) == 0
      , checkSolve (gen 5 :: Matrix (Mod 7 Z)) == 0
      , prodElements (konst (9:: Mod 10 I) (12::Int)) == product (replicate 12 (9:: Mod 10 I))
      , gm <> gm == konst 0 (3,3)
      , lgm <> lgm == konst 0 (3,3)
      , invershur tmm == luSolve' (luPacked' tmm) (ident (rows tmm))
      , luSolve' (luPacked' (tr $ ident 5 :: Matrix (I ./. 2))) (ident 5) == ident 5
      ]


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

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

Multidimensional root finding.

<http://www.gnu.org/software/gsl/manual/html_node/Multidimensional-Root_002dFinding.html>

The example in the GSL manual:

@import Numeric.GSL
import Numeric.LinearAlgebra(format)
import Text.Printf(printf)

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

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

main = do
    let (sol,path) = root Hybrids 1E-7 30 (rosenbrock 1 10) [-10,-5]
    print sol
    disp path

\> main
[1.0,1.0]
 0.000  -10.000  -5.000  11.000  -1050.000
 1.000   -3.976  24.827   4.976     90.203
 2.000   -3.976  24.827   4.976     90.203
 3.000   -3.976  24.827   4.976     90.203
 4.000   -1.274  -5.680   2.274    -73.018
 5.000   -1.274  -5.680   2.274    -73.018
 6.000    0.249   0.298   0.751      2.359
 7.000    0.249   0.298   0.751      2.359
 8.000    1.000   0.878  -0.000     -1.218
 9.000    1.000   0.989  -0.000     -0.108
10.000    1.000   1.000   0.000      0.000
@

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

module Numeric.GSL.Root (
    uniRoot, UniRootMethod(..),
    uniRootJ, UniRootMethodJ(..),
    root, RootMethod(..),
    rootJ, RootMethodJ(..),
) where

import Data.Packed.Internal
import Data.Packed.Matrix
import Numeric.GSL.Internal
import Foreign.Ptr(FunPtr, freeHaskellFunPtr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

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

data UniRootMethod = Bisection
                   | FalsePos
                   | Brent
                   deriving (Enum, Eq, Show, Bounded)

uniRoot :: UniRootMethod
        -> Double
        -> Int
        -> (Double -> Double)
        -> Double
        -> Double
        -> (Double, Matrix Double)
uniRoot method epsrel maxit fun xl xu = uniRootGen (fi (fromEnum method)) fun xl xu epsrel maxit

uniRootGen m f xl xu epsrel maxit = unsafePerformIO $ do
    fp <- mkDoublefun f
    rawpath <- createMIO maxit 4
                         (c_root m fp epsrel (fi maxit) xl xu)
                         "root"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (sol !! 1, path)

foreign import ccall safe "root"
    c_root:: CInt -> FunPtr (Double -> Double) -> Double -> CInt -> Double -> Double -> TM

-------------------------------------------------------------------------
data UniRootMethodJ = UNewton
                    | Secant
                    | Steffenson
                    deriving (Enum, Eq, Show, Bounded)

uniRootJ :: UniRootMethodJ
        -> Double
        -> Int
        -> (Double -> Double)
        -> (Double -> Double)
        -> Double
        -> (Double, Matrix Double)
uniRootJ method epsrel maxit fun dfun x = uniRootJGen (fi (fromEnum method)) fun
    dfun x epsrel maxit

uniRootJGen m f df x epsrel maxit = unsafePerformIO $ do
    fp <- mkDoublefun f
    dfp <- mkDoublefun df
    rawpath <- createMIO maxit 2
                         (c_rootj m fp dfp epsrel (fi maxit) x)
                         "rootj"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (sol !! 1, path)

foreign import ccall safe "rootj"
    c_rootj :: CInt -> FunPtr (Double -> Double) -> FunPtr (Double -> Double)
            -> Double -> CInt -> Double -> TM

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

data RootMethod = Hybrids
                | Hybrid
                | DNewton
                | Broyden
                deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional root finding using algorithms that do not require
-- any derivative information to be supplied by the user.
-- Any derivatives needed are approximated by finite differences.
root :: RootMethod
     -> Double                     -- ^ maximum residual
     -> Int                        -- ^ maximum number of iterations allowed
     -> ([Double] -> [Double])     -- ^ function to minimize
     -> [Double]                   -- ^ starting point
     -> ([Double], Matrix Double)  -- ^ solution vector and optimization path

root method epsabs maxit fun xinit = rootGen (fi (fromEnum method)) fun xinit epsabs maxit

rootGen m f xi epsabs maxit = unsafePerformIO $ do
    let xiv = fromList xi
        n   = dim xiv
    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
    rawpath <- vec xiv $ \xiv' ->
                   createMIO maxit (2*n+1)
                         (c_multiroot m fp epsabs (fi maxit) // xiv')
                         "multiroot"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    return (take n $ drop 1 sol, path)

foreign import ccall safe "multiroot"
    c_multiroot:: CInt -> FunPtr TVV -> Double -> CInt -> TVM

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

data RootMethodJ = HybridsJ
                 | HybridJ
                 | Newton
                 | GNewton
                deriving (Enum,Eq,Show,Bounded)

-- | Nonlinear multidimensional root finding using both the function and its derivatives.
rootJ :: RootMethodJ
      -> Double                     -- ^ maximum residual
      -> Int                        -- ^ maximum number of iterations allowed
      -> ([Double] -> [Double])     -- ^ function to minimize
      -> ([Double] -> [[Double]])   -- ^ Jacobian
      -> [Double]                   -- ^ starting point
      -> ([Double], Matrix Double)  -- ^ solution vector and optimization path

rootJ method epsabs maxit fun jac xinit = rootJGen (fi (fromEnum method)) fun jac xinit epsabs maxit

rootJGen m f jac xi epsabs maxit = unsafePerformIO $ do
    let xiv = fromList xi
        n   = dim xiv
    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
    jp <- mkVecMatfun (aux_vTom (checkdim2 n . fromLists . jac . toList))
    rawpath <- vec xiv $ \xiv' ->
                   createMIO maxit (2*n+1)
                         (c_multirootj m fp jp epsabs (fi maxit) // xiv')
                         "multiroot"
    let it = round (rawpath @@> (maxit-1,0))
        path = takeRows it rawpath
        [sol] = toLists $ dropRows (it-1) path
    freeHaskellFunPtr fp
    freeHaskellFunPtr jp
    return (take n $ drop 1 sol, path)

foreign import ccall safe "multirootj"
    c_multirootj:: CInt -> FunPtr TVV -> FunPtr TVM -> Double -> CInt -> TVM

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

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

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


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiWayIf   #-}
{-# LANGUAGE RankNTypes   #-}
{-# LANGUAGE ViewPatterns #-}
module Main where

import Control.Lens                       hiding (inside, transform)
import Control.Monad.Extra
import Control.Monad.ST
import Data.Array.Repa
import Data.Complex
import Data.Matrix                        (Matrix)
import Data.STRef
import GHC.Base                           (until)
import Graphics.Gloss.Interface.Pure.Game hiding (Vector, shift)
import Graphics.Gloss.Raster.Array
import Relude

import qualified Data.ByteString as B
import qualified Data.Matrix     as Mtx

data World = World { _v   :: Matrix Float
                   , _mtx :: Matrix Float
                   } deriving (Eq,Show)
makeLenses ''World

initialWorld :: World
initialWorld = World { _v = Mtx.fromLists [[-2],[-1]] -- FIXME no hardcode
                     , _mtx = Mtx.fromLists [[1/250,0],[0,1/250]]
                     }

type Vector a = Matrix a

data Settings = Settings { _windowSize :: (Int, Int)
                         , _offset     :: (Int, Int)
                         , _pointSize  :: (Int,Int)
                         , _threshold  :: forall a. Num a => a
                         , _fps        :: Int
                         , _iteration  :: Int
                         }
makeLenses ''Settings

settings :: Settings
settings = Settings { _windowSize = (750,500)
                    , _offset     = (250,250)
                    , _pointSize  = (1,1)
                    , _threshold  = 2
                    , _fps        = 60
                    , _iteration  = 100
                    }

window :: Display
window = InWindow "fractol" (settings^.windowSize) (settings^.offset)

toTup :: DIM2 -> (Int, Int)
toTup (Z :. y :. x) = (x,y)

toVec :: (Int, Int) -> Vector Float
toVec (!x, !y) = seq y $ seq x $ on f fromIntegral x y
  where
    f !x1 !x2 = Mtx.fromLists [[x1],[x2]]
{-# INLINE toVec #-}

transform :: World -> Vector Float -> Vector Float
transform !w !p = w^.mtx * p + w^.v
{-# INLINE transform #-}

mandel :: World -> Array D DIM2 Color
mandel !w = fromFunction (ix2 500 750) render
  where
    render (iter . transform w . toVec . toTup -> (len, n)) =
      let diff = len - settings^.threshold
          m = 2 * n
      in if diff >= 0
           then rgb' (0.9 * diff) (0.6 * diff) (0.1 * diff)
           else black

toComplex :: Vector Float -> Complex Float
toComplex !v = let [!x,!y] = Mtx.toList v
          in x :+ y
{-# INLINE toComplex #-}

iter :: Vector Float -> (Float,Int)
iter (toComplex -> !c) = let (!zn, !len, !n) = until condition endo start
                        in  (len, n)
  where
    start = (0,0,0)
    endo (!zn, !len, !n) = let new = zn * zn + c
                          in (new, magnitude new, succ n)
    {-# INLINE endo #-}
    condition (_, !len, !n) = len >= settings^.threshold
                                || n > settings^.iteration
    {-# INLINE condition #-}

{-# INLINE iter #-}

cast :: (Float, Float) -> (Int, Int)
cast (!x,!y) = on (,) floor x y
{-# INLINE cast #-}

shift :: Vector Float -> Vector Float
shift = (+ Mtx.fromLists [[375],[250]])
{-# INLINE shift #-}

eventHandler :: Event -> World -> World
eventHandler (EventKey (Char 'p') Up _ (shift . toVec . cast -> p)) w =
  let truePoint = transform w p
  in w & v -~ Mtx.scaleMatrix 0.1 truePoint
       & mtx %~ (Mtx.scaleMatrix 1.1 (Mtx.identity 2) *)
eventHandler (EventKey (Char 'n') Up _ (shift . toVec . cast -> p)) w =
  let truePoint = transform w p
  in w & v +~ Mtx.scaleMatrix 0.1 truePoint
       & mtx %~ (Mtx.scaleMatrix 0.9 (Mtx.identity 2) *)
eventHandler _ w = w

main :: IO ()
main = do
  let fractol = playArray window (settings^.pointSize) (settings^.fps)
  fractol initialWorld mandel eventHandler (const id)


{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}

module SimulatedAnnealing (
   State,
   energy,
   metric,
   perturb,
   Args(..),
   CoolingTemp(..),
   anneal
) where

import Metric (MetricSpace, Metric, metric)

import qualified Metric

import Prelude hiding (min)

import Control.Monad.Random (evalRand)
import Control.Monad.Random.Class (MonadRandom, getRandom)
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Int (Int32)
import Numeric.GSL.SimulatedAnnealing (SimulatedAnnealingParams(..), simanSolve)
import Numeric.LinearAlgebra.Data (Vector, (!))
import System.Random (mkStdGen, StdGen)

import Debug.Trace

class MetricSpace a => State a where
    energy :: a -> Double
    perturb :: MonadRandom m => Double -> a -> m a

data Args = Args {
    numTries :: Word,
    numItersPerTemp :: Word,
    maxStepSize :: Double,
    boltzmannK :: Double,
    coolingTemp :: CoolingTemp
}

data CoolingTemp = CoolingTemp {
    initial :: Double,
    rate :: Double,
    min :: Double
}

argsToParams :: Args -> SimulatedAnnealingParams
argsToParams Args {
    numTries,
    numItersPerTemp,
    maxStepSize,
    boltzmannK,
    coolingTemp = CoolingTemp {
        initial,
        rate,
        min
    }
} = SimulatedAnnealingParams {
    n_tries = fromIntegral numTries,
    iters_fixed_T = fromIntegral numItersPerTemp,
    step_size = maxStepSize,
    boltzmann_k = boltzmannK,
    cooling_t_initial = initial,
    cooling_mu_t = rate,
    cooling_t_min = min
}

data Annealer a = Annealer {
    seed :: Int,
    numRandoms :: Int,
    params :: SimulatedAnnealingParams,
    initialState :: a,
    energy' :: a -> Double,
    metric' :: a -> a -> Double,
    step :: Vector Double -> Double -> a -> a,
    maybeShow :: Maybe (a -> String)
}

rawAnneal :: Annealer a -> a
rawAnneal Annealer {
    seed,
    numRandoms,
    params,
    initialState,
    energy',
    metric',
    step,
    maybeShow
} = simanSolve
    seed
    numRandoms
    params
    initialState
    energy'
    metric'
    step
    maybeShow

anneal :: (State a, MonadRandom m) => a -> Args -> Maybe (a -> String) -> m a
anneal initialState args maybeShow = do
    seed <- getRandom
    Annealer {
        seed,
        numRandoms = 1,
        params = argsToParams args,
        initialState,
        energy' = energy,
        metric' = metric,
        step,
        maybeShow
    }
        & rawAnneal
        & return
      where
        -- It's too hard to thread the monad through the SA,
        -- b/c the FFI library isn't designed with monads.
        -- So I'm doing it manually here, seeding a new StdGen using mkStdGen
        -- with a seed from the MonadRandom.
        step vectors stepSize s = s'
          where
            seed = (vectors ! 0)
                * (fromIntegral (maxBound :: Int32))
                & truncate
            g = mkStdGen seed
            rand = perturb stepSize s
            s' = evalRand rand g

realAnneal :: (State a, MonadRandom m) => a -> Args -> m a
realAnneal initialState args = do
    undefined


{-# LANGUAGE BangPatterns #-}
module Main (main) where

import Data.Vector.Storable as V
import Numeric.FFT.Vector.Unnormalized as U
import Numeric.FFT.Vector.Invertible as I
import Numeric.FFT.Vector.Unitary as O
import Criterion.Main
import Data.Complex

main = do
    let mkSize planner k = do
            let !v = V.replicate k 17 :: Vector (Complex Double)
            let !p = plan planner k
            return $ bench (show k) $ whnf (\w -> execute p w) v
    let mkSize1 k = do
            let !v = V.replicate k 17 :: Vector (Complex Double)
            let !p = plan U.dft k
            return $ bench (show k) $ whnf (\w -> execute p w) v
    let mkSize2 k = do
            let !v = V.replicate k 17 :: Vector (Complex Double)
            let !p = plan O.dft k
            return $ bench (show k) $ whnf (\w -> execute p w) v

    let mkSize3 k = do
            let !v = V.replicate k 17 :: Vector (Complex Double)
            return $ bench (show k) $ whnf (\w -> I.run O.dft w) v

    let sizes = [64,256,2048,8192]
    let sizedGroup name f = fmap (bgroup name) $ Prelude.mapM f sizes
    benches <- sizedGroup "U.idft_1" (mkSize U.idft)
    benches1 <- sizedGroup "U.dft" mkSize1
    benches2 <- sizedGroup "O.dft" mkSize2
    benches3 <- sizedGroup "run_O.dft" mkSize3
    defaultMain $ [benches, benches1, benches2, benches3]


-- {-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE DataKinds, GADTs, TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}

module HBLAS.BLAS.Level1Spec (main, spec) where

import Data.Complex

import Numerical.HBLAS.MatrixTypes as Matrix
import Numerical.HBLAS.BLAS.Level1 as BLAS

import Test.Hspec

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  asumSpec
  axpySpec
  copySpec
  dotSpec
  sdotSpec
  dotcSpec
  dotuSpec
  nrm2Spec
  rotSpec
  rotgSpec
  rotmSpec
  rotmgSpec
  scalSpec
  swapSpec
  iamaxSpec
  iaminSpec

asumSpec :: Spec
asumSpec =
  context "?ASUM" $ do
    describe "SASUM" $ do
      it "vector of length 6 with incx 1" $ do
        vecTest1SASUM
      it "vector of length 12 with incx 2" $ do
        vecTest2SASUM

vecTest1SASUM :: IO ()
vecTest1SASUM = do
  vec <- Matrix.generateMutableDenseVector 6 (\idx -> [1 .. 6] !! idx)
  res <- BLAS.sasum vec
  res `shouldBe` (21.0 :: Float)

vecTest2SASUM :: IO ()
vecTest2SASUM = do
  vec <- Matrix.generateMutableDenseVectorWithStride 12 2 (\idx -> [1 .. 12] !! idx)
  res <- BLAS.sasum  vec
  res `shouldBe` (sum [1,3 .. 12] :: Float)

axpySpec :: Spec
axpySpec =
  context "?AXPY" $ do
    describe "SAXPY" $ do
      it "vectors of lengths 6 and 6 with both incx 1" $ do
        vecTest1SAXPY
      it "vectors of lengths 12 and 18 with incx 2 and 3" $ do
        vecTest2SAXPY

vecTest1SAXPY :: IO ()
vecTest1SAXPY = do
  input <- Matrix.generateMutableDenseVector 6 (\idx -> [1 .. 6] !! idx)
  output <- Matrix.generateMutableDenseVector 6 (\idx -> [2, 3, 4, 3, 5, 6] !! idx)
  BLAS.saxpy 6 (-1.0) input output
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector output
  resList `shouldBe` [1, 1, 1, -1, 0, 0]

vecTest2SAXPY :: IO ()
vecTest2SAXPY = do
  input <- Matrix.generateMutableDenseVectorWithStride 18 3 (\idx -> [1 .. 18] !! idx)
  output <- Matrix.generateMutableDenseVectorWithStride 12 2 (\idx -> [1 .. 12] !! idx)
  BLAS.saxpy 6 2.0 input output
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector output
  resList `shouldBe` [3, 2, 11, 4, 19, 6, 27, 8, 35, 10, 43, 12]

copySpec :: Spec
copySpec =
  context "?COPY" $ do
    describe "DCOPY" $ do
      it "vectors of lengths 6 and 6 with both incx 1" $ do
        vecTest1DCOPY
      it "vectors of lengths 6 and 9 with incx 2 and 3" $ do
        vecTest2DCOPY

vecTest1DCOPY :: IO ()
vecTest1DCOPY = do
  input <- Matrix.generateMutableDenseVector 6 (\idx -> [1 .. 6] !! idx)
  output <- Matrix.generateMutableDenseVector 6 (const 0.0)
  BLAS.dcopy 6 input output
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector output
  resList `shouldBe` [1, 2, 3, 4, 5, 6]

vecTest2DCOPY :: IO ()
vecTest2DCOPY = do
  input <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1 .. 6] !! idx)
  output <- Matrix.generateMutableDenseVectorWithStride 9 3 (const 0.0)
  BLAS.dcopy 3 input output
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector output
  resList `shouldBe` [1, 0, 0, 3, 0, 0, 5, 0, 0]

dotSpec :: Spec
dotSpec =
  context "?DOT" $ do
    describe "SDOT" $ do
      it "vectors of lengths 6 and 12 with incx 2 and 4" $ do
        vecTest1SDOT
    describe "DDOT" $ do
      it "vectors of lengths 12 and 6 with incx 2 and 1" $ do
        vecTest1DDOT

vecTest1SDOT :: IO ()
vecTest1SDOT = do
  left <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 12 4 (\idx -> [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0] !! idx)
  res <- sdot 3 left right
  res `shouldBe` ((1 + 15 + 45) :: Float)

vecTest1DDOT :: IO ()
vecTest1DDOT = do
  left <- Matrix.generateMutableDenseVectorWithStride 12 2 (((+) 1) . fromRational . toRational)
  right <- Matrix.generateMutableDenseVectorWithStride 6 1 (\idx -> [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] !! idx)
  res <- ddot 6 left right
  res `shouldBe` ((1 + 6 + 15 + 28 + 45 + 66) :: Double)

sdotSpec :: Spec
sdotSpec =
  context "?SDOT" $ do
    describe "SDSDOT" $ do
      it "vectors of lengths 6 and 12 with incx 2 and 4" $ do
        pending
    describe "DSDOT" $ do
      it "vectors of lengths 12 and 6 with incx 2 and 1" $ do
        vecTest1DSDOT

{-vecTest1SDSDOT :: IO ()-}
{-vecTest1SDSDOT = do-}
  {-left <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> ([1 .. 6] :: [Float]) !! idx) -- Float-}
  {-right <- Matrix.generateMutableDenseVectorWithStride 12 4 (\idx -> ([1 .. 12] :: [Float]) !! idx) -- Float-}
  {-res <- sdsdot 3 2.0 left right-}
  {-res `shouldBe` ((2 :: Float) + 1 + 15 + 45)-}

vecTest1DSDOT :: IO ()
vecTest1DSDOT = do
  left <- Matrix.generateMutableDenseVectorWithStride 12 2 (\idx -> ([1 .. 12] :: [Float]) !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 6 1 (\idx -> ([1 .. 6] :: [Float]) !! idx)
  res <- dsdot 6 left right
  res `shouldBe` ((1 + 6 + 15 + 28 + 45 + 66) :: Double)

dotcSpec :: Spec
dotcSpec =
  context "?DOTC" $ do
    describe "CDOTC" $ do
      it "vectors of length 6 and 9 with incx 2 and 3" $ do
        vecTest1CDOTC

vecTest1CDOTC :: IO ()
vecTest1CDOTC = do
  left <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [2:+3, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1)] !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 9 3 (\idx -> [1:+(-2), 1:+1, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1)] !! idx)
  cdotc 3 left right `shouldReturn` ((-2):+(-9) :: Complex Float)

dotuSpec :: Spec
dotuSpec =
  context "?DOTU" $ do
    describe "CDOTU" $ do
      it "vectors of length 6 and 9 with incx of 2 and 3" $ do
        vecTest1CDOTU

vecTest1CDOTU :: IO ()
vecTest1CDOTU = do
  left <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1:+1, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1)] !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 9 3 (\idx -> [1:+(-2), 1:+1, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1), 1:+1, 1:+(-1)] !! idx)
  --res <- Matrix.generateMutableValue (1:+1)
  res <- cdotu 3 left right

  res `shouldBe` 5:+1

nrm2Spec :: Spec
nrm2Spec =
  context "?NRM2" $ do
    describe "SNRM2" $ do
      it "vector of length 6 with incx of 1" $ do
        vecTest1SNRM2
    describe "DZNRM2" $ do
      it "vector of length 8 with incx of 2" $ do
        vecTest1DZNRM2

vecTest1SNRM2 :: IO ()
vecTest1SNRM2 = do
  input <- Matrix.generateMutableDenseVector 6 (\idx -> [1.0, -2.0, 3.0, -4.0, 5.0, -6.0] !! idx)
  res <- snrm2 6 input
  True `shouldBe` 1e-6 > (abs $ res - (sqrt $ sum $ fmap (\x->x^2) [1, 2, 3, 4, 5, 6]))

vecTest1DZNRM2 :: IO ()
vecTest1DZNRM2 = do
  input <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [1:+1, 1:+2, 2:+(-3), 2:+(-2), (-3):+1, (-3):+0, (-4):+2, (-4):+1] !! idx)
  res <- dznrm2 4 input
  True `shouldBe` 1e-12 > (abs $ res - (sqrt $ sum $ fmap (\x->x^2) [1, 1, 2, 3, 3, 1, 4, 2]))

rotSpec :: Spec
rotSpec =
  context "?ROT" $ do
    describe "SROT" $ do
      it "vectors of length 6 and 6 with incx of 2" $ do
        vecTest1SROT
    describe "DROT" $ do
      it "vectors of length 4 and 8 with incx of 1 and 2" $ do
        vecTest1DROT

vecTest1SROT :: IO ()
vecTest1SROT = do
  left <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [6.0, 5.0, 4.0, 3.0, 2.0, 1.0] !! idx)
  srot 3 left right (-1) 2
  resLeft <- Matrix.mutableVectorToList $ _bufferMutDenseVector left
  resRight <- Matrix.mutableVectorToList $ _bufferMutDenseVector right
  resLeft `shouldBe` [11.0, 2.0, 5.0, 4.0, -1.0, 6.0]
  resRight `shouldBe` [-8.0, 5.0, -10.0, 3.0, -12.0, 1.0]

vecTest1DROT :: IO ()
vecTest1DROT = do
  left <- Matrix.generateMutableDenseVectorWithStride 4 1 (\idx -> [1, 2, 3, 4] !! idx)
  right <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [8, 7, 6, 5, 4, 3, 2, 1] !! idx)
  drot 4 left right 0 (-2)
  resLeft <- Matrix.mutableVectorToList $ _bufferMutDenseVector left
  resRight <- Matrix.mutableVectorToList $ _bufferMutDenseVector right
  resLeft `shouldBe` [-16, -12, -8, -4]
  resRight `shouldBe` [2, 7, 4, 5, 6, 3, 8, 1]

rotgSpec :: Spec
rotgSpec =
  context "?ROTG" $ do
    describe "SROTG" $ do
      it "apply on 3 and 4" $ do
        pending
    describe "DROTG" $ do
      it "apply on 5.8 and 3.4" $ do
        pending

{-vecTest1SROTG :: IO ()-}
{-vecTest1SROTG = do-}
  {-a <- Matrix.generateMutableValue 3-}
  {-b <- Matrix.generateMutableValue 4-}
  {-c <- Matrix.generateMutableValue 0-}
  {-s <- Matrix.generateMutableValue 0-}
  {-srotg a b c s-}
  {-av <- Matrix.mutableValueToValue a-}
  {-bv <- Matrix.mutableValueToValue b-}
  {-cv <- Matrix.mutableValueToValue c-}
  {-sv <- Matrix.mutableValueToValue s-}
  {-av `shouldBe` 5-}
  {-True `shouldBe` 1e-6 > (abs $ bv - 1/0.6)-}
  {-cv `shouldBe` 0.6-}
  {-sv `shouldBe` 0.8-}

{-vecTest1DROTG :: IO ()-}
{-vecTest1DROTG = do-}
  {-a <- Matrix.generateMutableValue 5.8-}
  {-b <- Matrix.generateMutableValue 3.4-}
  {-c <- Matrix.generateMutableValue 0-}
  {-s <- Matrix.generateMutableValue 0-}
  {-drotg a b c s-}
  {-av <- Matrix.mutableValueToValue a-}
  {-bv <- Matrix.mutableValueToValue b-}
  {-cv <- Matrix.mutableValueToValue c-}
  {-sv <- Matrix.mutableValueToValue s-}
  {-True `shouldBe` 1e-12 > (abs $ av - sqrt(3.4^2 + 5.8^2))-}
  {-True `shouldBe` 1e-12 > (abs $ bv - 3.4 / sqrt(3.4^2 + 5.8^2))-}
  {-True `shouldBe` 1e-12 > (abs $ cv - 5.8 / sqrt(3.4^2 + 5.8^2))-}
  {-True `shouldBe` 1e-12 > (abs $ sv - 3.4 / sqrt(3.4^2 + 5.8^2))-}

rotmSpec :: Spec
rotmSpec =
  context "?ROTM" $ do
    it "vectors of length 4 and 8 with incx of 1 and 2, param starts with -1" $ do
      pending
    it "vectors of length 6 and 9 with incx of 2 and 3, param starts with 1" $ do
      pending

{-
vecTest1DROTM :: IO ()
vecTest1DROTM = do
  x <- Matrix.generateMutableDenseVectorWithStride 4 1 (\idx -> [1, 2, 3, 4] !! idx)
  y <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [8, 7, 6, 5, 4, 3, 2, 1] !! idx)
  param <- Matrix.generateMutableDenseVector 5 (\idx -> [-1, 0, -1, 1, 0] !! idx)
  drotm 4 x y param
  resX <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  resY <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resX `shouldBe` [8, 6, 4, 2]
  resY `shouldBe` [-1, 7, -2, 5, -3, 3, -4, 1]

vecTest1SROTM :: IO ()
vecTest1SROTM = do
  x <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  y <- Matrix.generateMutableDenseVectorWithStride 9 3 (\idx -> [9, 8, 7, 6, 5, 4, 3, 2, 1] !! idx)
  param <- Matrix.generateMutableDenseVector 5 (\idx -> [1, 1, 2, -2, 1] !! idx)
  srotm 3 x y param
  resX <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  resY <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resX `shouldBe` [10, 2, 9, 4, 8, 6]
  resY `shouldBe` [8, 8, 7, 3, 5, 4, -2, 2, 1]
-}

rotmgSpec :: Spec
rotmgSpec =
  context "?ROTMG" $ do
    describe "SROTMG" $ do
      it "todo" $ do
        pending

{-
vecTest1SROTMG :: IO ()
vecTest1SROTMG = do
  d1 <- Matrix.generateMutableValue 3
  d2 <- Matrix.generateMutableValue 6
  x <- Matrix.generateMutableValue 1
  let y = 1
  param <- Matrix.generateMutableDenseVector 5 (\idx -> [-1, 1, 1, -1, 1] !! idx)
  srotmg d1 d2 x y param
  paramR <- Matrix.mutableVectorToList $ _bufferMutDenseVector param
  updatedD1 <- Matrix.mutableValueToValue d1
  updatedD2 <- Matrix.mutableValueToValue d2
  updatedX <- Matrix.mutableValueToValue x
  paramR `shouldBe` [1, 0, 0.5, 0, 1]
  updatedD1 `shouldBe` 4
  updatedD2 `shouldBe` 2
  updatedX `shouldBe` 1.5
-}

scalSpec :: Spec
scalSpec =
  context "?SCAL" $ do
    describe "SSCAL" $ do
      it "vector of length 8 with incx 2" $ do
        vecTest1SSCAL
    describe "CSCAL" $ do
      it "vector of length 8 with incx 4" $ do
        vecTest1CSCAL
    describe "CSSCAL" $ do
      it "vector of length 8 with incx 1" $ do
        vecTest1CSSCAL

vecTest1SSCAL :: IO ()
vecTest1SSCAL = do
  x <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [1, 2, 3, 4, 5, 6, 7, 8] !! idx)
  sscal 4 (-2) x
  xRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  xRes `shouldBe` [-2, 2, -6, 4, -10, 6, -14, 8]

vecTest1CSCAL :: IO ()
vecTest1CSCAL = do
  x <- Matrix.generateMutableDenseVectorWithStride 8 4 (\idx -> [1:+1, 1:+2, 2:+(-3), 2:+(-2), (-3):+1, (-3):+0, (-4):+2, (-4):+1] !! idx)
  cscal 2 (2:+(-2)) x -- size 2?
  xRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  xRes `shouldBe` [4:+0, 1:+2, 2:+(-3), 2:+(-2), (-4):+8, (-3):+0, (-4):+2, (-4):+1]

vecTest1CSSCAL :: IO ()
vecTest1CSSCAL = do
  x <- Matrix.generateMutableDenseVector 8 (\idx -> [1:+1, 1:+2, 2:+(-3), 2:+(-2), (-3):+1, (-3):+0, (-4):+2, (-4):+1] !! idx)
  csscal 8 (-2) x
  xRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  xRes `shouldBe` [(-2):+(-2), (-2):+(-4), (-4):+6, (-4):+4, 6:+(-2), 6:+0, 8:+(-4), 8:+(-2)]

swapSpec :: Spec
swapSpec =
  context "?SWAP" $ do
    describe "SSWAP" $ do
      it "vectors of length 8 and 4 with incx 2 and 1" $ do
        vecTest1SSWAP
    describe "CSWAP" $ do
      it "vectors of length 9 and 6 with incx 3 and 2" $ do
        vecTest1CSWAP

vecTest1SSWAP :: IO ()
vecTest1SSWAP = do
  x <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [1, 2, 3, 4, 5, 6, 7, 8] !! idx)
  y <- Matrix.generateMutableDenseVectorWithStride 4 1 (\idx -> [-1, -2, -3, -4] !! idx)
  sswap 4 x y
  xRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  yRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  xRes `shouldBe` [-1, 2, -2, 4, -3, 6, -4, 8]
  yRes `shouldBe` [1, 3, 5, 7]

vecTest1CSWAP :: IO ()
vecTest1CSWAP = do
  x <- Matrix.generateMutableDenseVectorWithStride 9 3 (\idx -> [1:+1, 1:+2, 2:+(-3), 2:+(-2), (-3):+1, (-3):+0, (-4):+2, (-4):+1, 0:+9] !! idx)
  y <- Matrix.generateMutableDenseVectorWithStride 6 2 (\idx -> [1:+2, 1:+3, 3:+(-3), 2:+2, 3:+1, 3:+3] !! idx)
  cswap 3 x y
  xRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
  yRes <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  xRes `shouldBe` [1:+2, 1:+2, 2:+(-3), 3:+(-3), (-3):+1, (-3):+0, 3:+1, (-4):+1, 0:+9]
  yRes `shouldBe` [1:+1, 1:+3, 2:+(-2), 2:+2, (-4):+2, 3:+3]

iamaxSpec :: Spec
iamaxSpec =
  context "I?AMAX" $ do
    describe "ISAMAX" $ do
      it "vector of length 8 with incx 2" $ do
        vecTest1ISAMAX
    describe "ICAMAX" $ do
      it "vector of length 9 with incx 1" $ do
        vecTest1ICAMAX

vecTest1ISAMAX :: IO ()
vecTest1ISAMAX = do
  x <- Matrix.generateMutableDenseVectorWithStride 8 2 (\idx -> [1, 2, 3, 4, 5, 6, 7, 8] !! idx)
  idx <- isamax  x
  idx `shouldBe` 7

vecTest1ICAMAX :: IO ()
vecTest1ICAMAX = do
  x <- Matrix.generateMutableDenseVector 9 (\idx -> [1:+1, 1:+2, 2:+(-3), 2:+(-2), (-3):+1, (-3):+0, (-4):+2, (-4):+1, 0:+9] !! idx)
  idx <- icamax  x
  idx `shouldBe` 8

iaminSpec :: Spec
iaminSpec =
  context "I?AMIN" $ do
    describe "ISAMIN" $ do
      it "vector of length 8 with incx 2" $ do
        pending
    describe "ICAMIN" $ do
      it "vector of length 9 with incx 1" $ do
        pending

{-
vecTest1ISAMIN :: IO ()
vecTest1ISAMIN = do
  x <- Matrix.generateMutableDenseVector 8 (\idx -> [1, 2, 3, 4, -5, 6, 7, 8] !! idx)
  idx <- isamin 4 x 2
  idx `shouldBe` 2

vecTest1ICAMIN :: IO ()
vecTest1ICAMIN = do
  x <- Matrix.generateMutableDenseVector 9 (\idx -> [1:+2, 1:+2, (-2):+(-3), 2:+(-2), (-3):+1, (-2):+0, (-4):+2, (-4):+1, 0:+9] !! idx)
  idx <- icamin 9 x 1
  idx `shouldBe` 5
-}


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeInType          #-}
{-# LANGUAGE TypeOperators       #-}

-- |
-- Module      : Numeric.Hamilton
-- Description : Hamiltonian dynamics for physical systems on generalized
--               coordinates using automatic differentiation
-- Copyright   : (c) Justin Le 2016
-- License     : BSD-3
-- Maintainer  : justin@jle.im
-- Stability   : unstable
-- Portability : portable
--
-- Simulate physical systems on generalized/arbitrary coordinates using
-- Hamiltonian mechanics and automatic differentiation!
--
-- See the <https://github.com/mstksg/hamilton#readme> for more
-- information on usage!
--

module Numeric.Hamilton
  ( -- * Systems and states
    -- ** Systems
    System
  , mkSystem
  , mkSystem'
  , underlyingPos
    -- ** States
  , Config(..)
  , Phase(..)
  , toPhase
  , fromPhase
    -- * State functions
  , momenta
  , velocities
  , keC
  , keP
  , pe
  , lagrangian
  , hamiltonian
  , hamEqs
    -- * Simulating hamiltonian dynamics
    -- ** Over phase space
  , stepHam
  , evolveHam
  , evolveHam'
    -- ** Over configuration space
    -- | Convenience wrappers over the normal phase-space
    -- steppers/simulators that allow you to provide input and expect
    -- output in configuration space instead of in phase space.  Note that
    -- the simulation itself still runs in phase space, so these all
    -- require conversions to and from phase space under the hood.
  , stepHamC
  , evolveHamC
  , evolveHamC'
  ) where

import           Control.Monad
import           Data.Bifunctor
import           Data.Foldable
import           Data.Kind
import           Data.Maybe
import           Data.Proxy
import           Data.Type.Equality
import           GHC.Generics                        (Generic)
import           GHC.TypeLits
import           GHC.TypeLits.Compare
import           Numeric.AD
import           Numeric.GSL.ODE
import           Numeric.LinearAlgebra.Static        as H
import           Numeric.LinearAlgebra.Static.Vector
import qualified Data.Vector.Generic.Sized           as VG
import qualified Data.Vector.Sized                   as V
import qualified Numeric.LinearAlgebra               as LA

-- | Represents the full state of a system of @n@ generalized coordinates
-- in configuration space (informally, "positions and velocities")
--
-- A configuration space representaiton is more directly "physically
-- meaningful" and intuitive/understandable to humans than a phase space
-- representation.  However, it's much less mathematically ideal to work
-- with because of the lack of some neat underlying symmetries.
--
-- You can convert a @'Config' n@ into a @'Phase' n@ (convert from
-- configuration space to phase space) for a given system with 'toPhase'.
-- This allows you to state your system in configuration space and then
-- convert it to phase space before handing it off to the hamiltonian
-- machinery.
data Config :: Nat -> Type where
    Cfg :: { -- | The current values ("positions") of each of the @n@
             -- generalized coordinates
             cfgPositions :: !(R n)
             -- | The current rate of changes ("velocities") of each of the
             -- @n@ generalized coordinates
           , cfgVelocities :: !(R n)
           }
        -> Config n
  deriving (Generic)

deriving instance KnownNat n => Show (Config n)

-- | Represents the full state of a system of @n@ generalized coordinates
-- in phase space (informally, "positions and momentums").
--
-- Phase space representations are much nicer to work with mathematically
-- because of some neat underlying symmetries.  For one, positions and
-- momentums are "interchangeable" in a system; if you swap every
-- coordinate's positions with their momentums, and also swap them in the
-- equations of motions, you get the same system back.  This isn't the case
-- with configuration space representations.
--
-- A hamiltonian simulation basically describes the trajectory of each
-- coordinate through phase space, so this is the /state/ of the
-- simulation.  However, configuration space representations are much more
-- understandable to humans, so it might be useful to give an initial state
-- in configuration space using 'Config', and then convert it to a 'Phase'
-- with 'toPhase'.
data Phase :: Nat -> Type where
    Phs :: { -- | The current values ("positions") of each of the @n@
             -- generalized coordinates.
             phsPositions :: !(R n)
             -- | The current conjugate momenta ("momentums") to each of
             -- the @n@ generalized coordinates
           , phsMomenta :: !(R n)
           }
        -> Phase n
  deriving (Generic)

deriving instance KnownNat n => Show (Phase n)

-- | Represents a physical system in which physics happens.  A @'System'
-- m n@ is a system whose state described using @n@ generalized coordinates
-- (an "@n@-dimensional" system), where the underlying cartesian coordinate
-- space is @m@-dimensional.
--
-- For the most part, you are supposed to be able to ignore @m@.  @m@ is
-- only provided because it's useful when plotting/drawing the system with
-- a given state back in rectangular coordinates. (The only function that
-- use the @m@ at the moment is 'underlyingPos')
--
-- A @'System' m n@'s state is described using a @'Config' n@ (which
-- describes the system in configuration space) or a @'Phase' n@ (which
-- describes the system in phase space).
data System :: Nat -> Nat -> Type where
    Sys :: { _sysInertia       :: R m
           , _sysCoords        :: R n -> R m
           , _sysJacobian      :: R n -> L m n
           , _sysHessian       :: R n -> V.Vector n (L m n)
           , _sysPotential     :: R n -> Double
           , _sysPotentialGrad :: R n -> R n
           }
        -> System m n

-- | Converts the position of generalized coordinates of a system to the
-- coordinates of the system's underlying cartesian coordinate system.
-- Useful for plotting/drawing the system in cartesian space.
underlyingPos
    :: System m n
    -> R n
    -> R m
underlyingPos = _sysCoords

-- | The potential energy of a system, given the position in the
-- generalized coordinates of the system.
pe  :: System m n
    -> R n
    -> Double
pe = _sysPotential

vec2l
    :: (KnownNat m, KnownNat n)
    => V.Vector m (V.Vector n Double)
    -> L m n
vec2l = rowsL . fmap gvecR

-- | Create a system with @n@ generalized coordinates by describing its
-- coordinate space (by a function from the generalized coordinates to the
-- underlying cartesian coordinates), the inertia of each of those
-- underlying coordinates, and the pontential energy function.
--
-- The potential energy function is expressed in terms of the genearlized
-- coordinate space's positions.
mkSystem
    :: forall m n. (KnownNat m, KnownNat n)
    => R m      -- ^ The "inertia" of each of the @m@ coordinates
                -- in the underlying cartesian space of the system.  This
                -- should be mass for linear coordinates and rotational
                -- inertia for angular coordinates.
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
                -- ^ Conversion function to convert points in the
                -- generalized coordinate space to the underlying cartesian
                -- space of the system.
    -> (forall a. RealFloat a => V.Vector n a -> a)
                -- ^ The potential energy of the system as a function of
                -- the generalized coordinate space's positions.
    -> System m n
mkSystem m f u = Sys
    { _sysInertia       =                     m
    , _sysCoords        = gvecR             . f           . grVec
    , _sysJacobian      = tr   . vec2l      . jacobianT f . grVec
    , _sysHessian       = tr2  . fmap vec2l . hessianF f  . grVec
    , _sysPotential     =                     u           . grVec
    , _sysPotentialGrad = gvecR             . grad u      . grVec
    }
  where
    tr2 :: forall o. (KnownNat n, KnownNat o)
        => V.Vector m (L n o)
        -> V.Vector n (L m o)
    tr2 = fmap rowsL . traverse lRows
    {-# INLINE tr2 #-}

-- | Convenience wrapper over 'mkSystem' that allows you to specify the
-- potential energy function in terms of the underlying cartesian
-- coordinate space.
mkSystem'
    :: forall m n. (KnownNat m, KnownNat n)
    => R m      -- ^ The "inertia" of each of the @m@ coordinates
                -- in the underlying cartesian space of the system.  This
                -- should be mass for linear coordinates and rotational
                -- inertia for angular coordinates.
    -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a)
                -- ^ Conversion function to convert points in the
                -- generalized coordinate space to the underlying cartesian
                -- space of the system.
    -> (forall a. RealFloat a => V.Vector m a -> a)
                -- ^ The potential energy of the system as a function of
                -- the underlying cartesian coordinate space's positions.
    -> System m n
mkSystem' m f u = mkSystem m f (u . f)

-- | Compute the generalized momenta conjugate to each generalized
-- coordinate of a system by giving the configuration-space state of the
-- system.
--
-- Note that getting the momenta from a @'Phase' n@ involves just using
-- 'phsMomenta'.
momenta
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> R n
momenta Sys{..} Cfg{..} = tr j #> diag _sysInertia #> j #> cfgVelocities
  where
    j = _sysJacobian cfgPositions

-- | Convert a configuration-space representaiton of the state of the
-- system to a phase-space representation.
--
-- Useful because the hamiltonian simulations use 'Phase' as its working
-- state, but 'Config' is a much more human-understandable and intuitive
-- representation.  This allows you to state your starting state in
-- configuration space and convert to phase space for your simulation to
-- use.
toPhase
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Phase n
toPhase s = Phs <$> cfgPositions <*> momenta s

-- | The kinetic energy of a system, given the system's state in
-- configuration space.
keC :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Double
keC s = do
    vs <- cfgVelocities
    ps <- momenta s
    return $ (vs <.> ps) / 2

-- | The Lagrangian of a system (the difference between the kinetic energy
-- and the potential energy), given the system's state in configuration
-- space.
lagrangian
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Config n
    -> Double
lagrangian s = do
    t <- keC s
    u <- pe s . cfgPositions
    return (t - u)

-- | Compute the rate of change of each generalized coordinate by giving
-- the state of the system in phase space.
--
-- Note that getting the velocities from a @'Config' n@ involves just using
-- 'cfgVelocities'.
velocities
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> R n
velocities Sys{..} Phs{..} = inv jmj #> phsMomenta
  where
    j   = _sysJacobian phsPositions
    jmj = tr j H.<> diag _sysInertia H.<> j

-- | Invert 'toPhase' and convert a description of a system's state in
-- phase space to a description of the system's state in configuration
-- space.
--
-- Possibly useful for showing the phase space representation of a system's
-- state in a more human-readable/human-understandable way.
fromPhase
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Config n
fromPhase s = Cfg <$> phsPositions <*> velocities s

-- | The kinetic energy of a system, given the system's state in
-- phase space.
keP :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Double
keP s = do
    ps <- phsMomenta
    vs <- velocities s
    return $ (vs <.> ps) / 2

-- | The Hamiltonian of a system (the sum of kinetic energy and the
-- potential energy), given the system's state in phase space.
hamiltonian
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> Double
hamiltonian s = do
    t <- keP s
    u <- pe s . phsPositions
    return (t + u)

-- | The "hamiltonian equations" for a given system at a given state in
-- phase space.  Returns the rate of change of the positions and
-- conjugate momenta, which can be used to progress the simulation through
-- time.
--
-- Computed using the maths derived in
-- <https://blog.jle.im/entry/hamiltonian-dynamics-in-haskell.html>.
hamEqs
    :: (KnownNat m, KnownNat n)
    => System m n
    -> Phase n
    -> (R n, R n)
hamEqs Sys{..} Phs{..} = (dHdp, -dHdq)
  where
    mm   = diag _sysInertia
    j    = _sysJacobian phsPositions
    trj  = tr j
    jmj  = trj H.<> mm H.<> j
    ijmj = inv jmj
    dTdq = gvecR
         . flip fmap (_sysHessian phsPositions) $ \djdq ->
             -phsMomenta <.> ijmj #> trj #> mm #> djdq #> ijmj #> phsMomenta
    dHdp = ijmj #> phsMomenta
    dHdq = dTdq + _sysPotentialGrad phsPositions

-- | Step a system through phase space over over a single timestep.
stepHam
    :: forall m n. (KnownNat m, KnownNat n)
    => Double           -- ^ timestep to step through
    -> System m n       -- ^ system to simulate
    -> Phase n          -- ^ initial state, in phase space
    -> Phase n
stepHam r s p = evolveHam @m @n @2 s p (V.fromTuple (0, r))
                  `V.unsafeIndex` 1

-- | Evolve a system using a hamiltonian stepper, with the given initial
-- phase space state.
--
-- Desired solution times provided as a list instead of a sized 'V.Vector'.
-- The output list should be the same length as the input list.
evolveHam'
    :: forall m n. (KnownNat m, KnownNat n)
    => System m n  -- ^ system to simulate
    -> Phase n     -- ^ initial state, in phase space
    -> [Double]    -- ^ desired solution times
    -> [Phase n]
evolveHam' _ _ [] = []
evolveHam' s p0 ts = V.withSizedList (toList ts') $ \(v :: V.Vector s Double) ->
                       case Proxy @2 %<=? Proxy @s of
                         LE Refl -> (if l1 then tail else id)
                                  . toList
                                  $ evolveHam s p0 v
                         NLE{}   -> error "evolveHam': Internal error"
  where
    (l1, ts') = case ts of
      [x] -> (True , [0,x])
      _   -> (False, ts   )

-- | Evolve a system using a hamiltonian stepper, with the given initial
-- phase space state.
evolveHam
    :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s)
    => System m n           -- ^ system to simulate
    -> Phase n              -- ^ initial state, in phase space
    -> V.Vector s Double    -- ^ desired solution times
    -> V.Vector s (Phase n)
evolveHam s p0 ts = fmap toPs . fromJust . V.fromList . LA.toRows
                  $ odeSolveV RKf45 hi eps eps (const f) (fromPs p0) ts'
  where
    hi  = (V.unsafeIndex ts 1 - V.unsafeIndex ts 0) / 100
    eps = 1.49012e-08
    f :: LA.Vector Double -> LA.Vector Double
    f   = uncurry (\p m -> LA.vjoin [p,m])
        . join bimap extract . hamEqs s . toPs
    ts' = VG.fromSized . VG.convert $ ts
    n = fromInteger $ natVal (Proxy @n)
    fromPs :: Phase n -> LA.Vector Double
    fromPs p = LA.vjoin . map extract $ [phsPositions p, phsMomenta p]
    toPs :: LA.Vector Double -> Phase n
    toPs v = Phs pP pM
      where
        Just [pP, pM] = traverse create . LA.takesV [n, n] $ v

-- | A convenience wrapper for 'evolveHam'' that works on configuration
-- space states instead of phase space states.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the inputs
-- and outputs.
evolveHamC'
    :: forall m n. (KnownNat m, KnownNat n)
    => System m n       -- ^ system to simulate
    -> Config n         -- ^ initial state, in configuration space
    -> [Double]         -- ^ desired solution times
    -> [Config n]
evolveHamC' s c0 = fmap (fromPhase s) . evolveHam' s (toPhase s c0)

-- | A convenience wrapper for 'evolveHam' that works on configuration
-- space states instead of phase space states.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the inputs
-- and outputs.
evolveHamC
    :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s)
    => System m n           -- ^ system to simulate
    -> Config n             -- ^ initial state, in configuration space
    -> V.Vector s Double    -- ^ desired solution times
    -> V.Vector s (Config n)
evolveHamC s c0 = fmap (fromPhase s) . evolveHam s (toPhase s c0)

-- | Step a system through configuration space over over a single timestep.
--
-- Note that the simulation itself still runs in phase space; this function
-- just abstracts over converting to and from phase space for the input
-- and output.
stepHamC
    :: forall m n. (KnownNat m, KnownNat n)
    => Double           -- ^ timestep to step through
    -> System m n       -- ^ system to simulate
    -> Config n         -- ^ initial state, in phase space
    -> Config n
stepHamC r s = fromPhase s . stepHam r s . toPhase s


-- |
-- = Nonlinear transient computing
--
-- This module was developed as a part of author's PhD project:
-- https://www.researchgate.net/project/Theory-and-Modeling-of-Complex-Nonlinear-Delay-Dynamics-Applied-to-Neuromorphic-Computing
--

{-# LANGUAGE BangPatterns #-}

module RC.NTC
  ( new
  , learn
  , learnClassifier
  , predict
  , par0
  , NTCParameters (..)
  ) where

import           Numeric.LinearAlgebra
import           System.Random ( StdGen
                               , split
                               )
import qualified Data.List as List
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as VM
import           System.IO.Unsafe ( unsafePerformIO )
import qualified Learning
import qualified Numeric.DDE.Model as DDEModel

import           RC.NTC.Types
import qualified RC.NTC.Reservoir as Reservoir
import qualified RC.Helpers as H

-- | An untrained NTC network
new
  :: StdGen
  -> NTCParameters
  -> (Int, Int, Int)
  -- ^ Input dimension, network nodes, and output dimension
  -> NTC
new g par (ind, nodes, out) =
  let iwgen = _inputWeightsGenerator par
      iw = iwgen g (nodes, ind) (_inputWeightsRange par)
      ntc = NTC { _inputWeights = iw
                , _reservoir = _reservoirModel par
                , _outputWeights = Nothing
                , _par = par
                }
  in ntc

-- | Default NTC parameters
par0 :: NTCParameters
par0 = Par
  { _preprocess = id
  , _inputWeightsGenerator = H.randMatrix
  , _postprocess = H.addBiases  -- Usually `id` will work
  , _inputWeightsRange = undefined  -- To be manually set, e.g. (-1, 1)
  , _reservoirModel = Reservoir.genReservoir p
  }
  where
    filt' = DDEModel.BandpassFiltering {
              DDEModel._tau = 7.8125e-3
            , DDEModel._theta = recip 0.34375
            }

    p = DDEModel.RC { DDEModel._filt = filt'
                    , DDEModel._rho = 3.25
                    , DDEModel._fnl = H.hsigmoid (1.09375, 1.5, 0.0)
                    }

-- | Nonlinear transformation performed by an NTC network
forwardPass :: NTC  -- ^ NTC network
            -> Matrix Double  -- ^ Input information
            -> Matrix Double
forwardPass NTC { _par = Par { _preprocess = prep, _postprocess = post }
                , _inputWeights = iw
                , _reservoir = Reservoir res
                } !sample =
  let pipeline = post. res. (iw <>). prep
  in pipeline sample

-- | NTC training: learn the readout weights offline
learn
  :: NTC
  -> Int
  -- ^ Discard the first N points
  -> Matrix Double
  -- ^ Input matrix of features rows and observations columns
  -> Matrix Double
  -- ^ Desired output matrix of observations columns
  -> Either String NTC
learn ntc forgetPts inp out = ntc'
  where
    state' = forwardPass ntc inp ?? (All, Drop forgetPts)
    teacher' = out ?? (All, Drop forgetPts)
    ntc' = case Learning.learn' state' teacher' of
      Nothing -> Left "Cannot create a readout matrix"
      w -> Right $ ntc { _outputWeights = w }

_concatForwardPass :: NTC
                   -> Int
                   -- ^ Number of samples in the list (for
                   -- memory allocation)
                   -> [Matrix Double]
                   -- ^ List of samples
                   -> Matrix Double
_concatForwardPass ntc m (state0_:samples') = unsafePerformIO $ do
      -- Detect the postprocessing output dimension `nodes`.
      -- Alternatively, use Static shapes from hmatrix.
      let state0 = forwardPass ntc state0_
          nodes = rows state0
          state0' = H.flatten' state0

      -- Allocate memory for a mutable vector
      v <- VM.new (nodes * m)

      -- Copy the first computed state0'
      copy state0' v 0 0 nodes

      let foldA _ [] = return ()
          foldA f ((y, i):ys) = do
            let v0 = f y
            copy v0 v 0 (i * nodes) ((i + 1) * nodes)
            foldA f ys

      -- NB: Does compiler know (H.flatten'. H.unflatten') == id?
      -- If not, consider refactoring genReservoir and forwardPass
      foldA (H.flatten'. forwardPass ntc) $ zip samples' [1..m - 1]

      processed' <- V.unsafeFreeze v
      let processed = H.unflatten' nodes processed'
      return processed
  where
    copy !v0 !v !i0 !i !k
      | i < k = do
        VM.unsafeWrite v i (v0 V.! i0)
        copy v0 v (i0 + 1) (i + 1) k
      | otherwise = return ()

-- | NTC training specific to classification task.
-- The readout weights are learned offline.
--
-- Alternatively, one could use `learn` function. However,
-- to make sure the training samples do not mix in the reservoir RNN,
-- a significant padding of zeros would be needed. Although that
-- way directly corresponds to a physical experiment, that would be not
-- memory-efficient on a computer.
learnClassifier
  :: NTC
  -- ^ NTC network
  -> [Int]
  -- ^ Target labels
  -> Int
  -- ^ Number of inputs (for memory allocation)
  -> [Matrix Double]
  -- ^ Training inputs
  -> Either String (Learning.Classifier Int)
learnClassifier ntc labels m samples =
  case Learning.learn' state teacher' of
    Nothing -> Left "Cannot create a readout matrix"
    Just w -> let clf = Learning.winnerTakesAll w klasses. forwardPass ntc
           in Right (Learning.Classifier clf)
  where
    klasses = fromList. List.sort. List.nub $ labels
    klassesNo = V.length klasses
    teacher' = concatHor. map (flip (Learning.teacher klassesNo) 1) $ labels

    -- Alternatively, a streaming interface might be a solution
    -- https://hackage.haskell.org/package/pipes-4.3.7/docs/Pipes-Tutorial.html
    state = _concatForwardPass ntc m samples

-- | Horizontal matrix concatenation, alternative to fromBlocks
concatHor :: Element a => [Matrix a] -> Matrix a
concatHor ms@(m:_) = foldr (|||) (zeroCols (rows m)) ms
-- concatHor = concatMapHor id
{-# SPECIALIZE concatHor :: [Matrix Double] -> Matrix Double #-}

concatMapHor
  :: Element b =>
     (Matrix a -> Matrix b) -> [Matrix a] -> Matrix b
concatMapHor f ms@(m:_) = foldr (\a b -> f a ||| b) (zeroCols (rows m)) ms
{-# SPECIALIZE concatMapHor :: (Matrix Double -> Matrix Double) -> [Matrix Double] -> Matrix Double #-}

-- | Matrix with zero elements
zeroCols :: V.Storable a => Int -> Matrix a
zeroCols rows' = (rows'><0) []
{-# SPECIALIZE zeroCols :: Int -> Matrix Double #-}

-- | Run prediction using a "clean" (uninitialized) reservoir and then
-- forget the reservoir's state.
-- This can be used for both forecasting and classification tasks.
predict :: NTC  -- ^ Trained network
        -> Int  -- ^ Washout (forget) points
        -> Matrix Double
        -- ^ Input matrix where measurements are columns and features are rows
        -> Either String (Matrix Double)
        -- ^ Either error string or predicted output
predict ntc@NTC { _outputWeights = ow
                } forgetPts inp =
  case ow of
    Nothing -> Left "Please train the NTC first"
    Just w -> let y = forwardPass ntc inp
                  y2 = y ?? (All, Drop forgetPts)
                  prediction = w <> y2
              in Right prediction


module Q.Options.ImpliedVol.Normal where
import           Data.Default.Class (Default (..))
import           Numeric.IEEE (epsilon)
import           Numeric.Natural (Natural)
import           Numeric.RootFinding as R (RiddersParam (RiddersParam),
                                           Root (NotBracketed, Root, SearchFailed), Tolerance,
                                           ridders)
import           Q.Options.Bachelier (Bachelier (Bachelier), euOption)

import           Statistics.Distribution (cumulative, density)
import           Statistics.Distribution.Normal (standard)
import Q.Options

-- | Method to use to calculate the normal implied vol.
data Method =
    Jackel        -- ^ Jackel analytical formula approximation.
  | ChoKimKwak    -- ^ J. Choi, K kim, and M. Kwak (2009)
  -- | Numerical root finding. Currently Ridders is used.
  | RootFinding
        Natural                  -- ^ Maximum number of iterations.
        R.Tolerance              -- ^ Tolerance (relative or absolute)
        (Double, Double, Double) -- ^ Triple of @(low bound, initial
                                 --   guess, upper bound)@. If initial
                                 --   guess if out of bracket middle
                                 --   of bracket is taken as.

instance Default Method where
  def = Jackel

-- | Default method implementation of 'euImpliedVolWith' using 'Jackel'.
euImpliedVol :: OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVol = euImpliedVolWith def

-- | Calcualte the bachelier option implied vol of a european option.
--
-- If the options premium does not have time value @'hasTimeValue'@ return 0.
euImpliedVolWith :: Method -> OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVolWith m cp f k t r p
  | hasTimeValue cp f k p df = euImpliedVolWith' m cp f k t r p
  | otherwise                = Vol $ 0
  where df = discountFactor t r

euImpliedVolWith' :: Method-> OptionType-> Forward-> Strike-> YearFrac-> Rate-> Premium-> Vol
euImpliedVolWith' Jackel cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p)
  -- Case where interest rate is not 0 we need undiscount. The paper is written
  -- for the undiscounted Bachelier option prices.
  | r /= 0
    = euImpliedVol cp (Forward f) (Strike k) (YearFrac t) (Rate 0) (Premium (p/df))
  -- Case of ATM. Solve directly.
  | abs (k - f) <= epsilon = Vol $ p * sqrt2Pi / (sqrt t)
  -- Case of ITM option. Calcualte vol of the out of the money option with Put-Call-Parity.
  | phiStarTilde >= 0
    = euImpliedVol (succ cp) (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p')
  -- General case for an out of the money option.
  | otherwise  = let
      ẋ      = if phiStarTilde < c then
                 let g = 1 / (phiStarTilde - 0.5)
                     ξ = (0.032114372355 - (g**2)*(0.016969777977 - (g**2)*(2.6207332461E-3-(9.6066952861E-5)*g**2)))
                         /
                         (1-(g**2)*(0.6635646938 - (g**2)*(0.14528712196 - 0.010472855461*g**2)))
                 in g * (1 / sqrt2Pi + ξ*g**2)
               else
                 let h = sqrt $ (-log (-phiStarTilde))
                 in (9.4883409779-h*(9.6320903635-h*(0.58556997323 + 2.1464093351*h)))
                    /
                    (1-h*(0.65174820867 + h*(1.5120247828 + 6.6437847132E-5*h)))
      c       = (-0.001882039271)
      x       = ẋ + (3*q * ẋ * ẋ * (2 - q * ẋ * (2 + ẋ*ẋ)))
                    /
                    (6 + q*ẋ * ((-12) + ẋ *(6*q + ẋ * ((-6)*q*ẋ*(3+ẋ*ẋ)))))
      phiXBarTilde = (cumulative standard ẋ) + (density standard ẋ)/ẋ
      q       =  (phiXBarTilde-phiStarTilde)/ (density standard ẋ)
    in Vol $ (abs (k - f)) / (abs (x * sqrt t))
  where phiStarTilde = negate $ (abs (p - (max (theta * (f - k)) 0))) / (abs (k - f))
        theta        = if cp == Call then 1 else -1
        phiTilde     = (-theta) * p / (k - f)
        p'           = cpi * df * (f - k) + p
        cpi          = fromIntegral $ fromEnum cp --call put indicartor.
        df           = exp $ (-r) * t
        sqrt2Pi      = 2.506628274631000

euImpliedVolWith' ChoKimKwak cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p) =
  let df              = exp $ (-r) * t
      forwardPremium  = p / df
      straddlePremium = case cp of Call -> 2 * forwardPremium - (f - k)
                                   Put  -> 2 * forwardPremium + (f - k)
      nu'             = (f - k) / straddlePremium
      nu              = max (-1 + epsilon) (min nu' (1 - epsilon))
      eta             | abs nu < sqrtEpsilon = 1
                      | otherwise            = nu / (atanh nu)
      heta            = h eta
  in Vol $ sqrt (pi / (2 * t)) * straddlePremium * heta

euImpliedVolWith' (RootFinding maxIter tol bracket)  cp (Forward forward) k t r (Premium p) =
  let f vol        = p' - p  where
        (Premium p') = vPremium $ euOption b cp k
        b            = Bachelier t (Forward forward) r (Vol vol)
      (lb, _, ub) = bracket
      root = ridders (RiddersParam (fromEnum maxIter) tol) (lb, ub) f
  in case root of (Root vol)   -> Vol vol
                  NotBracketed -> error "not bracketed"
                  SearchFailed -> error "search failed"

sqrtEpsilon :: Double
sqrtEpsilon = sqrt epsilon

h :: Floating p => p -> p
h eta = sqrt(eta) * (num / den) where
  a0          = 3.994961687345134e-1
  a1          = 2.100960795068497e+1
  a2          = 4.980340217855084e+1
  a3          = 5.988761102690991e+2
  a4          = 1.848489695437094e+3
  a5          = 6.106322407867059e+3
  a6          = 2.493415285349361e+4
  a7          = 1.266458051348246e+4

  b0          = 1.000000000000000e+0
  b1          = 4.990534153589422e+1
  b2          = 3.093573936743112e+1
  b3          = 1.495105008310999e+3
  b4          = 1.323614537899738e+3
  b5          = 1.598919697679745e+4
  b6          = 2.392008891720782e+4
  b7          = 3.608817108375034e+3
  b8          = -2.067719486400926e+2
  b9          = 1.174240599306013e+1

  num = a0 + eta * (a1 + eta * (a2 + eta * (a3 + eta * (a4 + eta * (a5 + eta * (a6 + eta * a7))))))
  den = b0 + eta * (b1 + eta * (b2 + eta * (b3 + eta * (b4 + eta * (b5 + eta * (b6 + eta * (b7 + eta * (b8 + eta * b9))))))))


{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}

module Language.Grappa.Test.Inference.GradientDescent where

import           Control.Monad ( forM, forM_, replicateM )
import qualified Numeric.LinearAlgebra as La
import qualified Statistics.Distribution as S
import qualified Statistics.Distribution.StudentT as S
import qualified System.Random.MWC as MWC

import           Test.Hspec
import qualified Test.QuickCheck as Qc
import qualified Test.QuickCheck.Monadic as Qc
import qualified Test.QuickCheck.Property as Qc

import           Language.Grappa.Distribution
import           Language.Grappa.Rand.MWCRandM
import           Language.Grappa.Inference.GradientDescent

tests :: Spec
tests = do
  forM_ [1..8] $ \ x ->
    it ("should pass this QC test with k=" ++ show x) $
      prop_evalGradLogMVNormalDensity_correct x

-- | QuickCheck property asserting 'compareGrads' for given mv normal
-- dimension @k@.
prop_evalGradLogMVNormalDensity_correct :: Int -> Qc.Property
prop_evalGradLogMVNormalDensity_correct k =
  -- forAll' "k" Qc.arbitrary $ \(Qc.Positive (k :: Int)) ->
  forAll' "xList" (Qc.vector k) $ \xList ->
  forAll' "muList" (Qc.vector k) $ \muList ->
  forAll' "cList" (genCList k) $ \cList ->
  Qc.monadicIO $ do
    deltas <- Qc.run $ compareGrads k (i xList) (i $ muList ++ cList)
    Qc.monitor (Qc.counterexample $ "(Index, Delta): " ++ show (argMax deltas))
    Qc.assert $ all isSmall deltas
  where
    -- Return index of max entry, and max entry.
    argMax :: Ord a => [a] -> (Int, a)
    argMax (x:xs) = go 1 (0,x) xs
      where
        go _ p [] = p
        go j (i,x) (x':xs) = go (j+1) (if x < x' then (j, x') else (i, x)) xs

    i :: [Integer] -> [Double]
    i = map fromInteger
    -- Generate a list representation of a Cholesky factor. The
    -- strategy is to recursively generate row tails from the diagonal
    -- entry onwards, with non-zero diagonal entries and arbitrary
    -- other entries.
    genCList :: Int -> Qc.Gen [Integer]
    genCList k = do
      rows <- forM [k,k-1..1] $ \len -> do
        Qc.Positive diagEntry <- Qc.arbitrary
        rest <- Qc.vectorOf (len-1) Qc.arbitrary
        return $ diagEntry : rest
      return $ concat rows

    isSmall delta = abs delta <= 0.001

    -- A version of @Qc.forAll@ with named quantifiers. This makes the
    -- counterexamples much easier to read.
    forAll' name gen pf =
      -- The library provided @Qc.forAll@ that I adapted to create
      -- this @forAll'@ has a @Qc.again@ here. But I don't understand
      -- why I need it, and more importantly, I get an import error,
      -- so leaving it out ...
      --
      -- Qc.again $
      Qc.MkProperty $
      gen >>= \x ->
        Qc.unProperty $
        Qc.counterexample (name++": "++show x) (pf x)

-- | Check the QuickCheck props @numTrials@ many times.
--
-- For larger @k@ this fails, but I think it's due to numerical
-- instability. For small @k@, e.g. less than or equal to 5, we can
-- pass hundreds of thousands of tests. For large @k@, e.g. @k = 30@,
-- we fail pretty fast.
checkProps :: Int -> Int -> IO ()
checkProps k numTrials = do
  qc $ prop_evalGradLogMVNormalDensity_correct k
  where
    qc = Qc.quickCheckWith (Qc.stdArgs { Qc.maxSuccess = numTrials })

-- | For example:
--
-- > runGdTest (gd2 0.001) (gdTest3 4 [1,42])
runGdTest :: Gd Rand -> GdTest Rand -> IO (Double, [Double])
runGdTest gd GdTest{..} = do
  -- Use 'MWC.create' instead if you want repeatable/predictable
  -- results.
  g <- MWC.createSystemRandom
  lambda <- runRand g $ gd _gdParams
  return (_estimateError lambda, lambda)

----------------------------------------------------------------
-- * Tests
----------------------------------------------------------------

data GdTest m = GdTest
  { _gdParams      :: GdParams m
  , _estimateError :: [Double] -> Double
  }

-- | Test for objective function @f x = (x - 42)**2@, which is
-- minimized at @x = 42@.
gdTest1 :: SampleableIn m Normal => GdTest m
gdTest1 = GdTest
  { _gdParams = GdParams
    { _projectLambda = const id
    , _sampleLambda = defaultSampleLambda 1
    , _evalObj = \[x] -> return $ (x - center!!0)**2
    , _evalGradObj = \[x] -> return [2 * (x - center!!0)]
    }
  , _estimateError = \lambda -> ell1 $ lambda `subV` center
  }
  where
    center = [42]

-- | Test for objective function @f [x,y] = ([x,y] - [42,10^6])**2@,
-- which is minimized at @[42,10^6]@.
gdTest2 :: SampleableIn m Normal => GdTest m
gdTest2 = GdTest
  { _gdParams = GdParams
    { _projectLambda = const id
    , _sampleLambda = defaultSampleLambda 2
    , _evalObj = \[x,y] -> return $ sqrt ( (x - center!!0)**2
                                         + (y - center!!1)**2 )
    , _evalGradObj = \[x,y] -> return [ 2 * (x - center!!0)
                                      , 2 * (y - center!!1) ]
    }
  , _estimateError = \lambda -> ell1 $ lambda `subV` center
  }
  where
    center = [ 42, 1000000 ]

-- | Test for objective function @f x = (x - center)**k@ where
-- @center@ is an arbitrary vector. The dimension of @center@
-- determines the dimension of @x@. The objective is minimized at
-- @center@ for even powers @k@, and has no minimum for odd powers
-- @k@. You're on your own for fractional powers.
gdTest3 :: SampleableIn m Normal => Double -> [Double] -> GdTest m
gdTest3 k center = GdTest
  { _gdParams = GdParams
    { _projectLambda = const id
    , _sampleLambda = defaultSampleLambda (length center)
    , _evalObj = \lambda -> return $ (ell2 $ v lambda)**k
    , _evalGradObj = \lambda -> return
        [ k * (ell2 (v lambda))**(k-2) * vi
        | vi <- v lambda ]
    }
  , _estimateError = \lambda -> ell1 $ lambda `subV` center
  }
  where
    v lambda = lambda `subV` center

-- | Test for objective function
--
-- > f x = (x - center)**k -
-- >       amplitude * cos (2pi * |x - center| / period)
--
-- I.e. like 'gdTest3', but with an added noise factor of @- cos@ to
-- introduce infinitely many local minima. This is a better test of
-- the line search in 'gd2'. I expect that making the period and
-- amplitude large enough, and the center far enough from the origin,
-- would trick 'gd2' into a local minium, but I can't find parameters
-- that actually do this!
--
-- This objective is minimized at @center@ for even powers @k@, like
-- 'gdTest3'.
gdTest4 :: SampleableIn m Normal =>
  Double -> [Double] -> Double -> Double -> GdTest m
gdTest4 k center amplitude period = GdTest
  { _gdParams = GdParams
    { _projectLambda = const id
    , _sampleLambda = defaultSampleLambda (length center)
    , _evalObj = \lambda ->
        return $ (ell2 $ v lambda)**k -
                 amplitude * cos (2 * pi *
                                  (ell2 $ v lambda) / period)
    , _evalGradObj = \lambda -> return
        [ k * (ell2 (v lambda))**(k-2) * vi +
          sin (ell2 (v lambda)) *
            (2 * pi / period) *
            vi / ell2 (v lambda)
        | vi <- v lambda ]
    }
  , _estimateError = \lambda -> ell1 $ lambda `subV` center
  }
  where
    v lambda = lambda `subV` center

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


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Statistics.Sample.Histogram.MagnitudeSpec where

import qualified Data.Vector.Unboxed as U
import           Data.Monoid
import           Data.Functor.Identity
import           Statistics.Sample.Histogram.Magnitude
import           Test.Hspec
import           Test.QuickCheck
import           Test.QuickCheck.Checkers
import           Test.QuickCheck.Classes

main :: IO ()
main = hspec spec

spec :: Spec
spec = parallel . describe "Histogram" $ do

  it "calculates a meaningful histogram" $ do
    let x = foldHist 1 [-3, 1, 4, 6, 9, 9.9, 4.032, 9.2]
        result = U.fromList [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 2, 0, 1, 0, 0, 3]
    histBuckets x `shouldBe` result

  describe "calculates magnitude" $ do
    it "is -1" $ property $ \(MagN x) -> histMagnitude (fromList 1 [x]) == -1
    it "is 0" $ property $ \(Mag0 x) -> histMagnitude (fromList 1 [x]) == 0
    it "is 1" $ property $ \(Mag1 x) -> histMagnitude (fromList 1 [x]) == 1
    it "is 2" $ property $ \(Mag2 x) -> histMagnitude (fromList 1 [x]) == 2

  it "blends equal resolutions and magnitudes" $ do
    let x = foldHist 1 [1, 4, 6, 9]
        y = foldHist 1 [1, 6]
        result = U.fromList [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 1, 0, 2, 0, 0, 1]
    histBuckets (x <> y) `shouldBe` result

  it "scales non equal magnitudes to the higher magnitude" $ do
    let x = foldHist 1 [1, 4, 6, 9]
        y = foldHist 1 [134, 666]
        result = U.fromList [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 1, 0, 0, 0, 0, 1, 0, 0, 0]
    histMagnitude (x <> y) `shouldBe` 2
    histBuckets (x <> y) `shouldBe` result

  it "blurs non equal resolutions to the lower resolution" $ do
    let x = foldHist 1 [1, 4, 6, 9]
        y = foldHist 2 [1, 6.6]
        result = U.fromList [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 1, 0, 2, 0, 0, 1]
    histResolution (x <> y) `shouldBe` 1
    histBuckets (x <> y) `shouldBe` result

  it "blurs and scales non equal resolution and magnitude" $ do
    let x = foldHist 2 [1, 4, 6, 9]
        y = foldHist 3 [-17, 1, 66]
        vec = (U.replicate 100 0)
        result = U.accum (+) vec [(82, 1)]
              <> U.accum (+) vec [(1, 2), (4, 1), (6, 1), (9, 1), (66, 1)]
    histBuckets (x <> y) `shouldBe` result

  describe "produces keys for given magnitudes and scales" $ do
    it "2-1" $
      keys (fromList 2 [1]) `shouldBe` U.fromList (reverse [0, -0.1 .. -9.9 ] <> [0, 0.1 .. 9.9 :: Double])
    it "2-10" $
      keys (fromList 2 [10]) `shouldBe` U.fromList ([-99, -98 .. 0] <> [0, 1 .. 99 :: Double])
    it "1-9" $
      keys (fromList 1 [9]) `shouldBe` U.fromList ([-9, -8 .. 0] <> [0, 1 .. 9 :: Double])

  it "allows insertion" $ do
    let x = foldHist 1 [1, 4, 6, 9]
        result = U.fromList [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1]
    histBuckets (insert (-3) x) `shouldBe` result

  testBatch $ monoid (undefined :: Histogram Double)

  it "is commutable" $
    property $ isCommut ((<>) :: Histogram Double -> Histogram Double -> Histogram Double)

  it "has decreasing resolution on mappend" $
    property $ \(x :: Histogram Double, y :: Histogram Double) ->
      histResolution (x <> y) == min (histResolution x) (histResolution y)

  it "has increasing magnitude on mappend" $
    property $ \(x :: Histogram Double, y :: Histogram Double) ->
      histMagnitude (x <> y) == max (histMagnitude x) (histMagnitude y)

  it "is equivalent to mappend and insert" $
    property $ \(xs :: [Double]) ->
      foldHist 2 xs == foldMap (foldHist 2 . Identity) xs

testBatch :: TestBatch -> Spec
testBatch (batchName, tests) = describe ("laws for: " ++ batchName) $
    foldr (>>) (return ()) (map (uncurry it) tests)

instance (RealFrac a, Floating a, Arbitrary a) => Arbitrary (Histogram a) where
  arbitrary = do
    ys <- arbitrary
    x <- choose (0, 3)
    pure $ fromList (fromIntegral (x :: Int)) (getNonEmpty ys :: [a])

instance EqProp a => EqProp (Histogram a) where
  x =-= y = eq x y

newtype MagN = MagN Double
  deriving (Show, Ord, Eq, Num)

instance Arbitrary MagN where
  arbitrary = MagN <$> choose (0.1, 0.9)

newtype Mag0 = Mag0 Double
  deriving (Show, Ord, Eq, Num)

instance Arbitrary Mag0 where
  arbitrary = Mag0 <$> choose (1, 9.9)

newtype Mag1 = Mag1 Double
  deriving (Show, Ord, Eq, Num)

instance Arbitrary Mag1 where
  arbitrary = Mag1 <$> choose (10, 99.9)

newtype Mag2 = Mag2 Double
  deriving (Show, Ord, Eq, Num)

instance Arbitrary Mag2 where
  arbitrary = Mag2 <$> choose (100, 999.9)


-- Slope estimates

module Model.Slope
  ( estimateSlope
  , median
  , getSurroundingSlopes
  ) where

import           Data.Vector.Algorithms.Intro (partialSort)
import qualified Data.Vector.Unboxed          as V
import           Statistics.Test.Types
import           Statistics.Test.WilcoxonT    (wilcoxonMatchedPairTest)
import           Statistics.Types

import           Types.Indices

estimateSlope
  :: IVector Index1 Double
  -> IIntSet Index1
  -> Int
  -> Index1
  -> Double
estimateSlope ds badSegments radius i
  | V.length nonZeroSlopes < 5 = 0 -- p < 0.10 impossible with l.t. 5 slopes
  | pValue' < 0.10             = median slopes
  | otherwise                  = 0
  where
    slopes = V.fromList $ getSurroundingSlopes ds badSegments radius i
    nonZeroSlopes = V.filter (/= 0) slopes
    pairs = V.zip nonZeroSlopes (V.replicate (V.length nonZeroSlopes) 0)
    pValue' = pValue $ testSignificance
            $ wilcoxonMatchedPairTest SamplesDiffer pairs

-- Fails if vector is empty
median :: V.Vector Double -> Double
median v
  | odd n = v' V.! middleIndex
  | otherwise = mid (v' V.! pred middleIndex) (v' V.! middleIndex)
  where
    n = V.length v
    middle = n `quot` 2 + 1
    middleIndex = pred middle
    v' = V.modify (`partialSort` middle) v
    mid x y = (x+y)/2

getSurroundingSlopes
  :: IVector Index1 Double
  -> IIntSet Index1
  -> Int
  -> Index1
  -> [Double]
getSurroundingSlopes ds badSegments radius i =
  map (ivIndex ds) $ filter (not . flip iisMember badSegments) [i0..i1]
  where
    (i0, i1) = runIndexInterval $ iiBoundByIVector ds
             $ IndexInterval (iTranslate (-radius) i, iTranslate radius i)