{-# LANGUAGE ForeignFunctionInterface, GeneralizedNewtypeDeriving #-}

module Numerical.HBLAS.BLAS.FFI.Level3  where

import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex
import Numerical.HBLAS.BLAS.FFI

--------------------------------------------------------------------------------
------------------------------ | BLAS LEVEL 3 ROUTINES
--------------------------------------------------------------------------------
-----------------------  |  Level 3 ops are faster than Levels 1 or 2
--------------------------------------------------------------------------------

--void cblas_sgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);

--void cblas_dgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_cgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

-- |  Matrix mult for general dense matrices
type GemmFunFFI scale el = CBLAS_ORDERT ->   CBLAS_TRANSPOSET -> CBLAS_TRANSPOSET->
        CInt -> CInt -> CInt -> {- scal A * B -} scale  -> {- Matrix A-} Ptr el  -> CInt -> {- B -}  Ptr el -> CInt->
            scale -> {- C -}  Ptr el -> CInt -> IO ()

{- C := alpha*op( A )*op( B ) + beta*C ,  -}

-- matrix mult!
foreign import ccall unsafe "cblas_sgemm"
    cblas_sgemm_unsafe :: GemmFunFFI Float Float

foreign import ccall unsafe "cblas_dgemm"
    cblas_dgemm_unsafe :: GemmFunFFI Double Double

foreign import ccall unsafe "cblas_cgemm"
    cblas_cgemm_unsafe :: GemmFunFFI (Ptr(Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zgemm"
    cblas_zgemm_unsafe :: GemmFunFFI (Ptr (Complex Double)) (Complex Double)

-- safe ffi variant for large inputs
foreign import ccall "cblas_sgemm"
    cblas_sgemm_safe :: GemmFunFFI Float Float

foreign import ccall "cblas_dgemm"
    cblas_dgemm_safe :: GemmFunFFI Double Double

foreign import ccall "cblas_cgemm"
    cblas_cgemm_safe :: GemmFunFFI (Ptr(Complex Float)) (Complex Float)

foreign import ccall "cblas_zgemm"
    cblas_zgemm_safe :: GemmFunFFI (Ptr (Complex Double)) (Complex Double)

-----------------------------------------
----- |  Matrix mult for Symmetric Matrices
-----------------------------------------

--void cblas_ssymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);
--void cblas_dsymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_csymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zsymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

type SymmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT ->
     CInt->CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> scale ->Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_ssymm"
    cblas_ssymm_unsafe :: SymmFunFFI Float Float

foreign import ccall unsafe "cblas_dsymm"
    cblas_dsymm_unsafe :: SymmFunFFI Double Double

foreign import ccall unsafe "cblas_csymm"
    cblas_csymm_unsafe :: SymmFunFFI (Ptr (Complex Float )) (Complex Float)

foreign import ccall unsafe "cblas_zsymm"
    cblas_zsymm_unsafe :: SymmFunFFI (Ptr (Complex Double)) (Complex Double)

-- safe ffi variant,
foreign import ccall  "cblas_ssymm"
    cblas_ssymm_safe :: SymmFunFFI Float Float

foreign import ccall  "cblas_dsymm"
    cblas_dsymm_safe :: SymmFunFFI Double Double

foreign import ccall  "cblas_csymm"
    cblas_csymm_safe :: SymmFunFFI (Ptr (Complex Float )) (Complex Float)

foreign import ccall  "cblas_zsymm"
    cblas_zsymm_safe :: SymmFunFFI (Ptr (Complex Double)) (Complex Double)

-----------------------------------
--- |  symmetric rank k  matrix update, C := alpha*A*A' + beta*C
--- or C = alpha*A'*A + beta*C
------------------------------------

--void cblas_ssyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const float alpha, const float *A, const blasint lda, const float beta, float *C, const blasint ldc);
--void cblas_dsyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const double alpha, const double *A, const blasint lda, const double beta, double *C, const blasint ldc);
--void cblas_csyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const float *alpha, const float *A, const blasint lda, const float *beta, float *C, const blasint ldc);
--void cblas_zsyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const double *alpha, const double *A, const blasint lda, const double *beta, double *C, const blasint ldc);

type SyrkFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt  -> scale -> Ptr el -> CInt -> scale -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyrk"
    cblas_ssyrk_unsafe :: SyrkFunFFI Float Float
foreign import ccall unsafe "cblas_dsyrk"
    cblas_dsyrk_unsafe :: SyrkFunFFI Double Double
foreign import ccall unsafe "cblas_csyrk"
    cblas_csyrk_unsafe :: SyrkFunFFI (Ptr(Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zsyrk"
    cblas_zsyrk_unsafe :: SyrkFunFFI (Ptr(Complex Double)) (Complex Double)

foreign import ccall safe "cblas_ssyrk"
    cblas_ssyrk_safe :: SyrkFunFFI Float Float
foreign import ccall safe "cblas_dsyrk"
    cblas_dsyrk_safe :: SyrkFunFFI Double Double
foreign import ccall safe "cblas_csyrk"
    cblas_csyrk_safe :: SyrkFunFFI (Ptr(Complex Float)) (Complex Float)
foreign import ccall safe "cblas_zsyrk"
    cblas_zsyrk_safe :: SyrkFunFFI (Ptr(Complex Double)) (Complex Double)
----------------------
----- | Symmetric Rank 2k matrix update, C= alpha* A*B' + alpha* B*A' + beta * C
----- or C= alpha* A'*B + alpha* B'*A + beta * C
-------------------

--void cblas_ssyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);
--void cblas_dsyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_csyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zsyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
         --const blasint N, const blasint K, const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

type Syr2kFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET  ->
     CInt->CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt ->
     scale ->Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_ssyr2k"
    cblas_ssyr2k_unsafe :: Syr2kFunFFI Float Float
foreign import ccall unsafe "cblas_dsyr2k"
    cblas_dsyr2k_unsafe :: Syr2kFunFFI Double Double
foreign  import ccall unsafe "cblas_csyr2k"
    cblas_csyr2k_unsafe :: Syr2kFunFFI (Ptr (Complex Float)) (Complex Float)
foreign  import ccall unsafe "cblas_zsyr2k"
    cblas_zsyr2k_unsafe :: Syr2kFunFFI (Ptr (Complex Double)) (Complex Double)

foreign  import ccall safe "cblas_ssyr2k"
    cblas_ssyr2k_safe :: Syr2kFunFFI Float Float
foreign import ccall safe "cblas_dsyr2k"
    cblas_dsyr2k_safe :: Syr2kFunFFI Double Double
foreign  import ccall safe "cblas_csyr2k"
    cblas_csyr2k_safe :: Syr2kFunFFI (Ptr (Complex Float)) (Complex Float)
foreign  import ccall safe "cblas_zsyr2k"
    cblas_zsyr2k_safe :: Syr2kFunFFI (Ptr (Complex Double)) (Complex Double)

-------------------------------
--------  |  matrix matrix product for triangular matrices
------------------------------

type TrmmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign  import ccall unsafe "cblas_strmm"
    cblas_strmm_unsafe :: TrmmFunFFI Float Float
foreign  import ccall unsafe "cblas_dtrmm"
    cblas_dtrmm_unsafe :: TrmmFunFFI Double Double
foreign  import ccall unsafe "cblas_ctrmm"
    cblas_ctrmm_unsafe :: TrmmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall unsafe "cblas_ztrmm"
    cblas_ztrmm_unsafe :: TrmmFunFFI (Ptr (Complex Double )) (Complex Double)

foreign  import ccall safe "cblas_strmm"
    cblas_strmm_safe :: TrmmFunFFI Float Float
foreign  import ccall safe "cblas_dtrmm"
    cblas_dtrmm_safe :: TrmmFunFFI Double Double
foreign  import ccall safe "cblas_ctrmm"
    cblas_ctrmm_safe :: TrmmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall safe "cblas_ztrmm"
    cblas_ztrmm_safe :: TrmmFunFFI (Ptr (Complex Double )) (Complex Double)
--void cblas_strmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_dtrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);
--void cblas_ctrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float *alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_ztrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double *alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);

------------------------
--  |  triangular solvers
-----------------------

--
--TRSM solves  op(A)*X = alpha*B or  X*op(A) = alpha*B
--op(A) is one of op(A) = A, or op(A) = A', or op(A) = conjg(A').
-- A is a unit, or non-unit, upper or lower triangular matrix
----
type TrsmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign  import ccall unsafe "cblas_strsm"
    cblas_strsm_unsafe :: TrsmFunFFI Float Float
foreign  import ccall unsafe "cblas_dtrsm"
    cblas_dtrsm_unsafe :: TrsmFunFFI Double Double
foreign  import ccall unsafe "cblas_ctrsm"
    cblas_ctrsm_unsafe :: TrsmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall unsafe "cblas_ztrsm"
    cblas_ztrsm_unsafe :: TrsmFunFFI (Ptr (Complex Double )) (Complex Double)

foreign  import ccall safe "cblas_strsm"
    cblas_strsm_safe :: TrsmFunFFI Float Float
foreign  import ccall safe "cblas_dtrsm"
    cblas_dtrsm_safe :: TrsmFunFFI Double Double
foreign  import ccall safe "cblas_ctrsm"
    cblas_ctrsm_safe :: TrsmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall safe "cblas_ztrsm"
    cblas_ztrsm_safe :: TrsmFunFFI (Ptr (Complex Double )) (Complex Double)
--void cblas_strsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_dtrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);
--void cblas_ctrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float *alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_ztrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double *alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);

-------------------------
-- | hermitian matrix mult
------------------------

type HemmFunFFI  el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT ->
     CInt -> CInt -> Ptr el -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_chemm"
    cblas_chemm_unsafe :: HemmFunFFI (Complex Float)
foreign  import ccall unsafe "cblas_zhemm"
    cblas_zhemm_unsafe :: HemmFunFFI  (Complex Double)

foreign  import ccall safe "cblas_chemm"
    cblas_chemm_safe :: HemmFunFFI (Complex Float)
foreign  import ccall safe "cblas_zhemm"
    cblas_zhemm_safe :: HemmFunFFI  (Complex Double)

