{-# LANGUAGE CPP #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Eigen.Matrix (
    -- * Matrix type
    -- | Matrix aliases follows Eigen naming convention
    Matrix(..),
    MatrixXf,
    MatrixXd,
    MatrixXcf,
    MatrixXcd,
    I.Elem,
    I.CComplex,
    valid,
    -- * Matrix conversions
    fromList,
    toList,
    fromFlatList,
    toFlatList,
    generate,
    -- * Standard matrices and special cases
    empty,
    null,
    square,
    zero,
    ones,
    identity,
    constant,
    random,
    -- * Accessing matrix data
    cols,
    rows,
    dims,
    (!),
    coeff,
    unsafeCoeff,
    col,
    row,
    block,
    topRows,
    bottomRows,
    leftCols,
    rightCols,
    -- * Matrix properties
    sum,
    prod,
    mean,
    minCoeff,
    maxCoeff,
    trace,
    norm,
    squaredNorm,
    blueNorm,
    hypotNorm,
    determinant,
    -- * Generic reductions
    fold,
    fold',
    ifold,
    ifold',
    fold1,
    fold1',
    -- * Boolean reductions
    all,
    any,
    count,
    -- * Basic matrix algebra
    add,
    sub,
    mul,
    -- * Mapping over elements
    map,
    imap,
    filter,
    ifilter,
    -- * Matrix transformations
    diagonal,
    transpose,
    inverse,
    adjoint,
    conjugate,
    normalize,
    modify,
    convert,
    TriangularMode(..),
    triangularView,
    lowerTriangle,
    upperTriangle,
    -- * Matrix serialization
    encode,
    decode,
    -- * Mutable matrices
    thaw,
    freeze,
    unsafeThaw,
    unsafeFreeze,
    -- * Raw pointers
    unsafeWith,
) where

import qualified Prelude as P
import qualified Data.List as L
import Prelude hiding (null, sum, all, any, map, filter)
import Data.Tuple
import Data.Complex hiding (conjugate)
import Data.Binary hiding (encode, decode)
import qualified Data.Binary as B
import Foreign.Ptr
import Foreign.C.Types
import Foreign.C.String
import Foreign.Storable
import Foreign.Marshal.Alloc
import Text.Printf
import Control.Monad
import Control.Monad.ST
import Control.Monad.Primitive
#if __GLASGOW_HASKELL__ >= 710
#else
import Control.Applicative hiding (empty)
#endif
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Eigen.Internal as I
import qualified Data.Eigen.Matrix.Mutable as M
import qualified Data.ByteString.Lazy as BSL

-- | Matrix to be used in pure computations, uses column major memory layout, features copy-free FFI with C++ <http://eigen.tuxfamily.org Eigen> library.

data Matrix a b where
    Matrix :: I.Elem a b => !Int -> !Int -> !(VS.Vector b) -> Matrix a b

-- | Alias for single precision matrix
type MatrixXf = Matrix Float CFloat
-- | Alias for double precision matrix
type MatrixXd = Matrix Double CDouble
-- | Alias for single previsiom matrix of complex numbers
type MatrixXcf = Matrix (Complex Float) (I.CComplex CFloat)
-- | Alias for double prevision matrix of complex numbers
type MatrixXcd = Matrix (Complex Double) (I.CComplex CDouble)

-- | Pretty prints the matrix
instance (I.Elem a b, Show a) => Show (Matrix a b) where
    show m@(Matrix rows cols _) = concat [
        "Matrix ", show rows, "x", show cols,
        "\n", L.intercalate "\n" $ P.map (L.intercalate "\t" . P.map show) $ toList m, "\n"]

-- | Basic matrix math exposed through Num instance: @(*)@, @(+)@, @(-)@, `fromInteger`, `signum`, `abs`, `negate`
instance I.Elem a b => Num (Matrix a b) where
    (*) = mul
    (+) = add
    (-) = sub
    fromInteger = constant 1 1 . fromInteger
    signum = map signum
    abs = map abs
    negate = map negate

-- | Matrix binary serialization
instance I.Elem a b => Binary (Matrix a b) where
    put (Matrix rows cols vals) = do
        put $ I.magicCode (undefined :: b)
        put rows
        put cols
        put vals

    get = do
        get >>= (`when` fail "wrong matrix type") . (/= I.magicCode (undefined :: b))
        Matrix <$> get <*> get <*> get

-- | Encode the matrix as a lazy byte string
encode :: I.Elem a b => Matrix a b -> BSL.ByteString
encode = B.encode

-- | Decode matrix from the lazy byte string
decode :: I.Elem a b => BSL.ByteString -> Matrix a b
decode = B.decode

