{-# LANGUAGE BangPatterns , RankNTypes, GADTs, DataKinds #-}

{- | The 'Numerical.HBLAS.BLAS.Level2' module provides a fully general
yet type safe Level2 BLAS API.

When in doubt about the semantics of an operation,
consult your system's BLAS api documentation, or just read the documentation
for
<https://software.intel.com/sites/products/documentation/hpc/mkl/mklman/index.htm the Intel MKL BLAS distribution>

A few basic notes about how to invoke BLAS routines.

Many BLAS operations take one or more arguments of type 'Transpose'.
'Tranpose' has the following different constructors, which tell BLAS
routines what transformation to implicitly apply to an input matrix @mat@ with dimension @n x m@.

*  'NoTranspose' leaves the matrix @mat@ as is.

* 'Transpose' treats the @mat@ as being implicitly transposed, with dimension
    @m x n@. Entry @mat(i,j)@ being treated as actually being the entry
    @mat(j,i)@. For Real matrices this is also the matrix adjoint operation.
    ie @Tranpose(mat)(i,j)=mat(j,i)@

*  'ConjNoTranspose' will implicitly conjugate @mat@, which is a no op for Real ('Float' or 'Double') matrices, but for
'Complex Float' and 'Complex Double' matrices, a given matrix entry @mat(i,j)==x':+'y@
will be treated as actually being  @conjugate(mat)(i,j)=y':+'x@.

* 'ConjTranpose' will implicitly transpose and conjugate the input matrix.
ConjugateTranpose acts as matrix adjoint for both real and complex matrices.

The *gemm operations  work as follows (using 'sgemm' as an example):

* @'sgemm trLeft trRight alpha beta left right result'@, where @trLeft@ and @trRight@
are values of type 'Transpose' that respectively act on the matrices @left@ and @right@.

* the generalized matrix computation thusly formed can be viewed as being
@result = alpha * trLeft(left) * trRight(right) + beta * result@

the *gemv operations are akin to the *gemm operations, but with @right@ and @result@
being vectors rather than matrices.

the *trsv operations solve for @x@ in the equation @A x = y@ given @A@ and @y@.
The 'MatUpLo' argument determines if the matrix should be treated as upper or
lower triangular and 'MatDiag' determines if the triangular solver should treat
the diagonal of the matrix as being all 1's or not.  A general pattern of invocation
would be @'strsv' matuplo  tranposeMatA  matdiag  matrixA  xVector@.
A key detail to note is that the input vector is ALSO the result vector,
ie 'strsv' and friends updates the vector place.

-}

module Numerical.HBLAS.BLAS.Level2(
        sgbmv
        ,dgbmv
        ,cgbmv
        ,zgbmv

        ,sgemv
        ,dgemv
        ,cgemv
        ,zgemv

        ,sger
        ,dger
        ,cgerc
        ,zgerc
        ,cgeru
        ,zgeru

        ,chbmv
        ,zhbmv

        ,chemv
        ,zhemv

        ,cher
        ,zher
        ,cher2
        ,zher2

        ,chpmv
        ,zhpmv

        ,chpr
        ,zhpr
        ,chpr2
        ,zhpr2

        ,ssbmv
        ,dsbmv

        ,sspmv
        ,dspmv

        ,sspr
        ,dspr
        ,sspr2
        ,dspr2

        ,ssymv
        ,dsymv

        ,ssyr
        ,dsyr
        ,ssyr2
        ,dsyr2

        ,stbmv
        ,dtbmv
        ,ctbmv
        ,ztbmv

        ,stbsv
        ,dtbsv
        ,ctbsv
        ,ztbsv

        ,stpmv
        ,dtpmv
        ,ctpmv
        ,ztpmv

        ,stpsv
        ,dtpsv
        ,ctpsv
        ,ztpsv

        ,strmv
        ,dtrmv
        ,ctrmv
        ,ztrmv

        ,strsv
        ,dtrsv
        ,ctrsv
        ,ztrsv
            ) where

import Numerical.HBLAS.UtilsFFI
import Numerical.HBLAS.BLAS.FFI.Level2
import Numerical.HBLAS.BLAS.Internal.Level2
import Control.Monad.Primitive
import Data.Complex

sgbmv :: PrimMonad m => GbmvFun Float orient (PrimState m) m
sgbmv = gbmvAbstraction "sgbmv" cblas_sgbmv_safe cblas_sgbmv_unsafe (\x f -> f x)

dgbmv :: PrimMonad m => GbmvFun Double orient (PrimState m) m
dgbmv = gbmvAbstraction "dgbmv" cblas_dgbmv_safe cblas_dgbmv_unsafe (\x f -> f x)

cgbmv :: PrimMonad m => GbmvFun (Complex Float) orient (PrimState m) m
cgbmv = gbmvAbstraction "cgbmv" cblas_cgbmv_safe cblas_cgbmv_unsafe withRStorable_

zgbmv :: PrimMonad m => GbmvFun (Complex Double) orient (PrimState m) m
zgbmv = gbmvAbstraction "zgbmv" cblas_zgbmv_safe cblas_zgbmv_unsafe withRStorable_

sgemv :: PrimMonad m => GemvFun Float orient (PrimState m) m
sgemv = gemvAbstraction "sgemv" cblas_sgemv_safe cblas_sgemv_unsafe (flip ($))

dgemv :: PrimMonad m => GemvFun Double orient (PrimState m) m
dgemv = gemvAbstraction "dgemv" cblas_dgemv_safe cblas_dgemv_unsafe (flip ($))

cgemv :: PrimMonad m => GemvFun (Complex Float) orient (PrimState m) m
cgemv = gemvAbstraction "cgemv" cblas_cgemv_safe cblas_cgemv_unsafe withRStorable_

zgemv :: PrimMonad m => GemvFun (Complex Double) orient (PrimState m) m
zgemv = gemvAbstraction "zgemv" cblas_zgemv_safe cblas_zgemv_unsafe withRStorable_

sger :: PrimMonad m => GerFun Float orient (PrimState m) m
sger = gerAbstraction "sger" cblas_sger_safe cblas_sger_unsafe (\x f -> f x)

dger :: PrimMonad m => GerFun Double orient (PrimState m) m
dger = gerAbstraction "dger" cblas_dger_safe cblas_dger_unsafe (\x f -> f x)

cgerc :: PrimMonad m => GerFun (Complex Float) orient (PrimState m) m
cgerc = gerAbstraction "cgerc" cblas_cgerc_safe cblas_cgerc_unsafe withRStorable_

zgerc :: PrimMonad m => GerFun (Complex Double) orient (PrimState m) m
zgerc = gerAbstraction "zgerc" cblas_zgerc_safe cblas_zgerc_unsafe withRStorable_

cgeru :: PrimMonad m => GerFun (Complex Float) orient (PrimState m) m
cgeru = gerAbstraction "cgeru" cblas_cgeru_safe cblas_cgeru_unsafe withRStorable_

zgeru :: PrimMonad m => GerFun (Complex Double) orient (PrimState m) m
zgeru = gerAbstraction "zgeru" cblas_zgeru_safe cblas_zgeru_unsafe withRStorable_

chbmv :: PrimMonad m => HbmvFun (Complex Float) orient (PrimState m) m
chbmv = hbmvAbstraction "chbmv" cblas_chbmv_safe cblas_chbmv_unsafe withRStorable_

zhbmv :: PrimMonad m => HbmvFun (Complex Double) orient (PrimState m) m
zhbmv = hbmvAbstraction "zhbmv" cblas_zhbmv_safe cblas_zhbmv_unsafe withRStorable_

chemv :: PrimMonad m => HemvFun (Complex Float) orient (PrimState m) m
chemv = hemvAbstraction "chemv" cblas_chemv_safe cblas_chemv_unsafe withRStorable_

zhemv :: PrimMonad m => HemvFun (Complex Double) orient (PrimState m) m
zhemv = hemvAbstraction "zhemv" cblas_zhemv_safe cblas_zhemv_unsafe withRStorable_

cher :: PrimMonad m => HerFun Float (Complex Float) orient (PrimState m) m
cher = herAbstraction "cher" cblas_cher_safe cblas_cher_unsafe (\x f -> f x)

zher :: PrimMonad m => HerFun Double (Complex Double) orient (PrimState m) m
zher = herAbstraction "zher" cblas_zher_safe cblas_zher_unsafe (\x f -> f x)

cher2 :: PrimMonad m => Her2Fun (Complex Float) orient (PrimState m) m
cher2 = her2Abstraction "cher2" cblas_cher2_safe cblas_cher2_unsafe withRStorable_

zher2 :: PrimMonad m => Her2Fun (Complex Double) orient (PrimState m) m
zher2 = her2Abstraction "zher2" cblas_zher2_safe cblas_zher2_unsafe withRStorable_

chpmv :: PrimMonad m => HpmvFun (Complex Float) orient (PrimState m) m
chpmv = hpmvAbstraction "chpmv" cblas_chpmv_safe cblas_chpmv_unsafe withRStorable_

zhpmv :: PrimMonad m => HpmvFun (Complex Double) orient (PrimState m) m
zhpmv = hpmvAbstraction "zhpmv" cblas_zhpmv_safe cblas_zhpmv_unsafe withRStorable_

chpr :: PrimMonad m => HprFun Float (Complex Float) orient (PrimState m) m
chpr = hprAbstraction "chpr" cblas_chpr_safe cblas_chpr_unsafe (\x f -> f x)

zhpr :: PrimMonad m => HprFun Double (Complex Double) orient (PrimState m) m
zhpr = hprAbstraction "zhpr" cblas_zhpr_safe cblas_zhpr_unsafe (\x f -> f x)

chpr2 :: PrimMonad m => Hpr2Fun (Complex Float) orient (PrimState m) m
chpr2 = hpr2Abstraction "chpr2" cblas_chpr2_safe cblas_chpr2_unsafe withRStorable_

zhpr2 :: PrimMonad m => Hpr2Fun (Complex Double) orient (PrimState m) m
zhpr2 = hpr2Abstraction "zhpr2" cblas_zhpr2_safe cblas_zhpr2_unsafe withRStorable_

ssbmv :: PrimMonad m => SbmvFun Float orient (PrimState m) m
ssbmv = sbmvAbstraction "ssbmv" cblas_ssbmv_safe cblas_ssbmv_unsafe (\x f -> f x)

dsbmv :: PrimMonad m => SbmvFun Double orient (PrimState m) m
dsbmv = sbmvAbstraction "dsbmv" cblas_dsbmv_safe cblas_dsbmv_unsafe (\x f -> f x)

sspmv :: PrimMonad m => SpmvFun Float orient (PrimState m) m
sspmv = spmvAbstraction "sspmv" cblas_sspmv_safe cblas_sspmv_unsafe (\x f -> f x)

dspmv :: PrimMonad m => SpmvFun Double orient (PrimState m) m
dspmv = spmvAbstraction "dspmv" cblas_dspmv_safe cblas_dspmv_unsafe (\x f -> f x)

sspr :: PrimMonad m => SprFun Float orient (PrimState m) m
sspr = sprAbstraction "sspr" cblas_sspr_safe cblas_sspr_unsafe (\x f -> f x)

dspr :: PrimMonad m => SprFun Double orient (PrimState m) m
dspr = sprAbstraction "dspr" cblas_dspr_safe cblas_dspr_unsafe (\x f -> f x)

sspr2 :: PrimMonad m => Spr2Fun Float orient (PrimState m) m
sspr2 = spr2Abstraction "sspr2" cblas_sspr2_safe cblas_sspr2_unsafe (\x f -> f x)

dspr2 :: PrimMonad m => Spr2Fun Double orient (PrimState m) m
dspr2 = spr2Abstraction "dspr2" cblas_dspr2_safe cblas_dspr2_unsafe (\x f -> f x)

ssymv :: PrimMonad m => SymvFun Float orient (PrimState m) m
ssymv = symvAbstraction "ssymv" cblas_ssymv_safe cblas_ssymv_unsafe (\x f -> f x)

dsymv :: PrimMonad m => SymvFun Double orient (PrimState m) m
dsymv = symvAbstraction "dsymv" cblas_dsymv_safe cblas_dsymv_unsafe (\x f -> f x)

ssyr :: PrimMonad m => SyrFun Float orient (PrimState m) m
ssyr = syrAbstraction "ssyr" cblas_ssyr_safe cblas_ssyr_unsafe (\x f -> f x)

dsyr :: PrimMonad m => SyrFun Double orient (PrimState m) m
dsyr = syrAbstraction "dsyr" cblas_dsyr_safe cblas_dsyr_unsafe (\x f -> f x)

ssyr2 :: PrimMonad m => Syr2Fun Float orient (PrimState m) m
ssyr2 = syr2Abstraction "ssyr2" cblas_ssyr2_safe cblas_ssyr2_unsafe (\x f -> f x)

dsyr2 :: PrimMonad m => Syr2Fun Double orient (PrimState m) m
dsyr2 = syr2Abstraction "dsyr2" cblas_dsyr2_safe cblas_dsyr2_unsafe (\x f -> f x)

stbmv :: PrimMonad m => TbmvFun Float orient (PrimState m) m
stbmv = tbmvAbstraction "stbmv" cblas_stbmv_safe cblas_stbmv_unsafe

dtbmv :: PrimMonad m => TbmvFun Double orient (PrimState m) m
dtbmv = tbmvAbstraction "dtbmv" cblas_dtbmv_safe cblas_dtbmv_unsafe

ctbmv :: PrimMonad m => TbmvFun (Complex Float) orient (PrimState m) m
ctbmv = tbmvAbstraction "ctbmv" cblas_ctbmv_safe cblas_ctbmv_unsafe

ztbmv :: PrimMonad m => TbmvFun (Complex Double) orient (PrimState m) m
ztbmv = tbmvAbstraction "ztbmv" cblas_ztbmv_safe cblas_ztbmv_unsafe

stbsv :: PrimMonad m => TbsvFun Float orient (PrimState m) m
stbsv = tbsvAbstraction "stbsv" cblas_stbsv_safe cblas_stbsv_unsafe

dtbsv :: PrimMonad m => TbsvFun Double orient (PrimState m) m
dtbsv = tbsvAbstraction "dtbsv" cblas_dtbsv_safe cblas_dtbsv_unsafe

ctbsv :: PrimMonad m => TbsvFun (Complex Float) orient (PrimState m) m
ctbsv = tbsvAbstraction "ctbsv" cblas_ctbsv_safe cblas_ctbsv_unsafe

ztbsv :: PrimMonad m => TbsvFun (Complex Double) orient (PrimState m) m
ztbsv = tbsvAbstraction "ztbsv" cblas_ztbsv_safe cblas_ztbsv_unsafe

stpmv :: PrimMonad m => TpmvFun Float orient (PrimState m) m
stpmv = tpmvAbstraction "stpmv" cblas_stpmv_safe cblas_stpmv_unsafe

dtpmv :: PrimMonad m => TpmvFun Double orient (PrimState m) m
dtpmv = tpmvAbstraction "dtpmv" cblas_dtpmv_safe cblas_dtpmv_unsafe

ctpmv :: PrimMonad m => TpmvFun (Complex Float) orient (PrimState m) m
ctpmv = tpmvAbstraction "ctpmv" cblas_ctpmv_safe cblas_ctpmv_unsafe

ztpmv :: PrimMonad m => TpmvFun (Complex Double) orient (PrimState m) m
ztpmv = tpmvAbstraction "ztpmv" cblas_ztpmv_safe cblas_ztpmv_unsafe

stpsv :: PrimMonad m => TpsvFun Float orient (PrimState m) m
stpsv = tpsvAbstraction "stpsv" cblas_stpsv_safe cblas_stpsv_unsafe

dtpsv :: PrimMonad m => TpsvFun Double orient (PrimState m) m
dtpsv = tpsvAbstraction "dtpsv" cblas_dtpsv_safe cblas_dtpsv_unsafe

ctpsv :: PrimMonad m => TpsvFun (Complex Float) orient (PrimState m) m
ctpsv = tpsvAbstraction "ctpsv" cblas_ctpsv_safe cblas_ctpsv_unsafe

ztpsv :: PrimMonad m => TpsvFun (Complex Double) orient (PrimState m) m
ztpsv = tpsvAbstraction "ztpsv" cblas_ztpsv_safe cblas_ztpsv_unsafe

strmv :: PrimMonad m => TrmvFun Float orient (PrimState m) m
strmv = trmvAbstraction "strmv" cblas_strmv_safe cblas_strmv_unsafe

dtrmv :: PrimMonad m => TrmvFun Double orient (PrimState m) m
dtrmv = trmvAbstraction "dtrmv" cblas_dtrmv_safe cblas_dtrmv_unsafe

ctrmv :: PrimMonad m => TrmvFun (Complex Float) orient (PrimState m) m
ctrmv = trmvAbstraction "ctrmv" cblas_ctrmv_safe cblas_ctrmv_unsafe

ztrmv :: PrimMonad m => TrmvFun (Complex Double) orient (PrimState m) m
ztrmv = trmvAbstraction "ztrmv" cblas_ztrmv_safe cblas_ztrmv_unsafe

strsv :: PrimMonad m => TrsvFun Float orient (PrimState m) m
strsv = trsvAbstraction "strsv" cblas_strsv_safe cblas_strsv_unsafe

dtrsv :: PrimMonad m => TrsvFun Double orient (PrimState m) m
dtrsv = trsvAbstraction "dtrsv" cblas_dtrsv_safe cblas_dtrsv_unsafe

ctrsv :: PrimMonad m => TrsvFun (Complex Float) orient (PrimState m) m
ctrsv = trsvAbstraction "ctrsv" cblas_ctrsv_safe cblas_ctrsv_unsafe

ztrsv :: PrimMonad m => TrsvFun (Complex Double) orient (PrimState m) m
ztrsv = trsvAbstraction "ztrsv" cblas_ztrsv_safe cblas_ztrsv_unsafe


{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Test.QuickCheck.LinearAlgebra
-- Copyright  : Copyright (c) , Patrick Perry <patperry@stanford.edu>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability  : experimental
--
--
-- Test generators for linear algebra types.
--

module Test.QuickCheck.LinearAlgebra (
    -- Testable element types
    TestElem(..),

    -- * Generating random objects
    -- ** Dimensions
    dim,
    dim2,
    Dim(..),
    Dim2(..),

    -- ** Indices
    index,
    index2,
    Index(..),
    Index2(..),

    -- ** Elements
    elem,
    elems,
    -- realElem,
    -- realElems,

    -- ** Association lists
    assocs,
    assocs2,
    Assocs(..),
    Assocs2(..),
    -- bandedAssocs,

    -- ** Vectors
    vector,
    VectorPair(..),
    VectorTriple(..),
    VectorList(..),
    WeightedVectorList(..),
    NonEmptyVectorList(..),
    NonEmptyWeightedVectorList(..),

    --  ** Matrices
    matrix,
    MatrixPair(..),
    MatrixTriple(..),
    -- hermMatrix,
    -- triMatrix,

    --  Banded matrices
    -- bandwidths,
    -- banded,
    -- bandedWith,
    -- hermBanded,
    -- triBanded,

    -- BandedAssocs(..),

    ) where

import Prelude hiding ( elem )

-- import BLAS.Types( UpLoEnum(..), DiagEnum(..) )
import Control.Monad
import Data.Complex( magnitude )
import Data.Maybe( fromJust )

import Test.QuickCheck hiding ( vector )
import qualified Test.QuickCheck as QC

import Numeric.LinearAlgebra.Types
import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Matrix( Matrix )
import qualified Numeric.LinearAlgebra.Matrix as M
-- import Data.Matrix.Banded( Banded, maybeBandedFromMatrixStorage )
-- import Data.Matrix.Banded.ST( runSTBanded, unsafeThawBanded,
--     diagViewBanded )

class TestElem e where
    maybeToReal :: e -> Maybe Double
    toReal :: e -> Double
    norm1 :: e -> Double
    norm :: e -> Double
    conjugate :: e -> e

    toReal = fromJust . maybeToReal

instance (TestElem Double) where
    maybeToReal = Just
    toReal = id
    norm1 = abs
    norm = abs
    conjugate = id

instance (TestElem (Complex Double)) where
    maybeToReal (x :+ y) | y == 0    = Just x
                         | otherwise = Nothing
    norm1 (x :+ y) = abs x + abs y
    norm z = magnitude z
    conjugate (x :+ y) = x :+ (-y)

{-
-- | Element types that can be tested with QuickCheck properties.
class (Storable e, Arbitrary e, CoArbitrary e) => TestElem e where
    -- | Inicates whether or not the value should be used in tests.  For
    -- 'Double's, @isTestElem e@ is defined as
    -- @not (isNaN e || isInfinite e || isDenormalized e)@.
    isTestElem :: e -> Bool

instance TestElem Double where
    isTestElem e = not (isNaN e || isInfinite e || isDenormalized e)
    {-# INLINE isTestElem #-}

instance TestElem (Complex Double) where
    isTestElem (x :+ y) = isTestElem x && isTestElem y
    {-# INLINE isTestElem #-}
-}

minit:: Int ->Int
minit  n=  n `mod` 20

-- | Generate a random element.
elem :: (Arbitrary e) => Gen e
elem = arbitrary

-- | Generate a list of elements suitable for testing with.
elems :: (Arbitrary e) => Int -> Gen [e]
elems n = replicateM n elem

-- Generate a random element that has no imaginary part.
-- realElem :: (TestElem e) => Gen e
-- realElem = liftM fromReal elem

-- Generate a list of elements for testing that have no imaginary part.
-- realElems :: (TestElem e) => Int -> Gen [e]
-- realElems n = replicateM n realElem

-- | Get an appropriate dimension for a random vector
dim :: Gen Int
dim = sized $ \s -> do
    (NonNegative n) <- resize (s `div` 4) $ arbitrary
    return (minit n)

-- | Get an appropriate dimension for a random matrix
dim2 :: Gen (Int,Int)
dim2 = sized $ \s -> do
    m <- resize (s `div` 2) $ dim
    n <- resize (s `div` 2) $ dim
    return (m,n)

-- | A vector dimension.
newtype Dim = Dim Int
  deriving (Eq, Ord, Num, Integral, Real, Enum, Show, Read)

instance Arbitrary Dim where
    arbitrary = sized $ \s -> do
        (NonNegative n) <- resize (s `div` 4) $ arbitrary
        return $ Dim n

    shrink (Dim a) =
        [ Dim a' | (NonNegative a') <- shrink (NonNegative a) ]

-- | A matrix dimension.
newtype Dim2 = Dim2 (Int,Int)
  deriving (Eq, Show, Read)

instance Arbitrary Dim2 where
    arbitrary = do
        mn <- dim2
        return $ Dim2 mn

-- | Given a dimension generate a valid index.  The dimension must be positive.
index :: Int -> Gen Int
index n | n <= 0 =
            error $ "index " ++ (show n) ++ ":"
                  ++ " dimension must be positive (QuickCheck error)"
        | otherwise =
            choose (0,n-1)

-- | Given a matrix dimension generate a valid index.
index2 :: (Int,Int) -> Gen (Int,Int)
index2 (m,n) = do
    i <- index m
    j <- index n
    return (i,j)

-- | A dimension and a valid index for it.
data Index = Index Int Int deriving (Eq,Show)
instance Arbitrary Index where
    arbitrary = do
        n <- (1+) `fmap` dim
        i <- index n
        return $ Index n i

-- | A matrix dimension and a valid index for it.
data Index2 = Index2 (Int,Int) (Int,Int) deriving (Eq,Show)
instance Arbitrary Index2 where
    arbitrary = do
        (m',n') <- dim2
        let mn = (m' + 1, n' + 1)
        ij <- index2 mn
        return $ Index2 mn ij

-- | Generate an associations list for a vector of the given dimension.
assocs :: (Arbitrary e) => Int -> Gen [(Int,e)]
assocs n | n == 0    = return []
         | otherwise = do
    l <- choose(0, 2*n)
    is <- replicateM l $ index n
    es <- elems l
    return $ zip is es

-- | Generate an associations list for a matrix of the given shape.
assocs2 :: (Arbitrary e) => (Int,Int) -> Gen [((Int,Int),e)]
assocs2 (m,n) | m*n == 0  = return []
              | otherwise = do
    l <- choose(0, 2*m*n)
    is <- replicateM l $ index2 (m,n)
    es <- elems l
    return $ zip is es

-- | A dimension and an associations list.
data Assocs e = Assocs Int [(Int,e)] deriving (Eq,Show)

instance (Arbitrary e) => Arbitrary (Assocs e) where
    arbitrary = do
        n   <- dim
        ies <- assocs n
        return $ Assocs n ies

    shrink (Assocs n ies) =
        [ Assocs n' $ filter ((< n') . fst) ies
        | n' <- shrink n
        ] ++
        [ Assocs n ies'
        | ies' <- shrink ies
        ]

-- | A shape and an associations list.
data Assocs2 e = Assocs2 (Int,Int) [((Int,Int),e)] deriving (Eq,Show)
instance (Arbitrary e) => Arbitrary (Assocs2 e) where
    arbitrary = do
        mn   <- dim2
        ies  <- assocs2 mn
        return $ Assocs2 mn ies

-- | Generate a random vector of the given size.
vector :: (Arbitrary e, Storable e) => Int -> Gen (Vector e)
vector n = do
    es <- elems n
    return $ V.fromList n es

instance (Arbitrary e, Storable e) => Arbitrary (Vector e) where
    arbitrary = dim >>= vector

    shrink x =
        [ V.slice 0 n x
        | (NonNegative n) <- shrink (NonNegative $ V.dim x)
        ]

-- | Two vectors with the same dimension.
data VectorPair e f =
    VectorPair (Vector e) (Vector f) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f) =>
    Arbitrary (VectorPair e f) where
        arbitrary = do
            x <- arbitrary
            y <- vector (V.dim x)
            return $ VectorPair x y

        shrink (VectorPair x y) =
            [ VectorPair (V.slice 0 n' x) (V.slice 0 n' y)
            | n' <- shrink (V.dim x)
            ]