--void cblas_chemm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   CInt M,   CInt N,
--                   Float *alpha,   Float *A,   CInt lda,   Float *B,   CInt ldb,   Float *beta, Float *C,   CInt ldc);
--void cblas_zhemm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   CInt M,   CInt N,
--                   Double *alpha,   Double *A,   CInt lda,   Double *B,   CInt ldb,   Double *beta, Double *C,   CInt ldc);

type HerkFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> scale -> Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_cherk"
    cblas_cherk_unsafe :: HerkFunFFI Float (Complex Float)
foreign  import ccall unsafe "cblas_zherk"
    cblas_zherk_unsafe :: HerkFunFFI Double (Complex Double)

foreign  import ccall safe "cblas_cherk"
    cblas_cherk_safe :: HerkFunFFI Float (Complex Float)
foreign  import ccall safe "cblas_zherk"
    cblas_zherk_safe :: HerkFunFFI Double (Complex Double)
--void cblas_cherk(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                   Float alpha,   Float *A,   CInt lda,   Float beta, Float *C,   CInt ldc);
--void cblas_zherk(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                   Double alpha,   Double *A,   CInt lda,   Double beta, Double *C,   CInt ldc);

type Her2kFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt -> Ptr el  -> Ptr el -> CInt -> Ptr el -> CInt -> scale ->Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_cher2k"
    cblas_cher2k_unsafe :: Her2kFunFFI  Float  (Complex Float)
foreign  import ccall unsafe "cblas_zher2k"
    cblas_zher2k_unsafe :: Her2kFunFFI  Double  (Complex Double)

foreign  import ccall safe "cblas_cher2k"
    cblas_cher2k_safe :: Her2kFunFFI Float (Complex Float)
foreign  import ccall safe "cblas_zher2k"
    cblas_zher2k_safe :: Her2kFunFFI Double (Complex Double)
--void cblas_cher2k(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                    Float *alpha,   Float *A,   CInt lda,   Float *B,   CInt ldb,   Float beta, Float *C,   CInt ldc);
--void cblas_zher2k(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                    Double *alpha,   Double *A,   CInt lda,   Double *B,   CInt ldb,   Double beta, Double *C,   CInt ldc);

----void cblas_xerbla(CInt p, char *rout, char *form, ...);


{-# LANGUAGE RankNTypes #-}

module Test.Grenade.Layers.Internal.Reference where

import           Grenade.Types
import           Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Static as H
import           GHC.TypeLits

im2col :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2col nrows ncols srows scols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  im2colFit starts nrows ncols m

vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> [Matrix RealNum] -> Matrix RealNum
vid2col nrows ncols srows scols inputrows inputcols ms =
  let starts = fittingStarts inputrows nrows srows inputcols ncols scols
      subs   = fmap (im2colFit starts nrows ncols) ms
  in  foldl1 (|||) subs

im2colFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2colFit starts nrows ncols m =
  let imRows = fmap (\start -> flatten $ subMatrix start (nrows, ncols) m) starts
  in  fromRows imRows

col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> [Matrix RealNum]
col2vid nrows ncols srows scols drows dcols m =
  let starts = fittingStart (cols m) (nrows * ncols) (nrows * ncols)
      r      = rows m
      mats   = fmap (\s -> subMatrix (0,s) (r, nrows * ncols) m) starts
      colSts = fittingStarts drows nrows srows dcols ncols scols
  in  fmap (col2imfit colSts nrows ncols drows dcols) mats

col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im krows kcols srows scols drows dcols m =
  let rs       = map toList $ toColumns m
      rs'      = zip [0..] rs
      indicies = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      accums   = concatMap (\(offset, column) -> zipWith (comb offset) indicies column) rs'
  in accum (konst 0 (drows, dcols)) (+) accums
  where
    comb o (i, j) x = let w = (div (dcols - kcols) scols) + 1
                          (a, b) = divMod o w
                      in  ((i + srows * a, j + scols * b), x)

col2imfit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2imfit starts krows kcols drows dcols m =
  let indicies   = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      convs      = fmap (zip indicies . toList) . toRows $ m
      pairs      = zip convs starts
      accums     = concatMap (\(conv',(stx',sty')) -> fmap (\((ix,iy), val) -> ((ix + stx', iy + sty'), val)) conv') pairs
  in  accum (konst 0 (drows, dcols)) (+) accums

poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForward nrows ncols srows scols outputRows outputCols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols m

poolForwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum) -> f (Matrix RealNum)
poolForwardList nrows ncols srows scols inRows inCols outputRows outputCols ms =
  let starts = fittingStarts inRows nrows srows inCols ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols <$> ms

poolForwardFit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForwardFit starts nrows ncols _ outputCols m =
  let els    = fmap (\start -> maxElement $ subMatrix start (nrows, ncols) m) starts
  in  matrix outputCols els

poolBackward :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackward krows kcols srows scols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      starts     = fittingStarts inRows krows srows inCols kcols scols
  in  poolBackwardFit starts krows kcols inputMatrix gradientMatrix

poolBackwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum, Matrix RealNum) -> f (Matrix RealNum)
poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
  let starts     = fittingStarts inRows krows srows inCols kcols scols
  in  uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices

poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      inds       = fmap (\start -> maxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
      grads      = toList $ flatten gradientMatrix
      grads'     = zip3 starts grads inds
      accums     = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
  in  accum (konst 0 (inRows, inCols)) (+) accums

-- | These functions are not even remotely safe, but it's only called from the statically typed
--   commands, so we should be good ?!?!?
--   Returns the starting sub matrix locations which fit inside the larger matrix for the
--   convolution. Takes into account the stride and kernel size.
fittingStarts :: Int -> Int -> Int -> Int -> Int -> Int -> [(Int,Int)]
fittingStarts nrows kernelrows steprows ncols kernelcols stepcolsh =
  let rs = fittingStart nrows kernelrows steprows
      cs = fittingStart ncols kernelcols stepcolsh
      ls = sequence [rs, cs]
  in  fmap (\[a,b] -> (a,b)) ls

-- | Returns the starting sub vector which fit inside the larger vector for the
--   convolution. Takes into account the stride and kernel size.
fittingStart :: Int -> Int -> Int -> [Int]
fittingStart width kernel steps =
  let go left | left + kernel < width
              = left : go (left + steps)
              | left + kernel == width
              = [left]
              | otherwise
              = []
  in go 0

convBackProp :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int
             -> Int -> Int
             -> (Matrix RealNum, Matrix RealNum)
convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let fs       = [0..filters-1]
      hs       = [let h_start = h * strideRows in (h, h_start) | h <- [0..outRows - 1]]
      ws       = [let w_start = w * strideCols in (w, w_start) | w <- [0..outCols - 1]]
      fhws     = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      dX_accum = concatMap dxLoop fhws
      dX       = accum (konst 0 (channels * inRows, inCols)) (+) dX_accum
      dW_accum = concatMap dwLoop fhws
      dW       = accum (konst 0 (kernelRows * kernelCols * channels, filters)) (+) dW_accum
  in (dX, dW)
  where
    dxLoop (f, (h, h_start), (w, w_start))
      = [ ((c * inRows + h_start + i, w_start + j), (indexAtdOut f h w) * (indexAtKernel f c i j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    dwLoop (f, (h, h_start), (w, w_start))
      = [ ((c * kernelRows * kernelCols + i * kernelCols + j, f), (indexAtdOut f h w) * (indexAtIn c (h_start + i) (w_start + j))) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    indexAtIn       c x y = input `atIndex` (c * inRows + x, y)
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwards :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Int -> Int
             -> Int -> Int
             -> Matrix RealNum
convForwards input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
convForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in convForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols outRows outCols strideRows strideCols

convBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> (Matrix RealNum, Matrix RealNum)
convBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW)   = convBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'     = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX          = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW)

biasConvForwards :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Vector RealNum
                 -> Int -> Int
                 -> Int -> Int
                 -> Matrix RealNum
biasConvForwards input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), indexAtBias f + sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtBias i = bias `atIndex` i
    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

biasConvForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Vector RealNum
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
biasConvForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in biasConvForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols

biasConvBackProp :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int
                 -> Int -> Int
                 -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let (dX, dW)  = convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      db_accums = [ (f, sum [ indexAtdOut f i j | i <- [0..outRows-1], j <- [0..outCols-1]]) | f <- [0..filters-1]]
      db        = accum (konst 0 filters) (+) db_accums
  in (dX, dW, db)
  where
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)

biasConvBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int
                            -> Int -> Int
                            -> Int -> Int -> Int -> Int
                            -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums        = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r      = inRows + padt + padb
      padded_c      = inCols + padl + padr
      paddedInput   = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW, db) = biasConvBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'       = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX            = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW, db)

naiveFullyConnectedRunForwards :: forall i o. (KnownNat i, KnownNat o)
                               => H.L o i       -- Weights
                               -> H.R o         -- Biases
                               -> H.R i         -- Input
                               -> (H.R i, H.R o)  -- (Tape, Output)
naiveFullyConnectedRunForwards w b i = (i, b + (w H.#> i))

naiveFullyConnectedBackprop :: forall i o. (KnownNat i, KnownNat o)
                            => H.L o i             -- Weights
                            -> (H.R i, H.R o)        -- (Tape, Output)
                            -> (H.L o i, H.R o, H.R i) -- (NablaW, NablaB, derivatives)
naiveFullyConnectedBackprop w (tape, out) = (w', b', d)
  where
    b' = out
    w' = H.outer out tape
    d  = (H.tr w) H.#> out

-- Implementation reference: https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
naiveLRNForwards :: RealNum -> RealNum -> RealNum -> Int -> [[[RealNum]]] -> [[[RealNum]]]
naiveLRNForwards a b k n values = [
    [[ g ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f ch ro co = values!!ch!!ro!!co
    g ch ro co = (f ch ro co) / den
      where
        den  = den' ** b
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den' = k + a * summation

naiveLRNBackwards :: RealNum         -- a
                    -> RealNum       -- b
                    -> RealNum       -- k
                    -> Int          -- n
                    -> [[[RealNum]]] -- inputs
                    -> [[[RealNum]]] -- backpropagated error
                    -> [[[RealNum]]] -- error to propagate further
naiveLRNBackwards a b k n values errs = [
    [[ ng ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f  ch ro co = values!!ch!!ro!!co
    nf ch ro co = errs!!ch!!ro!!co
    c ch ro co = den
      where
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den = k + a * summation
    ng ch ro co = t1 - t2
      where
        t1 = (c ch ro co) ** (-b) * (nf ch ro co)
        t2 = 2 * b * a * (f ch ro co) * (c ch ro co) ** (-b - 1) * s
        s  = sum [(f q ro co) * (nf q ro co) | q <- [lower..upper]]

        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
module Q.Stochastic.Process
        where
import           Control.Monad
import           Control.Monad.State
import           Data.List             (foldl')
import           Data.RVar
import           Data.Random
import           Numeric.LinearAlgebra

rwalkState :: RVarT (State Double) Double
rwalkState = do
    prev <- lift get
    change  <- rvarT StdNormal

    let new = prev + change
    lift (put new)
    return new

type Time = Double

-- Dont know why this wasn't done.
-- Is there an easier way to do this where we either lift or return?
instance (Num a) => Num (RVarT m a) where
  (+) = liftM2 (+)
  (-) = liftM2 (-)
  (*) = liftM2 (*)
  abs = liftM abs
  signum = liftM signum
  fromInteger x = return $ fromInteger x

-- |Discretization of stochastic process over given interval
class (Num b) => Discretize d b where
  -- |Discretization of the drift process.
  dDrift  :: (StochasticProcess a b) => a -> d -> (Time, b) -> RVar b
  -- |Discretization of the diffusion process.
  dDiff   :: (StochasticProcess a b) => a -> d -> (Time, b) -> RVar b
  -- |dt used.
  dDt     :: (StochasticProcess a b) => a -> d -> (Time, b) -> Time

-- |A stochastic process of the form \(dX_t = \mu(X_t, t)dt + \sigma(S_t, t)dB_t \)
class (Num b) => StochasticProcess a b where
  -- |The process drift.
  pDrift  :: a -> (Time, b) -> RVar b
  -- |The process diffusion.
  pDiff   :: a -> (Time, b) -> RVar b

  -- |Evolve a process from a given state to a given time.
  pEvolve :: (Discretize d b) => a         -- ^The process
                             -> d         -- ^Discretization scheme
                             -> (Time, b) -- ^Initial state
                             -> Time      -- ^Target time t.
                             -> RVar b    -- ^\(dB_i\).
                             -> RVar b    -- ^\(X(t)\).
  pEvolve p disc s0@(t0, x0) t dw = do
    if t0 >= t then return x0 else do
      s'@(t', b') <- pEvolve' p disc s0 dw
      if t' >= t then return b' else pEvolve p disc s' t dw

  -- |Similar to evolve, but evolves the process with the discretization scheme \(dt\).
  pEvolve' :: (Discretize d b, Num b) => a -> d -> (Time, b) -> RVar b -> RVar (Time, b)
  pEvolve' process discr s@(t, b) dw = do
    let !newT = t + dDt process discr s
        !newX = do
               drift <- dDrift process discr s
               diff  <- dDiff process discr s
               dw' <- dw
               return $ b + drift + diff * dw'
        newX :: RVar b

    (newT,) <$>  newX

-- |Geometric Brownian motion
data GeometricBrownian = GeometricBrownian {
    gbDrift :: Double -- ^Drift
  , gbDiff  :: Double -- ^Vol
} deriving (Show)

instance StochasticProcess GeometricBrownian Double where
--  pDrift :: GeometricBrownian -> (Time, Double) -> RVar Double
  pDrift p (_, x) = return $ gbDrift p * x -- drift is prpotional to the spot.
  pDiff  p (_, x) = return $ gbDiff p  * x -- diffisuion is also prportional to the spot.

-- | Ito process
data ItoProcess = ItoProcess {
        ipDrift :: (Time, Double) -> Double,
        ipDiff  :: (Time, Double) -> Double
}


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ConstraintKinds #-}
module Frames.VegaLite.Regression
  (
    Frame2DRegressionScatterFit (..)
  , frameScatterWithFit
  , keyedLayeredFrameScatterWithFit
  , scatterWithFit
  , FitToPlot (..)
  , regressionCoefficientPlot
  , regressionCoefficientPlotMany
  ) where

import qualified Frames.VegaLite.Utils as FV
import qualified Frames.Transform as FT
import qualified Frames.Regression as FR
import qualified Math.Regression.Regression as RE

import qualified Control.Foldl          as FL
import qualified Data.Text              as T
import qualified Data.Vinyl             as V
import qualified Data.Vinyl.Functor     as V
import qualified Data.Vinyl.TypeLevel   as V
import qualified Frames                 as F
import qualified Frames.Melt            as F
import qualified Graphics.Vega.VegaLite as GV
import qualified Data.List              as List
import Text.Printf (printf)

import qualified Statistics.Types               as S

#if MIN_VERSION_hvega(0,4,0)
gvTitle :: Text -> GV.PropertySpec
gvTitle x = GV.title x []
#else
gvTitle :: Text -> (GV.VLProperty, GV.VLSpec)
gvTitle = GV.title
#endif

-- | Plot regression coefficents with error bars
-- | Flex version handles a foldable of results so we can, e.g.,
-- | 1. Compare across time or other variable in the data
-- | 2. Compare different fitting methods for the same data

-- TODO: Fix replicated y-axis ticks.  There since we need the difference between them (trailing "'"s) to provide y-offset
-- to the different results.  Otherwise they would overlap in y.  Maybe fix by switching to layers and only having Y-axis labels
-- on 0th?

-- TODO: Make this typed?  Using a regression result with typed parameters?  Power-to-weight?

regressionCoefficientPlot :: T.Text -> [T.Text] -> RE.RegressionResult Double -> S.CL Double -> GV.VegaLite
regressionCoefficientPlot title names r cl = regressionCoefficientPlotFlex False id title names (V.Identity ("",r)) cl

regressionCoefficientPlotMany :: Foldable f
                              => (k -> T.Text) -> T.Text -> [T.Text] -> f (k, RE.RegressionResult Double) -> S.CL Double -> GV.VegaLite
regressionCoefficientPlotMany = regressionCoefficientPlotFlex True

regressionCoefficientPlotFlex :: Foldable f
                              => Bool -> (k -> Text) -> T.Text -> [T.Text] -> f (k, RE.RegressionResult Double) -> S.CL Double -> GV.VegaLite
regressionCoefficientPlotFlex haveLegend printKey title names results cl =
  let toRow m (RE.NamedEstimate n e eci _) = [("Parameter",GV.Str (n <> T.replicate m "'")), ("Estimate",GV.Number e), ("Confidence",GV.Number eci)]
      addKey k l = ("Key", GV.Str $ printKey k) : l
      dataRowFold = FL.Fold (\(l,n) (k,regRes) -> (l ++ fmap (flip GV.dataRow [] . addKey k . toRow n) (RE.namedEstimates names regRes cl), n+1)) ([],0) (GV.dataFromRows [] . List.concat . fst)
      dat = FL.fold dataRowFold results
      xLabel = "Estimate (with " <> T.pack (printf "%2.0f" (100 * S.confidenceLevel cl)) <> "% confidence error bars)"
      estimateXEnc = GV.position GV.X [GV.PName "Estimate", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle xLabel]]
      estimateYEnc = GV.position GV.Y [GV.PName "Parameter", GV.PmType GV.Ordinal]
      handleLegend l = if haveLegend then l else GV.MLegend [] : l
      estimateColorEnc = GV.color $ handleLegend  [GV.MName "Key", GV.MmType GV.Nominal]
      estimateEnc = estimateXEnc . estimateYEnc . estimateColorEnc
      estLoCalc = "datum.Estimate - datum.Confidence/2"
      estHiCalc = "datum.Estimate + datum.Confidence/2"
      calcEstConf = GV.calculateAs estLoCalc "estLo" . GV.calculateAs estHiCalc "estHi"
      estConfLoEnc = GV.position GV.X [GV.PName "estLo", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle ""]]
      estConfHiEnc = GV.position GV.X2 [GV.PName "estHi", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle ""]]
      estConfEnc = estConfLoEnc . estConfHiEnc . estimateYEnc . estimateColorEnc
      estSpec = GV.asSpec [(GV.encoding . estimateEnc) [], GV.mark GV.Point []]
      confSpec = GV.asSpec [(GV.encoding . estConfEnc) [], GV.mark GV.Rule []]
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      vl = GV.toVegaLite
        [
          gvTitle title
        , (GV.transform . calcEstConf) []
        , GV.layer [estSpec, confSpec]
        , dat
        , configuration []
        ]
  in vl

--
{-
type family WeightElemOf (rs :: [(Symbol, Type)]) (w :: (Symbol, Type)) :: Constraint where
  WeightElemOf _ FR.Unweighted = ()
  WeightElemOf rs w  = F.ElemOf rs w
-}

frameRegressionError :: forall y wc as w rs. FR.FrameRegressionResult y wc as w rs -> Error rs
frameRegressionError (FR.FrameUnweightedRegressionResult _) = const 0
frameRegressionError (FR.FrameWeightedRegressionResult wf _) = (\r -> 1/(wf $ F.rgetField @w r))

type ScatterFitConstraints x y = ( F.ColumnHeaders '[x]
                                 , F.ColumnHeaders '[y]
                                 , FV.ToVLDataValue (F.ElField x)
                                 , FV.ToVLDataValue (F.ElField y)
                                 , V.KnownField y
                                 , Real (V.Snd y)
                                 )

type ScatterFitC1 rs x y = ( F.ElemOf (rs V.++ [YError,YFit,YFitError]) x
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) y
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YError
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YFit
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YFitError)

class Frame2DRegressionScatterFit rs a where
  regressionResultToError :: a -> Error rs
  regressionResultToFit :: Maybe T.Text -> a -> S.CL Double -> FitToPlot rs
--  scatterPlot :: (Foldable f, Functor f) => T.Text -> Maybe T.Text -> a -> Double -> f (F.Record rs) -> GV.VegaLite
  scatterPlotSpec :: (Foldable f, Functor f) => Maybe T.Text -> a -> S.CL Double -> f (F.Record rs) -> GV.VLSpec

frameScatterWithFit :: (Frame2DRegressionScatterFit rs a, Foldable f, Functor f)
                    => T.Text -> Maybe T.Text -> a -> S.CL Double -> f (F.Record rs) -> GV.VegaLite
frameScatterWithFit title fitNameM frr cl frame =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.layer [scatterPlotSpec fitNameM frr cl frame]
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]

keyedLayeredFrameScatterWithFit :: (Frame2DRegressionScatterFit rs a, Foldable f, Functor f, Foldable g, Functor g)
  => T.Text -> (k -> T.Text) -> g (k, a) -> S.CL Double -> f (F.Record rs) -> GV.VegaLite
keyedLayeredFrameScatterWithFit title keyText keyedFits cl dat =
  let toSpec (k, a) = scatterPlotSpec (Just $ keyText k) a cl dat
      specs = FL.fold FL.list (fmap toSpec keyedFits)
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
  in GV.toVegaLite [configuration [], GV.layer specs, gvTitle title]

instance ( V.KnownField x
         , Real (V.Snd x)
         , ScatterFitConstraints x y
         , ScatterFitC1 rs x y
         , F.ElemOf (rs V.++ [YError,YFit,YFitError]) x
         , F.ElemOf (rs V.++ [YError,YFit,YFitError]) y
         , F.ElemOf rs x
         , F.ElemOf rs y
         , V.KnownField w
         , rs F.⊆ rs
         ) => Frame2DRegressionScatterFit rs (FR.FrameRegressionResult y 'True '[x] w rs) where
  regressionResultToError = frameRegressionError
  regressionResultToFit fitNameM frr cl =
    let label = fromMaybe "fit" fitNameM
        dof = RE.degreesOfFreedom (FR.regressionResult frr)
        predF = RE.predictFromEstimateAtConfidence dof frr cl
    in FitToPlot label predF

--  scatterPlot title fitNameM frr ci frame =
--    scatterWithFit @x @y @rs @rs (regressionResultToError frr) (regressionResultToFit fitNameM frr ci) Nothing title frame

  scatterPlotSpec fitNameM frr cl frame =
      scatterWithFitSpec @x @y @rs @rs (regressionResultToError frr) (regressionResultToFit fitNameM frr cl) Nothing frame

-- in this case we have y = a(x1) + b(x2) so we plot (y/x1) vs (a + b(x1/x2))
type YOverX1 = "y_over_x1" F.:-> Double
type X2OverX1 = "x2_over_x1" F.:-> Double

instance ( F.ElemOf rs x1
         , F.ElemOf rs x2
         , F.ElemOf rs y
         , V.KnownField y
         , V.KnownField x1
         , V.KnownField x2
         , V.KnownField w
         , Real (V.Snd x1)
         , Real (V.Snd x2)
         , Real (V.Snd y)
         , F.ElemOf [x1,x2] x1
         , F.ElemOf [x1,x2] x2
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) x1
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) x2
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) y
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) X2OverX1
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) YOverX1
         , rs F.⊆ (rs V.++  [YOverX1, X2OverX1])
         , ScatterFitC1 (rs V.++ [YOverX1, X2OverX1]) YOverX1 X2OverX1
         ) => Frame2DRegressionScatterFit rs (FR.FrameRegressionResult y 'False '[x1,x2] w rs) where
  regressionResultToError frr =
    let x1 = realToFrac . F.rgetField @x1
    in (\r -> frameRegressionError frr r/x1 r)

  regressionResultToFit fitNameM frr cl =
    let label = fromMaybe "fit" fitNameM
        dof = RE.degreesOfFreedom (FR.regressionResult frr)
        predF r =
          let x1 = realToFrac $ F.rgetField @x1 r
              (y,dy) = RE.predictFromEstimateAtConfidence dof frr cl r
          in (y/x1, dy/x1)
    in FitToPlot label predF

  scatterPlotSpec fitNameM frr cl frame =
    let mut :: F.Record rs -> F.Record [YOverX1, X2OverX1]
        mut r =
          let y = F.rgetField @y r
              x1 = F.rgetField @x1 r
              x2 = F.rgetField @x2 r
          in (realToFrac y/realToFrac x1) F.&: (realToFrac x2/realToFrac x1) F.&: V.RNil
        mutData = fmap (FT.mutate mut) frame
        yName = FV.colName @y
        x1Name = FV.colName @x1
        x2Name = FV.colName @x2
        xLabel = x2Name <> "/" <> x1Name
        yLabel = yName <> "/" <> x1Name
    in scatterWithFitSpec @X2OverX1 @YOverX1 @(rs V.++ [YOverX1,X2OverX1]) @rs (regressionResultToError @rs frr) (regressionResultToFit @rs fitNameM frr cl) (Just (xLabel,yLabel)) mutData