-- | Empty 0x0 matrix
{-# INLINE empty #-}
empty :: I.Elem a b => Matrix a b
empty = Matrix 0 0 VS.empty

-- | Is matrix empty?
{-# INLINE null #-}
null :: I.Elem a b => Matrix a b -> Bool
null (Matrix rows cols _) = rows == 0 && cols == 0

-- | Is matrix square?
{-# INLINE square #-}
square :: I.Elem a b => Matrix a b -> Bool
square (Matrix rows cols _) = rows == cols

-- | Matrix where all coeffs are filled with given value
{-# INLINE constant #-}
constant :: I.Elem a b => Int -> Int -> a -> Matrix a b
constant rows cols val = Matrix rows cols $ VS.replicate (rows * cols) (I.cast val)

-- | Matrix where all coeff are 0
{-# INLINE zero #-}
zero :: I.Elem a b => Int -> Int -> Matrix a b
zero rows cols = constant rows cols 0

-- | Matrix where all coeff are 1
{-# INLINE ones #-}
ones :: I.Elem a b => Int -> Int -> Matrix a b
ones rows cols = constant rows cols 1

-- | The identity matrix (not necessarily square).
identity :: I.Elem a b => Int -> Int -> Matrix a b
identity rows cols = I.performIO $ do
    m <- M.new rows cols
    I.call $ M.unsafeWith m I.identity
    unsafeFreeze m

-- | The random matrix of a given size
random :: I.Elem a b => Int -> Int -> IO (Matrix a b)
random rows cols = do
    m <- M.new rows cols
    I.call $ M.unsafeWith m I.random
    unsafeFreeze m

-- | Number of rows for the matrix
{-# INLINE rows #-}
rows :: I.Elem a b => Matrix a b -> Int
rows (Matrix rows _ _) = rows

-- | Number of columns for the matrix
{-# INLINE cols #-}
cols :: I.Elem a b => Matrix a b -> Int
cols (Matrix _ cols _) = cols

-- | Mtrix size as (rows, cols) pair
{-# INLINE dims #-}
dims :: I.Elem a b => Matrix a b -> (Int, Int)
dims (Matrix rows cols _) = (rows, cols)

-- | Matrix coefficient at specific row and col
{-# INLINE (!) #-}
(!) :: I.Elem a b => Matrix a b -> (Int, Int) -> a
(!) m (row,col) = coeff row col m

-- | Matrix coefficient at specific row and col
{-# INLINE coeff #-}
coeff :: I.Elem a b => Int -> Int -> Matrix a b -> a
coeff row col m@(Matrix rows cols _)
    | not (valid m) = error "matrix is not valid"
    | row < 0 || row >= rows = error $ printf "Matrix.coeff: row %d is out of bounds [0..%d)" row rows
    | col < 0 || col >= cols = error $ printf "Matrix.coeff: col %d is out of bounds [0..%d)" col cols
    | otherwise = unsafeCoeff row col m

-- | Unsafe version of coeff function. No bounds check performed so SEGFAULT possible
{-# INLINE unsafeCoeff #-}
unsafeCoeff :: I.Elem a b => Int -> Int -> Matrix a b -> a
unsafeCoeff row col (Matrix rows _ vals) = I.cast $ VS.unsafeIndex vals $ col * rows + row

-- | List of coefficients for the given col
{-# INLINE col #-}
col :: I.Elem a b => Int -> Matrix a b -> [a]
col c m@(Matrix rows _ _) = [coeff r c m | r <- [0..pred rows]]

-- | List of coefficients for the given row
{-# INLINE row #-}
row :: I.Elem a b => Int -> Matrix a b -> [a]
row r m@(Matrix _ cols _) = [coeff r c m | c <- [0..pred cols]]

-- | Extract rectangular block from matrix defined by startRow startCol blockRows blockCols
block :: I.Elem a b => Int -> Int -> Int -> Int -> Matrix a b -> Matrix a b
block startRow startCol blockRows blockCols m =
    generate blockRows blockCols $ \row col ->
        coeff (startRow + row) (startCol + col) m

-- | Verify matrix dimensions and memory layout
{-# INLINE valid #-}
valid :: I.Elem a b => Matrix a b -> Bool
valid (Matrix rows cols vals) = rows >= 0 && cols >= 0 && VS.length vals == rows * cols

-- | The maximum coefficient of the matrix
{-# INLINE maxCoeff #-}
maxCoeff :: (I.Elem a b, Ord a) => Matrix a b -> a
maxCoeff = fold1' max

-- | The minimum coefficient of the matrix
{-# INLINE minCoeff #-}
minCoeff :: (I.Elem a b, Ord a) => Matrix a b -> a
minCoeff = fold1' min

-- | Top @N@ rows of matrix
{-# INLINE topRows #-}
topRows :: I.Elem a b => Int -> Matrix a b -> Matrix a b
topRows n m@(Matrix _ cols _) = block 0 0 n cols m

-- | Bottom @N@ rows of matrix
{-# INLINE bottomRows #-}
bottomRows :: I.Elem a b => Int -> Matrix a b -> Matrix a b
bottomRows n m@(Matrix rows cols _) = block (rows - n) 0 n cols m

-- | Left @N@ columns of matrix
{-# INLINE leftCols #-}
leftCols :: I.Elem a b => Int -> Matrix a b -> Matrix a b
leftCols n m@(Matrix rows _ _) = block 0 0 rows n m

-- | Right @N@ columns of matrix
{-# INLINE rightCols #-}
rightCols :: I.Elem a b => Int -> Matrix a b -> Matrix a b
rightCols n m@(Matrix rows cols _) = block 0 (cols - n) rows n m

-- | Construct matrix from a list of rows, column count is detected as maximum row length. Missing values are filled with 0
fromList :: I.Elem a b => [[a]] -> Matrix a b
fromList list = Matrix rows cols vals where
    rows = length list
    cols = L.foldl' max 0 $ P.map length list
    vals = VS.create $ do
        vm <- VSM.replicate (rows * cols) (I.cast (0 `asTypeOf` (head (head list))))
        forM_ (zip [0..] list) $ \(row, vals) ->
            forM_ (zip [0..] vals) $ \(col, val) ->
                VSM.write vm (col * rows + row) (I.cast val)
        return vm

-- | Convert matrix to a list of rows
toList :: I.Elem a b => Matrix a b -> [[a]]
toList m@(Matrix rows cols vals)
    | not (valid m) = error "matrix is not valid"
    | otherwise = [[I.cast $ vals `VS.unsafeIndex` (col * rows + row) | col <- [0..pred cols]] | row <- [0..pred rows]]

-- | Build matrix of given dimensions and values from given list split on rows. Invalid list length results in error.
fromFlatList :: I.Elem a b => Int -> Int -> [a] -> Matrix a b
fromFlatList rows cols list
    | not (rows * cols == (length list)) = error $ concat ["cannot construct ", show rows, "x", show cols, " matrix from ", show $ length list, " values"]
    | otherwise = Matrix rows cols vals where
        vals = VS.create $ do
            vm <- VSM.replicate (rows * cols) (I.cast (0 `asTypeOf` (head list)))
            forM_ (zip [(col * rows + row) | row <- [0..pred rows], col <- [0..pred cols]] list) $ \(idx, val) ->
                VSM.write vm idx (I.cast val)
            return vm

-- | Convert matrix to a list by concatenating rows
toFlatList :: I.Elem a b => Matrix a b -> [a]
toFlatList m@(Matrix rows cols vals)
    | not (valid m) = error "matrix is not valid"
    | otherwise = [I.cast $ vals `VS.unsafeIndex` (col * rows + row) | row <- [0..pred rows], col <- [0..pred cols]]

-- | [generate rows cols (λ row col -> val)]
--
-- Create matrix using generator function @λ row col -> val@
--
generate :: I.Elem a b => Int -> Int -> (Int -> Int -> a) -> Matrix a b
generate rows cols f = Matrix rows cols $ VS.create $ do
    vals <- VSM.new (rows * cols)
    forM_ [0..pred rows] $ \row ->
        forM_ [0..pred cols] $ \col ->
            VSM.write vals (col * rows + row) (I.cast $ f row col)
    return vals

-- | The sum of all coefficients of the matrix
sum :: I.Elem a b => Matrix a b -> a
sum = _prop I.sum

-- | The product of all coefficients of the matrix
prod :: I.Elem a b => Matrix a b -> a
prod = _prop I.prod

-- | The mean of all coefficients of the matrix
mean :: I.Elem a b => Matrix a b -> a
mean = _prop I.mean

-- | The trace of a matrix is the sum of the diagonal coefficients and can also be computed as sum (diagonal m)
trace :: I.Elem a b => Matrix a b -> a
trace = _prop I.trace

-- | Applied to a predicate and a matrix, all determines if all elements of the matrix satisfies the predicate
all :: I.Elem a b => (a -> Bool) -> Matrix a b -> Bool
all f = VS.all (f . I.cast) . _vals

-- | Applied to a predicate and a matrix, any determines if any element of the matrix satisfies the predicate
any :: I.Elem a b => (a -> Bool) -> Matrix a b -> Bool
any f = VS.any (f . I.cast) . _vals

-- | Returns the number of coefficients in a given matrix that evaluate to true
count :: I.Elem a b => (a -> Bool) -> Matrix a b -> Int
count f = VS.foldl' (\n x -> if f (I.cast x) then succ n else n) 0 . _vals

{-| For vectors, the l2 norm, and for matrices the Frobenius norm.
    In both cases, it consists in the square root of the sum of the square of all the matrix entries.
    For vectors, this is also equals to the square root of the dot product of this with itself.
-}
norm :: I.Elem a b => Matrix a b -> a
norm = _prop I.norm

-- | For vectors, the squared l2 norm, and for matrices the Frobenius norm. In both cases, it consists in the sum of the square of all the matrix entries. For vectors, this is also equals to the dot product of this with itself.
squaredNorm :: I.Elem a b => Matrix a b -> a
squaredNorm = _prop I.squaredNorm

-- | The l2 norm of the matrix using the Blue's algorithm. A Portable Fortran Program to Find the Euclidean Norm of a Vector, ACM TOMS, Vol 4, Issue 1, 1978.
blueNorm :: I.Elem a b => Matrix a b -> a
blueNorm = _prop I.blueNorm

-- | The l2 norm of the matrix avoiding undeflow and overflow. This version use a concatenation of hypot calls, and it is very slow.
hypotNorm :: I.Elem a b => Matrix a b -> a
hypotNorm = _prop I.hypotNorm

-- | The determinant of the matrix
determinant :: I.Elem a b => Matrix a b -> a
determinant m
    | square m = _prop I.determinant m
    | otherwise = error "Matrix.determinant: non-square matrix"

-- | Adding two matrices by adding the corresponding entries together. You can use @(+)@ function as well.
add :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
add m1 m2
    | dims m1 == dims m2 = _binop const I.add m1 m2
    | otherwise = error "Matrix.add: matrices should have the same size"

-- | Subtracting two matrices by subtracting the corresponding entries together. You can use @(-)@ function as well.
sub :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
sub m1 m2
    | dims m1 == dims m2 = _binop const I.sub m1 m2
    | otherwise = error "Matrix.add: matrices should have the same size"

-- | Matrix multiplication. You can use @(*)@ function as well.
mul :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
mul m1 m2
    | cols m1 == rows m2 = _binop (\(rows, _) (_, cols) -> (rows, cols)) I.mul m1 m2
    | otherwise = error "Matrix.mul: number of columns for lhs matrix should be the same as number of rows for rhs matrix"

{- | Apply a given function to each element of the matrix.

Here is an example how to implement scalar matrix multiplication:

>>> let a = fromList [[1,2],[3,4]] :: MatrixXf

>>> a
Matrix 2x2
1.0 2.0
3.0 4.0

>>> map (*10) a
Matrix 2x2
10.0    20.0
30.0    40.0

-}
map :: I.Elem a b => (a -> a) -> Matrix a b -> Matrix a b
map f (Matrix rows cols vals) = Matrix rows cols (VS.map (I.cast . f . I.cast) vals)

{- | Apply a given function to each element of the matrix.

Here is an example how upper triangular matrix can be implemented:

>>> let a = fromList [[1,2,3],[4,5,6],[7,8,9]] :: MatrixXf

>>> a
Matrix 3x3
1.0 2.0 3.0
4.0 5.0 6.0
7.0 8.0 9.0

>>> imap (\row col val -> if row <= col then val else 0) a
Matrix 3x3
1.0 2.0 3.0
0.0 5.0 6.0
0.0 0.0 9.0

-}

imap :: I.Elem a b => (Int -> Int -> a -> a) -> Matrix a b -> Matrix a b
imap f (Matrix rows cols vals) = Matrix rows cols (VS.imap (\n -> let (c, r) = divMod n rows in I.cast . f r c . I.cast) vals)

data TriangularMode
    -- | View matrix as a lower triangular matrix.
    = Lower
    -- | View matrix as an upper triangular matrix.
    | Upper
    -- | View matrix as a lower triangular matrix with zeros on the diagonal.
    | StrictlyLower
    -- | View matrix as an upper triangular matrix with zeros on the diagonal.
    | StrictlyUpper
    -- | View matrix as a lower triangular matrix with ones on the diagonal.
    | UnitLower
    -- | View matrix as an upper triangular matrix with ones on the diagonal.
    | UnitUpper deriving (Eq, Enum, Show, Read)

-- | Triangular view extracted from the current matrix
triangularView :: I.Elem a b => TriangularMode -> Matrix a b -> Matrix a b
triangularView Lower         = imap $ \row col val -> case compare row col of { LT -> 0; _ -> val }
triangularView Upper         = imap $ \row col val -> case compare row col of { GT -> 0; _ -> val }
triangularView StrictlyLower = imap $ \row col val -> case compare row col of { GT -> val; _ -> 0 }
triangularView StrictlyUpper = imap $ \row col val -> case compare row col of { LT -> val; _ -> 0 }
triangularView UnitLower     = imap $ \row col val -> case compare row col of { GT -> val; LT -> 0; EQ -> 1 }
triangularView UnitUpper     = imap $ \row col val -> case compare row col of { LT -> val; GT -> 0; EQ -> 1 }

-- | Lower trinagle of the matrix. Shortcut for @triangularView Lower@
lowerTriangle :: I.Elem a b => Matrix a b -> Matrix a b
lowerTriangle = triangularView Lower

-- | Upper trinagle of the matrix. Shortcut for @triangularView Upper@
upperTriangle :: I.Elem a b => Matrix a b -> Matrix a b
upperTriangle = triangularView Upper

-- | Filter elements in the matrix. Filtered elements will be replaced by 0
filter :: I.Elem a b => (a -> Bool) -> Matrix a b -> Matrix a b
filter f = map (\x -> if f x then x else 0)

-- | Filter elements in the matrix. Filtered elements will be replaced by 0
ifilter :: I.Elem a b => (Int -> Int -> a -> Bool) -> Matrix a b -> Matrix a b
ifilter f = imap (\r c x -> if f r c x then x else 0)

-- | Reduce matrix using user provided function applied to each element.
fold :: I.Elem a b => (c -> a -> c) -> c -> Matrix a b -> c
fold f a (Matrix _ _ vals) = VS.foldl (\a x -> f a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element. This is strict version of 'fold'
fold' :: I.Elem a b => (c -> a -> c) -> c -> Matrix a b -> c
fold' f a (Matrix _ _ vals) = VS.foldl' (\a x -> f a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element and it's index
ifold :: I.Elem a b => (Int -> Int -> c -> a -> c) -> c -> Matrix a b -> c
ifold f a (Matrix rows _ vals) = VS.ifoldl (\a n x -> let (c,r) = divMod n rows in f r c a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element and it's index. This is strict version of 'ifold'
ifold' :: I.Elem a b => (Int -> Int -> c -> a -> c) -> c -> Matrix a b -> c
ifold' f a (Matrix rows _ vals) = VS.ifoldl' (\a n x -> let (c,r) = divMod n rows in f r c a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element.
fold1 :: I.Elem a b => (a -> a -> a) -> Matrix a b -> a
fold1 f = foldl1 f . P.map I.cast . VS.toList . _vals

-- | Reduce matrix using user provided function applied to each element. This is strict version of 'fold'
fold1' :: I.Elem a b => (a -> a -> a) -> Matrix a b -> a
fold1' f = L.foldl1' f . P.map I.cast . VS.toList . _vals

-- | Diagonal of the matrix
diagonal :: I.Elem a b => Matrix a b -> Matrix a b
diagonal = _unop (\(rows, cols) -> (min rows cols, 1)) I.diagonal

{- | Inverse of the matrix

For small fixed sizes up to 4x4, this method uses cofactors. In the general case, this method uses PartialPivLU decomposition
-}
inverse :: I.Elem a b => Matrix a b -> Matrix a b
inverse m
    | square m = _unop id I.inverse m
    | otherwise = error "Matrix.inverse: non-square matrix"

-- | Adjoint of the matrix
adjoint :: I.Elem a b => Matrix a b -> Matrix a b
adjoint = _unop swap I.adjoint

-- | Transpose of the matrix
transpose :: I.Elem a b => Matrix a b -> Matrix a b
transpose = _unop swap I.transpose

-- | Conjugate of the matrix
conjugate :: I.Elem a b => Matrix a b -> Matrix a b
conjugate = _unop id I.conjugate

-- | Nomalize the matrix by deviding it on its 'norm'
normalize :: I.Elem a b => Matrix a b -> Matrix a b
normalize (Matrix rows cols vals) = I.performIO $ do
    vals <- VS.thaw vals
    VSM.unsafeWith vals $ \p ->
        I.call $ I.normalize p (I.cast rows) (I.cast cols)
    Matrix rows cols <$> VS.unsafeFreeze vals

-- | Apply a destructive operation to a matrix. The operation will be performed in place if it is safe to do so and will modify a copy of the matrix otherwise.
modify :: I.Elem a b => (forall s. M.MMatrix a b s -> ST s ()) -> Matrix a b -> Matrix a b
modify f (Matrix rows cols vals) = Matrix rows cols (VS.modify (f . M.MMatrix rows cols) vals)

-- | Convert matrix to different type using user provided element converter
convert :: (I.Elem a b, I.Elem c d) => (a -> c) -> Matrix a b -> Matrix c d
convert f (Matrix rows cols vals) = Matrix rows cols $ VS.map (I.cast . f . I.cast) vals

-- | Yield an immutable copy of the mutable matrix
freeze :: I.Elem a b => PrimMonad m => M.MMatrix a b (PrimState m) -> m (Matrix a b)
freeze (M.MMatrix mrows mcols mvals) = VS.freeze mvals >>= return . Matrix mrows mcols

-- | Yield a mutable copy of the immutable matrix
thaw :: I.Elem a b => PrimMonad m => Matrix a b -> m (M.MMatrix a b (PrimState m))
thaw (Matrix rows cols vals) = VS.thaw vals >>= return . M.MMatrix rows cols

-- | Unsafe convert a mutable matrix to an immutable one without copying. The mutable matrix may not be used after this operation.
unsafeFreeze :: I.Elem a b => PrimMonad m => M.MMatrix a b (PrimState m) -> m (Matrix a b)
unsafeFreeze (M.MMatrix mrows mcols mvals) = VS.unsafeFreeze mvals >>= return . Matrix mrows mcols

-- | Unsafely convert an immutable matrix to a mutable one without copying. The immutable matrix may not be used after this operation.
unsafeThaw :: I.Elem a b => PrimMonad m => Matrix a b -> m (M.MMatrix a b (PrimState m))
unsafeThaw (Matrix rows cols vals) = VS.unsafeThaw vals >>= return . M.MMatrix rows cols

-- | Pass a pointer to the matrix's data to the IO action. The data may not be modified through the pointer.
unsafeWith :: I.Elem a b => Matrix a b -> (Ptr b -> CInt -> CInt -> IO c) -> IO c
unsafeWith m@(Matrix rows cols vals) f
    | not (valid m) = fail "Matrix.unsafeWith: matrix layout is invalid"
    | otherwise = VS.unsafeWith vals $ \p -> f p (I.cast rows) (I.cast cols)

{-# INLINE _prop #-}
_prop :: I.Elem a b => (Ptr b -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> a
_prop f m = I.cast $ I.performIO $ alloca $ \p -> do
    I.call $ unsafeWith m (f p)
    peek p

{-# INLINE _binop #-}
_binop :: I.Elem a b => ((Int, Int) -> (Int, Int) -> (Int, Int)) -> (Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> Matrix a b -> Matrix a b
_binop f g m1 m2 = I.performIO $ do
    m0 <- uncurry M.new $ f (dims m1) (dims m2)
    M.unsafeWith m0 $ \vals0 rows0 cols0 ->
        unsafeWith m1 $ \vals1 rows1 cols1 ->
            unsafeWith m2 $ \vals2 rows2 cols2 ->
                I.call $ g
                    vals0 rows0 cols0
                    vals1 rows1 cols1
                    vals2 rows2 cols2
    unsafeFreeze m0

{-# INLINE _unop #-}
_unop :: I.Elem a b => ((Int,Int) -> (Int,Int)) -> (Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> Matrix a b
_unop f g m1 = I.performIO $ do
    m0 <- uncurry M.new $ f (dims m1)
    M.unsafeWith m0 $ \vals0 rows0 cols0 ->
        unsafeWith m1 $ \vals1 rows1 cols1 ->
            I.call $ g
                vals0 rows0 cols0
                vals1 rows1 cols1
    unsafeFreeze m0

{-# INLINE _vals #-}
_vals :: I.Elem a b => Matrix a b -> VS.Vector b
_vals (Matrix _ _ vals) = vals


{-# LANGUAGE QuasiQuotes #-}
module Myelin.PyNN.SpiNNakerSpec where

import Control.Lens
import Control.Monad.State.Lazy
import Control.Monad.Trans.Except
import Data.Aeson
import Data.ByteString.Lazy.Char8
import Data.Either
import qualified Data.Map.Strict as Map
import Data.String.Interpolate
import Numeric.LinearAlgebra
import Text.Regex as Regex

import Myelin.Model
import Myelin.Neuron
import Myelin.PyNN.PyNN
import Myelin.SNN

import Test.Hspec

main :: IO ()
main = hspec spec

eval :: PyNNState a -> Either String a
eval s = evalState (runExceptT s) emptyPyNNModel

exec :: PyNNState a -> PyNNModel
exec s = execState (runExceptT s) emptyPyNNModel

spec :: Spec
spec = do
  describe "PyNN backend" $ do
    it "can translate an entire SNN model to a SpiNNaker script" $ do
      let input = Population 2 if_cond_exp "p0" 0
      let hidden = Population 4 if_cond_exp "p1" 1
      let output = Population 3 if_cond_exp "p2" 2
      let dictWithType = unpack $ encode if_cond_exp
      let dict = Regex.subRegex typeRegex dictWithType ""
      let effect1 = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1)))
      let effect2 = Static Excitatory (AllToAll (Biases [0.0, 1.1, 2.2]) (WeightGenerator (GaussianRandom 2 1)))
      let edges = [ DenseProjection effect1 input hidden, DenseProjection effect2 hidden output ]
      let network = Network 0 [input] [hidden] edges [output]
      let task = Task SpiNNaker network 50
      let preample = PyNNPreample "# some config"
      let code = [i|import numpy as np
import volrpynn.spiNNaker as v
import pyNN.spiNNaker as pynn

# some config

p0 = pynn.Population(2, pynn.IF_cond_exp(**#{dict}))
p1 = pynn.Population(4, pynn.IF_cond_exp(**#{dict}))
p2 = pynn.Population(3, pynn.IF_cond_exp(**#{dict}))
layer0 = v.Dense(p0, p1, weights=1.0, biases=0.0)
layer1 = v.Dense(p1, p2, weights=np.random.normal(2.0, 1.0, (4, 3)), biases=np.array([0.0,1.1,2.2]))
l_decode = v.Decode(p2)
model = v.Model(layer0, layer1, l_decode)

optimiser = v.GradientDescentOptimiser(0.1, simulation_time=50.0)
if __name__ == "__main__":
    v.Main(model).train(optimiser)
|]
      translate task preample `shouldBe` Right code


{-# LANGUAGE CPP                     #-}
{-# LANGUAGE FlexibleContexts        #-}
{-# LANGUAGE TypeApplications        #-}
{-# LANGUAGE ScopedTypeVariables     #-}
{-# LANGUAGE ConstraintKinds         #-}
{-# LANGUAGE TypeFamilies            #-}
{-# LANGUAGE TypeOperators           #-}
{-# LANGUAGE DataKinds               #-}

-- For OkLC as a class
{-# LANGUAGE UndecidableInstances    #-}
{-# LANGUAGE FlexibleInstances       #-}
{-# LANGUAGE MultiParamTypeClasses   #-}

{-# OPTIONS_GHC -Wall #-}

{-# OPTIONS -Wno-type-defaults #-}

{-# OPTIONS_GHC -Wno-missing-signatures #-}
{-# OPTIONS_GHC -Wno-unused-imports #-}

{-# OPTIONS_GHC -fsimpl-tick-factor=500 #-}
{-# OPTIONS_GHC -dsuppress-idinfo #-}
{-# OPTIONS_GHC -fdicts-strict #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Main where

import Prelude hiding (unzip,zip,zipWith) -- (id,(.),curry,uncurry)
import qualified Prelude as P

import Data.Monoid (Sum(..))
import Data.Foldable (fold)
import Control.Applicative (liftA2)
import Control.Arrow (second)
import Control.Monad ((<=<))
import Data.List (unfoldr)  -- TEMP
import Data.Complex (Complex)
import GHC.Float (int2Double)

import qualified ConCat.AltCat as A
import ConCat.AltCat
 (toCcc,toCcc',unCcc,unCcc',conceal,(:**:)(..),Ok,Ok2,U2,equal)

import ConCat.Rebox

import ConCat.Circuit (GenBuses,(:>))
import ConCat.Syntactic (Syn,render)
import ConCat.RunCircuit (run)

type EC = Syn :**: (:>)

runU2 :: U2 a b -> IO ()
runU2 = print
{-# INLINE runU2 #-}

type GO a b = (GenBuses a, Ok2 (:>) a b)

runSyn :: Syn a b -> IO ()
runSyn syn = putStrLn ('\n' : render syn)
{-# INLINE runSyn #-}

runSynCirc :: GO a b => String -> EC a b -> IO ()
runSynCirc nm (syn :**: circ) = runSyn syn >> runCirc nm circ
{-# INLINE runSynCirc #-}

runCirc :: GO a b => String -> (a :> b) -> IO ()
runCirc nm circ = run nm [] circ
{-# INLINE runCirc #-}

add5 :: Double -> Double
add5 x = x + 5

decide :: Int -> Bool
decide 4 = True
decide _ = False

fix :: (a -> a) -> a
fix f = let {x = f x} in x

fac :: Int -> Int
fac 1 = 1
fac n = n*fac(n-1)

main :: IO ()
--main = print "hello world!"

main = sequence_ [
   putChar '\n' -- return ()

--  -- Circuit graphs
  , runSynCirc "add"         $ toCcc $ (+) @Double
  , runSynCirc "add5"        $ toCcc add5
  , runSynCirc "add-uncurry" $ toCcc $ uncurry ((+) @Double)
  , runSynCirc "decide"      $ toCcc $ decide
--  , runSynCirc "fac"         $ toCcc (fix (\rec n -> if n == 0 then 1 else n * rec (n-1)) :: Int -> Int)

--  , runSynCirc "fac"         $ toCcc $ fac
--  , runSynCirc "dup"         $ toCcc $ A.dup @(->) @Int
--  , runSynCirc "fst"         $ toCcc $ fst @R @R
--  , runSynCirc "twice"       $ toCcc $ twice @R
--  , runSynCirc "sqr"         $ toCcc $ sqr @R
--  , runSynCirc "complex-mul" $ toCcc $ uncurry ((*) @C)
--  , runSynCirc "magSqr"      $ toCcc $ magSqr @R
--  , runSynCirc "cosSinProd"  $ toCcc $ cosSinProd @R
--  , runSynCirc "xp3y"        $ toCcc $ \ (x,y) -> x + 3 * y :: R
--  , runSynCirc "horner"      $ toCcc $ horner @R [1,3,5]
--  , runSynCirc "cos-2xx"     $ toCcc $ \ x -> cos (2 * x * x) :: R
--
  ]


module HScheme.Data where

import Control.Monad.Except
import Data.Complex ( Complex )
import Text.ParserCombinators.Parsec ( ParseError )
import qualified Data.Vector as V

data Value = Atom String
    | List [Value]
    | DottedList [Value] Value
    | Vector (V.Vector Value)
    | Number Integer
    | Float Double
    | Ratio Rational
    | Complex (Complex Double)
    | String String
    | Character Char
    | Bool Bool
    deriving (Eq)

instance Show Value where
    show (String s) = "\"" ++ s ++ "\""
    show (Character c) = "#\\" ++ [c]
    show (Atom name) = name
    show (Number n) = show n
    show (Float n) = show n
    show (Ratio n) = show n
    show (Complex n) = show n
    show (Bool True) = "#t"
    show (Bool False) = "#f"
    show (List xs) = "(" ++ unwordsList xs ++ ")"
    show (DottedList head tail) = "(" ++ unwordsList head ++ " . " ++ show tail ++ ")"
    show (Vector vs) = "#(" ++ unwordsList (V.toList vs) ++ ")"

unwordsList :: [Value] -> String
unwordsList = unwords . map show

{-
    -- Error handling --
-}

data Error = NumArgs Integer [Value]
    | TypeMismatch String Value
    | Parser ParseError
    | BadSpecialForm String Value
    | NotFunction String String
    | UnboundVar String String
    | Default String

instance Show Error where
    show (NumArgs expected found)       = "Expected " ++ show expected
                                        ++ " args; found values " ++ unwordsList found
    show (TypeMismatch expected found)  = "Invalid Type: expected " ++ expected
                                        ++ ", found " ++ show found
    show (Parser err)                   = "Parse error at " ++ show err
    show (BadSpecialForm message form)  = message ++ ": " ++ show form
    show (NotFunction message func)     = message ++ ": " ++ show func
    show (UnboundVar message varname)   = message ++ ": " ++ varname
    show (Default message)              = message

type ThrowsErr = Either Error

trapErr :: (MonadError a m, Show a) => m String -> m String
trapErr action = catchError action (return . show)

extractValue :: ThrowsErr a -> a
extractValue (Right val) = val


module Statistics.Classification (
    module Classification
  ) where

import Statistics.Classification.ConfusionMatrix as Classification
import Statistics.Classification.ROC as Classification
import Statistics.Classification.Types as Classification


module Vote
  ( processVote,
    reprocessVotes,
    encryptBallot,
    decryptBallot,
  )
where

import Conduit (foldlC, mapC)
import Control.Lens
import Control.Monad.Trans.Reader (ReaderT)
import Data.Binary.Get (getWord32be, runGet)
import Data.Binary.Put (putWord32be, runPut)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Data.Time (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime, utcTimeToPOSIXSeconds)
import Database.Persist.Sql (SqlBackend)
import Import hiding (toList)
import Model.IsaacVersion
import Numeric.LinearAlgebra
import Ranks
import qualified Web.ClientSession as WS

processVote :: (MonadIO m) => IsaacVersion -> Int -> Int -> UTCTime -> Text -> Text -> ReaderT SqlBackend m ()
processVote ver w l timestamp voter rawBallot = do
  Entity w' _ <- getBy404 $ UniqueItem ver w
  Entity l' _ <- getBy404 $ UniqueItem ver l
  insert400_ $ Ballot rawBallot timestamp
  insert400_ $ Vote ver w' l' timestamp voter
  return ()

reprocessVotes :: (MonadIO m, MonadResource m) => ReaderT SqlBackend m ()
reprocessVotes = do
  items <- M.fromList . map (\(Entity a b) -> (a, b)) <$> selectList [] []
  bms <- runConduit $ selectSource [] [] .| mapC entityVal .| beatMatrix
  recalculatePairs items bms
  now <- liftIO getCurrentTime
  deleteWhere [BallotTimestamp <=. (- validTime) `addUTCTime` now]
  return ()

type Matchup = (Key Item, Key Item)

type BeatMatrix = M.Map Matchup Int

beatMatrix :: Monad m => ConduitT Vote o m (M.Map IsaacVersion BeatMatrix)
beatMatrix = foldlC upd M.empty
  where
    upd bm v = bm & at ver . non M.empty . at (w, l) . non 0 +~ 1
      where
        w = v ^. voteWinner
        l = v ^. voteLoser
        ver = v ^. voteVersion

ranks :: BeatMatrix -> [Key Item] -> [(Key Item, R)]
ranks bm items =
  zip items . toList $ ilsrPairwise bm' 0.01
  where
    n = length items
    itemToIdx = M.fromList (zip items [0 ..])
    bm' = assoc (n, n) 0 (convR <$> M.toList bm)
    convR ((w, l), s) = ((itemToIdx M.! w, itemToIdx M.! l), fromIntegral s)

recalculatePairs :: (MonadIO m, MonadResource m) => M.Map (Key Item) Item -> M.Map IsaacVersion BeatMatrix -> ReaderT SqlBackend m ()
recalculatePairs items bms =
  ifor_ bms $ \ver bm ->
    for_ (ranks bm (itemsFor ver items)) $ \(itemId, rating) ->
      update
        itemId
        [ ItemRating =. rating,
          ItemVotes =. sumOf (ifolded . ifiltered (\(i1, i2) _ -> i1 == itemId || i2 == itemId)) bm
        ]
  where
    itemsFor ver = findIndicesOf ifolded (\i -> i ^. itemVersion == ver)

encodeBallot :: UTCTime -> Int -> Int -> B.ByteString
encodeBallot expiry winner loser = toStrict . runPut $ do
  putWord32be (truncate . utcTimeToPOSIXSeconds $ expiry)
  putWord32be (fromIntegral winner)
  putWord32be (fromIntegral loser)

decodeBallot :: B.ByteString -> (UTCTime, Int, Int)
decodeBallot = go . fromStrict
  where
    go = runGet $ do
      expiry <- getWord32be
      winner <- getWord32be
      loser <- getWord32be
      return
        ( posixSecondsToUTCTime (fromIntegral expiry),
          fromIntegral winner,
          fromIntegral loser
        )

-- 1 hour
validTime :: NominalDiffTime
validTime = 3600

encryptBallot :: UTCTime -> Int -> Int -> Handler Text
encryptBallot now winner loser = do
  key <- getsYesod appBallotKey
  out <- liftIO $ WS.encryptIO key (encodeBallot (validTime `addUTCTime` now) winner loser)
  return . T.pack . BC.unpack $ out

decryptBallot :: UTCTime -> Text -> Handler (Int, Int)
decryptBallot now b = do
  key <- getsYesod appBallotKey
  (expiry, winner, loser) <-
    case decodeBallot <$> (WS.decrypt key . BC.pack . T.unpack) b of
      Just r -> pure r
      Nothing -> error "invalid ballot!"
  if expiry < now
    then error "expired ballot!"
    else return (winner, loser)


{-# LANGUAGE OverloadedStrings #-}

module JsonParser where

import           Data.Aeson              hiding ( Series )
import           Data.Aeson.Types               ( parseMaybe )
import qualified Data.ByteString.Lazy          as BSL
import           Data.Complex                   ( Complex((:+)) )

import           Form

data Vec = Vec {
    number :: Int,
    real   :: Double,
    imag   :: Double
} deriving (Show)

instance ToJSON Vec where
    toJSON (Vec number real imag) =
        object ["number" .= number, "real" .= real, "imag" .= imag]

instance FromJSON Vec where
    parseJSON = withObject "Vec"
        $ \v -> Vec <$> v .: "number" <*> v .: "real" <*> v .: "imag"

newtype Series = Series {
    series :: [Vec]
} deriving (Show)

instance ToJSON Series where
    toJSON (Series series) = object ["series" .= series]

instance FromJSON Series where
    parseJSON = withObject "Series" $ \s -> Series <$> s .: "series"

instance Semigroup Series where
    (Series xs) <> (Series ys) = Series (xs <> ys)

instance Monoid Series where
    mempty = Series []

convert :: Vec -> Form
convert (Vec n r i) = form (r :+ i) n

getForm :: Series -> [Form]
getForm (Series xs) = map convert xs

jsonToCoord :: FilePath -> IO [(Int, Int)]
jsonToCoord path = formsToCoords <$> jsonToForm path

jsonToForm :: FilePath -> IO [Form]
jsonToForm path = do
    json <- BSL.readFile path
    case decode json of
        Just js -> return $ getForm js
        Nothing -> error $ "invalid json file: " ++ path

jsonToCoords :: [FilePath] -> IO [(Int, Int)]
jsonToCoords = foldr (mappend . jsonToCoord) mempty


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Grenade.Layers.FullyConnected (
    FullyConnected (..)
  , FullyConnected' (..)
  , randomFullyConnected
  ) where

import           Control.DeepSeq
import           Control.Monad.Primitive        (PrimBase, PrimState)
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           System.Random.MWC              hiding (create)
import           Data.List                      (foldl1')
import           Data.Either                    (fromRight)
import           Data.Proxy
import           Data.Serialize

import qualified Numeric.LinearAlgebra          as LA
import           Numeric.LinearAlgebra.Static

import           Grenade.Core
import           Grenade.Layers.Internal.Update
import           Grenade.Onnx
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore
import           Grenade.Types

import           Lens.Micro

-- | A basic fully connected (or inner product) neural network layer.
data FullyConnected i o = FullyConnected
                        !(FullyConnected' i o)   -- Neuron weights
                        !(ListStore (FullyConnected' i o))   -- momentum store
                        deriving (Generic)

instance NFData (FullyConnected i o) where
  rnf (FullyConnected w store) = rnf w `seq` rnf store

data FullyConnected' i o = FullyConnected'
                         !(R o)   -- Bias
                         !(L o i) -- Activations
                        deriving (Generic)

instance NFData (FullyConnected' i o) where
  rnf (FullyConnected' b w) = rnf b `seq` rnf w

instance Show (FullyConnected i o) where
  show FullyConnected {} = "FullyConnected"

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => UpdateLayer (FullyConnected i o) where
  type Gradient (FullyConnected i o) = (FullyConnected' i o)
  type MomentumStore (FullyConnected i o) = ListStore (FullyConnected' i o)
  runUpdate opt@OptSGD{} x@(FullyConnected (FullyConnected' oldBias oldActivations) store) (FullyConnected' biasGradient activationGradient) =
    let (FullyConnected' oldBiasMomentum oldMomentum) = getData opt x store
        VectorResultSGD newBias newBiasMomentum = descendVector opt (VectorValuesSGD oldBias biasGradient oldBiasMomentum)
        MatrixResultSGD newActivations newMomentum = descendMatrix opt (MatrixValuesSGD oldActivations activationGradient oldMomentum)
        newStore = setData opt x store (FullyConnected' newBiasMomentum newMomentum)
     in FullyConnected (FullyConnected' newBias newActivations) newStore
  runUpdate opt@OptAdam{} x@(FullyConnected (FullyConnected' oldBias oldActivations) store) (FullyConnected' biasGradient activationGradient) =
    let [FullyConnected' oldMBias oldMActivations, FullyConnected' oldVBias oldVActivations] = getData opt x store
        VectorResultAdam newBias newMBias newVBias = descendVector opt (VectorValuesAdam (getStep store) oldBias biasGradient oldMBias oldVBias)
        MatrixResultAdam newActivations newMActivations newVActivations = descendMatrix opt (MatrixValuesAdam (getStep store) oldActivations activationGradient oldMActivations oldVActivations)
        newStore = setData opt x store [FullyConnected' newMBias newMActivations, FullyConnected' newVBias newVActivations]
    in FullyConnected (FullyConnected' newBias newActivations) newStore

  reduceGradient grads = FullyConnected' (dvmap (/l) bs) (dmmap (/l) as)
    where
      FullyConnected' bs as = foldl1' (\(FullyConnected' bs as) (FullyConnected' bs' as') -> FullyConnected' (bs + bs') (as + as')) grads
      l = fromIntegral $ length grads :: RealNum

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => LayerOptimizerData (FullyConnected i o) (Optimizer 'SGD) where
  type MomentumDataType (FullyConnected i o) (Optimizer 'SGD) = FullyConnected' i o
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = FullyConnected' (konst 0) (konst 0)

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => LayerOptimizerData (FullyConnected i o) (Optimizer 'Adam) where
  type MomentumDataType (FullyConnected i o) (Optimizer 'Adam) = FullyConnected' i o
  type MomentumExpOptResult (FullyConnected i o) (Optimizer 'Adam) = [FullyConnected' i o]
  getData = getListStore
  setData = setListStore
  newData _ _ = FullyConnected' (konst 0) (konst 0)

instance (KnownNat i, KnownNat o) => FoldableGradient (FullyConnected' i o) where
  mapGradient f (FullyConnected' bias activations) = FullyConnected' (dvmap f bias) (dmmap f activations)
  squaredSums (FullyConnected' bias activations) = [sumV . squareV $ bias, sumM . squareM $ activations]

instance (KnownNat i, KnownNat o, KnownNat (i * o)) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
  type Tape (FullyConnected i o) ('D1 i) ('D1 o) = R i
  -- Do a matrix vector multiplication and return the result.
  runForwards (FullyConnected (FullyConnected' wB wN) _) (S1D v) = (v, S1D (wB + wN #> v))

  -- Run a backpropogation step for a full connected layer.
  runBackwards (FullyConnected (FullyConnected' _ wN) _) x (S1D dEdy) =
          let wB'  = dEdy
              mm'  = dEdy `outer` x
              -- calcluate derivatives for next step
              dWs  = tr wN #> dEdy
          in  (FullyConnected' wB' mm', S1D dWs)

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected i o) where
  put (FullyConnected w ms) = put w >> put ms
  get = FullyConnected <$> get <*> get

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected' i o) where
  put (FullyConnected' b w) = do
    putListOf put . LA.toList . extract $ b
    putListOf put . LA.toList . LA.flatten . extract $ w
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy i)
      b     <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
      k     <- maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
      return $ FullyConnected' b k

instance (KnownNat i, KnownNat o, KnownNat (i*o)) => RandomLayer (FullyConnected i o) where
  createRandomWith = randomFullyConnected

randomFullyConnected :: forall m i o . (PrimBase m, KnownNat i, KnownNat o, KnownNat (i*o))
                     => WeightInitMethod -> Gen (PrimState m) -> m (FullyConnected i o)
randomFullyConnected m gen = do
  wN <- getRandomMatrix i o m gen
  wB <- getRandomVector i o m gen
  return $ FullyConnected (FullyConnected' wB wN) mkListStore
  where i = natVal (Proxy :: Proxy i)
        o = natVal (Proxy :: Proxy o)

instance OnnxOperator (FullyConnected i o) where
  onnxOpTypeNames _ = ["Gemm"]

instance (KnownNat i, KnownNat o) => OnnxLoadable (FullyConnected i o) where
  loadOnnxNode inits node = case (node ^. #input) of
    [_, b, c] -> do
      -- FIXME: Proper attribute checking
      -- node `doesNotHaveAttribute` "alpha"
      -- node `doesNotHaveAttribute` "transA"
      -- node `doesNotHaveAttribute` "transB"

      let beta = fromRight 1 (readFloatAttributeToRealNum "beta" node)
      loadedB <- readInitializerMatrix inits b
      loadedC <- readInitializerVector inits c

      return $ FullyConnected (FullyConnected' loadedC (dmmap (*beta) loadedB)) mkListStore
    _         -> onnxIncorrectNumberOfInputs


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module Math.HiddenMarkovModel.Private where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Distribution (State(State))

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.LinearAlgebra.Util as LinAlg
import qualified Numeric.Container as NC
import qualified Data.Packed.Development as Dev
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Numeric.Container ((<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import Control.DeepSeq (NFData, rnf)
import Foreign.Storable (Storable)

import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import qualified Data.Semigroup as Sg
import qualified Data.List as List
import Data.Traversable (Traversable, mapAccumL)
import Data.Tuple.HT (mapPair, mapFst, mapSnd, swap)

{- |
A Hidden Markov model consists of a number of (hidden) states
and a set of emissions.
There is a vector for the initial probability of each state
and a matrix containing the probability for switching
from one state to another one.
The 'distribution' field points to probability distributions
that associate every state with emissions of different probability.
Famous distribution instances are discrete and Gaussian distributions.
See "Math.HiddenMarkovModel.Distribution" for details.

The transition matrix is transposed
with respect to popular HMM descriptions.
But I think this is the natural orientation, because this way
you can write \"transition matrix times probability column vector\".

The type has two type parameters,
although the one for the distribution would be enough.
However, replacing @prob@ by @Distr.Probability distr@
would prohibit the derived Show and Read instances.
-}
data T distr prob =
   Cons {
      initial :: Vector prob,
      transition :: Matrix prob,
      distribution :: distr
   }
   deriving (Show, Read)

instance
   (NFData distr, NFData prob, Storable prob) =>
      NFData (T distr prob) where
   rnf hmm = rnf (initial hmm, transition hmm, distribution hmm)

emission ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> emission -> Vector prob
emission  =  Distr.emissionProb . distribution

forward ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> prob
forward hmm = NC.sumElements . NonEmpty.last . alpha hmm

alpha ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   NonEmpty.T f emission -> NonEmpty.T f (Vector prob)
alpha hmm (NonEmpty.Cons x xs) =
   NonEmpty.scanl
      (\alphai xi -> NC.mul (emission hmm xi) (transition hmm <> alphai))
      (NC.mul (emission hmm x) (initial hmm))
      xs

backward ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> prob
backward hmm (NonEmpty.Cons x xs) =
   NC.sumElements $
   NC.mul (initial hmm) $
   NC.mul (emission hmm x) $
   NonEmpty.head $ beta hmm xs

beta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   f emission -> NonEmpty.T f (Vector prob)
beta hmm =
   NonEmpty.scanr
      (\xi betai -> NC.mul (emission hmm xi) betai <> transition hmm)
      (NC.constant 1 (NC.dim $ initial hmm))

alphaBeta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   NonEmpty.T f emission ->
   (prob, NonEmpty.T f (Vector prob), NonEmpty.T f (Vector prob))
alphaBeta hmm xs =
   let alphas = alpha hmm xs
       betas = beta hmm $ NonEmpty.tail xs
       recipLikelihood = recip $ NC.sumElements $ NonEmpty.last alphas
   in  (recipLikelihood, alphas, betas)

xiFromAlphaBeta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> prob ->
   NonEmpty.T [] emission ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob) ->
   [Matrix prob]
xiFromAlphaBeta hmm recipLikelihood xs alphas betas =
   zipWith3
      (\x alpha0 beta1 ->
         NC.scale recipLikelihood $
         NC.mul
            (NC.outer (NC.mul (emission hmm x) beta1) alpha0)
            (transition hmm))
      (NonEmpty.tail xs)
      (NonEmpty.init alphas)
      (NonEmpty.tail betas)

zetaFromXi ::
   (NC.Product prob) => T distr prob -> [Matrix prob] -> [Vector prob]
zetaFromXi hmm xis =
   map (NC.constant 1 (Matrix.rows $ transition hmm) <>) xis

zetaFromAlphaBeta ::
   (NC.Container Vector prob) =>
   prob ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob)
zetaFromAlphaBeta recipLikelihood alphas betas =
   fmap (NC.scale recipLikelihood) $
   NonEmptyC.zipWith NC.mul alphas betas

{- |
In constrast to Math.HiddenMarkovModel.reveal
this does not normalize the vector.
This is slightly simpler but for long sequences
the product of probabilities might be smaller
than the smallest representable number.
-}
reveal ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> NonEmpty.T f State
reveal hmm (NonEmpty.Cons x xs) =
   fmap State $
   uncurry (NonEmpty.scanr Dev.at') $
   mapFst NC.maxIndex $
   mapAccumL
      (\alphai xi ->
         swap $ mapSnd (NC.mul (emission hmm xi)) $
         matrixMaxMul (transition hmm) alphai)
      (NC.mul (emission hmm x) (initial hmm)) xs

matrixMaxMul ::
   (NC.Container Vector a) =>
   Matrix a -> Vector a -> (Vector Int, Vector a)
matrixMaxMul m v =
   mapPair (Vector.fromList, Vector.fromList) $ unzip $
   map ((\x -> (NC.maxIndex x, NC.maxElement x)) . NC.mul v) $
   Matrix.toRows m

{- |
A trained model is a temporary form of a Hidden Markov model
that we need during the training on multiple training sequences.
It allows to collect knowledge over many sequences with 'mergeTrained',
even with mixed supervised and unsupervised training.
You finish the training by converting the trained model
back to a plain modul using 'finishTraining'.

You can create a trained model in three ways:

* supervised training using an emission sequence with associated states,

* unsupervised training using an emission sequence and an existing Hidden Markov Model,

* derive it from state sequence patterns, cf. "Math.HiddenMarkovModel.Pattern".
-}
data Trained distr prob =
   Trained {
      trainedInitial :: Vector prob,
      trainedTransition :: Matrix prob,
      trainedDistribution :: distr
   }
   deriving (Show, Read)

instance
   (NFData distr, NFData prob, Storable prob) =>
      NFData (Trained distr prob) where
   rnf hmm =
      rnf (trainedInitial hmm, trainedTransition hmm, trainedDistribution hmm)

sumTransitions ::
   (NC.Container Vector e) =>
   T distr e -> [Matrix e] -> Matrix e
sumTransitions hmm =
   List.foldl' NC.add (NC.konst 0 $ LinAlg.size $ transition hmm)

{- |
Baum-Welch algorithm
-}
trainUnsupervised ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> NonEmpty.T [] emission -> Trained tdistr prob
trainUnsupervised hmm xs =
   let (recipLikelihood, alphas, betas) = alphaBeta hmm xs
       zetas = zetaFromAlphaBeta recipLikelihood alphas betas

   in  Trained {
          trainedInitial = NonEmpty.head zetas,
          trainedTransition =
             sumTransitions hmm $
             xiFromAlphaBeta hmm recipLikelihood xs alphas betas,
          trainedDistribution =
             Distr.accumulateEmissions $ map (zip (NonEmpty.flatten xs)) $
             List.transpose $ map Vector.toList $ NonEmpty.flatten zetas
       }

mergeTrained ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
   Trained tdistr prob -> Trained tdistr prob -> Trained tdistr prob
mergeTrained hmm0 hmm1 =
   Trained {
      trainedInitial = NC.add (trainedInitial hmm0) (trainedInitial hmm1),
      trainedTransition =
         NC.add (trainedTransition hmm0) (trainedTransition hmm1),
      trainedDistribution =
         Distr.combine
            (trainedDistribution hmm0) (trainedDistribution hmm1)
   }

instance
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
      Sg.Semigroup (Trained tdistr prob) where
   (<>) = mergeTrained

toCells ::
   (Distr.CSV distr, Algo.Field prob, Show prob) =>
   T distr prob -> [[String]]
toCells hmm =
   (HMMCSV.cellsFromVector $ initial hmm) :
   (HMMCSV.cellsFromMatrix $ transition hmm) ++
   [] :
   (Distr.toCells $ distribution hmm)

parseCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   HMMCSV.CSVParser (T distr prob)
parseCSV = do
   v <- HMMCSV.parseNonEmptyVectorCells
   m <- HMMCSV.parseSquareMatrixCells $ Vector.dim v
   HMMCSV.skipEmptyRow
   distr <- Distr.parseCells $ Vector.dim v
   return $ Cons {
      initial = v,
      transition = m,
      distribution = distr
   }


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Main where

import qualified AnalyticsLaop as LM
import qualified AnalyticsHMatrix as HM
import qualified AnalyticsMatrix as DM
import qualified LAoP.Matrix.Type as LM
import qualified Numeric.LinearAlgebra as HM
import qualified Numeric.LinearAlgebra.Data as HM
import qualified Numeric.LinearAlgebra.HMatrix as HM
import qualified Data.Matrix as DM
import GHC.TypeLits
import Data.Proxy
import Criterion.Main
import Test.QuickCheck

randomBDLaop ::
             forall jn en .
             ( KnownNat jn,
               KnownNat en,
               LM.FromLists Int (LM.FromNat (jn + 1)) (LM.FromNat 3),
               LM.FromLists Int (LM.FromNat (jn + 1)) (),
               LM.FromLists Int (LM.FromNat (en + 1)) (LM.FromNat 5),
               LM.FromLists Int (LM.FromNat (en + 1)) (LM.FromNat 3),
               LM.FromLists Int (LM.FromNat (en + 1)) (LM.FromNat 2)
             ) => Gen (LM.BD jn en)
randomBDLaop = do
  let jn' = fromInteger (natVal (Proxy :: Proxy jn)) + 1
      en' = fromInteger (natVal (Proxy :: Proxy en)) + 1
  descs <- vectorOf (jn' * 3) (elements [0, 1])
  salary <- vectorOf jn' (choose (0, 10000))
  code <- vectorOf (jn' * 3) (elements [0, 1])
  let jobs = LM.JT {
      LM.jDesc   = LM.fromLists (buildList descs jn'),
      LM.jSalary = LM.fromLists (buildList salary jn'),
      LM.jCode   = LM.fromLists (buildList code jn')
          }
  name <- vectorOf (en' * 5) (elements [0, 1])
  job <- vectorOf (en' * 3) (elements [0, 1])
  country <- vectorOf (en' * 2) (elements [0, 1])
  branch <- vectorOf (en' * 2) (elements [0, 1])
  let employees = LM.ET {
      LM.eName    = LM.fromLists (buildList name en'),
      LM.eJob     = LM.fromLists (buildList job en'),
      LM.eCountry = LM.fromLists (buildList country en'),
      LM.eBranch  = LM.fromLists (buildList branch en')
          }
  return LM.BD { LM.jobsTable = jobs, LM.employeesTable = employees }

randomBDHMatrix ::
             forall jn en .
             ( KnownNat jn,
               KnownNat en
             ) => Gen (HM.BD jn en)
randomBDHMatrix = do
  let jn' = fromInteger (natVal (Proxy :: Proxy jn)) + 1
      en' = fromInteger (natVal (Proxy :: Proxy en)) + 1
  descs <- vectorOf (jn' * 3) (elements [0, 1])
  salary <- vectorOf jn' (choose (0, 10000))
  code <- vectorOf (jn' * 3) (elements [0, 1])
  let jobs = HM.JT {
      HM.jDesc   = HM.fromLists (buildList descs jn'),
      HM.jSalary = HM.fromLists (buildList salary jn'),
      HM.jCode   = HM.fromLists (buildList code jn')
          }
  name <- vectorOf (en' * 5) (elements [0, 1])
  job <- vectorOf (en' * 3) (elements [0, 1])
  country <- vectorOf (en' * 2) (elements [0, 1])
  branch <- vectorOf (en' * 2) (elements [0, 1])
  let employees = HM.ET {
      HM.eName    = HM.fromLists (buildList name en'),
      HM.eJob     = HM.fromLists (buildList job en'),
      HM.eCountry = HM.fromLists (buildList country en'),
      HM.eBranch  = HM.fromLists (buildList branch en')
          }
  return HM.BD { HM.jobsTable = jobs, HM.employeesTable = employees }

randomBDMatrix ::
             forall jn en .
             ( KnownNat jn,
               KnownNat en
             ) => Gen (DM.BD jn en)
randomBDMatrix = do
  let jn' = fromInteger (natVal (Proxy :: Proxy jn)) + 1
      en' = fromInteger (natVal (Proxy :: Proxy en)) + 1
  descs <- vectorOf (jn' * 3) (elements [0, 1])
  salary <- vectorOf jn' (choose (0, 10000))
  code <- vectorOf (jn' * 3) (elements [0, 1])
  let jobs = DM.JT {
      DM.jDesc   = DM.fromLists (buildList descs jn'),
      DM.jSalary = DM.fromLists (buildList salary jn'),
      DM.jCode   = DM.fromLists (buildList code jn')
          }
  name <- vectorOf (en' * 5) (elements [0, 1])
  job <- vectorOf (en' * 3) (elements [0, 1])
  country <- vectorOf (en' * 2) (elements [0, 1])
  branch <- vectorOf (en' * 2) (elements [0, 1])
  let employees = DM.ET {
      DM.eName    = DM.fromLists (buildList name en'),
      DM.eJob     = DM.fromLists (buildList job en'),
      DM.eCountry = DM.fromLists (buildList country en'),
      DM.eBranch  = DM.fromLists (buildList branch en')
          }
  return DM.BD { DM.jobsTable = jobs, DM.employeesTable = employees }

buildList [] _ = []
buildList l r  = take r l : buildList (drop r l) r

setupEnv1 = do
  -- HMatrix
  bd1 <- generate (resize 1 (randomBDHMatrix @999 @999))
  -- Matrix
  bd2 <- generate (resize 1 (randomBDMatrix @999 @999))
  -- Laop
  bd3 <- generate (resize 1 (randomBDLaop @999 @999))
  return (bd1, bd2, bd3)

benchmark :: IO ()
benchmark = do
  print "Starting benchmarks..."
  defaultMain [
   env setupEnv1 $ \ ~(bd1, bd2, bd3) -> bgroup "DB query" [
   bgroup "1000 entries" [
     bench "hmatrix" $ nf HM.query bd1
   , bench "matrix" $ nf DM.query bd2
   , bench "laop" $ nf LM.query bd3
   ] ] ]

main :: IO ()
main = benchmark


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Dynamic.Layers.Deconvolution
  ( SpecDeconvolution (..)
  , specDeconvolution2DInput
  , specDeconvolution3DInput
  , deconvolution
  ) where

import           Control.Monad.Primitive             (PrimBase, PrimState)
import           Data.Constraint                     (Dict (..))
import           Data.Proxy
import           Data.Reflection                     (reifyNat)
import           Data.Singletons
import           Data.Singletons.Prelude.Num         ((%*))
import           Data.Singletons.TypeLits            hiding (natVal)
import           GHC.TypeLits

import           Numeric.LinearAlgebra.Static        hiding (build, toRows, (&),
                                                      (|||), size)
import           System.Random.MWC                   (Gen)
import           Unsafe.Coerce                       (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic.Internal.Build
import           Grenade.Dynamic.Specification
import           Grenade.Dynamic.Network
import           Grenade.Layers.Deconvolution
import           Grenade.Utils.ListStore

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

instance (KnownNat channels, KnownNat filters, KnownNat kernelRows, KnownNat kernelColumns, KnownNat strideRows, KnownNat strideColumns) =>
         FromDynamicLayer (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  fromDynamicLayer inp _ _ =
    SpecNetLayer $
    SpecDeconvolution
      (tripleFromSomeShape inp)
      (natVal (Proxy :: Proxy channels))
      (natVal (Proxy :: Proxy filters))
      (natVal (Proxy :: Proxy kernelRows))
      (natVal (Proxy :: Proxy kernelColumns))
      (natVal (Proxy :: Proxy strideRows))
      (natVal (Proxy :: Proxy strideColumns))

instance ToDynamicLayer SpecDeconvolution where
  toDynamicLayer  = toDynamicLayer'

toDynamicLayer' :: (PrimBase m) => WeightInitMethod -> Gen (PrimState m) -> SpecDeconvolution -> m SpecNetwork
toDynamicLayer' _ _ (SpecDeconvolution inp@(_, 1, 1) _ _ _ _ _ _) = error $ "1D input to a deconvolutional layer is not permited! you specified: " ++ show inp
toDynamicLayer' wInit gen (SpecDeconvolution (rows, cols, depth) ch fil kerRows kerCols strRows strCols) =
    reifyNat ch $ \(pxCh :: (KnownNat channels) => Proxy channels) ->
    reifyNat fil $ \(pxFil :: (KnownNat filters) => Proxy filters) ->
    reifyNat kerRows $ \(pxKerRows :: (KnownNat kernelRows) => Proxy kernelRows) ->
    reifyNat kerCols $ \(pxKerCols :: (KnownNat kernelCols) => Proxy kernelCols) ->
    reifyNat strRows $ \(_ :: (KnownNat strideRows) => Proxy strideRows) ->
    reifyNat strCols $ \(_ :: (KnownNat strideCols) => Proxy strideCols) ->
    reifyNat rows $ \(pxRows :: (KnownNat rows) => Proxy rows) ->
    reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
    reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
    reifyNat ((rows - 1) * strRows + kerRows) $ \(pxOutRows :: (KnownNat outRows) => Proxy outRows) ->
    reifyNat ((cols - 1) * strCols + kerCols) $ \(_ :: (KnownNat outCols) => Proxy outCols) ->
    case ( (singByProxy pxKerRows %* singByProxy pxKerCols) %* singByProxy pxFil
         , (singByProxy pxKerRows %* singByProxy pxKerCols) %* singByProxy pxCh -- this is the input: i = (kernelRows * kernelCols) * channels)
         , singByProxy pxCh %* ((singByProxy pxKerRows %* singByProxy pxKerCols) %* singByProxy pxFil)
         , singByProxy pxOutRows %* singByProxy pxFil -- 'D3 representation
         , singByProxy pxRows %* singByProxy pxCh -- 'D3 representation
         ) of
      (SNat, SNat, SNat, SNat, SNat) | ch == 1 && fil == 1 && depth == 0 ->
        case (unsafeCoerce (Dict :: Dict()) :: Dict (channels ~ 1, filters ~ 1, ((rows - 1) * strideRows) ~ (outRows - kernelRows), ((cols - 1) * strideCols) ~ (outCols - kernelCols)) ) of
          Dict -> do
            (layer  :: Deconvolution 1 1 kernelRows kernelCols strideRows strideCols) <- createRandomWith wInit gen
            return $ SpecLayer layer (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 outRows outCols))
      (SNat, SNat, SNat, SNat, SNat) | ch == 1 ->
        case (unsafeCoerce (Dict :: Dict()) :: Dict (channels ~ 1, ((rows - 1) * strideRows) ~ (outRows - kernelRows), ((cols - 1) * strideCols) ~ (outCols - kernelCols)) ) of
          Dict -> do
            (layer  :: Deconvolution 1 filters kernelRows kernelCols strideRows strideCols) <- createRandomWith wInit gen
            return $ SpecLayer layer (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D3 outRows outCols filters))
      (SNat, SNat, SNat, SNat, SNat) | fil == 1 ->
        case (unsafeCoerce (Dict :: Dict()) :: Dict (filters ~ 1, ((rows - 1) * strideRows) ~ (outRows - kernelRows), ((cols - 1) * strideCols) ~ (outCols - kernelCols)) ) of
          Dict -> do
            (layer  :: Deconvolution channels 1 kernelRows kernelCols strideRows strideCols) <- createRandomWith wInit gen
            return $ SpecLayer layer (sing :: Sing ('D3 rows cols channels)) ( sing :: Sing ('D2 outRows outCols))
      (SNat, SNat, SNat, SNat, SNat) ->
        case (unsafeCoerce (Dict :: Dict()) :: Dict (((rows - 1) * strideRows) ~ (outRows - kernelRows), ((cols - 1) * strideCols) ~ (outCols - kernelCols)) ) of
          Dict -> do
            (layer :: Deconvolution channels filters kernelRows kernelCols strideRows strideCols) <- createRandomWith wInit gen
            return $ SpecLayer layer (sing :: Sing ('D3 rows cols channels)) ( sing :: Sing ('D3 outRows outCols filters))

-- | Creates a specification for a deconvolutional layer with 2D input to the layer. If channels and filters are both 1 then the output is 2D otherwise it is 3D. The output sizes are `out = (in - 1) *
-- stride + kernel`, for rows and cols and the depth is filters for 3D output.
specDeconvolution2DInput ::
     (Integer, Integer) -- ^ Number of input rows.
  -> Integer -- ^ Number of channels, for the first layer this could be RGB for instance.
  -> Integer -- ^ Number of filters, this is the number of channels output by the layer.
  -> Integer -- ^ The number of rows in the kernel filter
  -> Integer -- ^ The number of column in the kernel filter
  -> Integer -- ^ The row stride of the deconvolution filter
  -> Integer -- ^ The cols stride of the deconvolution filter
  -> SpecNet
specDeconvolution2DInput (rows, cols) = specDeconvolution3DInput (rows, cols, 1)

-- | Creates a specification for a deconvolutional layer with 3D input to the layer. If the filter is 1 then the output is 2D, otherwise it is 3D. The output sizes are `out = (in - 1) * stride +
-- kernel`, for rows and cols and the depth is filters for 3D output.
specDeconvolution3DInput ::
     (Integer, Integer, Integer) -- ^ Input to layer (rows, cols, depths). Use 1 if not used or the function @specDeconvolution1DInput@ and @specDeconvolution2DInput@.
  -> Integer -- ^ Number of channels, for the first layer this could be RGB for instance.
  -> Integer -- ^ Number of filters, this is the number of channels output by the layer.
  -> Integer -- ^ The number of rows in the kernel filter
  -> Integer -- ^ The number of column in the kernel filter
  -> Integer -- ^ The row stride of the deconvolution filter
  -> Integer -- ^ The cols stride of the deconvolution filter
  -> SpecNet
specDeconvolution3DInput inp channels filters kernelRows kernelCols strideRows strideCols =
  SpecNetLayer $ SpecDeconvolution inp channels filters kernelRows kernelCols strideRows strideCols

-- | A deconvolution layer. 2D and 3D input/output only!
deconvolution ::
     Integer -- ^ Number of channels, for the first layer this could be RGB for instance.
  -> Integer -- ^ Number of filters, this is the number of channels output by the layer.
  -> Integer -- ^ The number of rows in the kernel filter
  -> Integer -- ^ The number of column in the kernel filter
  -> Integer -- ^ The row stride of the deconvolution filter
  -> Integer -- ^ The cols stride of the deconvolution filter
  -> BuildM ()
deconvolution channels filters kernelRows kernelCols strideRows strideCols = do
  inp@(r, c, _) <- buildRequireLastLayerOut IsNot1D
  let outRows = (r - 1) * strideRows + kernelRows
      outCols = (c - 1) * strideCols + kernelCols
  buildAddSpec $ SpecNetLayer $ SpecDeconvolution inp channels filters kernelRows kernelCols strideRows strideCols
  buildSetLastLayer (outRows, outCols, filters)

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

instance (KnownNat strideCols, KnownNat strideRows, KnownNat kernelCols, KnownNat kernelRows, KnownNat filters, KnownNat channels, KnownNat ((kernelRows * kernelCols) * filters)) =>
         GNum (Deconvolution channels filters kernelRows kernelCols strideRows strideCols) where
  n |* (Deconvolution w store) = Deconvolution (dmmap (fromRational n *) w) (n |* store)
  (Deconvolution w1 store1) |+ (Deconvolution w2 store2) = Deconvolution (w1 + w2) (store1 |+ store2)
  gFromRational r = Deconvolution (fromRational r) mkListStore

instance (KnownNat strideCols, KnownNat strideRows, KnownNat kernelCols, KnownNat kernelRows, KnownNat filters, KnownNat channels, KnownNat ((kernelRows * kernelCols) * filters)) =>
         GNum (Deconvolution' channels filters kernelRows kernelCols strideRows strideCols) where
  n |* (Deconvolution' g) = Deconvolution' (dmmap (fromRational n *) g)
  (Deconvolution' g) |+ (Deconvolution' g2) = Deconvolution' (g + g2)
  gFromRational r = Deconvolution' (fromRational r)


{-# LANGUAGE GADTs, RankNTypes, MultiParamTypeClasses, FlexibleInstances, UndecidableInstances, KindSignatures, ScopedTypeVariables, ConstraintKinds, TemplateHaskell, GeneralizedNewtypeDeriving, TypeFamilies #-}

module Numeric.Trainee.Types (
	Gradee(..),
	Ow(..),
	params, forwardPass, Learnee(..),
	Loss, Regularization, HasNorm(..),
	Sample(..), (⇢),
	Samples, samples,

	module Numeric.Trainee.Params
	) where

import Prelude hiding (id, (.))
import Prelude.Unicode

import Control.Category
import Control.DeepSeq
import Control.Lens
import qualified Data.Vector as V
import Numeric.LinearAlgebra (Normed(norm_1), R, Vector)

import Numeric.Trainee.Params

newtype Gradee a b = Gradee {
	runGradee ∷ Lens' a b }

instance Category Gradee where
	id = Gradee id
	Gradee f . Gradee g = Gradee (g . f)

class Category a ⇒ Ow a where
	first ∷ a b c → a (b, d) (c, d)
	first = flip stars id

	second ∷ a b c → a (d, b) (d, c)
	second = stars id

	stars ∷ a b c → a b' c' → a (b, b') (c, c')
	vstars ∷ a b c → a (V.Vector b) (V.Vector c)

instance Ow Gradee where
	stars (Gradee f) (Gradee g) = Gradee $ lens g' s' where
		g' (x, y) = (view f x, view g y)
		s' (x, y) (x', y') = (set f x' x, set g y' y)

	vstars (Gradee f) = Gradee $ lens g' s' where
		g' = V.map (view f)
		s' = V.zipWith (flip (set f))

data Learnee a b = Learnee {
	_params ∷ Params,
	_forwardPass ∷ Params → a → (b, b → (Params, a)) }

makeLenses ''Learnee

instance NFData (Learnee a b) where
	rnf (Learnee ws _) = rnf ws

instance Show (Learnee a b) where
	show (Learnee ws _) = show ws

instance Category Learnee where
	id = Learnee (Params NoParams) fn where
		fn _ x = (x, const (Params NoParams, x))
	Learnee rws g . Learnee lws f = lws `deepseq` rws `deepseq` Learnee (Params (lws, rws)) (h ∘ castParams) where
		h (lws', rws') x = x `seq` y `seq` lws' `deepseq` rws' `deepseq` (z, up) where
			(y, f') = f lws' x
			(z, g') = g rws' y
			up dz = dz `seq` dy `seq` lws'' `deepseq` rws'' `deepseq` (Params (lws'', rws''), dx) where
				(rws'', dy) = g' dz
				(lws'', dx) = f' dy

type Loss b = b → b → (b, b)

type Regularization b = b → (b, b)

class HasNorm a where
	type Norm a
	norm ∷ a → Norm a

instance HasNorm Float where
	type Norm Float = Float
	norm = id

instance HasNorm Double where
	type Norm Double = Double
	norm = id

instance Normed (Vector a) ⇒ HasNorm (Vector a) where
	type Norm (Vector a) = R
	norm = norm_1

data Sample a b = Sample {
	sampleInput ∷ a,
	sampleOutput ∷ b }
		deriving (Eq, Ord)

(⇢) ∷ a → b → Sample a b
(⇢) = Sample

instance Bifunctor Sample where
	bimap f g (Sample x y) = Sample (f x) (g y)

instance (Show a, Show b) ⇒ Show (Sample a b) where
	show (Sample xs ys) = show xs ++ " => " ++ show ys

instance (NFData a, NFData b) ⇒ NFData (Sample a b) where
	rnf (Sample i o) = rnf i `seq` rnf o

type Samples a b = V.Vector (Sample a b)

samples ∷ [Sample a b] → Samples a b
samples = V.fromList


{-# LANGUAGE FlexibleContexts #-}

module Lib
  ( someFunc,
  )
where

import Control.Exception (throw)
import Data.Char (digitToInt, isHexDigit, toLower, toUpper)
import Data.Complex (Complex ((:+)))
import Data.Foldable (Foldable (foldl'))
import qualified Data.Functor
import Data.Ratio (Ratio, denominator, numerator)
import Foreign.C (isValidErrno)
import GHC.Arr (Array, listArray)
import GHC.Float (rationalToDouble)
import GHC.Real ((%))
import GHC.Unicode (isHexDigit)
import Numeric (readFloat, readHex, readOct)
import System.Environment (getArgs)
import Text.Parsec (ParseError, alphaNum, anyChar, digit, optional, sepBy, sepEndBy, skipMany)
import Text.Parsec.Char (digit)
import Text.ParserCombinators.Parsec
  ( Parser,
    char,
    digit,
    endBy,
    hexDigit,
    letter,
    many,
    many1,
    noneOf,
    notFollowedBy,
    octDigit,
    oneOf,
    parse,
    sepBy,
    skipMany1,
    space,
    string,
    try,
    (<?>),
    (<|>),
  )
import Text.ParserCombinators.Parsec.Char (digit)
import Text.Show (Show)
import Control.Monad.Except (throwError, catchError)

someFunc :: IO ()
someFunc = do
  (expr : _) <- getArgs
  let evaled = fmap show  $ readExpr expr >>= eval
  putStrLn $ extractValue $ trapError evaled

trapError :: ThrowsError String -> ThrowsError String
trapError action = catchError action (return . show)

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

readExpr :: String -> ThrowsError LispVal
readExpr input = case parse parseExpr "lisp" input of
  Left err -> throwError $ Parser err
  Right val -> return val

spaces :: Parser ()
spaces = skipMany1 space

spaces1 :: Parser ()
spaces1 = skipMany space

data LispVal
  = Atom String
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | String String
  | Bool Bool
  | Character Char
  | Number Integer
  | Float Double
  | Ratio Rational
  | Complex (Complex Double)
  | Vector (Array Int LispVal)

instance Show LispVal where show = showVal

showVal :: LispVal -> String
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name) = name
showVal (Number contents) = show contents
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Character c) = "\'" ++ show c ++ "\'"
showVal (Float f) = show f
showVal (Ratio r) = show r
showVal (Complex c) = show c
showVal (List contents) = "(" ++ unwordsList contents ++ ")"
showVal (DottedList head tail) =
  "(" ++ unwordsList head ++ " . "
    ++ showVal tail
    ++ ")"
showVal (Vector array) = show array

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

parseString :: Parser LispVal
parseString = do
  char '"'
  x <- many $ escapedChar <|> noneOf "\\\""
  char '"'
  return $ String x

escapedChar :: Parser Char
escapedChar = do
  char '\\'
  oneOf "\\\"\n\r\t"

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many (letter <|> digit <|> symbol)
  let atom = first : rest
  return $ case atom of
    "#t" -> Bool True
    "#f" -> Bool False
    _ -> Atom atom

parseNumber :: Parser LispVal
parseNumber =
  parseDecimal1
    <|> parseDecimal2
    <|> parseHex
    <|> parseOct
    <|> parseBin

parseDecimal1 :: Parser LispVal
parseDecimal1 = Number . read <$> many1 digit

parseDecimal2 :: Parser LispVal
parseDecimal2 = do
  try $ string "#d"
  x <- many1 digit
  return $ (Number . read) x

parseHex :: Parser LispVal
parseHex = do
  try $ string "#x"
  x <- many1 hexDigit
  return $ Number $ hex2dig x

parseOct :: Parser LispVal
parseOct = do
  try $ string "#o"
  x <- many1 octDigit
  return $ Number $ oct2dig x

parseBin :: Parser LispVal
parseBin = do
  try $ string "#b"
  x <- many1 (oneOf "10")
  return $ Number $ bin2dig x

oct2dig :: (Eq a, Num a) => String -> a
oct2dig x = fst $ head $ readOct x

hex2dig :: (Eq a, Num a) => String -> a
hex2dig x = fst $ head $ readHex x

bin2dig :: [Char] -> Integer
bin2dig = foldl' (\acc x -> acc * 2 + (toInteger . digitToInt) x) 0

parseBool :: Parser LispVal
parseBool = do
  char '#'
  (char 't' >> return (Bool True))
    <|> (char 'f' >> return (Bool False))

parseChar :: Parser LispVal
parseChar = try parseCharacter <|> parseCharacterName

parseCharacter :: Parser LispVal
parseCharacter =
  do
    string "#\\"
    x <- anyChar
    notFollowedBy alphaNum
    return $ Character x

parseCharacterName :: Parser LispVal
parseCharacterName =
  do
    string "#\\"
    val <- caseInsensitiveString "newline" <|> caseInsensitiveString "space"
    return $
      Character $ case val of
        "space" -> ' '
        "newline" -> '\n'
        _ -> ' '

caseInsensitiveChar :: Char -> Parser Char
caseInsensitiveChar c = char (toLower c) <|> char (toUpper c)

caseInsensitiveString :: String -> Parser String
caseInsensitiveString s = mapM caseInsensitiveChar s <?> "\"" ++ s ++ "\""

parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  char '.'
  y <- many1 digit
  return $ Float (fst . head $ readFloat $ x ++ "." ++ y)

parseRatio :: Parser LispVal
parseRatio = do
  x <- many1 digit
  char '/'
  y <- many1 digit
  return $ Ratio $ read x % read y

parseComplex :: Parser LispVal
parseComplex = do
  x <- try parseFloat <|> parseNumber
  char '+'
  y <- try parseFloat <|> parseNumber
  char 'i'
  return $ Complex $ toDouble x :+ toDouble y

toDouble :: LispVal -> Double
toDouble (Float f) = realToFrac f
toDouble (Number n) = fromIntegral n
toDouble (Ratio r) = rationalToDouble (numerator r) (denominator r)
toDouble _ = error "not implement"

-- parse list
parseList :: Parser LispVal
parseList = List <$> sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList = do
  head <- endBy parseExpr spaces
  tail <- char '.' >> spaces >> parseExpr
  return $ DottedList head tail

parseListAll :: Parser LispVal
parseListAll = do
  char '('
  spaces1
  x <- try parseList <|> parseDottedList
  spaces1
  char ')'
  return x

parseListAllWithoutTry :: Parser LispVal
parseListAllWithoutTry = do
  char '('
  spaces1
  head <- sepEndBy parseExpr spaces
  do
    char '.' >> spaces
    tail <- parseExpr
    spaces1 >> char ')'
    return $ DottedList head tail
    <|> (char ')' >> return (List head))

parseQuoted :: Parser LispVal
parseQuoted = do
  char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseBackQuoted :: Parser LispVal
parseBackQuoted = do
  char '`'
  x <- parseExpr
  return $ List [Atom "quasiquote", x]

parseUnquote :: Parser LispVal
parseUnquote = do
  try $ char ','
  x <- parseExpr
  return $ List [Atom "unquote", x]

parseUnquoteSplicing :: Parser LispVal
parseUnquoteSplicing = do
  try $ string ",@"
  x <- parseExpr
  return $ List [Atom "unquote-splicing", x]

parseVector :: Parser LispVal
parseVector = do
  string "#("
  arrayVal <- sepBy parseExpr spaces
  char ')'
  return $ Vector $ listArray (0, length arrayVal - 1) arrayVal

parseNumber1 :: Parser LispVal
parseNumber1 = do
  num <- many1 digit
  return $ (Number . read) num

parseNumber2 :: Parser LispVal
parseNumber2 = many1 digit Data.Functor.<&> (Number . read)

parseExpr :: Parser LispVal
parseExpr =
  parseAtom
    <|> parseString
    <|> try parseComplex
    <|> try parseFloat
    <|> try parseRatio
    <|> try parseNumber -- we need the 'try' because these can all
    <|> try parseBool -- start with  the '#' char
    <|> try parseChar
    <|> parseQuoted
    <|> parseBackQuoted
    <|> parseUnquoteSplicing
    <|> parseUnquote
    <|> try parseVector
    <|> parseListAllWithoutTry

eval :: LispVal -> ThrowsError LispVal
eval val@(String _) = return val
eval val@(Number _) = return  val
eval val@(Bool _) = return val
eval val@(Character _) = return val
eval val@(Float _) = return val
eval val@(Ratio _) = return val
eval val@(Complex _) = return val
eval val@(Vector _) = return val
eval val@(DottedList _ _) = return val
eval (List [Atom "quote", val]) = return  val
eval (List [Atom "quasiquote", val]) = return val
eval (List (Atom func : args)) =  mapM  eval args >>= apply func
eval badForm = throwError $ BadSpecialForm "Unrecognized special form" badForm

apply :: String -> [LispVal] -> ThrowsError LispVal
apply func args =
  maybe
    (throwError $ NotFunction "Unrecognized primitive function args" func)
    ($ args)
    (lookup func primitives)

{-
boolean? --Boolean? returns #t if obj is either #t or #f and returns #f otherwise.
pair? --Pair? returns #t if obj is a pair, and otherwise returns #f.
null? --Returns #t if obj is the empty list, otherwise returns #f.
list? --Returns #t if obj is a list, otherwise returns #f.
symbol? --Returns #t if obj is a symbol, otherwise returns #f.
char? --Returns #t if obj is a character, otherwise returns #f.
string? --Returns #t if obj is a string, otherwise returns #f.
vector? --Returns #t if obj is a vector, otherwise returns #f.

number?
complex?
real?
rational?
integer?
-}
primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives =
  [ ("+", numericBinop (+)),
    ("-", numericBinop (-)),
    ("*", numericBinop (*)),
    ("/", numericBinop div),
    ("mod", numericBinop mod),
    ("quotient", numericBinop quot),
    ("remainder", numericBinop rem),
    ("boolean?", unaryOp booleanOp),
    ("pair?", unaryOp pairOp),
    ("null?", unaryOp nullOp),
    ("list?", unaryOp listOp),
    ("symbol?", unaryOp symbolOp),
    ("char?", unaryOp charOp),
    ("string?", unaryOp stringOp),
    ("vector?", unaryOp vectorOp),
    ("integer?", unaryOp integerOp),
    ("rational?", unaryOp rationalOp),
    ("real?", unaryOp realOp),
    ("complex?", unaryOp complexOp),
    ("number?", unaryOp numberOp),
    -- symbol handler
    ("symbol->string", unaryOp symbol2stringOp),
    ("string->symbol", unaryOp string2symbolOp)
  ]

string2symbolOp :: LispVal -> LispVal
string2symbolOp (String s) = Atom s
string2symbolOp n = error $ show n ++ " is not string."

symbol2stringOp :: LispVal -> LispVal
symbol2stringOp (Atom a) = String a
symbol2stringOp n = error $ show n ++ " is not atom."

integerOp :: LispVal -> LispVal
integerOp (Number _) = Bool True
integerOp _ = Bool False

rationalOp :: LispVal -> LispVal
rationalOp (Number _) = Bool True
rationalOp (Ratio _) = Bool True
rationalOp _ = Bool False

realOp :: LispVal -> LispVal
realOp (Number _) = Bool True
realOp (Ratio _) = Bool True
realOp (Float _) = Bool True
realOp _ = Bool False

complexOp :: LispVal -> LispVal
complexOp (Number _) = Bool True
complexOp (Ratio _) = Bool True
complexOp (Float _) = Bool True
complexOp (Complex _) = Bool True
complexOp _ = Bool False

numberOp :: LispVal -> LispVal
numberOp = complexOp

vectorOp :: LispVal -> LispVal
vectorOp (Vector _) = Bool True
vectorOp _ = Bool False

stringOp :: LispVal -> LispVal
stringOp (String _) = Bool True
stringOp _ = Bool False

charOp :: LispVal -> LispVal
charOp (Character _) = Bool True
charOp _ = Bool False

symbolOp :: LispVal -> LispVal
symbolOp (Atom _) = Bool True
symbolOp _ = Bool False

listOp :: LispVal -> LispVal
listOp (List _) = Bool True
listOp _ = Bool False

nullOp :: LispVal -> LispVal
nullOp (List []) = Bool True
nullOp _ = Bool False

pairOp :: LispVal -> LispVal
pairOp (List (a : _)) = Bool True
pairOp (DottedList _ _) = Bool True
pairOp _ = Bool False

booleanOp :: LispVal -> LispVal
booleanOp (Bool _) = Bool True
booleanOp _ = Bool False

unaryOp :: (LispVal -> LispVal) -> [LispVal] -> ThrowsError LispVal
unaryOp f [v] = return $ f v
unaryOp _ _ = error "only support one argument."

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal] -> ThrowsError LispVal
numericBinop op [] = throwError $ NumArgs 2 []
numericBinop op singleVal@[_] = throwError $ NumArgs 2 singleVal
numericBinop op params =  mapM unpackNum params >>= return . Number . foldl1 op

unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum notNum = throwError $ TypeMismatch "number" notNum

-- error checking and exceptions

data LispError
  = NumArgs Integer [LispVal]
  | TypeMismatch String LispVal
  | Parser ParseError
  | BadSpecialForm String LispVal
  | NotFunction String String
  | UnboundVar String String
  | Default String

showError :: LispError -> String
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (NumArgs expected found) =
  "Expected " ++ show expected
    ++ " args; found values "
    ++ unwordsList found
showError (TypeMismatch expected found) =
  "Invalid type: expected " ++ expected
    ++ ", found "
    ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr
showError (Default str) = show str

instance Show LispError where show = showError

type ThrowsError = Either LispError

extractValue :: ThrowsError a -> a
extractValue (Right val) = val


{-# LANGUAGE BangPatterns #-}
module Polestar.Core.Eval where
import Polestar.Core.Type
import Polestar.Core.TypeCheck
import Data.Complex
import GHC.Float (expm1, log1p)

termSubstD :: Int -> Term -> Int -> Term -> Term
termSubstD !depth s !i t = case t of
  TmAbs name ty body -> TmAbs name (typeShift (-1) i ty) (termSubstD (depth + 1) s (i + 1) body)
  TmRef j | j == i -> termShift depth 0 s
          | j > i -> TmRef (j - 1)
          | otherwise -> t
  TmApp u v -> TmApp (termSubstD depth s i u) (termSubstD depth s i v)
  TmLet name def body -> TmLet name (termSubstD depth s i def) (termSubstD (depth + 1) s (i + 1) body)
  TmTypedLet name ty def body -> TmTypedLet name ty (termSubstD depth s i def) (termSubstD (depth + 1) s (i + 1) body)
  TmIf cond then_ else_ -> TmIf (termSubstD depth s i cond) (termSubstD depth s i then_) (termSubstD depth s i else_)
  TmPrim _ -> t
  TmTuple components -> TmTuple $ termSubstD depth s i <$> components
  TmProj tuple j -> TmProj (termSubstD depth s i tuple) j
  TmIterate n init succ -> TmIterate (termSubstD depth s i n) (termSubstD depth s i init) (termSubstD depth s i succ)

-- replaces occurrences of TRef j (j > i) with TRef (j-1), and TRef i with the given term
termSubst = termSubstD 0

natFromValue :: Term -> Either String Integer
natFromValue (TmPrim (PVNat x)) = return x
natFromValue x = Left ("type error (expected a natural number, but got " ++ show x ++ ")")

intFromValue :: Term -> Either String Integer
intFromValue (TmPrim (PVInt x)) = return x
intFromValue x = Left ("type error (expected an integer, but got " ++ show x ++ ")")

realFromValue :: Term -> Either String Double
realFromValue (TmPrim (PVReal x)) = return x
realFromValue x = Left ("type error (expected a real number, but got " ++ show x ++ ")")

boolFromValue :: Term -> Either String Bool
boolFromValue (TmPrim (PVBool x)) = return x
boolFromValue x = Left ("type error (expected a boolean, but got " ++ show x ++ ")")

applyBuiltinUnary :: UnaryFn -> Term -> Either String Term
applyBuiltinUnary f v = case f of
  UNegateInt -> (TmPrim . PVInt) <$> intFromValue v
  UNegateReal -> (TmPrim . PVReal) <$> realFromValue v
  UNatToInt -> (TmPrim . PVInt) <$> natFromValue v
  UIntToReal -> (TmPrim . PVReal . fromInteger) <$> intFromValue v
  UIntToNat -> (TmPrim . PVNat . max 0) <$> intFromValue v

applyBuiltinBinary :: BinaryFn -> Term -> Term -> Either String Term
applyBuiltinBinary f u v = case f of
  UAddNat -> (TmPrim . PVNat) <$> ((+) <$> natFromValue u <*> natFromValue v)
  UMulNat -> (TmPrim . PVNat) <$> ((*) <$> natFromValue u <*> natFromValue v)
  UTSubNat -> (TmPrim . PVNat . max 0) <$> ((-) <$> natFromValue u <*> natFromValue v)
  UPowNatNat -> (TmPrim . PVNat) <$> ((^) <$> natFromValue u <*> natFromValue v)
  UEqualNat -> (TmPrim . PVBool) <$> ((==) <$> natFromValue u <*> natFromValue v)
  ULessThanNat -> (TmPrim . PVBool) <$> ((<) <$> natFromValue u <*> natFromValue v)
  ULessEqualNat -> (TmPrim . PVBool) <$> ((<=) <$> natFromValue u <*> natFromValue v)
  UAddInt -> (TmPrim . PVInt) <$> ((+) <$> intFromValue u <*> intFromValue v)
  USubInt -> (TmPrim . PVInt) <$> ((-) <$> intFromValue u <*> intFromValue v)
  UMulInt -> (TmPrim . PVInt) <$> ((*) <$> intFromValue u <*> intFromValue v)
  UPowIntNat -> (TmPrim . PVInt) <$> ((^) <$> intFromValue u <*> natFromValue v)
  UEqualInt -> (TmPrim . PVBool) <$> ((==) <$> intFromValue u <*> intFromValue v)
  ULessThanInt -> (TmPrim . PVBool) <$> ((<) <$> intFromValue u <*> intFromValue v)
  ULessEqualInt -> (TmPrim . PVBool) <$> ((<=) <$> intFromValue u <*> intFromValue v)
  UAddReal -> (TmPrim . PVReal) <$> ((+) <$> realFromValue u <*> realFromValue v)
  USubReal -> (TmPrim . PVReal) <$> ((-) <$> realFromValue u <*> realFromValue v)
  UMulReal -> (TmPrim . PVReal) <$> ((*) <$> realFromValue u <*> realFromValue v)
  UDivReal -> (TmPrim . PVReal) <$> ((/) <$> realFromValue u <*> realFromValue v)
  UPowRealReal -> (TmPrim . PVReal) <$> ((**) <$> realFromValue u <*> realFromValue v)
  UEqualReal -> (TmPrim . PVBool) <$> ((==) <$> realFromValue u <*> realFromValue v)
  ULessThanReal -> (TmPrim . PVBool) <$> ((<) <$> realFromValue u <*> realFromValue v)
  ULessEqualReal -> (TmPrim . PVBool) <$> ((<=) <$> realFromValue u <*> realFromValue v)

data ValueBinding = ValueBind !Term
                  | TypeBind
                  deriving (Eq,Show)

getValueFromContext :: [ValueBinding] -> Int -> Term
getValueFromContext ctx i
  | i < length ctx = case ctx !! i of
                       ValueBind x -> x
                       b -> error ("TRef: expected a variable binding, found " ++ show b)
  | otherwise = error "TRef: index out of bounds"

eval1 :: [ValueBinding] -> Term -> Either String Term
eval1 ctx t = case t of
  TmPrim _ -> return t
  TmAbs _ _ _ -> return t
  TmRef i -> return $ getValueFromContext ctx i
  TmApp u v
    | isValue u && isValue v -> case u of
        TmAbs _name _ty body -> return $ termSubst v 0 body -- no type checking here
        TmPrim (PVUnary f) -> applyBuiltinUnary f v
        TmPrim (PVBinary _) -> return t -- partial application
        TmApp (TmPrim (PVBinary f)) u' -> applyBuiltinBinary f u' v
        _ -> Left "invalid function application (expected function type)"
    | isValue u -> TmApp u <$> (eval1 ctx v)
    | otherwise -> TmApp <$> (eval1 ctx u) <*> pure v
  TmLet name def body
    | isValue def -> return $ termSubst def 0 body
    | otherwise -> TmLet name <$> eval1 ctx def <*> pure body
  TmTypedLet name ty def body
    | isValue def -> return $ termSubst def 0 body
    | otherwise -> TmTypedLet name ty <$> eval1 ctx def <*> pure body
  TmIf cond then_ else_
    | isValue cond -> case cond of
        TmPrim (PVBool True) -> return then_  -- no type checking here
        TmPrim (PVBool False) -> return else_ -- no type checking here
        _ -> Left "if-then-else: condition must be boolean"
    | otherwise -> TmIf <$> (eval1 ctx cond) <*> pure then_ <*> pure else_
  TmTuple components -> case span isValue components of
    (_,[]) -> return t
    (v,w:ws) -> eval1 ctx w >>= \w' -> return (TmTuple (v ++ (w':ws)))
  TmProj tuple j
    | isValue tuple -> case tuple of
        TmTuple components | length components > j -> return $ components !! j
                           | otherwise -> Left "tuple too short"
        _ -> Left "projection: not a tuple"
    | otherwise -> TmProj <$> eval1 ctx tuple <*> pure j
  TmIterate n init succ
    | isValue n && isValue init && isValue succ -> do
        n' <- natFromValue n
        if n' == 0
          then return $ init
          else return $ TmIterate (TmPrim (PVNat (n' - 1))) (TmApp succ init) succ
    | isValue n && isValue init -> TmIterate n init <$> eval1 ctx succ
    | isValue n -> TmIterate n <$> eval1 ctx init <*> pure succ
    | otherwise -> TmIterate <$> eval1 ctx n <*> pure init <*> pure succ

eval :: [ValueBinding] -> Term -> Either String Term
eval ctx t | isValue t = return t
           | otherwise = eval1 ctx t >>= eval ctx


-- |
-- Module      : Cartesian.Internal.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created October 31 2015

-- TODO | - Use TemplateHaskell (?)
--        - Strictness
--        - Performance, inlining

-- SPEC | -
--        -

------------------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
------------------------------------------------------------------------------------------------------------------------------------------------------
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE FlexibleInstances      #-}

------------------------------------------------------------------------------------------------------------------------------------------------------
-- API
------------------------------------------------------------------------------------------------------------------------------------------------------
module Cartesian.Internal.Types where

------------------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
------------------------------------------------------------------------------------------------------------------------------------------------------
import Control.Lens (Simple, Lens, lens)

import Data.Complex (Complex(..))

import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4

------------------------------------------------------------------------------------------------------------------------------------------------------
-- Types
------------------------------------------------------------------------------------------------------------------------------------------------------

-- Synonyms ------------------------------------------------------------------------------------------------------------------------------------------

-- TODO: Add some aliased lenses for these aliased types (?)

-- | A lens focusing on a single [vector-]component in a BoundingBox
type BoxLens v v' f f' = Lens (BoundingBox (v f)) (BoundingBox (v' f')) f f'

-- | An axis represented as (begin, length)
type Axis a = (a, a)

-- | A vector where each component represents a single axis (cf. 'Axis')
type Axes v a = v (Axis a)

-- |
-- type Domain

-- |
-- TODO: Rename (eg. 'Shape') (?)
type Polygon m v f = m (v f)

-- | Coordinate system wrappers
newtype Normalised v = Normalised { absolute   :: v } --
newtype Absolute   v = Absoloute  { normalised :: v } --

-- Types ---------------------------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Anchors (eg. C, N, S, E W and combinations thereof, perhaps represented as relative Vectors)
-- TODO: Define some standard instances (eg. Functor, Applicative)
data BoundingBox v = BoundingBox { cornerOf :: v, sizeOf :: v } deriving (Show, Eq)

-- |
-- TODO: Use record (eg. from, to) (?)
data Line v = Line v v deriving (Show, Eq)

-- |
data Linear f = Linear { interceptOf :: f, slopeOf :: f } deriving (Show, Eq)

-- |
data Inclusivity r = Inclusive r | Exclusive r -- TODO: Rename (?)
data Interval r    = Interval (Inclusivity r) (Inclusivity r)

-- |
-- TODO: Use existing type instead (?)
-- data Side = SideLeft | SideRight | SideTop | SideBottom

-- Classes -------------------------------------------------------------------------------------------------------------------------------------------

-- TODO: How do you generate lenses for non-record types (?)
class HasX a f | a -> f where { x :: Simple Lens a f }
class HasY a f | a -> f where { y :: Simple Lens a f }
class HasZ a f | a -> f where { z :: Simple Lens a f }

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

-- TODO: Generate these instances with TemplateHaskell (?)

instance HasX (V1 f) f where
  x = lens (\(V1 x') -> x') (\_ x' -> V1 x')

instance HasX (V2 f) f where
  x = lens (\(V2 x' _) -> x') (\(V2 _ y') x' -> V2 x' y')

instance HasY (V2 f) f where
  y = lens (\(V2 _ y') -> y') (\(V2 x' _) y' -> V2 x' y')

instance HasX (V3 f) f where
  x = lens (\(V3 x' _ _) -> x') (\(V3 _ y' z') x' -> V3 x' y' z')

instance HasY (V3 f) f where
  y = lens (\(V3 _ y' _) -> y') (\(V3 x' _ z') y' -> V3 x' y' z')

instance HasZ (V3 f) f where
  z = lens (\(V3 _ _ z') -> z') (\(V3 x' y' _) z' -> V3 x' y' z')

instance HasX (V4 f) f where
  x = lens (\(V4 x' _ _ _) -> x') (\(V4 _ y' z' w') x' -> V4 x' y' z' w')

instance HasY (V4 f) f where
  y = lens (\(V4 _ y' _ _) -> y') (\(V4 x' _ z' w') y' -> V4 x' y' z' w')

instance HasZ (V4 f) f where
  z = lens (\(V4 _ _ z' _) -> z') (\(V4 x' y' _ w') z' -> V4 x' y' z' w')

instance HasX (Complex f) f where
  x = lens (\(x':+_) -> x') (\(_:+y') x' -> x':+y')

instance HasY (Complex f) f where
  y = lens (\(_':+y') -> y') (\(x':+_) y' -> x':+y')