-- | Three vectors with the same dimension.
data VectorTriple e f g =
    VectorTriple (Vector e) (Vector f) (Vector g) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f,
          Arbitrary g, Storable g) =>
    Arbitrary (VectorTriple e f g) where
        arbitrary = do
            x <- arbitrary
            y <- vector (V.dim x)
            z <- vector (V.dim x)
            return $ VectorTriple x y z

-- | A nonempty list of vectors with the same dimension.
data NonEmptyVectorList e = NonEmptyVectorList Int [Vector e] deriving (Eq, Show)
instance (Arbitrary e, Storable e) => Arbitrary (NonEmptyVectorList e) where
    arbitrary = do
        x <- arbitrary
        n <- choose (0,20)
        let p = V.dim x
        xs <- replicateM n $ vector p
        return $ NonEmptyVectorList p $ x:xs

-- | A nonempty list of (weight, vector) pairs, with the weights all non-negative
-- and the vectors all having the same dimension.
data NonEmptyWeightedVectorList e = NonEmptyWeightedVectorList Int [(e, Vector e)]
    deriving (Eq, Show)
instance (Arbitrary e, Storable e, Num e) => Arbitrary (NonEmptyWeightedVectorList e) where
    arbitrary = do
        (NonEmptyVectorList p xs) <- arbitrary
        ws <- replicateM (length xs) $ fmap abs arbitrary
        return $ NonEmptyWeightedVectorList p $ zip ws xs