{-
frame2DRegressionScatter :: forall x y rs a f. (Foldable f, Frame2DRegressionScatterFit rs x y a)
  => T.Text -> Maybe T.Text -> a -> Double -> f (F.Record rs) -> GV.VegaLite
frame2DRegressionScatter title fitNameM frr ci frame =
-}

--
type YError = "yError" F.:-> Double
type YFit = "yFit" F.:-> Double
type YFitError = "yFitError" F.:->Double

type Error rs = (F.Record rs -> Double)
data FitToPlot rs = FitToPlot { fitLabel :: T.Text, fitFunction :: F.Record rs -> (Double, Double) }

-- | 2D Scatter of Data with calculated error and fit function.
-- | Use TypeApplications to specify x and y columns for scatter and then provide calculated error and fit.
-- | Since both calculations use the record itself as the domain, you can put the error and fit into the record and just use
-- | field selection.  But this allows more flexibility and doesn't require adding things to the input frame record just for the plot.

scatterWithFit :: forall x y rs as f. ( ScatterFitConstraints x y
                                      , as F.⊆ rs
                                      , ScatterFitC1 rs x y
                                      , Foldable f)
               => Error as -> FitToPlot as -> Maybe (T.Text, T.Text) -> Text -> f (F.Record rs) -> GV.VegaLite
scatterWithFit err fit axisLabelsM title frame =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.specification $ scatterWithFitSpec @x @y @rs @as @f err fit axisLabelsM frame
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]

scatterWithFitSpec :: forall x y rs as f. ( ScatterFitConstraints x y
                                          , as F.⊆ rs
                                          , ScatterFitC1 rs x y
                                          , Foldable f)
               => Error as -> FitToPlot as -> Maybe (T.Text, T.Text) -> f (F.Record rs) -> GV.VLSpec
scatterWithFitSpec err fit axisLabelsM frame =
  let mut :: F.Record rs -> F.Record '[YError, YFit, YFitError]
      mut r =
        let a = F.rcast r
            (f,fe) = fitFunction fit a
        in err a F.&: f F.&: fe F.&: V.RNil
      vegaDat = FV.recordsToVLData (F.rcast @[x,y,YError,YFit,YFitError] . FT.mutate mut) frame
  in scatterWithFitSpec' @x @y @YError @YFit @YFitError axisLabelsM (fitLabel fit) vegaDat

{-
scatterWithFit' :: forall x y ye fy fye. ( F.ColumnHeaders '[x]
                                         , F.ColumnHeaders '[y]
                                         , F.ColumnHeaders '[ye]
                                         , F.ColumnHeaders '[fy]
                                         , F.ColumnHeaders '[fye])
  => Maybe (T.Text, T.Text) -> Text -> Text -> GV.Data -> GV.VegaLite
scatterWithFit' axisLabelsM title fitLbl dat =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.specification $ scatterWithFitSpec' @x @y @ye @fy @fye axisLabelsM fitLbl dat
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]
-}
-- TODO: Add xErrors as well, in scatter and in fit
scatterWithFitSpec' :: forall x y ye fy fye. ( F.ColumnHeaders '[x]
                                             , F.ColumnHeaders '[y]
                                             , F.ColumnHeaders '[ye]
                                             , F.ColumnHeaders '[fy]
                                             , F.ColumnHeaders '[fye])
  => Maybe (T.Text, T.Text) -> Text -> GV.Data -> GV.VLSpec
scatterWithFitSpec' axisLabelsM fitLbl dat =
-- create 4 new cols so we can use rules/areas for errors
  let yLoCalc yName yErrName = "datum." <> yName <> " - (datum." <> yErrName <> ")/2"
      yHiCalc yName yErrName = "datum." <> yName <> " + (datum." <> yErrName <> ")/2"
      calcs = GV.calculateAs (yLoCalc (FV.colName @y) (FV.colName @ye)) "yLo"
              . GV.calculateAs (yHiCalc (FV.colName @y) (FV.colName @ye)) "yHi"
              . GV.calculateAs (yLoCalc (FV.colName @fy) (FV.colName @fye)) "fyLo"
              . GV.calculateAs (yHiCalc (FV.colName @fy) (FV.colName @fye)) "fyHi"
              . GV.calculateAs ("\"" <> fitLbl <> "\"")  "fitLabel"
      xLabel = maybe (FV.colName @x) fst axisLabelsM
      yLabel = maybe (FV.colName @y) snd axisLabelsM
      xEnc = GV.position GV.X [FV.pName @x, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle xLabel]]
      yEnc = GV.position GV.Y [FV.pName @y, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle yLabel]]
      yLEnc = GV.position GV.Y [GV.PName "yLo", GV.PmType GV.Quantitative]
      yHEnc = GV.position GV.Y2 [GV.PName "yHi", GV.PmType GV.Quantitative]
--      yErrorEnc = GV.position GV.YError [FV.pName @ye, GV.PmType GV.Quantitative]
      yFitEnc t = GV.position GV.Y [FV.pName @fy, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle t]]
      yFitLEnc = GV.position GV.Y [GV.PName "fyLo", GV.PmType GV.Quantitative]
      yFitHEnc = GV.position GV.Y2 [GV.PName "fyHi", GV.PmType GV.Quantitative]
      colorEnc = GV.color [GV.MName "fitLabel", GV.MmType GV.Nominal]
  --    yFitErrorEnc = GV.position GV.YError [FV.pName @fye, GV.PmType GV.Quantitative]
      scatterEnc = xEnc . yEnc
      scatterBarEnc = xEnc . yLEnc . yHEnc
      fitEnc t = xEnc . yFitEnc t . colorEnc
      fitBandEnc = xEnc . yFitLEnc . yFitHEnc . colorEnc

      selectScalesS = GV.select ("scalesS" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesSE = GV.select("scalesSE" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesF = GV.select ("scalesF" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesFE = GV.select ("scalesFE" <> fitLbl) GV.Interval [GV.BindScales]
      scatterSpec = GV.asSpec
        [
          (GV.encoding . scatterEnc) []
        , GV.mark GV.Point []
        , (GV.selection . selectScalesS) []
        ]
      scatterBarSpec = GV.asSpec
        [
          (GV.encoding . scatterBarEnc) []
        , GV.mark GV.Rule []
        , (GV.selection . selectScalesSE) []
        ]
      fitLineSpec = GV.asSpec
        [
          (GV.encoding . fitEnc fitLbl) []
        , GV.mark GV.Line []
        , (GV.selection . selectScalesF) []
        ]
      fitBandSpec = GV.asSpec
        [
          (GV.encoding . fitBandEnc) []
        , GV.mark GV.Area [GV.MOpacity 0.5, GV.MFillOpacity 0.3]
        , (GV.selection . selectScalesFE) []
        ]
      layers = GV.layer [ scatterSpec, scatterBarSpec, fitLineSpec, fitBandSpec]
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      spec =
        GV.asSpec
        [
          (GV.transform . calcs) []
        , layers
        , dat
        , configuration []
        ]
  in spec


{-# LANGUAGE DataKinds, TypeFamilies #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
-- | Vector-based (meaning that dual numbers for gradient computation
-- consider vectors, not scalars, as the primitive differentiable type)
-- implementation of fully connected neutral network for classification
-- of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistFcnnVector where

import Prelude

import           Control.Exception (assert)
import qualified Data.Array.DynamicS as OT
import qualified Data.Vector.Generic as V
import           GHC.Exts (inline)
import           Numeric.LinearAlgebra (Vector)

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, var1)
import HordeAd.Tool.MnistData

sumTrainableInputsV
  :: IsScalar d r
  => DualNumber d (Vector r) -> Int -> DualNumberVariables d r -> DualNumber d r
sumTrainableInputsV x offset variables =
  let v = var1 variables offset
  in v <.>! x

sumTrainableInputsL
  :: forall d r. IsScalar d r
  => DualNumber d (Vector r) -> Int -> DualNumberVariables d r -> Int
  -> DualNumber d (Vector r)
sumTrainableInputsL x offset variables width =
  let f :: Int -> DualNumber d r
      f i = sumTrainableInputsV x (offset + i) variables
  in seq1 $ V.generate width f

sumConstantDataV
  :: IsScalar d r
  => Vector r -> Int -> DualNumberVariables d r -> DualNumber d r
sumConstantDataV x offset variables =
  let v = var1 variables offset
  in v <.>!! x

sumConstantDataL
  :: forall d r. IsScalar d r
  => Vector r -> Int -> DualNumberVariables d r -> Int
  -> DualNumber d (Vector r)
sumConstantDataL x offset variables width =
  let f :: Int -> DualNumber d r
      f i = sumConstantDataV x (offset + i) variables
  in seq1 $ V.generate width f

fcnnMnistLen1 :: Int -> Int -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
fcnnMnistLen1 widthHidden widthHidden2 =
  ( 0
  , replicate widthHidden sizeMnistGlyph ++ [widthHidden]
    ++ replicate widthHidden2 widthHidden ++ [widthHidden2]
    ++ replicate sizeMnistLabel widthHidden2 ++ [sizeMnistLabel]
  , []
  , []
  )

-- | Fully connected neural network for the MNIST digit classification task.
-- There are two hidden layers and both use the same activation function.
-- The output layer uses a different activation function.
-- The widths of the hidden layers are @widthHidden@ and @widthHidden2@
-- and from these, the @len*@ functions compute the number and dimensions
-- of scalars (none in this case) and vectors of dual number parameters
-- (variables) to be given to the program.
fcnnMnist1 :: forall d r m. DualMonad d r m
           => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> Int
           -> Int
           -> Vector r
           -> DualNumberVariables d r
           -> m (DualNumber d (Vector r))
fcnnMnist1 factivationHidden factivationOutput widthHidden widthHidden2
          input variables = do
  let !_A = assert (sizeMnistGlyph == V.length input) ()
  let hiddenLayer1 = sumConstantDataL input 0 variables widthHidden
                     + var1 variables widthHidden  -- bias
  nonlinearLayer1 <- factivationHidden hiddenLayer1
  let offsetMiddle = widthHidden + 1
      hiddenLayer2 = sumTrainableInputsL nonlinearLayer1 offsetMiddle
                                         variables widthHidden2
                     + var1 variables (offsetMiddle + widthHidden2)  -- bias
  nonlinearLayer2 <- factivationHidden hiddenLayer2
  let offsetOutput = offsetMiddle + widthHidden2 + 1
      outputLayer = sumTrainableInputsL nonlinearLayer2 offsetOutput
                                        variables sizeMnistLabel
                    + var1 variables (offsetOutput + sizeMnistLabel)  -- bias
  factivationOutput outputLayer

-- | The neural network applied to concrete activation functions
-- and composed with the appropriate loss function.
fcnnMnistLoss1
  :: DualMonad d r m
  => Int -> Int -> MnistData r -> DualNumberVariables d r
  -> m (DualNumber d r)
fcnnMnistLoss1 widthHidden widthHidden2 (input, target) variables = do
  result <- inline fcnnMnist1 logisticAct softMaxActV
                              widthHidden widthHidden2 input variables
  lossCrossEntropyV target result

-- | A function testing the neural network given testing set of inputs
-- and the trained parameters.
fcnnMnistTest1
  :: forall r. IsScalar 'DModeGradient r
  => Int -> Int -> [MnistData r] -> (Domain0 r, Domain1 r) -> r
fcnnMnistTest1 widthHidden widthHidden2 inputs (params0, params1) =
  let matchesLabels :: MnistData r -> Bool
      matchesLabels (glyph, label) =
        let nn = inline fcnnMnist1 logisticAct softMaxActV
                                        widthHidden widthHidden2 glyph
            value = primalValue nn (params0, params1, V.empty, V.empty)
        in V.maxIndex value == V.maxIndex label
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)


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

module Language.Hakaru.Distribution where

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

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

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

bern :: Double -> Dist Bool
bern p = Dist {logDensity = (\ (Discrete x) -> log (if x then p else 1 - p)),
               distSample = (\ g -> do t <- MWC.uniformR (0,1) g
                                       return $ Discrete (t <= p))}

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

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

marsaglia :: (MWC.Variate a, Ord a, Floating a, PrimMonad m) => PRNG m -> m (a, a)
marsaglia g = do -- "Marsaglia polar method"
  x <- MWC.uniformR (-1,1) g
  y <- MWC.uniformR (-1,1) g
  let s = x * x + y * y
      q = sqrt ((-2) * log s / s)
  if 1 >= s && s > 0 then return (x * q, y * q) else marsaglia g

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

chooseIndex :: (PrimMonad m) => [Double] -> PRNG m -> m Int
chooseIndex probs g = do
  p <- MWC.uniform g
  return $ fromMaybe (error ("chooseIndex: failure p=" ++ show p))
           (findIndex (p <=) (scanl1 (+) probs))

normal_rng :: (Real a, Floating a, MWC.Variate a, PrimMonad m) =>
              a -> a -> PRNG m -> m a
normal_rng mu sd g | sd > 0 = do (x, _) <- marsaglia g
                                 return (mu + sd * x)
normal_rng _ _ _ = error "normal: invalid parameters"

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

normal :: Double -> Double -> Dist Double
normal mu sd = Dist {logDensity = normalLogDensity mu sd . fromLebesgue,
                     distSample = (\g -> liftM Lebesgue $ normal_rng mu sd g)}

categoricalLogDensity :: (Eq b, Floating a) => [(b, a)] -> b -> a
categoricalLogDensity list x = log $ fromMaybe 0 (lookup x list)

categoricalSample :: (Num b, Ord b, PrimMonad m, MWC.Variate b) =>
    [(t,b)] -> PRNG m -> m t
categoricalSample list g = do
  let total = sum $ map snd list
  p <- MWC.uniformR (0, total) g
  let sumList = scanl1 (\acc (a, b) -> (a, b + snd(acc))) list
      elem' = fst $ head $ filter (\(_,p0) -> p <= p0) sumList
  return elem'

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

lnFact :: Int -> Double
lnFact = logFactorial

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

         make_poisson :: (PrimMonad m) => PRNG m -> m Int
         make_poisson g = do u <- MWC.uniformR (-0.5,0.5) g
                             v <- MWC.uniformR (0,1) g
                             let us = 0.5 - abs u
                                 k = floor $ (2*a / us + b)*u + lambda + 0.43
                             case () of
                               () | us >= 0.07 && v <= vr -> return k
                               () | k < 0 -> make_poisson g
                               () | us <= 0.013 && v > us -> make_poisson g
                               () | accept_region us v k -> return k
                               _  -> make_poisson g

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

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

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

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

gamma :: Double -> Double -> Dist Double
gamma shape scl = Dist {logDensity = gammaLogDensity shape scl . fromLebesgue,
                        distSample = (\g -> liftM Lebesgue $ gamma_rng shape scl g)}

beta_rng :: (PrimMonad m) => Double -> Double -> PRNG m -> m Double
beta_rng a b g | a <= 1.0 && b <= 1.0 = do
                 u <- MWC.uniformR (0.0, 1.0) g
                 v <- MWC.uniformR (0.0, 1.0) g
                 let x = u ** (recip a)
                     y = v ** (recip b)
                 case (x+y) <= 1.0 of
                   True -> return $ x / (x + y)
                   False -> beta_rng a b g
beta_rng a b g = do ga <- gamma_rng a 1 g
                    gb <- gamma_rng b 1 g
                    return $ ga / (ga + gb)

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

beta :: Double -> Double -> Dist Double
beta a b = Dist {logDensity = betaLogDensity a b . fromLebesgue,
                 distSample = (\g -> liftM Lebesgue $ beta_rng a b g)}

laplace_rng :: (PrimMonad m) => Double -> Double -> PRNG m -> m Double
laplace_rng mu sd g = MWC.uniformR (0.0, 1.0) g >>= sample
   where sample u = return $ case u < 0.5 of
                               True  -> mu + sd * log (u + u)
                               False -> mu - sd * log (2.0 - u - u)

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

laplace :: Double -> Double -> Dist Double
laplace mu sd = Dist {logDensity = laplaceLogDensity mu sd . fromLebesgue,
                      distSample = (\g -> liftM Lebesgue $ laplace_rng mu sd g)}

-- Consider having dirichlet return Vector
-- Note: This is actually symmetric dirichlet
dirichlet_rng :: (PrimMonad m) => Int ->  Double -> PRNG m -> m [Double]
dirichlet_rng n' a g' = liftM normalize $ gammas g' n'
  where gammas _ 0 = return ([], 0)
        gammas g n = do (xs, total) <- gammas g (n-1)
                        x <- gamma_rng a 1 g
                        return ((x : xs), x+total)
        normalize (b, total) = map (/ total) b

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


module LatticeSymmetries.ComplexRational
  ( ComplexRational (..),
    ConvertibleToComplexDouble (..),
    realPart,
    imagPart,
    conjugate,
    magnitudeSquared,
    Cscalar,
  )
where

import Data.Complex (Complex (..))
import Foreign.C.Types (CDouble (..))
import Text.PrettyPrint.ANSI.Leijen (Pretty (..))
import qualified Text.PrettyPrint.ANSI.Leijen as Pretty

data ComplexRational = ComplexRational {-# UNPACK #-} !Rational {-# UNPACK #-} !Rational
  deriving stock (Eq, Show)

prettyRational :: Rational -> Pretty.Doc
prettyRational x
  | realToFrac (fromRational x :: Double) == x = Pretty.double (fromRational x :: Double)
  | otherwise = Pretty.rational x

instance Pretty ComplexRational where
  pretty (ComplexRational r i)
    | i == 0 = prettyRational r
    | otherwise = Pretty.parens $ prettyRational r <> Pretty.text " + " <> prettyRational i <> "𝕚"

realPart :: ComplexRational -> Rational
realPart (ComplexRational r _) = r
{-# INLINE realPart #-}

imagPart :: ComplexRational -> Rational
imagPart (ComplexRational _ i) = i
{-# INLINE imagPart #-}

conjugate :: ComplexRational -> ComplexRational
conjugate (ComplexRational r i) = (ComplexRational r (-i))
{-# INLINE conjugate #-}

magnitudeSquared :: ComplexRational -> Rational
magnitudeSquared (ComplexRational r i) = r * r + i * i
{-# INLINE magnitudeSquared #-}

instance Num ComplexRational where
  {-# INLINE (+) #-}
  {-# INLINE (-) #-}
  {-# INLINE (*) #-}
  {-# INLINE fromInteger #-}
  (ComplexRational r i) + (ComplexRational r' i') = ComplexRational (r + r') (i + i')
  (ComplexRational r i) - (ComplexRational r' i') = ComplexRational (r - r') (i - i')
  (ComplexRational r i) * (ComplexRational r' i') = ComplexRational (r * r' - i * i') (r * i' + i * r')
  negate (ComplexRational r i) = ComplexRational (-r) (-i)
  abs _ = error "Num instance of ComplexRational does not implement abs"
  signum _ = error "Num instance of ComplexRational does not implement signum"
  fromInteger n = ComplexRational (fromInteger n) 0

instance Fractional ComplexRational where
  {-# INLINE (/) #-}
  {-# INLINE fromRational #-}
  (ComplexRational r i) / (ComplexRational r' i') =
    ComplexRational ((r * r' + i * i') / d) ((-r * i' + i * r') / d)
    where
      d = r' * r' + i' * i'
  fromRational a = ComplexRational a 0

-- class Fractional a => ComplexFloating a where
--   toComplexDouble :: a -> Complex Double
--   fromComplexDouble :: Complex Double -> a

type Cscalar = Complex CDouble

class Fractional a => ConvertibleToComplexDouble a where
  toComplexDouble :: a -> Cscalar
  fromComplexDouble :: Cscalar -> a

instance ConvertibleToComplexDouble ComplexRational where
  toComplexDouble (ComplexRational r i) = (fromRational r) :+ (fromRational i)
  fromComplexDouble (r :+ i) = ComplexRational (toRational r) (toRational i)

instance ConvertibleToComplexDouble (Complex CDouble) where
  toComplexDouble = id
  fromComplexDouble = id

instance ConvertibleToComplexDouble (Complex Double) where
  toComplexDouble = coerce
  fromComplexDouble = coerce


{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}

module Lib
    ( someFunc
    ) where

import Numeric.GSL.Fitting
import Numeric.LinearAlgebra
import Numeric.AD
import Numeric.AD.Internal.Reverse
import Data.Reflection

someFunc :: IO ()
someFunc = putStrLn "someFunc"

type FitRes = ([(Double, Double)], Matrix Double)
type FitData = [([Double],([Double],Double))]

-- type ManualJacob = ([Double] -> [Double] -> [[Double]])

data JacobType = AutoJacob | ManualJacob ([Double] -> [Double] -> [[Double]])

data FitParams = FitParams { jacob :: JacobType
                           , guess :: [Double]
                           , iter :: Int
                           , absTol :: Double
                           , relTol :: Double
                           }

defaultParams = FitParams { jacob = AutoJacob,
                            guess = repeat 1,
                            iter = 1000,
                            absTol = 1E-20,
                            relTol = 1E-20 }

     -- ([Double] -> [Double] -> [Double]) model
     -- ([Double] -> [Double] -> [[Double]]) model der
-- Only Doubles since hmatrix requires it
fit :: [Double]
     -> [Double]
     -> [Double]
     -> (forall a . Floating a => [a]->[a]->[a])
     -> JacobType
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, sanitizedJacob jacobian) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma
    sanitizedJacob jacob = case jacob of AutoJacob -> mkJac model
                                         ManualJacob f -> f

{-
fit :: [Double]
     -> [Double]
     -> [Double]
     -> (forall a . Floating a => [a]->[a]->[a])
     -> ([Double] -> [Double] -> [[Double]])
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, mkJac model) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma
-}

{-
fit :: [Double]
     -> [Double]
     -> [Double]
     -> ([Double] -> [Double] -> [Double])
     -> ([Double] -> [Double] -> [[Double]])
     -> [Double]->  Maybe FitRes

fit xs ys sigma model jacobian guess
    | xeqy && ( seqx || seqone)  = Just $ fitModelScaled 1E-20 1E-20 500 (model, jacobian) fitdata guess
    | otherwise  = Nothing
    where
    lx = length xs
    ly = length ys
    ls = length sigma
    xeqy = lx == ly
    seqx = ls == lx
    seqone = ls == 1
    fitdata = unsafeFormatData xs ys sigma -}

unsafeFormatData :: [Double] -> [Double] -> [Double] -> FitData
unsafeFormatData xs ys sigma
    |  length sigma == length xs = zip (return <$> xs) $ zip (return <$> ys) sigma
    |  length sigma == 1 = zip (return <$> xs) $ zip (return <$> ys) (repeat (head sigma))

--jac :: (Num a, Traversable f, Functor g) =>
  --     (forall s. Reifies s Tape => f (Reverse s a) -> [Reverse s a] -> g (Reverse s a))
    --   -> f a
      -- -> g (f a)

jac2 :: (Num a, Traversable f1, Traversable f2, Functor g) =>
       (forall s. Reifies s Tape => f1 (Reverse s a) -> f2 (Reverse s a) -> g (Reverse s a))
       -> f1 a
       -> f2 a

       -> g (f1 a)

jac2 f p x = jacobian ((flip f) (fmap auto x)) p

jac3 :: (Floating a, Traversable f1, Traversable f2, Functor g) =>
       (forall s. Reifies s Tape => f1 (Reverse s a) -> f2 (Reverse s a) -> g (Reverse s a))
       -> f1 a
       -> f2 a

       -> g (f1 a)

jac3 f p x = jacobian ((flip f) (fmap auto x)) p

mkJac :: (Num a) => (forall s. Reifies s Tape => [Reverse s a] -> [Reverse s a] -> [Reverse s a])
       -> [a]
       -> [a]

       -> [[a]]

mkJac f p x = jacobian ((flip f) (fmap auto x)) p

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

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

-- to generate data
expModel2 [a,lambda,b] t = a * exp (-lambda * t)+b

expModelDer2 [a,lambda,b] t = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]
-- to go from unbox to box:
-- a = (\p -> \x -> return (expModel2 p (x !! 0)))::Floating a => [a] -> [a] -> [a]

xs = [a | a <- [0..100]]
sigma = [0.1]
ys = (map (expModel2 [5,0.1,1]) xs)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE InstanceSigs           #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE NamedFieldPuns         #-}
{-# LANGUAGE QuantifiedConstraints  #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}

module Q.MonteCarlo where
import           Control.Monad.State
import           Data.RVar
import           Q.Stochastic.Discretize
import           Q.Stochastic.Process
import           Control.Monad
import           Q.ContingentClaim
import Data.Random
import Q.Time
import Data.Time
import Statistics.Distribution (cumulative)
import Statistics.Distribution.Normal (standard)
import Q.ContingentClaim.Options
import Q.Types

type Path b = [(Time, b)]

-- |Summary type class aggregates all priced values of paths
class (PathPricer p)  => Summary m p | m->p where
  -- | Updates summary with given priced pathes
  sSummarize      :: m -> [p] -> m

  -- | Defines a metric, i.e. calculate distance between 2 summaries
  sNorm           :: m -> m -> Double

-- | Path generator is a stochastic path generator
class PathGenerator m where
  pgMkNew         :: m->IO m
  pgGenerate      :: Integer -> m -> Path b

-- | Path pricer provides a price for given path
class PathPricer m where
  ppPrice :: m -> Path b -> m

type MonteCarlo s a = StateT [(Time, s)] RVar a

-- | Generate a single trajectory stopping at each provided time.
trajectory :: forall a b d. (StochasticProcess a b, Discretize d b) =>
             d        -- ^ Discretization scheme
           -> a        -- ^ The stochastic process
           -> b        -- ^ \(S(0)\)
           -> [Time]   -- ^ Stopping points \(\{t_i\}_i^n \) where \(t_i > 0\)
           -> [RVar b] -- ^ \(dW\)s. One for each stopping point.
           -> RVar [b] -- ^ \(S(0) \cup \{S(t_i)\}_i^n \)
trajectory disc p s0 times dws = reverse <$> evalStateT (onePath times dws) initState' where
  initState' :: [(Time, b)]
  initState' = [(0, s0)]

  onePath :: [Time] -> [RVar b] -> MonteCarlo b [b]
  onePath [] _ = do
    s <- get
    return $ map snd s
  onePath (t1:tn) (dw1:dws) = do
    s <- get
    let t0 = head s
    b <- lift $ pEvolve p disc t0 t1 dw1
    put $ (t1, b) : s
    onePath tn dws

-- | Generate multiple trajectories. See 'trajectory'
trajectories:: forall a b d. (StochasticProcess a b, Discretize d b) =>
             Int        -- ^Num of trajectories
           -> d          -- ^Discretization scheme
           -> a          -- ^The stochastic process
           -> b          -- ^\(S(0)\)
           -> [Time]     -- ^Stopping points \(\{t_i\}_i^n \) where \(t_i > 0\)
           -> [RVar b]   -- ^\(dW\)s. One for each stopping point.
           -> RVar [[b]] -- ^\(S(0) \cup \{S(t_i)\}_i^n \)
trajectories n disc p initState times dws = replicateM n $ trajectory disc p initState times dws

observationTimes :: ContingentClaim a -> [Day]
observationTimes = undefined

class Model a b | a -> b where
  discountFactor :: a -> YearFrac -> YearFrac -> RVar Rate
  evolve   :: a -> YearFrac -> StateT (YearFrac, b) RVar Double


{-# LANGUAGE TemplateHaskell #-}

module RBM where

import           Control.Lens          hiding ((<.>), (|>))
import           Control.Monad.State
import           Data.Foldable
import           Numeric.LinearAlgebra
import           System.Random

type Input = Vector R

toBool :: Functor f => f Double -> f Bool
toBool = fmap (> 0.5)

sign :: (Ord a, Num a, Num t) => a -> t
sign x
  | x >= 0 = 1
  | otherwise = -1

sign01 :: (Ord a, Num a, Num t) => a -> t
sign01 x
  | x >= 0 = 1
  | otherwise = 0

data RBM = RBM { _weights :: Matrix R, _vbias :: Vector R, _hbias :: Vector R } deriving (Show)
makeLenses ''RBM

nrVisible :: Getter RBM Int
nrVisible = to (rows . view weights)

nrHidden :: Getter RBM Int
nrHidden = to (cols . view weights)

nullrbm :: Int -> Int -> RBM
nullrbm visible hidden = RBM ((hidden >< visible) zeros) (visible |> zeros) (hidden |> zeros)
  where
    zeros = repeat 0.0

-- rngrbm :: Int -> Int -> Double -> State StdGen RBM
rngrbm visible hidden rnghi = do
  let nulled = nullrbm visible hidden
      addv = flip addUniformNoise  rnghi
      addm = flip addUniformNoiseM rnghi
  vis <- addv $ nulled ^. vbias
  hid <- addv $ nulled ^. hbias
  wei <- addm $ nulled ^. weights
  return $ nulled & vbias .~ vis & hbias .~ hid & weights .~ wei

sigmoid :: Double -> Double
sigmoid x = 1.0 / (1.0 + exp(-x))

sigmoidvec :: Vector R -> Vector R
sigmoidvec x = 1 / (1 + cmap exp (-x))

nRandoms :: Int -> State StdGen [Double]
nRandoms n = do
  replicateM n getone
  where
    getone = do
      rng <- get
      let (nr, rng') = randomR (0.0, 1.0) rng
      put rng'
      return nr

energy :: RBM -> Vector Double -> Vector Double -> Double
energy rbm v h = - a - b - c
  where
    a = v <.> (rbm ^. vbias)
    b = h <.> (rbm ^. hbias)
    c = (h <# (rbm ^. weights)) <.> v

-- | Requires binary data, i.e. v binary, for mathematical correctness
freeEnergyBin :: RBM -> Vector Double -> Double
freeEnergyBin rbm v = let
  b = rbm ^. vbias
  c = rbm ^. hbias
  w = rbm ^. weights
  lhs = b <.> v
  rhs = sumElements $ 1 + exp (c + w #> v)
  in -lhs -rhs

addUniformNoise :: Vector R -> R -> State StdGen (Vector R)
addUniformNoise v hi = do
  rngs <- nRandoms $ size v
  let rngv = (size v |> rngs) * 2 - 1
  return $ scalar hi * rngv + v

addUniformNoiseM :: Matrix R -> R -> State StdGen (Matrix R)
addUniformNoiseM m hi = do
  let (i,j) = size m
  rngs <- nRandoms $ i * j
  let rngm = ((i><j) rngs) * 2 - 1
  return $ scalar hi * rngm + m

-- | random parameters that deeplearning.net used for initializing weights
goodNoiseParam :: Floating a => RBM -> a
goodNoiseParam rbm = 4 * sqrt (6 / fromIntegral (rbm ^. nrHidden + rbm ^. nrVisible))

-- | Probability of hidden nodes firing given sample v
pHid :: RBM -> Input -> Vector R
pHid rbm v = sigmoidvec $ b' + (w #> v)
  where
    b' = rbm ^. hbias
    w = rbm ^. weights

-- | Probability of visible nodes firing given sample h
pVis :: RBM -> Vector R -> Vector R
pVis rbm h = sigmoidvec $ v' + (h <# w)
  where
    v' = rbm ^. vbias
    w = rbm ^. weights

testfire :: Vector R -> State StdGen (Vector R)
testfire probs = do
  rngs <- nRandoms $ size probs
  return $ cmap sign01 $ probs - vector rngs

inpbook :: [Input]
inpbook = vector <$>
  [[1,1,1,0,0,0],[1,0,1,0,0,0],[1,1,1,0,0,0],[0,0,1,1,1,0], [0,0,1,1,0,0],[0,0,1,1,1,0]]

-- | The algorithm part of the book is sort of half on-line, just
-- generally broken and using weeird notation if they actually mean
-- something that would end up as something like a weight. The code,
-- instead, does a full batch thingy without a single mention of it,
-- since the code would work either way. Anyhow, a weight update
-- scheme that might work for one-input-at-a-time is available at
-- http://image.diku.dk/igel/paper/AItRBM-proof.pdf , which this
-- follows. All three weight diffs for this one input, as (pos - neg),
-- is returned.  I.e: dw (Matrix), dv (vector), dh (vector). No
-- assumptions on h, but that algo seems to imply that it maybe should
-- stay as a probability, and not as sampled. v' is newer than v, etc.
-- Should probably divide this by number of inputs (?).
compCDs :: Vector R -> Vector R -> Vector R -> Vector R -> (Matrix R, Vector R, Vector R)
compCDs v h v' h' = (dw, dv, dh)
  where
    dw = h `outer` v - h' `outer` v'
    dv = v - v'
    dh = h - h'

-- | gibbs starting on v, h from input unused (structured this way for ease of folding)
gibbsv :: RBM -> (Vector R, Vector R) -> State StdGen (Vector R, Vector R)
gibbsv rbm (v, _) = do
  h' <- testfire $ pHid rbm v
  v' <- testfire $ pVis rbm h'
  return (v', h')

getDiffs :: RBM -> Input -> State StdGen (Matrix R, Vector R, Vector R)
getDiffs rbm inp = do
  let v = inp
  h <- testfire $ pHid rbm v
  (v', h') <- foldM (\vh _ -> gibbsv rbm vh) (v, h) [1..1]
  return $ compCDs v h v' h'

train :: R -> RBM -> [Input] -> State StdGen RBM
train eta rbm inputs = do
  diffs <- mapM (getDiffs rbm) inputs
  let (dw, dv, dh) = foldr (\(a, b, c) (a', b', c') -> (a+a', b+b', c+c')) (0,0,0) diffs
      dims = fromIntegral $ length inputs
      scale = eta / dims
      newrbm = rbm & weights %~ (+ (dw * scalar scale))
                   & vbias %~ (+ (dv * scalar scale))
                   & hbias %~ (+ (dh * scalar scale))
  return newrbm

fwd :: RBM -> Input -> State StdGen (Vector R)
fwd rbm inp = do
  -- (v', h') <- gibbsv rbm (inp, inp)
  h' <- testfire $ pHid rbm inp
  return $ pVis rbm h'

test :: IO ()
test = do
  rng <- newStdGen
  let tren = train 0.01
      test = vector [0,0,0,1,1,0]
  print . flip evalState rng $ do
      t0 <- rngrbm 6 4 0.01

      t1 <- tren t0 inpbook
      t2 <- foldM tren t1 (replicate 4000 inpbook)
      t3 <- tren t2 inpbook
      out <- fwd t3 (head inpbook)
      outtest <- fwd t3 test
      return (out, outtest)

-- | We can see in the output of test that the probabilities after
-- forwarding the test vector on the one where it is active is very high,
-- often >0.7, which is completely unlike for (head inpbook), where
-- they are <0.2 (positions 4,5, that is). Similarly the first of the
-- training inputs fires higher probabilities on its activations, and
-- all this very consistently. If we keep gibbs-ing, the probabilities
-- tend towards the relative presence in the training set, which
-- asserts the generative-ness of it. Pos 3 is ~1 and pos 6 ~0 in all
-- cases, which is reasonable for our data.
--
-- Doing it batch-wise would probably be faster, doing persistent CD
-- likewise. Absolutely none of the improvements from hinton's very
-- reasonable practical guide are implemented, and it would improve
-- everything manyfold. (Also, the fold is pointless on the singleton
-- list, but it really seems to be the case that sampling once
-- performs the best)
--
-- As it is though, it seems to do some kind of learning!


--
-- PoolLayer : pooling layer
--

module CNN.PoolLayer (
  poolMax
, depoolMax
, reversePooling
) where

import Numeric.LinearAlgebra

import CNN.Algebra
import CNN.Image
import CNN.LayerType

type Pix = (Double, Double)

{- |
poolMax

  In : pooling size (x = y)
       image

  OUT: updated image
       position of max values

>>> let im = [fromLists [[1.0,2.0,3.0,4.0],[8.0,7.0,6.0,5.0],[1.0,3.0,5.0,7.0],[2.0,4.0,6.0,8.0]], fromLists [[1.0,2.0,3.0,4.0],[2.0,2.0,3.0,4.0],[3.0,3.0,3.0,4.0],[4.0,4.0,4.0,4.0]]]
>>> poolMax 2 im
[[(2><2)
 [ 8.0, 6.0
 , 4.0, 8.0 ],(2><2)
 [ 2.0, 4.0
 , 4.0, 4.0 ]],[(2><2)
 [ 2.0, 2.0
 , 3.0, 3.0 ],(2><2)
 [ 1.0, 1.0
 , 2.0, 1.0 ]]]

-}

poolMax :: Int -> Image -> [Image]
poolMax s im = [os, is]
  where
    pl = head im
    x  = cols pl `div` s
    y  = rows pl `div` s
    ps = [(i*s, j*s) | i <- [0..(x-1)], j <- [0..(y-1)]]
    (os, is) = unzip $ map (toPlain x y . maxPix s ps) im

toPlain :: Int -> Int -> [Pix] -> (Plain, Plain)
toPlain x y pls = ((x><y) op, (x><y) ip)
  where
    (op, ip) = unzip pls

maxPix :: Int -> [(Int, Int)] -> Plain -> [Pix]
maxPix s ps is = map (max' . toPix) ps
  where
    toPix :: (Int, Int) -> [Pix]
    toPix p = zip (concat $ toLists $ subMatrix p (s, s) is) [0.0..]

max' :: [Pix] -> Pix
max' [] = error "empty list!"
max' [x] = x
max' (x:xs) = maximum' x (max' xs)

maximum' :: Pix -> Pix -> Pix
maximum' a@(v1, _) b@(v2, _) = if v1 < v2 then b else a

-- back prop

{- |
depoolMax

>>> let im = [fromLists [[0.0,2.0],[3.0,1.0]]]
>>> let dl = [fromLists [[0.1,0.2],[0.3,0.4]]]
>>> depoolMax 2 im dl
([(4><4)
 [ 0.1, 0.0, 0.0, 0.0
 , 0.0, 0.0, 0.2, 0.0
 , 0.0, 0.0, 0.0, 0.4
 , 0.0, 0.3, 0.0, 0.0 ]],Nothing)

-}

depoolMax :: Int -> Image -> Delta -> (Delta, Maybe Layer)
depoolMax s im d = (zipWith (concatWith (expand s 0.0)) im d, Nothing)
  where
    concatWith :: (Double -> Double -> Matrix R) ->  Matrix R -> Matrix R
               -> Matrix R
    concatWith f i d = fromBlocks $ zipWith (zipWith f) is ds
      where
        is = toLists i
        ds = toLists d

{- |
expand

  IN : size of pooling
       filling value
       positions
       delta values

>>> expand 2 0.0 0 3
(2><2)
 [ 3.0, 0.0
 , 0.0, 0.0 ]
>>> expand 2 0.0 2 4
(2><2)
 [ 0.0, 0.0
 , 4.0, 0.0 ]
>>> expand 3 0.0 5 3
(3><3)
 [ 0.0, 0.0, 0.0
 , 0.0, 0.0, 3.0
 , 0.0, 0.0, 0.0 ]
>>> expand 3 0.0 8 3
(3><3)
 [ 0.0, 0.0, 0.0
 , 0.0, 0.0, 0.0
 , 0.0, 0.0, 3.0 ]

-}

expand :: Int -> Double -> Double -> Double -> Matrix R
expand s r p d = (s><s) $ (replicate p1 r ++ [d] ++ replicate p2 r)
  where
    p1 = truncate p
    p2 = max 0 (s * s - p1 - 1)

-- reverse

reversePooling :: Int -> Layer
reversePooling = MaxPoolLayer


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}

module Main where

import Data.Array
import Data.Complex
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate)
import Data.Map (empty, fromList, union)
import Data.Maybe (fromJust)
import Data.STRef.Strict
import qualified Data.Set as Set
import HashedExpression.Derivative.Partial
import Graphics.EasyPlot
import HashedExpression
import HashedExpression.Derivative
import HashedExpression.Interp
import HashedExpression.Operation
import qualified HashedExpression.Operation
import HashedExpression.Prettify
import Data.String.Interpolate
import Prelude hiding ((^))

main :: IO ()
main = do
  let x = variable "x"
  let y = variable "y"
  let f = sum [x^2 , 2*x , 1] * (x^2 + 2 * y + 1)
  print $ allEntries f


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

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

Solution of ordinary differential equation (ODE) initial value problems.

<http://www.gnu.org/software/gsl/manual/html_node/Ordinary-Differential-Equations.html>

A simple example:

@import Numeric.GSL
import Numeric.LinearAlgebra
import Graphics.Plot

xdot t [x,v] = [v, -0.95*x - 0.1*v]

ts = linspace 100 (0,20 :: Double)

sol = odeSolve xdot [10,0] ts

main = mplot (ts : toColumns sol)@

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

module Numeric.GSL.ODE (
    odeSolve, odeSolveV, ODEMethod(..), Jacobian
) where

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

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

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

type Jacobian = Double -> Vector Double -> Matrix Double

-- | Stepping functions
data ODEMethod = RK2 -- ^ Embedded Runge-Kutta (2, 3) method.
               | RK4 -- ^ 4th order (classical) Runge-Kutta. The error estimate is obtained by halving the step-size. For more efficient estimate of the error, use the embedded methods.
               | RKf45 -- ^ Embedded Runge-Kutta-Fehlberg (4, 5) method. This method is a good general-purpose integrator.
               | RKck -- ^ Embedded Runge-Kutta Cash-Karp (4, 5) method.
               | RK8pd -- ^ Embedded Runge-Kutta Prince-Dormand (8,9) method.
               | RK2imp Jacobian -- ^ Implicit 2nd order Runge-Kutta at Gaussian points.
               | RK4imp Jacobian -- ^ Implicit 4th order Runge-Kutta at Gaussian points.
               | BSimp Jacobian -- ^ Implicit Bulirsch-Stoer method of Bader and Deuflhard. The method is generally suitable for stiff problems.
               | RK1imp Jacobian -- ^ Implicit Gaussian first order Runge-Kutta. Also known as implicit Euler or backward Euler method. Error estimation is carried out by the step doubling method.
               | MSAdams -- ^ A variable-coefficient linear multistep Adams method in Nordsieck form. This stepper uses explicit Adams-Bashforth (predictor) and implicit Adams-Moulton (corrector) methods in P(EC)^m functional iteration mode. Method order varies dynamically between 1 and 12.
               | MSBDF Jacobian -- ^ A variable-coefficient linear multistep backward differentiation formula (BDF) method in Nordsieck form. This stepper uses the explicit BDF formula as predictor and implicit BDF formula as corrector. A modified Newton iteration method is used to solve the system of non-linear equations. Method order varies dynamically between 1 and 5. The method is generally suitable for stiff problems.

-- | A version of 'odeSolveV' with reasonable default parameters and system of equations defined using lists.
odeSolve
    :: (Double -> [Double] -> [Double])        -- ^ xdot(t,x)
    -> [Double]        -- ^ initial conditions
    -> Vector Double   -- ^ desired solution times
    -> Matrix Double   -- ^ solution
odeSolve xdot xi ts = odeSolveV RKf45 hi epsAbs epsRel (l2v xdot) (fromList xi) ts
    where hi = (ts@>1 - ts@>0)/100
          epsAbs = 1.49012e-08
          epsRel = 1.49012e-08
          l2v f = \t -> fromList  . f t . toList

-- | Evolution of the system with adaptive step-size control.
odeSolveV
    :: ODEMethod
    -> Double -- ^ initial step size
    -> Double -- ^ absolute tolerance for the state vector
    -> Double -- ^ relative tolerance for the state vector
    -> (Double -> Vector Double -> Vector Double)   -- ^ xdot(t,x)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveV RK2 = odeSolveV' 0 Nothing
odeSolveV RK4 = odeSolveV' 1 Nothing
odeSolveV RKf45 = odeSolveV' 2 Nothing
odeSolveV RKck = odeSolveV' 3 Nothing
odeSolveV RK8pd = odeSolveV' 4 Nothing
odeSolveV (RK2imp jac) = odeSolveV' 5 (Just jac)
odeSolveV (RK4imp jac) = odeSolveV' 6 (Just jac)
odeSolveV (BSimp jac) = odeSolveV' 7 (Just jac)
odeSolveV (RK1imp jac) = odeSolveV' 8 (Just jac)
odeSolveV MSAdams = odeSolveV' 9 Nothing
odeSolveV (MSBDF jac) = odeSolveV' 10 (Just jac)

odeSolveV'
    :: CInt
    -> Maybe (Double -> Vector Double -> Matrix Double)   -- ^ optional jacobian
    -> Double -- ^ initial step size
    -> Double -- ^ absolute tolerance for the state vector
    -> Double -- ^ relative tolerance for the state vector
    -> (Double -> Vector Double -> Vector Double)   -- ^ xdot(t,x)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveV' method mbjac h epsAbs epsRel f  xiv ts = unsafePerformIO $ do
    let n   = dim xiv
    fp <- mkDoubleVecVecfun (\t -> aux_vTov (checkdim1 n . f t))
    jp <- case mbjac of
        Just jac -> mkDoubleVecMatfun (\t -> aux_vTom (checkdim2 n . jac t))
        Nothing  -> return nullFunPtr
    sol <- vec xiv $ \xiv' ->
            vec (checkTimes ts) $ \ts' ->
             createMIO (dim ts) n
              (ode_c (method) h epsAbs epsRel fp jp // xiv' // ts' )
              "ode"
    freeHaskellFunPtr fp
    return sol

foreign import ccall safe "ode"
    ode_c :: CInt -> Double -> Double -> Double -> FunPtr (Double -> TVV) -> FunPtr (Double -> TVM) -> TVVM

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

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

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

checkTimes ts | dim ts > 1 && all (>0) (zipWith subtract ts' (tail ts')) = ts
              | otherwise = error "odeSolve requires increasing times"
    where ts' = toList ts