-- | A list of vectors with the same dimension.
data VectorList e = VectorList Int [Vector e] deriving (Eq, Show)
instance (Arbitrary e, Storable e) => Arbitrary (VectorList e) where
    arbitrary = do
        (NonEmptyVectorList p (_:xs)) <- arbitrary
        return $ VectorList p xs

-- | A list of (weight, vector) pairs, with the weights all non-negative
-- and the vectors all having the same dimension.
data WeightedVectorList e = WeightedVectorList Int [(e, Vector e)]
    deriving (Eq, Show)
instance (Arbitrary e, Storable e, Num e) => Arbitrary (WeightedVectorList e) where
    arbitrary = do
        (NonEmptyWeightedVectorList p (_:wxs)) <- arbitrary
        return $ WeightedVectorList p wxs

-- | Generate a random matrix of the given size.
matrix :: (Arbitrary e, Storable e) => (Int,Int) -> Gen (Matrix e)
matrix (m,n) =
    oneof [ raw, sub ]
  where
    raw = do
        es <- elems $ m * n
        return $ M.fromList (m,n) es
    sub = do
        m' <- choose (m, 2*m)
        es <- elems $ m' * n
        return $ M.slice (0,0) (m,n) (M.fromList (m',n) es)

instance (Arbitrary e, Storable e) => Arbitrary (Matrix e) where
    arbitrary = dim2 >>= matrix

-- | Two matrices with the same dimension.
data MatrixPair e f =
    MatrixPair (Matrix e) (Matrix f) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f) =>
    Arbitrary (MatrixPair e f) where
        arbitrary = do
            x <- arbitrary
            y <- matrix (M.dim x)
            return $ MatrixPair x y

-- | Three matrices with the same dimension.
data MatrixTriple e f g =
    MatrixTriple (Matrix e) (Matrix f) (Matrix g) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f,
          Arbitrary g, Storable g) =>
    Arbitrary (MatrixTriple e f g) where
        arbitrary = do
            x <- arbitrary
            y <- matrix (M.dim x)
            z <- matrix (M.dim x)
            return $ MatrixTriple x y z

instance Arbitrary Trans where
    arbitrary = elements [ NoTrans, Trans, ConjTrans ]
{-

-- | Generate a triangular dense matrix.
triMatrix :: (TestElem e) => (Int,Int) -> Gen (Tri Matrix e)
triMatrix (m,n) = do
    a <- matrix (m,n)
    u <- QC.elements [ Lower, Upper  ]
    d <- QC.elements [ Unit, NonUnit ]
    return $ Tri u d a

-- | Generate a Hermitian dense matrix.
hermMatrix :: (TestElem e) => Int -> Gen (Herm Matrix e)
hermMatrix n = do
    a <- matrix (n,n)
    d <- realElems n
    let a' = runSTMatrix $ do
                 ma <- unsafeThawMatrix a
                 setElems (diagView ma 0) d
                 return ma
    u <- QC.elements [ Lower, Upper ]
    return $ Herm u a'

rawMatrix :: (TestElem e) => (Int,Int) -> Gen (Matrix e)
rawMatrix (m,n) = do
    es <- elems (m*n)
    return $ M.fromList (m,n) es

subMatrix :: (TestElem e) => (Int,Int) -> Gen (SubMatrix e)
subMatrix (m,n) =
    oneof [ rawSubMatrix (m,n)
          , rawSubMatrix (n,m) >>= \(SubMatrix a (i,j) (m',n')) ->
                return $ SubMatrix (herm a) (j,i) (n',m')
          ]

rawSubMatrix :: (TestElem e) => (Int,Int) -> Gen (SubMatrix e)
rawSubMatrix (m,n) = do
    i <- choose (0,5)
    j <- choose (0,5)
    e <- choose (0,5)
    f <- choose (0,5)
    x <- rawMatrix (i+m+e, j+n+f)
    return $ SubMatrix x (i,j) (m,n)

instance (TestElem e) => Arbitrary (SubMatrix e) where
    arbitrary = do
        (m,n) <- shape
        (SubMatrix a ij mn) <- subMatrix (m,n)
        return $ SubMatrix a ij mn

-- | Generate valid bandwidth for a given matrix dimension size
bandwidth :: Int -> Gen Int
bandwidth n = if n == 0 then return 0 else choose (0,n-1)

-- | Generate valid bandwidths for the given matrix shape.
bandwidths :: (Int,Int) -> Gen (Int,Int)
bandwidths (m,n) = liftM2 (,) (bandwidth m) (bandwidth n)

-- | Generate a random banded matrix of the given shape.
banded :: (TestElem e) => (Int,Int) -> Gen (Banded e)
banded mn = do
    lu <- bandwidths mn
    bandedWith lu mn

-- | Generate a random banded matrix with the given bandwidths.
bandedWith :: (TestElem e)
           => (Int,Int) -> (Int,Int) -> Gen (Banded e)
bandedWith lu mn = frequency [ (3, rawBanded mn lu)
                             , (2, hermedBanded mn lu)
                             ]

-- | Generate a triangular banded matrix.
triBanded :: (TestElem e) => Int -> Gen (Tri Banded e)
triBanded n = do
    a <- banded (n,n)
    u <- QC.elements [ Lower, Upper  ]
    d <- QC.elements [ Unit, NonUnit ]
    return $ Tri u d a

-- | Generate a Hermitian banded matrix.
hermBanded :: (TestElem e) => Int -> Gen (Herm Banded e)
hermBanded n = do
    a <- banded (n,n)
    d <- realElems n
    let a' = runSTBanded $ do
                 ma <- unsafeThawBanded a
                 setElems (diagViewBanded ma 0) d
                 return ma
    u <- QC.elements [ Lower, Upper ]
    return $ Herm u a'

rawBanded :: (TestElem e) =>
    (Int,Int) -> (Int,Int) -> Gen (Banded e)
rawBanded (m,n) (kl,ku) =
    let bw = kl+ku+1
    in do
        a <- frequency [ (2, rawMatrix (bw,n))
                       , (1, rawSubMatrix (bw,n) >>= \(SubMatrix b ij _) ->
                                 return $ submatrix b ij (bw,n))
                       ]
        return $ fromJust (maybeBandedFromMatrixStorage (m,n) (kl,ku) a)

hermedBanded :: (TestElem e) =>
    (Int,Int) -> (Int,Int) -> Gen (Banded e)
hermedBanded (m,n) (kl,ku) = do
    x <- rawBanded (n,m) (ku,kl)
    return $ herm x

-- | Generate an associations list for a banded matrix of the given shape
-- and bandwidths.
bandedAssocs :: (TestElem e) => (Int,Int) -> (Int,Int) -> Gen [((Int,Int),e)]
bandedAssocs (m,n) (kl,ku) | m*n == 0  = return []
                           | otherwise = do
    (Nat l) <- arbitrary
    ijs     <- replicateM l $ index2 (kl+1+ku,n)
    let ijs' = mapMaybe (\(i,j) -> let i' = i - j - ku in
                                   if 0 <= i' && i' < m then Just (i',j)
                                                        else Nothing    ) ijs
    es      <- replicateM l elem
    return $ zip ijs' es

-- | A shape, bandwidths, and an associations list.
data BandedAssocs e = BandedAssocs (Int,Int) (Int,Int) [((Int,Int),e)] deriving (Eq,Show)
instance (TestElem e) => Arbitrary (BandedAssocs e) where
    arbitrary = do
        mn  <- shape
        bw  <- bandwidths mn
        ies <- bandedAssocs mn bw
        return $ BandedAssocs mn bw ies
-}


{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE DeriveGeneric        #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ViewPatterns         #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

import           Control.DeepSeq
import           Criterion.Main
import           Criterion.Types
import           Data.Char
import           Data.Functor.Identity
import           Data.Time
import           GHC.Generics                 (Generic)
import           GHC.TypeLits
import           Lens.Micro
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.Backprop.Class
import           Numeric.LinearAlgebra.Static
import           System.Directory
import qualified Data.Vector                  as V
import qualified Numeric.LinearAlgebra        as HM
import qualified System.Random.MWC            as MWC

type family HKD f a where
    HKD Identity a = a
    HKD f        a = f a

data Layer' i o f =
    Layer { _lWeights :: !(HKD f (L o i))
          , _lBiases  :: !(HKD f (R o))
          }
  deriving (Generic)

type Layer i o = Layer' i o Identity

deriving instance (KnownNat i, KnownNat o) => Show (Layer i o)
instance NFData (Layer i o)

makeLenses ''Layer'

data Network' i h1 h2 o f =
    Net { _nLayer1 :: !(HKD f (Layer i  h1))
        , _nLayer2 :: !(HKD f (Layer h1 h2))
        , _nLayer3 :: !(HKD f (Layer h2 o ))
        }
  deriving (Generic)

type Network i h1 h2 o = Network' i h1 h2 o Identity

deriving instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Show (Network i h1 h2 o)
instance NFData (Network i h1 h2 o)

makeLenses ''Network'

main :: IO ()
main = do
    g     <- MWC.initialize
           . V.fromList
           . map (fromIntegral . ord)
           $ "hello world"
    test0 <- MWC.uniformR @(R 784, R 10) ((0,0),(1,1)) g
    net0  <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g
    t     <- getZonedTime
    let tstr = formatTime defaultTimeLocale "%Y%m%d-%H%M%S" t
    createDirectoryIfMissing True "bench-results"
    defaultMainWith defaultConfig
          { reportFile = Just $ "bench-results/mnist-bench_" ++ tstr ++ ".html"
          , timeLimit  = 10
          } [
        bgroup "gradient"
          [ let runTest x y     = gradNetManual x y net0
            in  bench "manual"  $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (netErr x y) net0
            in  bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (netErrHKD x y) net0
            in  bench "bp-hkd"  $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (\n' -> netErrHybrid n' y x) net0
            in  bench "hybrid"  $ nf (uncurry runTest) test0
          ]
      , bgroup "descent"
          [ let runTest x y     = trainStepManual 0.02 x y net0
            in  bench "manual"  $ nf (uncurry runTest) test0
          , let runTest x y     = trainStep 0.02 x y net0
            in  bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y     = trainStepHKD 0.02 x y net0
            in  bench "bp-hkd"  $ nf (uncurry runTest) test0
          , let runTest x y     = trainStepHybrid 0.02 x y net0
            in  bench "hybrid"  $ nf (uncurry runTest) test0
          ]
      , bgroup "run"
          [ let runTest         = runNetManual net0
            in  bench "manual"  $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetwork` x) net0
            in  bench "bp-lens" $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetworkHKD` x) net0
            in  bench "bp-hkd"  $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetHybrid` x) net0
            in  bench "hybrid"  $ nf runTest (fst test0)
          ]
      ]

-- ------------------------------
-- - "Backprop" Lens Mode       -
-- ------------------------------

runLayer
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (Layer i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLayer l x = (l ^^. lWeights) #>! x + (l ^^. lBiases)
{-# INLINE runLayer #-}

softMax :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s (R n)
softMax x = konst' (1 / sumElements' expx) * expx
  where
    expx = exp x
{-# INLINE softMax #-}

runNetwork
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetwork n = softMax
             . runLayer (n ^^. nLayer3)
             . logistic
             . runLayer (n ^^. nLayer2)
             . logistic
             . runLayer (n ^^. nLayer1)
             . auto
{-# INLINE runNetwork #-}

crossEntropy
    :: (KnownNat n, Reifies s W)
    => R n
    -> BVar s (R n)
    -> BVar s Double
crossEntropy t r = negate $ log r <.>! auto t
{-# INLINE crossEntropy #-}

netErr
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => R i
    -> R o
    -> BVar s (Network i h1 h2 o)
    -> BVar s Double
netErr x t n = crossEntropy t (runNetwork n x)
{-# INLINE netErr #-}

trainStep
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStep r !x !t !n = n - realToFrac r * gradBP (netErr x t) n
{-# INLINE trainStep #-}

-- ------------------------------
-- - "Backprop" HKD Mode        -
-- ------------------------------

runLayerHKD
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (Layer i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLayerHKD (splitBV->Layer w b) x = w #>! x + b
{-# INLINE runLayerHKD #-}

runNetworkHKD
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetworkHKD (splitBV->Net l1 l2 l3) = softMax
                                      . runLayerHKD l3
                                      . logistic
                                      . runLayerHKD l2
                                      . logistic
                                      . runLayerHKD l1
                                      . auto
{-# INLINE runNetworkHKD #-}

netErrHKD
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => R i
    -> R o
    -> BVar s (Network i h1 h2 o)
    -> BVar s Double
netErrHKD x t n = crossEntropy t (runNetworkHKD n x)
{-# INLINE netErrHKD #-}

trainStepHKD
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepHKD r !x !t !n = n - realToFrac r * gradBP (netErrHKD x t) n
{-# INLINE trainStepHKD #-}

-- ------------------------------
-- - "Manual" Mode              -
-- ------------------------------

runLayerManual
    :: (KnownNat i, KnownNat o)
    => Layer i o
    -> R i
    -> R o
runLayerManual l x = (l ^. lWeights) #> x + (l ^. lBiases)
{-# INLINE runLayerManual #-}

softMaxManual :: KnownNat n => R n -> R n
softMaxManual x = konst (1 / sumElements expx) * expx
  where
    expx = exp x
{-# INLINE softMaxManual #-}

runNetManual
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Network i h1 h2 o
    -> R i
    -> R o
runNetManual n = softMaxManual
               . runLayerManual (n ^. nLayer3)
               . logistic
               . runLayerManual (n ^. nLayer2)
               . logistic
               . runLayerManual (n ^. nLayer1)
{-# INLINE runNetManual #-}

gradNetManual
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
gradNetManual x t (Net (Layer w1 b1) (Layer w2 b2) (Layer w3 b3)) =
    let y1 = w1 #> x
        z1 = y1 + b1
        x2 = logistic z1
        y2 = w2 #> x2
        z2 = y2 + b2
        x3 = logistic z2
        y3 = w3 #> x3
        z3 = y3 + b3
        o0 = exp z3
        o1 = HM.sumElements (extract o0)
        o2 = o0 / konst o1
        -- o3 = - (log o2 <.> t)
        dEdO3 = 1
        dEdO2 = dEdO3 * (- t / o2)
        dEdO1 = - (dEdO2 <.> o0) / (o1 ** 2)
        dEdO0 = konst dEdO1 + dEdO2 / konst o1
        dEdZ3 = dEdO0 * o0
        dEdY3 = dEdZ3
        dEdX3 = tr w3 #> dEdY3
        dEdZ2 = dEdX3 * (x3 * (1 - x3))
        dEdY2 = dEdZ2
        dEdX2 = tr w2 #> dEdY2
        dEdZ1 = dEdX2 * (x2 * (1 - x2))
        dEdY1 = dEdZ1
        dEdB3 = dEdZ3
        dEdW3 = dEdY3 `outer` x3
        dEdB2 = dEdZ2
        dEdW2 = dEdY2 `outer` x2
        dEdB1 = dEdZ1
        dEdW1 = dEdY1 `outer` x
    in  Net (Layer dEdW1 dEdB1) (Layer dEdW2 dEdB2) (Layer dEdW3 dEdB3)
{-# INLINE gradNetManual #-}

trainStepManual
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepManual r !x !t !n =
    let gN = gradNetManual x t n
    in  n - (realToFrac r * gN)

-- ------------------------------
-- - "Hybrid" Mode              -
-- ------------------------------

layerOp :: (KnownNat i, KnownNat o) => Op '[Layer i o, R i] (R o)
layerOp = op2 $ \(Layer w b) x ->
    ( w #> x + b
    , \g -> (Layer (g `outer` x) g, tr w #> g)
    )
{-# INLINE layerOp #-}

logisticOp
    :: Floating a
    => Op '[a] a
logisticOp = op1 $ \x ->
    let lx = logistic x
    in  (lx, \g -> lx * (1 - lx) * g)
{-# INLINE logisticOp #-}

softMaxOp
    :: KnownNat n
    => Op '[R n] (R n)
softMaxOp = op1 $ \x ->
    let expx   = exp x
        tot    = sumElements expx
        invtot = 1 / tot
        res    = konst invtot * expx
    in  ( res
        , \g -> res - konst (invtot ** 2) * exp (2 * x) * g
        )
{-# INLINE softMaxOp #-}

softMaxCrossEntropyOp
    :: KnownNat n
    => R n
    -> Op '[R n] Double
softMaxCrossEntropyOp targ = op1 $ \x ->
    let expx   = exp x
        sm     = konst (1 / sumElements expx) * expx
        ce     = negate $ log sm <.> targ
    in  ( ce
        , \g -> (sm - targ) * konst g
        )
{-# INLINE softMaxCrossEntropyOp #-}

runNetHybrid
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetHybrid n = liftOp1 softMaxOp
               . liftOp2 layerOp (n ^^. nLayer3)
               . liftOp1 logisticOp
               . liftOp2 layerOp (n ^^. nLayer2)
               . liftOp1 logisticOp
               . liftOp2 layerOp (n ^^. nLayer1)
               . auto
{-# INLINE runNetHybrid #-}

netErrHybrid
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R o
    -> R i
    -> BVar s Double
netErrHybrid n t = liftOp1 (softMaxCrossEntropyOp t)
                 . liftOp2 layerOp (n ^^. nLayer3)
                 . liftOp1 logisticOp
                 . liftOp2 layerOp (n ^^. nLayer2)
                 . liftOp1 logisticOp
                 . liftOp2 layerOp (n ^^. nLayer1)
                 . auto
{-# INLINE netErrHybrid #-}

trainStepHybrid
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepHybrid r !x !t !n =
    let gN = gradBP (\n' -> netErrHybrid n' t x) n
    in  n - (realToFrac r * gN)
{-# INLINE trainStepHybrid #-}

-- ------------------------------
-- - Operations                 -
-- ------------------------------

infixr 8 #>!
(#>!)
    :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (L m n)
    -> BVar s (R n)
    -> BVar s (R m)
(#>!) = liftOp2 . op2 $ \m v ->
  ( m #> v, \g -> (g `outer` v, tr m #> g) )
{-# INLINE (#>!) #-}

infixr 8 <.>!
(<.>!)
    :: (KnownNat n, Reifies s W)
    => BVar s (R n)
    -> BVar s (R n)
    -> BVar s Double
(<.>!) = liftOp2 . op2 $ \x y ->
  ( x <.> y, \g -> (konst g * y, x * konst g)
  )
{-# INLINE (<.>!) #-}

konst'
    :: (KnownNat n, Reifies s W)
    => BVar s Double
    -> BVar s (R n)
konst' = liftOp1 . op1 $ \c -> (konst c, HM.sumElements . extract)
{-# INLINE konst' #-}

sumElements :: KnownNat n => R n -> Double
sumElements = HM.sumElements . extract
{-# INLINE sumElements #-}

sumElements'
    :: (KnownNat n, Reifies s W)
    => BVar s (R n)
    -> BVar s Double
sumElements' = liftOp1 . op1 $ \x -> (sumElements x, konst)
{-# INLINE sumElements' #-}

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))
{-# INLINE logistic #-}

-- ------------------------------
-- - Instances                  -
-- ------------------------------

instance (KnownNat i, KnownNat o) => Num (Layer i o) where
    Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2)
    Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2)
    Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2)
    abs    (Layer w b)        = Layer (abs    w) (abs    b)
    signum (Layer w b)        = Layer (signum w) (signum b)
    negate (Layer w b)        = Layer (negate w) (negate b)
    fromInteger x             = Layer (fromInteger x) (fromInteger x)

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Num (Network i h1 h2 o) where
    Net a b c + Net d e f = Net (a + d) (b + e) (c + f)
    Net a b c - Net d e f = Net (a - d) (b - e) (c - f)
    Net a b c * Net d e f = Net (a * d) (b * e) (c * f)
    abs    (Net a b c)    = Net (abs    a) (abs    b) (abs    c)
    signum (Net a b c)    = Net (signum a) (signum b) (signum c)
    negate (Net a b c)    = Net (negate a) (negate b) (negate c)
    fromInteger x         = Net (fromInteger x) (fromInteger x) (fromInteger x)

instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where
    Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2)
    recip (Layer w b)         = Layer (recip w) (recip b)
    fromRational x            = Layer (fromRational x) (fromRational x)

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Fractional (Network i h1 h2 o) where
    Net a b c / Net d e f = Net (a / d) (b / e) (c / f)
    recip (Net a b c)     = Net (recip a) (recip b) (recip c)
    fromRational x        = Net (fromRational x) (fromRational x) (fromRational x)

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

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

instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where
    uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => MWC.Variate (Network i h1 h2 o) where
    uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance Backprop (R n) where
    zero = zeroNum
    add  = addNum
    one  = oneNum

instance (KnownNat n, KnownNat m) => Backprop (L m n) where
    zero = zeroNum
    add  = addNum
    one  = oneNum

instance (KnownNat i, KnownNat o) => Backprop (Layer i o)
instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Backprop (Network i h1 h2 o)


{-# LANGUAGE RecordWildCards #-}

module Data.BenchAna.Stats(stats) where

import qualified Data.Vector.Unboxed as V
import qualified Statistics.Sample as S
import qualified Statistics.Function as S

import Data.BenchAna.Types

stats :: BenchResults -> BenchStats
stats = fmap singleStats

singleStats :: V.Vector Int -> Stats
singleStats vec = Stats { .. }
    where statsWithOutliers = singleStats' vec'
          statsNoOutliers = singleStats' $ V.init $ V.tail $ S.sort vec'
          vec' = V.map fromIntegral vec

singleStats' :: V.Vector Double -> StatsPair
singleStats' vec = StatsPair { .. }
    where mean = S.mean vec
          stddev = S.stdDev vec


{-# LANGUAGE RankNTypes, BangPatterns, GADTs #-}
{-# OPTIONS -Wall #-}

module Language.Hakaru.Distribution where

import System.Random
import Language.Hakaru.Mixture
import Language.Hakaru.Types
import Data.Ix
import Data.Maybe (fromMaybe)
import Data.List (findIndex, foldl')
import Numeric.SpecFunctions
import qualified Data.Map.Strict as M
import qualified Data.Number.LogFloat as LF

mapFst :: (t -> s) -> (t, u) -> (s, u)
mapFst f (a,b) = (f a, b)

dirac :: (Eq a) => a -> Dist a
dirac theta = Dist {logDensity = (\ (Discrete x) -> if x == theta then 0 else log 0),
                    distSample = (\ g -> (Discrete theta,g))}

bern :: Double -> Dist Bool
bern p = Dist {logDensity = (\ (Discrete x) -> log (if x then p else 1 - p)),
               distSample = (\ g -> case randomR (0, 1) g of
                                  (t, g') -> (Discrete $ t <= p, g'))}

uniform :: Double -> Double -> Dist Double
uniform lo hi =
    let uniformLogDensity lo' hi' x | lo' <= x && x <= hi' = log (recip (hi' - lo'))
        uniformLogDensity _ _ _ = log 0
    in Dist {logDensity = (\ (Lebesgue x) -> uniformLogDensity lo hi x),
             distSample = (\ g -> mapFst Lebesgue $ randomR (lo, hi) g)}

uniformD :: (Ix a, Random a) => a -> a -> Dist a
uniformD lo hi =
    let uniformLogDensity lo' hi' x | lo' <= x && x <= hi' = log density
        uniformLogDensity _ _ _ = log 0
        density = recip (fromInteger (toInteger (rangeSize (lo,hi))))
    in Dist {logDensity = (\ (Discrete x) -> uniformLogDensity lo hi x),
             distSample = (\ g -> mapFst Discrete $ randomR (lo, hi) g)}

marsaglia :: (RandomGen g, Random a, Ord a, Floating a) => g -> ((a, a), g)
marsaglia g0 = -- "Marsaglia polar method"
  let (x, g1) = randomR (-1,1) g0
      (y, g ) = randomR (-1,1) g1
      s       = x * x + y * y
      q       = sqrt ((-2) * log s / s)
  in if 1 >= s && s > 0 then ((x * q, y * q), g) else marsaglia g

choose :: (RandomGen g) => Mixture k -> g -> (k, Prob, g)
choose (Mixture m) g0 =
  let peak = maximum (M.elems m)
      unMix = M.map (LF.fromLogFloat . (/peak)) m
      total = M.foldl' (+) (0::Double) unMix
      (p, g) = randomR (0, total) g0
      f !k !v b !p0 = let p1 = p0 + v in if p <= p1 then k else b p1
      err p0 = error ("choose: failure p0=" ++ show p0 ++
                      " total=" ++ show total ++
                      " size=" ++ show (M.size m))
  in (M.foldrWithKey f err unMix 0, LF.logFloat total * peak, g)

chooseIndex :: (RandomGen g) => [Double] -> g -> (Int, g)
chooseIndex probs g0 =
  let (p, g) = random g0
      k = fromMaybe (error ("chooseIndex: failure p=" ++ show p))
                    (findIndex (p <=) (scanl1 (+) probs))
  in (k, g)

normal_rng :: (Real a, Floating a, Random a, RandomGen g) =>
              a -> a -> g -> (a, g)
normal_rng mu sd g | sd > 0 = case marsaglia g of
                                ((x, _), g1) -> (mu + sd * x, g1)
normal_rng _ _ _ = error "normal: invalid parameters"

normalLogDensity :: Floating a => a -> a -> a -> a
normalLogDensity mu sd x = (-tau * square (x - mu)
                            + log (tau / pi / 2)) / 2
  where square y = y * y
        tau = 1 / square sd

normal :: Double -> Double -> Dist Double
normal mu sd = Dist {logDensity = normalLogDensity mu sd . fromLebesgue,
                     distSample = mapFst Lebesgue . normal_rng mu sd}

categoricalLogDensity :: (Eq b, Floating a) => [(b, a)] -> b -> a
categoricalLogDensity list x = log $ fromMaybe 0 (lookup x list)
categoricalSample :: (Num b, Ord b, RandomGen g, Random b) =>
    [(t,b)] -> g -> (t, g)
categoricalSample list g = (elem', g1)
    where
      (p, g1) = randomR (0, total) g
      elem' = fst $ head $ filter (\(_,p0) -> p <= p0) sumList
      sumList = scanl1 (\acc (a, b) -> (a, b + snd(acc))) list
      total = sum $ map snd list

categorical :: Eq a => [(a,Double)] -> Dist a
categorical list = Dist {logDensity = categoricalLogDensity list . fromDiscrete,
                         distSample = mapFst Discrete . categoricalSample list}

lnFact :: Integer -> Double
lnFact = logFactorial

-- Makes use of Atkinson's algorithm as described in:
-- Monte Carlo Statistical Methods pg. 55
--
-- Further discussion at:
-- http://www.johndcook.com/blog/2010/06/14/generating-poisson-random-values/
poisson_rng :: (RandomGen g) => Double -> g -> (Integer, g)
poisson_rng lambda g0 = make_poisson g0
   where smu = sqrt lambda
         b  = 0.931 + 2.53*smu
         a  = -0.059 + 0.02483*b
         vr = 0.9277 - 3.6224/(b - 2)
         arep  = 1.1239 + 1.1368/(b-3.4)
         lnlam = log lambda

         make_poisson :: (RandomGen g) => g -> (Integer,g)
         make_poisson g = let (u, g1) = randomR (-0.5,0.5) g
                              (v, g2) = randomR (0,1) g1
                              us = 0.5 - abs u
                              k = floor $ (2*a / us + b)*u + lambda + 0.43 in
                          case () of
                            () | us >= 0.07 && v <= vr -> (k, g2)
                            () | k < 0 -> make_poisson g2
                            () | us <= 0.013 && v > us -> make_poisson g2
                            () | accept_region us v k -> (k, g2)
                            _  -> make_poisson g2

         accept_region :: Double -> Double -> Integer -> Bool
         accept_region us v k = log (v * arep / (a/(us*us)+b)) <=
                                -lambda + (fromIntegral k)*lnlam - lnFact k

poisson :: Double -> Dist Integer
poisson l =
    let poissonLogDensity l' x | l' > 0 && x> 0 = (fromIntegral x)*(log l') - lnFact x - l'
        poissonLogDensity l' x | x==0 = -l'
        poissonLogDensity _ _ = log 0
    in Dist {logDensity = poissonLogDensity l . fromDiscrete,
             distSample = mapFst Discrete . poisson_rng l}

-- Direct implementation of  "A Simple Method for Generating Gamma Variables"
-- by George Marsaglia and Wai Wan Tsang.
gamma_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
gamma_rng shape _   _ | shape <= 0.0  = error "gamma: got a negative shape paramater"
gamma_rng _     scl _ | scl <= 0.0  = error "gamma: got a negative scale paramater"
gamma_rng shape scl g | shape <  1.0  = (gvar2, g2)
                      where (gvar1, g1) = gamma_rng (shape + 1) scl g
                            (w,  g2) = randomR (0,1) g1
                            gvar2 = scl * gvar1 * (w ** recip shape)
gamma_rng shape scl g =
    let d = shape - 1/3
        c = recip $ sqrt $ 9*d
        -- Algorithm recommends inlining normal generator
        n = normal_rng 1 c
        (v, g2) = until (\y -> fst y > 0.0) (\ (_, g') -> normal_rng 1 c g') (n g)
        x = (v - 1) / c
        sqr = x * x
        v3 = v * v * v
        (u, g3) = randomR (0.0, 1.0) g2
        accept  = u < 1.0 - 0.0331*(sqr*sqr) || log u < 0.5*sqr + d*(1.0 - v3 + log v3)
    in case accept of
         True -> (scl*d*v3, g3)
         False -> gamma_rng shape scl g3

gammaLogDensity :: Double -> Double -> Double -> Double
gammaLogDensity shape scl x | x>= 0 && shape > 0 && scl > 0 =
     scl * log shape - scl * x + (shape - 1) * log x - logGamma shape
gammaLogDensity _ _ _ = log 0

gamma :: Double -> Double -> Dist Double
gamma shape scl = Dist {logDensity = gammaLogDensity shape scl . fromLebesgue,
                          distSample = mapFst Lebesgue . gamma_rng shape scl}

beta_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
beta_rng a b g | a <= 1.0 && b <= 1.0 =
                 let (u, g1) = randomR (0.0, 1.0) g
                     (v, g2) = randomR (0.0, 1.0) g1
                     x = u ** (recip a)
                     y = v ** (recip b)
                 in  case (x+y) <= 1.0 of
                       True -> (x / (x + y), g2)
                       False -> beta_rng a b g2
beta_rng a b g = let (ga, g1) = gamma_rng a 1 g
                     (gb, g2) = gamma_rng b 1 g1
                 in (ga / (ga + gb), g2)

betaLogDensity :: Double -> Double -> Double -> Double
betaLogDensity _ _ x | x < 0 || x > 1 = error "beta: value must be between 0 and 1"
betaLogDensity a b _ | a <= 0 || b <= 0 = error "beta: parameters must be positve"
betaLogDensity a b x = (logGamma (a + b)
                        - logGamma a
                        - logGamma b
                        + (a - 1) * log x
                        + (b - 1) * log (1 - x))

beta :: Double -> Double -> Dist Double
beta a b = Dist {logDensity = betaLogDensity a b . fromLebesgue,
                 distSample = mapFst Lebesgue . beta_rng a b}

laplace_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
laplace_rng mu sd g = sample (randomR (0.0, 1.0) g)
   where sample (u, g1) = case u < 0.5 of
                            True  -> (mu + sd * log (u + u), g1)
                            False -> (mu - sd * log (2.0 - u - u), g1)

laplaceLogDensity :: Floating a => a -> a -> a -> a
laplaceLogDensity mu sd x = - log (2 * sd) - abs (x - mu) / sd

laplace :: Double -> Double -> Dist Double
laplace mu sd = Dist {logDensity = laplaceLogDensity mu sd . fromLebesgue,
                      distSample = mapFst Lebesgue . laplace_rng mu sd}

-- Consider having dirichlet return Vector
-- Note: This is acutally symmetric dirichlet
dirichlet_rng :: (RandomGen g) => Int ->  Double -> g -> ([Double], g)
dirichlet_rng n' a g' = normalize (gammas g' n')
  where gammas g 0 = ([], 0, g)
        gammas g n = let (xs, total, g1) = gammas g (n-1)
                         ( x, g2) = gamma_rng a 1 g1
                     in ((x : xs), x+total, g2)
        normalize (b, total, h) = (map (/ total) b, h)

dirichletLogDensity :: [Double] -> [Double] -> Double
dirichletLogDensity a x | all (> 0) x = sum' (zipWith logTerm a x) + logGamma (sum a)
  where sum' = foldl' (+) 0
        logTerm b y = (b-1) * log y - logGamma b
dirichletLogDensity _ _ = error "dirichlet: all values must be between 0 and 1"


module Section_2_5_2 where

import Numeric.LinearAlgebra

import Section_2_3_3 (and', nand, or')

xor :: Double -> Double -> Double
xor x1 x2 =
  let s1 = nand x1 x2
      s2 = or' x1 x2
      y = and' s1 s2
   in y


module LogisticRegressionExercise where

import qualified Graphics.Gnuplot.Advanced as GP

import qualified Graphics.Gnuplot.Graph as Graph

import qualified Graphics.Gnuplot.Plot.ThreeDimensional as Plot3D
import qualified Graphics.Gnuplot.Graph.ThreeDimensional as Graph3D

import qualified Graphics.Gnuplot.Plot.TwoDimensional as Plot2D
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D
import qualified Graphics.Gnuplot.Terminal.Default as DefaultTerm
import Graphics.Gnuplot.Plot.TwoDimensional (linearScale, )

import qualified Data.ByteString.Lazy as B
import Data.Csv.HMatrix
import Data.Csv.Incremental
import HVX
import LogisticRegression
import Numeric.LinearAlgebra
import Prelude hiding ((<>))
import System.Directory
import Util

-- replace this with the path to your data.
exerciseTwoDataOne :: IO String
exerciseTwoDataOne = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex2data1.txt")

exerciseTwoDataTwo :: IO String
exerciseTwoDataTwo = do
  home <- getHomeDirectory
  return (home ++ "/code/haskell/CourseraMachineLearning/app/ex2data2.txt")

costFunction :: Matrix Double -> Matrix Double -> Matrix Double -> (Double, Matrix Double)
costFunction theta x y = ((sumElements cost) / m, grad)
  where sig = sigmoid $ x <> theta
        m = fromIntegral $ fst (size y)
        n = snd (size x)
	-- Cost function for logistic regression. Faithful formula translation.
        cost = (((-1) * (tr y)) <> (log sig)) - ((1 - (tr y)) <> (log (1 - sig)))
	-- Gradient. Faithful translation of the formula.
        grad = (columnWise sumElements (x * repmat (sig - y) 1 n)) / (scalar m)

plotData :: Matrix Double -> IO (Plot2D.T Double Double)
plotData completeMatrix = do
  let matrixRows = toRows completeMatrix
  -- select all rows whose third column is 1.
  let pos = fromRows $ filter (((==) 1) . (flip atIndex 2)) matrixRows
  -- do the same for those whose third column is 0.
  let neg = fromRows $ filter (((==) 0) . (flip atIndex 2)) matrixRows
  -- Helper function: Convert column vectors to lists for plotting.
  let getPointLists p = map toList $ toColumns (p ?? (All, Take 2))
  let [xPos, yPos] = getPointLists pos
  let [xNeg, yNeg] = getPointLists neg
  return $ plotPoints xPos yPos `mappend` plotPoints xNeg yNeg

plotDecisionBoundary :: Matrix Double -> Matrix Double -> IO (Plot2D.T Double Double)
plotDecisionBoundary theta x = do
  let theta' = flatten theta
  -- We only need two points to plot a line.
  -- So we pick the smallest and largest.
  let lowerPoint = minimum ((toList . flatten) $ x ¿ [1]) - 2
  let upperPoint = maximum ((toList . flatten) $ x ¿ [1]) + 2
  let plotX = [lowerPoint, upperPoint]
  -- A helper function to calculate the y values given by our theta.
  let f val = (-1) / (atIndex theta' 2) * ((atIndex theta' 1) * val + (atIndex theta' 0))
  let plotY = [f lowerPoint, f upperPoint]
  return $ plotLine plotX plotY

exerciseTwo = do
  dataOne <- exerciseTwoDataOne
  contents <- B.readFile dataOne
  let completeMatrix = decodeMatrix NoHeader contents
  let x'' = completeMatrix ?? (All {- Rows -}, Take 2 {- Columns -})
  let y = completeMatrix ¿ [2]
  let (m, n) = size x''
  let x = ones m 1 ||| x''
  let initialTheta = zeros (n + 1) 1
  let (cost, grad) = costFunction initialTheta x y
  putStrLn $ "Cost at initial theta (zeros): " ++ (show cost)
  putStrLn $ "Expected cost (approx): 0.693"
  putStrLn $ "Gradient at initial theta (zeros): " ++ (show grad)
  putStrLn $ "Expected gradients (approx):\n -0.1000\n -12.0092\n -11.2628\n"
  -- Compute and display cost and gradient with non-zero theta.
  -- the learning rate was achieved through empirical tests.
  -- Gives theta closest to fminunc.
  -- TODO(mchav): Find haskell optimisation library.
  let (_, theta) = sgd x y initialTheta costFunction 0.0041 1000000
  decisionBoundary <- plotDecisionBoundary theta x
  points <- plotData completeMatrix
  GP.plotSync DefaultTerm.cons $ decisionBoundary `mappend` points


module Main where

import Numeric.GSL.Integration
import Complex
import System.IO
import Data.Fixed
import Data.Array
import Data.Maybe
import Control.Monad
import Data.Packed.Matrix
import Data.Packed.Vector
import Numeric.LinearAlgebra.Algorithms
import qualified Data.IntMap as M

data Ctree a b = End (a,b)
               | Branche (a,b) [Ctree a b]
    deriving (Show)

-- Given N boxes, return all possible combinations of P particles in these boxes.

c n p = g 1 0
    where g i a | i == n     = [ End (n, p - a) ]
                | a == p - 1 = [ End (i, 1) ] ++ g (i + 1) a
                | otherwise  = [ if r == p - a
                                 then End (i,r)
                                 else Branche (i,r) $ g (i+1) (a + r) |
                                                         i <- [i..n-1],
                                                         r <- [1..p-a] ]
                                                            ++ [ End (n, p-a) ]

-- Convert combinations of Ctree do list format

cl [] = []
cl ((End p):l) = [p]:cl l
cl ((Branche p b):l) = (map (p:) $ cl b) ++ cl l

-- Create newstates aplying A+(i)A+(j)A(k)A(l) operator on state 'a'

newState i j k l n a = do cp i =<< cp j =<< dp k =<< dp l a
    where
        cp i a = if a!i < n
                 then Just $ a//[(i,a!i + 1)]
                 else Nothing

        dp i a = if a!i > 0
                 then Just $ a//[(i,a!i - 1)]
                 else Nothing

-- List all states that can have operator A+(i)A+(j)A(k)A(l) applied

possibleStates i j k l nb np = filter (isJust . newState i j k l np) . aStates nb

-- Autovalue of operator A+(i)A+(j)A(k)A(l)

value i j k l n a = sqrt $ ((a!i + 1)*(a!j + 1)*a!k*a!l) / ((prod $ elems a')*(prod $ elems a))
    where a' = fromJust $ newState i j k l n a
          prod x = prod' x 1
          prod' [] a = a
          prod' (x:xs) a | x == 0 = prod' xs a
                         | otherwise = prod' xs a*x

sumValues i j k l n = sum . map (value i j k l n)

-- Convert state from List to Array type

aStates nb = map (listArray (1,nb) (repeat 0) //)

-- Total quasi-momentum of base state 'a'

qt a = qt' a 0 (min + 1)
    where qt' a t i | i == max  = mod (t + (max - 1)*a!max) max
                    | otherwise = qt' a (t + (i-1)*a!i) (i + 1)
          (min,max) = bounds a

-- Group states by quasi-momentum

groupQt states = M.fromListWith (++) $ map (\x -> (qt x, [x])) states

-- Print grouped states

printGroups states = do mapM_ printG states
    where printG (n,s) = do putStrLn $ (show n) ++ ":"
                            mapM_ (putStrLn.show.elems) s
                            putStrLn ""

na = array (0, 4) [ (0, 1.81772),
                    (1, 1.75048),
                    (2, 0.998368),
                    (3, 0.111417),
                    (4, 1.19828) ]

a = array (0,4) [ (0, 1.94699),
                  (1, 1.98301),
                  (2, 2.29932),
                  (3, 2.49687),
                  (4, 2.22725) ]

n = 5 :: Int
nd = 5 :: Double

bloch phi w a b = exp(0 :+ a*phi)*exp(0 :+ (-wa*phy))*(cos(b*phy)*sin(b*pi/nd)*cos(wa*pi/nd) :+ sin(b*phy)*cos(b*pi/nd)*sin(wa*pi/nd))
    where phy = -pi/nd + mod' (phi + pi/nd) (2*pi/nd)
          wa = a - w

cbloch phi w a b = conjugate $ bloch phi w a b

toInt phi w i j k l = 1/((na!i)*(na!j)*(na!k)*(na!l))*
                            cbloch phi w (fi i) (a!i) * cbloch phi w (fi j) (a!j) *
                            bloch phi w (fi k) (a!k) * bloch phi w (fi l) (a!l)
      where fi = fromIntegral

-- Calculate matrix for each quasi-momentum

matFrom :: Int -> Int -> Array Int (Array Int Int) -> Double -> Array (Int,Int) Double
matFrom nb np gs c = matrix
    where matrix = array ((1,1),(size,size)) [((i,j), f i j) | i <- [1..size], j <- [1..size]]
          f i j | i > j =  matrix!(j,i)
                | otherwise = energy i j + c*(sum $ map (\x -> sq x (gs!j)*(im!x)) [ (oi,oj,ok,ol) | oi <- [1..nb],
                                                                                oj <- [1..nb],
                                                                                ok <- [1..nb], ol <- [1..nb],
                                                                                newState oi oj ok ol np (gs!j) == Just (gs!i) ])
          sq (i,j,k,l) s = sqrt $ (fi $ (s!i) + 1)*(fi $ (s!j) + 1)*(fi $ s!k)*(fi $ s!l)
          energy i j | i == j = e (gs!j) 1 0
                     | otherwise = 0
          e s i t | i == nb = t*(fi $ s!i)
                  | otherwise = e s (i+1) (t + (a!(i-1))*(fi $ s!i))
          size = snd $ bounds gs
          im = intMatrix nb
          fi = fromIntegral

intMatrix n = array ((1,1,1,1),(n,n,n,n)) [ ((i,j,k,l), integrate (i,j,k,l) n) | i <- [1..n], j <- [1..n], k <- [1..n], l <- [1..n] ]

integrate (i,j,k,l) n | mod (i+j-k-l) n == 0 = fst $ integrateQAGS 1E-6 1000 (\x -> realPart $ toInt x 0.4 (i-1) (j-1) (k-1) (l-1)) 0 (2*pi)
                      | otherwise = 0

-- Create Matrix

genMatrix nb np k cst = matFrom nb np (listArray (1,length $ qmGroup M.! k) $ qmGroup M.! k) cst
    where qmGroup = groupQt $ aStates nb $ cl $ c nb np
          qmKeys = M.keys qmGroup

-- Print eigenvalues and eigenstates to all states, group by quasi-momentum

printStates nb np cst file = do
    let qmGroup = groupQt $ aStates nb $ cl $ c nb np
        qmKeys = M.keys qmGroup

        calEig n x = eig $ (n><n) $ elems x
        printEig (s,v) = do forM_ (zip (toList s) (toLists v)) $ \(e,v) -> do
                                 hPutStrLn file $ "Eigenvalue: " ++ (show e) ++ "\n"
                                 mapM_ (\n -> hPutStrLn file $ (show n) ++ " ") v
                                 hPutStrLn file "\n"

    forM_ qmKeys $ \k -> do
        let kSize = length $ qmGroup M.! k
        hPutStrLn file $ "quasi-momentum: " ++ (show k) ++ "\n"
        printEig $ calEig (kSize) $ matFrom nb np (listArray (1, kSize) $ qmGroup M.! k) cst

main = do
    file <- openFile "output" WriteMode
    printStates 5 10 1 file
    hClose file


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

-- The unused variables are left as documentation
{-# OPTIONS -Wno-unused-local-binds #-}

{-|
Module      : ControlSystems.DynamicSystems.Conversion
Description : Converts continuous time systems to discrete time, using static vectors.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module ControlSystems.DynamicSystems.Conversions.Static (
  c2d
) where

import GHC.TypeLits
import Numeric.LinearAlgebra.Static

-- | Convert a continuous time linear differential system to a discrete form,
-- using state space models, using type level naturals.
--
-- Specifically note this implementation against
-- `ControlSystems.DynamicSystems.Conversions.c2d`, which takes many more
-- code to handle the complexities of creating a block matrix of the right
-- size and then splitting it apart again.
c2d :: forall m n. (KnownNat m, KnownNat n) =>
       Double         -- ^ The sampling time
    -> L n n          -- ^ The state transition matrix, often called A
    -> L n m          -- ^ The state input matrix, often called B
    -> (L n n, L n m) -- ^ The discretized A & B system
c2d t a b = discreteSystem
  where
    block = a ||| b  -- [ A B ]
              ===    -- [ 0 0 ]
               0

    -- Calculate the discrete time block matrix, which is the following
    -- ⌈ A_d B_d ⌉
    -- ⌊  0   I  ⌋
    exp_block = expm (konst t * block)

    -- Finally we can extract the correct submatrices assuming we pull out
    -- the right pieces. Note that we do not need to pass the indices here
    -- directly; the type system ensures that we pull out the correct A_d
    -- and B_d parts of the exp_block matrix.
    discreteSystem@(a_d, b_d) = splitCols $ fst (splitRows exp_block)


{-# LANGUAGE CPP             #-}
{-# LANGUAGE TemplateHaskell #-}

{-|
Module      : Main
Description : Main module running the GMM
Copyright   : (c) Julian Kopka Larsen, 2015
Stability   : experimental

-}
module Main (
    main
) where

import           Control.Monad          (forM, unless)
import           Data.List              (transpose, (\\))
import           FileParsing
import           GHC.Arr (range)
import           Graphics.EasyPlot
import           Numeric.LinearAlgebra
-- import           Numeric.Statistics.PCA
import           System.CPUTime
import           System.Exit            (exitFailure)
import           System.IO              (readFile)
import           Text.Printf
import           System.Random (mkStdGen)
import Debug.Trace

import           MCMC
import           Partition
import           Math
import           Distributions (lNormW, dlNormW,lnormalInvWishartSS)

list2Touple (a:b:_) = (a,b)
list2Touple _ = (0,0)

toTouples :: Num a => [[a]] -> [(a,a)]
toTouples = map list2Touple

--to2D = toTouples . toLists
plotClusters :: X -> Partition -> [Graph2D Double Double]
plotClusters x p = map toplot $ range (0,((k p)-1))
                      where toplot i = Data2D [Title "d"] [] (toTouples $ map toList $ select x p i)

qtr x = trace ("value: " ++ show x) x

data State = State (Partition,[Sstat]) deriving Show

instance Sampleable State X where
    condMove x g (State (p,s)) = State (newPar, update)
                  where newPar = move g p
                        update = ss' x p s newPar

    llikelihood _ (State (_,s)) = sum $ map (\(i,m,v) -> lnormalInvWishartSS i m v) s
    llikDiff _ (State (_,s')) (State (_,s))
                = f s' - f s
                where f s = sum $ map (\(i,m,v) -> lnormalInvWishartSS i m v) s

ss :: X -> Partition -> [Sstat]
ss x part = zip3 csizes scatters means
            where csizes = componentSizes part
                  scatters = map scatterMatrix xs
                  means = map meanv xs
                  xs = groups x part

ss' :: X -> Partition -> [Sstat] -> Partition -> [Sstat]
ss' x par s npar = map go $ zip3 (p par) s (p npar)
        where go (c1, s, c2)
               | c1 == c2 = s
               | otherwise = updateSstat x c1 s c2

ssComponent :: X -> [Int] -> Sstat
ssComponent x []= (0, konst 0 (d,d), konst 0 d)
                where  d = dim $ head x
ssComponent x c = (length c, scatterMatrix xs, meanv xs)
                where  xs = group x c

updateSstat :: X -> [Int] -> Sstat -> [Int] -> Sstat
updateSstat x c1 s c2 = (s <-> removed) <+> added
                where removed =  (ssComponent x (c1\\c2))
                      added =   (ssComponent x (c2\\c1))

gmmElement :: X -> Int -> Int -> Int -> Partition
gmmElement x seed k num = p
          where State (p,_) = getElement gen x start num
                gen = mkStdGen seed
                start = State (startP, ss x startP)
                startP = naiveFromNk n k
                n = length x

time :: IO t -> IO t
time a = do
    start <- getCPUTime
    v <- a
    end   <- getCPUTime
    let diff = (fromIntegral (end - start)) / (10^12)
    printf "Computation time: %0.3f sec\n" (diff :: Double)
    return v

-- Main
main = do
    --contents <- readFile "../../Data/AccidentData.csv"
    --contents <- readFile "../../Data/2Clusters.csv"
    --contents <- readFile "../../Data/synthetic.6clust.csv"
    contents <- readFile "../../Data/synth2c2d.csv"

    let num_Components = 2
        num_Samples = 4000
        stdData = p2NormList contents
        result = gmmElement stdData 1 num_Components
        p a = plot X11 $ plotClusters stdData a

    --Plot Data
    putStrLn "Starting..."
    time $ p (result num_Samples)
    putStrLn "Done."


{-# LANGUAGE QuantifiedConstraints,
             TupleSections,
             DefaultSignatures,
             FlexibleInstances,
             FlexibleContexts,
             UndecidableInstances,
             ConstraintKinds,
             MultiParamTypeClasses,
             TypeFamilies #-}

module Tensor where
import Data.Complex
import Helper
import EitherTrans
import Control.Monad.Trans.State
import Data.Monoid
import Data.List
import Data.Function
import System.Random
import System.Random.Stateful

-- orphan instance
-- sort complex numbers by their magnitudes
instance RealFloat a => Ord (Complex a) where
    compare a b = compare (magnitude a) (magnitude b)

instance (UniformRange a) => UniformRange (Complex a) where
    uniformRM (lr :+ li, hr :+ hi) g = do
        mr <- uniformRM (lr, hr) g
        mi <- uniformRM (li, hi) g
        return $ mr :+ mi

instance (Uniform a) => Uniform (Complex a) where
    uniformM g = do
        r <- uniformM g
        i <- uniformM g
        return $ r :+ i

{-
    The definition of tensor product should follow conventional tensor product properties.
    Tensor product should form a monoid upto isomorphism
-}

-- higer-kinded definition
-- Polymorphic constraints

class Cmplx c where
    conj :: c -> c
    conj = error "please define conjugate function"

instance (Num a) => Cmplx (Complex a) where
    conj = conjugate

class (Monoid (t a), Cmplx a) => Tensor t a where
    (|*|) :: t a -> t a -> t a
    (|*|) = (<>)
    tconcat :: [t a] -> t a
    tconcat = mconcat
    norm :: Floating a => t a -> t a
    fromList :: Integral b => [a] -> b -> t a
    rank :: Integral b => t a -> b
    isSquare :: t a -> Bool
    conjTr :: t a -> t a
    inner :: t a -> t a -> t a

class (Tensor q a, Num (q a)) => QuantumRegister q a where

    densityOp :: q a
        -> q a

    densityOp x = x * conjTr x

    -- Builds a transformation tensor to apply input gate to selected qubit
    build :: (Integral b, Integral d) => b -- ^ Number of qubits this tensor applies to
        -> [(d, q a)]-- ^ Index / Gate pair (Pre-sorted)
        -> q a -- ^ Default Gate
        -> q a
    build i xs d = tconcat $ go i (sortOnFst xs) d 0
        where
            go i2 [] d2 n = replicate (fromIntegral i2 - n) d2
            go i2 (x:xs) d2 n = replicate (fromIntegral (fst x) - n) d2 ++ (snd x : go i2 xs d2 (fromIntegral $ fst x + 1))

    getProb :: Int -- ^ Index of the qubit
        -> q a -- ^ Qubits
        -> Either String (a, a) -- ^ Resulting Probability
    default getProb :: (Foldable q, Floating a) => Int -> q a -> Either String (a, a) -- ^ Resulting Probability
    getProb i q
        | i < 0 || i >= rank q = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid Quantum Register"
        | otherwise = do
            a <- applyGate mask0 i q
            b <- applyGate mask1 i q
            return (f a, f b)
        where
            f = getSum . foldMap (Sum . (\x -> conj x * x))

    initQubit :: a -- ^ Constant for |0>
        -> a -- ^ Constant for |1>
        -> Either String (q a) -- ^ Resulting Qubit

    initQubit0 :: q a -- ^ Initialize |0> qubit at z-direction

    initNumQubit0 :: (Integral b) => b -- ^ Number of qubits to initialize
        -> Either String (q a) -- ^ Initialized quantum register

    toQuantumRegister :: [a] -- ^ State vector array to convert
        -> Either String (q a) -- ^ Converted Quantum Register

    collapse :: Int -- ^ Index of the qubit
        -> Bool -- ^ Collapsed State
        -> q a -- ^ Qubits to collapse
        -> Either String (q a) -- ^ Resulting Tensor

    default collapse :: (Floating a) => Int -> Bool -> q a -> Either String (q a)
    collapse i s q
        | i < 0 || i >= rank q = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid state vector"
        | s = norm <$> applyGate mask1 i q
        | otherwise = norm <$> applyGate mask0 i q

    measure :: (Monad m, StatefulGen g m, UniformRange a) => Int -- ^ Index of the qubit
        -> g -- ^ generator
        -> StateT (q a) (EitherT String m) Bool -- ^ Resulting State Transformer

    default measure :: (Monad m, Num a, Ord a, StatefulGen g m, UniformRange a) => Int -- ^ Index of the qubit
        -> g -- ^ generator
        -> StateT (q a) (EitherT String m) Bool -- ^ Resulting State Transformer
    measure i g = StateT go
        where
            go q
                | i < 0 || i >= rank q = EitherT $ return $ Left "Index out of bounds"
                | not (isSquare q) = EitherT $ return $ Left "Invalid state vector"
                | otherwise = EitherT $ do
                    p <- uniformRM (0, 1) g
                    return $ do
                        t <- getProb i q
                        if p < fst t
                        then (False, ) <$> collapse i False q
                        else (True, ) <$> collapse i True q

    applyGate :: q a -- ^ Gate tensor
        -> Int -- ^ Index
        -> q a -- ^ Quantum Registers
        -> Either String (q a)
    applyGate a i b
        | not (isSquare a) = Left "Invalid Gate Tensor"
        | not (isSquare b) = Left "Invalid Statevector"
        | i < 0 || i >= rank b = Left "Index out of bounds"
        | otherwise = Right $ build (rank b) [(i, a)] pauliId * b

    applyGateAll :: q a -- ^ Gate tensor
        -> q a -- ^ Quantum Registers
        -> Either String (q a)

    applyGateAll a b
        | not (isSquare a) = Left "Invalid Gate Tensor"
        | not (isSquare b) = Left "Invalid Statevector"
        | otherwise = Right $ build (rank b) [] a * b

    applyControl :: q a -- ^ Gate tensor
        -> Int -- ^ Control qubit index
        -> Int -- ^ Apply qubit index
        -> q a -- ^ Quantum Registers
        -> Either String (q a)

    applyControl a ctl i b
        | not (isSquare a) = Left "Invalid gate tensor"
        | not (isSquare b) = Left "Invalid state vector"
        | i < 0 || i >= rank b = Left "Application qubit index out of bounds"
        | ctl < 0 || ctl >= rank b = Left "Control qubit index out of bounds"
        | ctl == i = Left "Control and application qubit cannot be the same"
        | otherwise = Right $ (build (rank b) [(ctl, mask0)] pauliId
            + build (rank b) (sortBy (compare `on` fst) [(ctl, mask1), (i, a)]) pauliId) * b

    illegalPeek :: q a
        -> String

    subSystem :: [Int] -> q a -> Either String (q a)

    isEntangled :: [Int] -> q a -> Either String Bool

    mask0 :: q a
    mask1 :: q a
    pauliId :: q a

class (QuantumRegister q a) => Gates q a where
    pauliX :: q a
    pauliY :: q a
    pauliZ :: q a
    hadamard :: q a

-- Num-only declaration

-- class (forall a. Num a => Monoid (t a)) => Tensor t where
--     (|*|) :: (Num a) => t a -> t a -> t a
--     (|*|) = (<>)
--     norm :: Floating a => t a -> t a
--     fromList :: Integral b => [a] -> b -> t a
--     rank :: Integral b => t a -> b
--     isSquare :: t a -> Bool

-- class (Tensor q, forall a. Num a => Num (q a)) => QuantumRegister q where
--     -- Builds a transformation tensor to apply input gate to selected qubit
--     build :: (RealFloat a, Integral b, Integral c) => b -- ^ Number of qubits this tensor applies to
--         -> [(c, q (Complex a))]-- ^ Index / Gate pair (Pre-sorted)
--         -> q (Complex a) -- ^ Default Gate
--         -> q (Complex a)
--     build i xs d = mconcat $ go i (sortOnFst xs) d 0
--         where
--             go i2 [] d2 n = replicate (fromIntegral i2 - n) d2
--             go i2 (x:xs) d2 n = replicate (fromIntegral (fst x) - n) d2 ++ (snd x : go i2 xs d2 (fromIntegral $ fst x + 1))

--     getProb :: (RealFloat a) => Int -- ^ Index of the qubit
--         -> q (Complex a) -- ^ Qubits
--         -> Either String (a, a) -- ^ Resulting Probability

--     initQubit :: (RealFloat a) => Complex a -- ^ Constant for |0>
--         -> Complex a -- ^ Constant for |1>
--         -> Either String (q (Complex a)) -- ^ Resulting Qubit

--     initQubit0 :: (RealFloat a) => q (Complex a) -- ^ Initialize |0> qubit at z-direction

--     initNumQubit0 :: (RealFloat a, Integral b) => b -- ^ Number of qubits to initialize
--         -> Either String (q (Complex a)) -- ^ Initialized quantum register

--     toQuantumRegister :: (RealFloat a) => [Complex a] -- ^ State vector array to convert
--         -> Either String (q (Complex a)) -- ^ Converted Quantum Register

--     collapse :: (RealFloat a) => Int -- ^ Index of the qubit
--         -> Complex a -- ^ Collapsed State
--         -> q (Complex a) -- ^ Qubits to collapse
--         -> EitherT String IO (q (Complex a)) -- ^ Resulting Tensor

--     measure :: (RealFloat a, Monad m) => Int -- ^ Index of the qubit
--         -> StateT (q (Complex a)) (EitherT String m) (Complex a)
--     -- -> q (Complex a) -- ^ Qubits
--     -- -> Either String ( m (Complex a, q (Complex a))) -- ^ Resulting measurement and the collapsed qubits

--     pauliX :: (RealFloat a) => q (Complex a)
--     pauliY :: (RealFloat a) => q (Complex a)
--     pauliZ :: (RealFloat a) => q (Complex a)
--     pauliId :: (RealFloat a) => q (Complex a)
--     hadamard :: (RealFloat a) => q (Complex a)
--     mask0 :: (RealFloat a) => q (Complex a)
--     mask1 :: (RealFloat a) => q (Complex a)


{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, GeneralizedNewtypeDeriving #-}

-- | This module defines the Syntax of the Quantum IO Monad, which is an embedded
-- language for writing quantum computations.
module QIO.QioSyn where

import Data.Monoid as Monoid
import Data.Complex
import Control.Applicative (Applicative(..))
import Control.Monad       (liftM, ap)

-- | For Real numbers, we simply use the built in Double type
type RR = Double

-- | For Complex numbers, we use the built in Complex numbers, over our Real
-- number type (i.e. Double)
type CC = Complex RR

-- | The amplitude of a complex number is the magnitude squared.
amp :: CC -> RR
amp k = (magnitude k)*(magnitude k)

-- | The type of Qubits in QIO are simply integer references.
newtype Qbit = Qbit Int deriving (Num, Enum, Eq, Ord)

-- | A rotation is in essence a two-by-two complex valued matrix
type Rotation = ((Bool,Bool) -> CC)

-- | The underlying data type of a U unitary operation
data U = UReturn | Rot Qbit Rotation U
       | Swap Qbit Qbit U | Cond Qbit (Bool -> U) U | Ulet Bool (Qbit -> U) U

-- | The underlying data type of a QIO Computation
data QIO a = QReturn a | MkQbit Bool (Qbit -> QIO a) | ApplyU U (QIO a)
           | Meas Qbit (Bool -> QIO a)

-- | The type "U" forms a Monoid
instance Semigroup U where
    --mempty = UReturn
    UReturn <> u = u
    (Rot x a u) <> u' = Rot x a (u <> u')
    (Swap x y u) <> u' = Swap x y (u <> u')
    (Cond x br u') <> u'' = Cond x br (u' <> u'')
    (Ulet b f u) <> u' = Ulet b f (u <> u')

instance Monoid  U where
  mempty = UReturn

-- | Apply the given rotation to the given qubit
rot :: Qbit -> Rotation -> U
rot x r = Rot x r UReturn

-- | Swap the state of the two given qubits
swap :: Qbit -> Qbit -> U
swap x y = Swap x y UReturn

-- | Apply the conditional unitary, depending on the value of the given qubit
cond :: Qbit -> (Bool -> U) -> U
cond x br = Cond x br UReturn

-- | Introduce an Ancilla qubit in the given state, for use in the sub-unitary
ulet :: Bool -> (Qbit -> U) -> U
ulet b ux = Ulet b ux UReturn

-- | Returns the inverse (or reverse) of the given unitary operation
urev :: U -> U
urev UReturn = UReturn
urev (Rot x r u) = urev u <> rot x (rrev r)
urev (Swap x y u) = urev u <> swap x y
urev (Cond x br u) = urev u <> cond x (urev.br)
urev (Ulet b xu u) = urev u <> ulet b (urev.xu)

-- | Apply a not rotation to the given qubit
unot :: Qbit -> U
unot x = rot x rnot

-- | Apply a hadamard rotation to the given qubit
uhad :: Qbit -> U
uhad x = rot x rhad

-- | Apply a phase rotation (of the given angle) to the given qubit
uphase :: Qbit -> RR -> U
uphase x r = rot x (rphase r)

instance Functor QIO where
    fmap = liftM

instance Applicative QIO  where
    pure  = QReturn
    (<*>) = ap

-- | The "QIO" type forms a Monad
instance Monad QIO where
    return = pure
    (QReturn a) >>= f = f a
    (MkQbit b g) >>= f = MkQbit b (\ x -> g x >>= f)
    (ApplyU u q) >>= f = ApplyU u (q >>= f)
    (Meas x g) >>= f = Meas x (\ b -> g b >>= f)

-- | Initialise a qubit in the given state (adding it to the overall quantum state)
mkQbit :: Bool -> QIO Qbit
mkQbit b = MkQbit b return

-- | Apply the given unitary operation to the current quantum state
applyU :: U -> QIO ()
applyU u =  ApplyU u (return ())

-- | Measure the given qubit, and return the measurement outcome (note that this
-- operation may affect the overall quantum state, as a measurement is destructive)
measQbit :: Qbit -> QIO Bool
measQbit x = Meas x return

-- | The identity rotation
rid :: Rotation
rid (x,y) = if x==y then 1 else 0

-- | The not rotation
rnot :: Rotation
rnot (x,y) = if x==y then 0 else 1

-- | The hadamard rotation
rhad :: Rotation
rhad (x,y) = if x && y then -h else h where h = (1/sqrt 2)

-- | The phase rotation
rphase :: RR -> Rotation
rphase _ (False,False)  = 1
rphase r (True,True)    = exp(0:+r)
rphase _ (_,_)          = 0

-- | Returns the inverse (or reverse) of the given rotation
rrev :: Rotation -> Rotation
rrev r (False,True)   = conjugate (r (True,False))
rrev r (True,False)   = conjugate (r (False,True))
rrev r xy             = conjugate (r xy)

-- | Rotations can be compared for equality.
-- They are equal if the define the same matrix.
instance Eq Rotation where
    f == g =    (f (False,False)  == g (False,False))
             && (f (False,True)   == g (False,True))
             && (f (True,False)   == g (True,False))
             && (f (True,True)    == g (True,True))
    f /= g =    (f (False,False)  /= g (False,False))
             || (f (False,True)   /= g (False,True))
             || (f (True,False)   /= g (True,False))
             || (f (True,True)    /= g (True,True))

-- | We can display a qubit reference
instance Show Qbit where
    show (Qbit q) = "(Qbit:" ++ show q ++ ")"

-- | We can display the matrix representation of a rotation
instance Show Rotation where
    show f = "(" ++ (show (f (False,False))) ++ "," ++ (show (f (False,True))) ++ "," ++ (show (f (True,False))) ++ "," ++ (show (f (True,True))) ++ ")"

-- | We can display a representation of a unitary
instance Show U where
    show u = show' u 0 (-1)

-- | A helper function for the show instance of U
show' :: U -> Int -> Int -> String
show' (UReturn) x fv = ""
show' (Rot q a u) x fv = spaces x ++ "Rotate " ++ show q ++ " by " ++ show a ++ ".\n" ++ show' u x fv
show' (Swap q1 q2 u) x fv = spaces x ++ "Swap " ++ show q1 ++ " and " ++ show q2 ++ ".\n" ++ show' u x fv
show' (Cond q f u) x fv = spaces x ++ "Cond (if " ++ show q ++ " then \n" ++ spaces (x+1) ++ "(\n" ++ show' (f True) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ spaces x ++ "else \n" ++ spaces (x+1) ++ "(\n" ++ show' (f False) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ show' u x fv
show' (Ulet b f u) x fv = spaces x ++ "Ulet " ++ show b ++ " (\\" ++ show (Qbit fv) ++ "->\n " ++ show' (f (Qbit fv)) x (fv-1) ++ ")\n" ++ show' u x fv

-- | A helper function that returns a string of 4\x\ spaces.
spaces :: Int -> String
spaces 0 = ""
spaces n = if (n < 0) then error "spaces: negative argument"
                      else "    " ++ spaces (n-1)


{-# LANGUAGE DataKinds, KindSignatures, TypeOperators #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- | Parsing and pre-processing MNIST data.
module HordeAd.Tool.MnistData where

import Prelude

import           Codec.Compression.GZip (decompress)
import           Control.Arrow (first)
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import qualified Data.ByteString.Lazy as LBS
import           Data.IDX
import           Data.List (sortOn)
import           Data.Maybe (fromMaybe)
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Unboxed
import           GHC.TypeLits (KnownNat, Nat)
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM
import           System.IO (IOMode (ReadMode), withBinaryFile)
import           System.Random

-- until stylish-haskell accepts NoStarIsType
import qualified GHC.TypeLits

type SizeMnistWidth = 28 :: Nat
type SizeMnistHeight = SizeMnistWidth

sizeMnistWidth :: Int
sizeMnistWidth = valueOf @SizeMnistWidth

type SizeMnistGlyph = SizeMnistWidth GHC.TypeLits.* SizeMnistHeight

sizeMnistGlyph :: Int
sizeMnistGlyph = valueOf @SizeMnistGlyph

type SizeMnistLabel = 10 :: Nat

sizeMnistLabel :: Int
sizeMnistLabel = valueOf @SizeMnistLabel

type LengthTestData = 10000 :: Nat

-- Actually, a better representation, supported by @Data.IDX@,
-- is an integer label and a picture (the same vector as below).
-- Then we'd use @lossCrossEntropy@ that picks a component according
-- to the label instead of performing a dot product with scaling.
-- This results in much smaller Delta expressions.
-- Our library makes this easy to express and gradients compute fine.
-- OTOH, methods with only matrix operations and graphs can't handle that.
-- However, the goal of the exercise it to implement the same
-- neural net that backprop uses for comparative benchmarks.
-- Also, loss computation is not the bottleneck and the more general
-- mechanism that admits non-discrete target labels fuses nicely
-- with softMax. This also seems to be the standard or at least
-- a simple default in tutorial.
type MnistData r = (Vector r, Vector r)

type MnistData2 r = (Matrix r, Vector r)

type MnistDataS r =
  ( OS.Array '[SizeMnistHeight, SizeMnistWidth] r
  , OS.Array '[SizeMnistLabel] r )

type MnistDataBatchS batch_size r =
  ( OS.Array '[batch_size, SizeMnistHeight, SizeMnistWidth] r
  , OS.Array '[batch_size, SizeMnistLabel] r )

shapeBatch :: Numeric r => MnistData r -> MnistDataS r
shapeBatch (input, target) = (OS.fromVector input, OS.fromVector target)

packBatch :: forall batch_size r. (Numeric r, KnownNat batch_size)
          => [MnistDataS r] -> MnistDataBatchS batch_size r
packBatch l =
  let (inputs, targets) = unzip l
  in (OS.ravel $ OSB.fromList inputs, OS.ravel $ OSB.fromList targets)

readMnistData :: LBS.ByteString -> LBS.ByteString -> [MnistData Double]
readMnistData glyphsBS labelsBS =
  let glyphs = fromMaybe (error "wrong MNIST glyphs file")
               $ decodeIDX glyphsBS
      labels = fromMaybe (error "wrong MNIST labels file")
               $ decodeIDXLabels labelsBS
      intData = fromMaybe (error "can't decode MNIST file into integers")
                $ labeledIntData labels glyphs
      f :: (Int, Data.Vector.Unboxed.Vector Int) -> MnistData Double
      -- Copied from library backprop to enable comparison of results.
      -- I have no idea how this is different from @labeledDoubleData@, etc.
      f (labN, v) =
        ( V.map (\r -> fromIntegral r / 255) $ V.convert v
        , V.generate sizeMnistLabel (\i -> if i == labN then 1 else 0) )
  in map f intData

trainGlyphsPath, trainLabelsPath, testGlyphsPath, testLabelsPath :: FilePath
trainGlyphsPath = "samplesData/train-images-idx3-ubyte.gz"
trainLabelsPath = "samplesData/train-labels-idx1-ubyte.gz"
testGlyphsPath  = "samplesData/t10k-images-idx3-ubyte.gz"
testLabelsPath  = "samplesData/t10k-labels-idx1-ubyte.gz"

loadMnistData :: FilePath -> FilePath -> IO [MnistData Double]
loadMnistData glyphsPath labelsPath =
  withBinaryFile glyphsPath ReadMode $ \glyphsHandle ->
    withBinaryFile labelsPath ReadMode $ \labelsHandle -> do
      glyphsContents <- LBS.hGetContents glyphsHandle
      labelsContents <- LBS.hGetContents labelsHandle
      return $! readMnistData (decompress glyphsContents)
                              (decompress labelsContents)

loadMnistData2 :: FilePath -> FilePath -> IO [MnistData2 Double]
loadMnistData2 glyphsPath labelsPath = do
  ds <- loadMnistData glyphsPath labelsPath
  return $! map (first $ HM.reshape sizeMnistWidth) ds

-- Good enough for QuickCheck, so good enough for me.
shuffle :: RandomGen g => g -> [a] -> [a]
shuffle g l =
  let rnds = randoms g :: [Int]
  in map fst $ sortOn snd $ zip l rnds

chunksOf :: Int -> [e] -> [[e]]
chunksOf n = go where
  go [] = []
  go l = let (chunk, rest) = splitAt n l
         in chunk : go rest


import Test.Tasty
import Test.Tasty.HUnit

import qualified Data.Vector.Storable.Mutable as MV
import qualified Data.Vector.Storable as V
import Data.Vector.Storable ((!))
import Numeric.LinearAlgebra.HArpack
import Control.Monad
import Foreign.C.Types
import System.IO.Unsafe
import Data.Complex

import Data.List
import Data.Ord
import Data.Maybe
import Data.Function (on)

main = defaultMain tests

tests :: TestTree
tests = testGroup "Tests" [unitTests]

unitTests = testGroup "Unit tests"
  [ testCase "Diagonal Matrix LM"   $ diagMatLMCheck   @?= True
  , testCase "Diagonal Matrix SM"   $ diagMatSMCheck   @?= True
  , testCase "Non-Symmetric Sanity" $ complexEigsCheck @?= True ]

diagMat :: SparseMatrix
diagMat = SparseMatrix 1000
            [((n-1,n-1),fromIntegral n) | n <- [1 .. 1000]]

complexEigs :: SparseMatrix
complexEigs = SparseMatrix 7
  [((0,0),3)
  ,((0,1),-2)
  ,((1,0),4)
  ,((1,1),-1)
  ,((2,2),2)
  ,((3,3),1)
  ,((4,4),0.5)
  ,((5,5),0.25)
  ,((6,6),0.125)]

diagMatLMCheck :: Bool
diagMatLMCheck =
  let RReal ePairs' = eigs diagMat (Symmetric SLM) 6
      ePairs = decreasingR ePairs'
      evals = map fst ePairs
      trueEvals = [1000,999..]
      diff = map abs $ zipWith (-) evals trueEvals
  in all (< 1e-10) diff

decreasingC = sortBy (compare `on` (\x -> - realPart (abs (fst x))))

decreasingR = sortBy (compare `on` (\x -> - abs (fst x)))

diagMatSMCheck :: Bool
diagMatSMCheck =
  let RComplex ePairs = eigs diagMat (NonSymmetric NSM) 6
      evals = map fst ePairs
      trueEvals = map (:+ 0) [1,2..]
      diff = map (realPart . abs) $ zipWith (-) evals trueEvals
  in all (< 1e-10) diff

complexEigsCheck :: Bool
complexEigsCheck =
  let RComplex ePairs' = eigs complexEigs (NonSymmetric NLM) 5
      ePairs = decreasingC ePairs'
      evals = map fst ePairs
      trueEvals = [1 :+ 2, 1 :+ (-2), 2, 1, 0.5]
      diff = map (realPart . abs) $ zipWith (-) evals trueEvals
  in all (< 1e-10) diff


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
-- |
-- Monoids for calculating various statistics in constant space. This
-- module contains algorithms that should be generally avoided unless
-- there's specific reason to use them.
module Statistics.Monoid.Extra (
    -- * Mean
    WelfordMean(..)
  , asWelfordMean
  , MeanKahan(..)
  , asMeanKahan
    -- $references
  ) where

import Data.Semigroup               (Semigroup(..))
import Data.Monoid                  (Monoid(..))
import Data.Data                    (Typeable,Data)
import Data.Vector.Unboxed          (Unbox)
import Data.Vector.Unboxed.Deriving (derivingUnbox)
import Numeric.Sum
import GHC.Generics                 (Generic)

import Statistics.Monoid.Class
-- COMPAT
import qualified Data.Vector.Generic         -- Needed for GHC7.4
import qualified Data.Vector.Generic.Mutable -- Needed for GHC7.4

----------------------------------------------------------------
-- Mean
----------------------------------------------------------------

-- | Incremental calculation of mean. Sum of elements is calculated
--   using compensated Kahan summation. It's provided only for sake of
--   completeness. 'Statistics.Monoid.Numeric.KBNSum' should be used
--   instead.
data MeanKahan = MeanKahan !Int !KahanSum
             deriving (Show,Eq,Typeable,Data,Generic)

asMeanKahan :: MeanKahan -> MeanKahan
asMeanKahan = id

instance Semigroup MeanKahan where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Monoid MeanKahan where
  mempty = MeanKahan 0 mempty
  MeanKahan 0  _  `mappend` m               = m
  m               `mappend` MeanKahan 0  _  = m
  MeanKahan n1 s1 `mappend` MeanKahan n2 s2 = MeanKahan (n1+n2) (s1 `mappend` s2)

instance Real a => StatMonoid MeanKahan a where
  addValue (MeanKahan n m) x = MeanKahan (n+1) (addValue m x)

instance CalcCount MeanKahan where
  calcCount (MeanKahan n _) = n
instance CalcMean MeanKahan where
  calcMean (MeanKahan 0 _) = Nothing
  calcMean (MeanKahan n s) = Just (kahan s / fromIntegral n)

-- | Incremental calculation of mean. Note that this algorithm doesn't
--   offer better numeric precision than plain summation. Its only
--   advantage is protection against double overflow:
--
--   > λ> calcMean $ asMeanKBN     $ reduceSample (replicate 100 1e308)
--   > Just NaN
--   > λ> calcMean $ asWelfordMean $ reduceSample (replicate 100 1e308)
--   > Just 1.0e308
--
--   Unless this feature is needed 'Statistics.Monoid.Numeric.KBNSum'
--   should be used. Algorithm is due to Welford [Welford1962]
data WelfordMean = WelfordMean !Int    -- Number of entries
                               !Double -- Current mean
  deriving (Show,Eq,Typeable,Data,Generic)

-- | Type restricted 'id'
asWelfordMean :: WelfordMean -> WelfordMean
asWelfordMean = id

instance Semigroup WelfordMean where
  (<>) = mappend
  {-# INLINE (<>) #-}

instance Monoid WelfordMean where
  mempty = WelfordMean 0 0
  mappend (WelfordMean 0 _) m = m
  mappend m (WelfordMean 0 _) = m
  mappend (WelfordMean n x) (WelfordMean k y)
    = WelfordMean (n + k) ((x*n' + y*k') / (n' + k'))
    where
      n' = fromIntegral n
      k' = fromIntegral k
  {-# INLINE mempty  #-}
  {-# INLINE mappend #-}

-- | \[ s_n = s_{n-1} + \frac{x_n - s_{n-1}}{n} \]
instance Real a => StatMonoid WelfordMean a where
  addValue (WelfordMean n m) !x
    = WelfordMean n' (m + (realToFrac x - m) / fromIntegral n')
    where
      n' = n+1
  {-# INLINE addValue #-}

instance CalcCount WelfordMean where
  calcCount (WelfordMean n _) = n
instance CalcMean WelfordMean where
  calcMean (WelfordMean 0 _) = Nothing
  calcMean (WelfordMean _ m) = Just m

----------------------------------------------------------------
-- Unboxed instances
----------------------------------------------------------------

derivingUnbox "MeanKahan"
  [t| MeanKahan -> (Int,Double,Double) |]
  [| \(MeanKahan a (KahanSum b c)) -> (a,b,c)   |]
  [| \(a,b,c) -> MeanKahan a (KahanSum b c) |]

derivingUnbox "WelfordMean"
  [t| WelfordMean -> (Int,Double) |]
  [| \(WelfordMean a b) -> (a,b)  |]
  [| \(a,b) -> WelfordMean a b    |]

-- $references
--
-- * [Welford1962] Welford, B.P. (1962) Note on a method for
--   calculating corrected sums of squares and
--   products. /Technometrics/
--   4(3):419-420. <http://www.jstor.org/stable/1266577>