{-# LANGUAGE ForeignFunctionInterface, GeneralizedNewtypeDeriving #-}

module Numerical.HBLAS.BLAS.FFI.Level2  where

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

type GbmvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_TRANSPOSET -> CInt -> CInt
    -> CInt -> CInt -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sgbmv"
    cblas_sgbmv_unsafe :: GbmvFunFFI Float Float
foreign import ccall unsafe "cblas_dgbmv"
    cblas_dgbmv_unsafe :: GbmvFunFFI Double Double
foreign import ccall unsafe "cblas_cgbmv"
    cblas_cgbmv_unsafe :: GbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgbmv"
    cblas_zgbmv_unsafe :: GbmvFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall "cblas_sgbmv"
    cblas_sgbmv_safe :: GbmvFunFFI Float Float
foreign import ccall "cblas_dgbmv"
    cblas_dgbmv_safe :: GbmvFunFFI Double Double
foreign import ccall "cblas_cgbmv"
    cblas_cgbmv_safe :: GbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall "cblas_zgbmv"
    cblas_zgbmv_safe :: GbmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_sgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Float alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Double alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);
--void cblas_cgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Float *alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zgbmv(  enum CBLAS_ORDER order,   enum CBLAS_TRANSPOSE TransA,   CInt M,   CInt N,
--                   CInt KL,   CInt KU,   Double *alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

{-
matrix vector product for general matrices
 perform one of the matrix-vector operations   y :=
      alpha*A*x + beta*y, or y := alpha*A'*x + beta*y,
-}

type GemvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_TRANSPOSET -> CInt -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sgemv"
    cblas_sgemv_unsafe :: GemvFunFFI Float Float
foreign import ccall safe   "cblas_sgemv"
    cblas_sgemv_safe   :: GemvFunFFI Float Float

foreign import ccall unsafe "cblas_dgemv"
    cblas_dgemv_unsafe :: GemvFunFFI Double Double
foreign import ccall safe   "cblas_dgemv"
    cblas_dgemv_safe   :: GemvFunFFI Double Double

foreign import ccall unsafe "cblas_cgemv"
    cblas_cgemv_unsafe :: GemvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_cgemv"
    cblas_cgemv_safe   :: GemvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zgemv"
    cblas_zgemv_unsafe :: GemvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zgemv"
    cblas_zgemv_safe   :: GemvFunFFI (Ptr (Complex Double)) (Complex Double)

--void cblas_sgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Float alpha,   Float  *a,   CInt lda,    Float  *x,   CInt incx,    Float beta,  Float  *y,   CInt incy);
--void cblas_dgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Double alpha,   Double  *a,   CInt lda,    Double  *x,   CInt incx,    Double beta,  Double  *y,   CInt incy);
--void cblas_cgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Float *alpha,   Float  *a,   CInt lda,    Float  *x,   CInt incx,    Float *beta,  Float  *y,   CInt incy);
--void cblas_zgemv(  enum CBLAS_ORDER order,    enum CBLAS_TRANSPOSE trans,    CInt m,   CInt n,
--           Double *alpha,   Double  *a,   CInt lda,    Double  *x,   CInt incx,    Double *beta,  Double  *y,   CInt incy);

-- perform the rank 1 operation   A := alpha*x*y' + A,

type GerxFunFFI scale el = CBLAS_ORDERT -> CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_sger" cblas_sger_unsafe ::
        GerxFunFFI Float Float
foreign import ccall safe   "cblas_sger" cblas_sger_safe   ::
        GerxFunFFI Float Float

foreign import ccall unsafe "cblas_dger" cblas_dger_unsafe ::
        GerxFunFFI Double Double
foreign import ccall safe   "cblas_dger" cblas_dger_safe   ::
        GerxFunFFI Double Double

foreign import ccall unsafe "cblas_cgerc" cblas_cgerc_unsafe ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgerc" cblas_zgerc_unsafe ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall safe   "cblas_cgerc" cblas_cgerc_safe   ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_zgerc" cblas_zgerc_safe   ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall unsafe "cblas_cgeru" cblas_cgeru_unsafe ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zgeru" cblas_zgeru_unsafe ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

foreign import ccall safe   "cblas_cgeru" cblas_cgeru_safe   ::
        GerxFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_zgeru" cblas_zgeru_safe   ::
        GerxFunFFI (Ptr (Complex Double)) (Complex Double)

--void cblas_sger (  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float   alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_dger (  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double  alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);
--void cblas_cgeru(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float  *alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_cgerc(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Float  *alpha,   Float  *X,   CInt incX,   Float  *Y,   CInt incY, Float  *A,   CInt lda);
--void cblas_zgeru(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);
--void cblas_zgerc(  enum CBLAS_ORDER order,   CInt M,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);

type HbmvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chbmv"
    cblas_chbmv_unsafe :: HbmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_chbmv"
    cblas_chbmv_safe   :: HbmvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhbmv"
    cblas_zhbmv_unsafe :: HbmvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zhbmv"
    cblas_zhbmv_safe   :: HbmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,
--           Float *alpha,   Float *A,   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,
--           Double *alpha,   Double *A,   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HemvFunFFI sc el =
       CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
    -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chemv"
    cblas_chemv_unsafe :: HemvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe   "cblas_chemv"
    cblas_chemv_safe   :: HemvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhemv"
    cblas_zhemv_unsafe :: HemvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe   "cblas_zhemv"
    cblas_zhemv_safe   :: HemvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chemv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhemv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HerFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_cher"
    cblas_cher_unsafe :: HerFunFFI Float (Complex Float)
foreign import ccall safe "cblas_cher"
    cblas_cher_safe :: HerFunFFI Float (Complex Float)

foreign import ccall unsafe "cblas_zher"
    cblas_zher_unsafe :: HerFunFFI Double (Complex Double)
foreign import ccall safe "cblas_zher"
    cblas_zher_safe :: HerFunFFI Double (Complex Double)
--void cblas_cher(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A,   CInt lda);
--void cblas_zher(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *A,   CInt lda);

type Her2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_cher2"
    cblas_cher2_unsafe :: Her2FunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_cher2"
    cblas_cher2_safe :: Her2FunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zher2"
    cblas_zher2_unsafe :: Her2FunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zher2"
    cblas_zher2_safe :: Her2FunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_cher2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *X,   CInt incX,
--                  Float *Y,   CInt incY, Float *A,   CInt lda);
--void cblas_zher2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *X,   CInt incX,
--                  Double *Y,   CInt incY, Double *A,   CInt lda);

type HpmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_chpmv"
    cblas_chpmv_unsafe :: HpmvFunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_chpmv"
    cblas_chpmv_safe :: HpmvFunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhpmv"
    cblas_zhpmv_unsafe :: HpmvFunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zhpmv"
    cblas_zhpmv_safe :: HpmvFunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,
--           Float *alpha,   Float *Ap,   Float *X,   CInt incX,   Float *beta, Float *Y,   CInt incY);
--void cblas_zhpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,
--           Double *alpha,   Double *Ap,   Double *X,   CInt incX,   Double *beta, Double *Y,   CInt incY);

type HprFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_chpr"
    cblas_chpr_unsafe :: HprFunFFI Float (Complex Float)
foreign import ccall safe "cblas_chpr"
    cblas_chpr_safe :: HprFunFFI Float (Complex Float)

foreign import ccall unsafe "cblas_zhpr"
    cblas_zhpr_unsafe :: HprFunFFI Double (Complex Double)
foreign import ccall safe "cblas_zhpr"
    cblas_zhpr_safe :: HprFunFFI Double (Complex Double)
--void cblas_chpr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A);
--void cblas_zhpr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,  CInt incX, Double *A);

type Hpr2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_chpr2"
    cblas_chpr2_unsafe :: Hpr2FunFFI (Ptr (Complex Float)) (Complex Float)
foreign import ccall safe "cblas_chpr2"
    cblas_chpr2_safe :: Hpr2FunFFI (Ptr (Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zhpr2"
    cblas_zhpr2_unsafe :: Hpr2FunFFI (Ptr (Complex Double)) (Complex Double)
foreign import ccall safe "cblas_zhpr2"
    cblas_zhpr2_safe :: Hpr2FunFFI (Ptr (Complex Double)) (Complex Double)
--void cblas_chpr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float *alpha,   Float *X,   CInt incX,   Float *Y,   CInt incY, Float *Ap);
--void cblas_zhpr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double *alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *Ap);

type SbmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssbmv"
    cblas_ssbmv_unsafe :: SbmvFunFFI Float Float
foreign import ccall safe "cblas_ssbmv"
    cblas_ssbmv_safe :: SbmvFunFFI Float Float

foreign import ccall unsafe "cblas_dsbmv"
    cblas_dsbmv_unsafe :: SbmvFunFFI Double Double
foreign import ccall safe "cblas_dsbmv"
    cblas_dsbmv_safe :: SbmvFunFFI Double Double
--void cblas_ssbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,   Float alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dsbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   CInt K,   Double alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

---------------
--- | packed symmetric matrix * vector product  y:= alpha * Av  + beta * y
---------------
type SpmvFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> Ptr el -> CInt -> sc -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_sspmv"
    cblas_sspmv_unsafe :: SpmvFunFFI Float Float
foreign import ccall safe "cblas_sspmv"
    cblas_sspmv_safe :: SpmvFunFFI Float Float

foreign import ccall unsafe "cblas_dspmv"
    cblas_dspmv_unsafe :: SpmvFunFFI Double Double
foreign import ccall safe "cblas_dspmv"
    cblas_dspmv_safe :: SpmvFunFFI Double Double
--void cblas_sspmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *Ap,
--                   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dspmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *Ap,
--                   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

type SprFunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_sspr"
    cblas_sspr_unsafe :: SprFunFFI Float Float
foreign import ccall safe "cblas_sspr"
    cblas_sspr_safe :: SprFunFFI Float Float

foreign import ccall unsafe "cblas_dspr"
    cblas_dspr_unsafe :: SprFunFFI Double Double
foreign import ccall safe "cblas_dspr"
    cblas_dspr_safe :: SprFunFFI Double Double
--void cblas_sspr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *Ap);
--void cblas_dspr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *Ap);

type Spr2FunFFI sc el =
      CBLAS_ORDERT -> CBLAS_UPLOT -> CInt
      -> sc -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> IO ()
foreign import ccall unsafe "cblas_sspr2"
    cblas_sspr2_unsafe :: Spr2FunFFI Float Float
foreign import ccall safe "cblas_sspr2"
    cblas_sspr2_safe :: Spr2FunFFI Float Float

foreign import ccall unsafe "cblas_dspr2"
    cblas_dspr2_unsafe :: Spr2FunFFI Double Double
foreign import ccall safe "cblas_dspr2"
    cblas_dspr2_safe :: Spr2FunFFI Double Double
--void cblas_sspr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX,   Float *Y,   CInt incY, Float *A);
--void cblas_dspr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX,   Double *Y,   CInt incY, Double *A);

----------------------------------
---- |  (unpacked) symmetric matrix vector product    x:=Av, writes result x into v
---------------------------------

type SymvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el ->  CInt ->
                        Ptr el -> CInt -> el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssymv"
    cblas_ssymv_unsafe :: SymvFunFFI Float
foreign import ccall safe "cblas_ssymv"
    cblas_ssymv_safe :: SymvFunFFI Float

foreign import ccall unsafe "cblas_dsymv"
    cblas_dsymv_unsafe :: SymvFunFFI Double
foreign import ccall safe "cblas_dsymv"
    cblas_dsymv_safe :: SymvFunFFI Double
--void cblas_ssymv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *A,
--                   CInt lda,   Float *X,   CInt incX,   Float beta, Float *Y,   CInt incY);
--void cblas_dsymv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *A,
--                   CInt lda,   Double *X,   CInt incX,   Double beta, Double *Y,   CInt incY);

type SyrFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyr"
      cblas_ssyr_unsafe :: SyrFunFFI Float
foreign import ccall safe "cblas_ssyr"
      cblas_ssyr_safe :: SyrFunFFI Float

foreign import ccall unsafe "cblas_dsyr"
      cblas_dsyr_unsafe :: SyrFunFFI Double
foreign import ccall safe "cblas_dsyr"
      cblas_dsyr_safe :: SyrFunFFI Double
--void cblas_ssyr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Float alpha,   Float *X,   CInt incX, Float *A,   CInt lda);
--void cblas_dsyr(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,   CInt incX, Double *A,   CInt lda);

type Syr2FunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CInt -> el -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyr2"
      cblas_ssyr2_unsafe :: Syr2FunFFI Float
foreign import ccall safe "cblas_ssyr2"
      cblas_ssyr2_safe :: Syr2FunFFI Float

foreign import ccall unsafe "cblas_dsyr2"
      cblas_dsyr2_unsafe :: Syr2FunFFI Double
foreign import ccall safe "cblas_dsyr2"
      cblas_dsyr2_safe :: Syr2FunFFI Double
--void cblas_ssyr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,  CInt N,   Float alpha,   Float *X,
--                  CInt incX,   Float *Y,   CInt incY, Float *A,   CInt lda);
--void cblas_dsyr2(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   CInt N,   Double alpha,   Double *X,
--                  CInt incX,   Double *Y,   CInt incY, Double *A,   CInt lda);

type TbmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stbmv"
      cblas_stbmv_unsafe :: TbmvFunFFI Float
foreign import ccall safe "cblas_stbmv"
      cblas_stbmv_safe :: TbmvFunFFI Float

foreign import ccall unsafe "cblas_dtbmv"
      cblas_dtbmv_unsafe :: TbmvFunFFI Double
foreign import ccall safe "cblas_dtbmv"
      cblas_dtbmv_safe :: TbmvFunFFI Double

foreign import ccall unsafe "cblas_ctbmv"
      cblas_ctbmv_unsafe :: TbmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctbmv"
      cblas_ctbmv_safe :: TbmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztbmv"
      cblas_ztbmv_unsafe :: TbmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztbmv"
      cblas_ztbmv_safe :: TbmvFunFFI (Complex Double)
--void cblas_stbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztbmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);

----------------
--- | solves  Ax=v where A is k+1 banded triangular matrix, and x and
----------------
type TbsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stbsv"
      cblas_stbsv_unsafe :: TbsvFunFFI Float
foreign import ccall safe "cblas_stbsv"
      cblas_stbsv_safe :: TbsvFunFFI Float

foreign import ccall unsafe "cblas_dtbsv"
      cblas_dtbsv_unsafe :: TbsvFunFFI Double
foreign import ccall safe "cblas_dtbsv"
      cblas_dtbsv_safe :: TbsvFunFFI Double

foreign import ccall unsafe "cblas_ctbsv"
      cblas_ctbsv_unsafe :: TbsvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctbsv"
      cblas_ctbsv_safe :: TbsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztbsv"
      cblas_ztbsv_unsafe :: TbsvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztbsv"
      cblas_ztbsv_safe :: TbsvFunFFI (Complex Double)
--void cblas_stbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztbsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   CInt K,   Double *A,   CInt lda, Double *X,   CInt incX);

-------------------------------------------------------------------------
-- | matrix vector product Av, writes result into v, where A is a packed triangular nxn matrix
-------------------------------------------------------------------------
type TpmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stpmv"
      cblas_stpmv_unsafe :: TpmvFunFFI Float
foreign import ccall safe "cblas_stpmv"
      cblas_stpmv_safe :: TpmvFunFFI Float

foreign import ccall unsafe "cblas_dtpmv"
      cblas_dtpmv_unsafe :: TpmvFunFFI Double
foreign import ccall safe "cblas_dtpmv"
      cblas_dtpmv_safe :: TpmvFunFFI Double

foreign import ccall unsafe "cblas_ctpmv"
      cblas_ctpmv_unsafe :: TpmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctpmv"
      cblas_ctpmv_safe :: TpmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztpmv"
      cblas_ztpmv_unsafe :: TpmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztpmv"
      cblas_ztpmv_safe :: TpmvFunFFI (Complex Double)
--void cblas_stpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_dtpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);
--void cblas_ctpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_ztpmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);

--------------------------------------------------
---  | solve  Ax=v where A is a nxn packed triangular matrix, v vector input, writes the solution into x.
--------------------------------------------------
type TpsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_stpsv"
      cblas_stpsv_unsafe :: TpsvFunFFI Float
foreign import ccall safe "cblas_stpsv"
      cblas_stpsv_safe :: TpsvFunFFI Float

foreign import ccall unsafe "cblas_dtpsv"
      cblas_dtpsv_unsafe :: TpsvFunFFI Double
foreign import ccall safe "cblas_dtpsv"
      cblas_dtpsv_safe :: TpsvFunFFI Double

foreign import ccall unsafe "cblas_ctpsv"
      cblas_ctpsv_unsafe :: TpsvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctpsv"
      cblas_ctpsv_safe :: TpsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztpsv"
      cblas_ztpsv_unsafe :: TpsvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztpsv"
      cblas_ztpsv_safe :: TpsvFunFFI (Complex Double)
--void cblas_stpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_dtpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);
--void cblas_ctpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Float *Ap, Float *X,   CInt incX);
--void cblas_ztpsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,
--                   CInt N,   Double *Ap, Double *X,   CInt incX);

type TrmvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_strmv"
      cblas_strmv_unsafe :: TrmvFunFFI Float
foreign import ccall safe "cblas_strmv"
      cblas_strmv_safe :: TrmvFunFFI Float

foreign import ccall unsafe "cblas_dtrmv"
      cblas_dtrmv_unsafe :: TrmvFunFFI Double
foreign import ccall safe "cblas_dtrmv"
      cblas_dtrmv_safe :: TrmvFunFFI Double

foreign import ccall unsafe "cblas_ctrmv"
      cblas_ctrmv_unsafe :: TrmvFunFFI (Complex Float)
foreign import ccall safe "cblas_ctrmv"
      cblas_ctrmv_safe :: TrmvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztrmv"
      cblas_ztrmv_unsafe :: TrmvFunFFI (Complex Double)
foreign import ccall safe "cblas_ztrmv"
      cblas_ztrmv_safe :: TrmvFunFFI (Complex Double)
--void cblas_strmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztrmv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);

--STRSV - solve one of the systems of equations   A*x = b, or A'*x = b, where A is a (non)unit upper(/lower) triangular matrix
type TrsvFunFFI el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
                      CInt -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_strsv"
  cblas_strsv_unsafe :: TrsvFunFFI Float
foreign import ccall safe   "cblas_strsv"
  cblas_strsv_safe   :: TrsvFunFFI Float

foreign import ccall unsafe "cblas_dtrsv"
  cblas_dtrsv_unsafe :: TrsvFunFFI Double
foreign import ccall safe   "cblas_dtrsv"
  cblas_dtrsv_safe   :: TrsvFunFFI Double

foreign import ccall unsafe "cblas_ctrsv"
  cblas_ctrsv_unsafe :: TrsvFunFFI (Complex Float)
foreign import ccall safe   "cblas_ctrsv"
  cblas_ctrsv_safe   :: TrsvFunFFI (Complex Float)

foreign import ccall unsafe "cblas_ztrsv"
  cblas_ztrsv_unsafe :: TrsvFunFFI (Complex Double)
foreign import ccall safe   "cblas_ztrsv"
  cblas_ztrsv_safe   :: TrsvFunFFI (Complex Double)

--void cblas_strsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_dtrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);
--void cblas_ctrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Float *A,   CInt lda, Float *X,   CInt incX);
--void cblas_ztrsv(  enum CBLAS_ORDER order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,   enum CBLAS_DIAG Diag,   CInt N,   Double *A,   CInt lda, Double *X,   CInt incX);


module Lambda.Parse where

-- Local Imports
import Lambda.Types
-- Global Imports
import Control.Monad
import Control.Monad.Except
import System.Environment
import Text.ParserCombinators.Parsec hiding (spaces)
import Numeric
import Data.Ratio
import Data.Complex

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

spaces :: Parser ()
spaces = skipMany1 space

escapedChars :: Parser Char
escapedChars = do char '\\'
                  x <- oneOf "\\\"nrt"
                  return $ case x of
                             '\\' -> x
                             '"'  -> x
                             'n'  -> '\n'
                             'r'  -> '\r'
                             't'  -> '\t'

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

parseNumber :: Parser LispVal
parseNumber = many1 digit >>= return . Number . read

parseDecimal1 :: Parser LispVal
parseDecimal1 = many1 digit >>= (return . Number . read)

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

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

hex2dig x = fst $ readHex x !! 0

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

oct2dig x = fst $ readOct x !! 0

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

bin2dig  = bin2dig' 0
bin2dig' digint ""     = digint
bin2dig' digint (x:xs) = let old = 2 * digint + (if x == '0' then 0 else 1) in
                               bin2dig' old xs

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

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

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

toDouble :: LispVal -> Double
toDouble(Float f)  = realToFrac f
toDouble(Number n) = fromIntegral n

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

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

parseCharacter :: Parser LispVal
parseCharacter = do
  try $ string "#\\"
  value <- try (string "newline" <|> string "space")
          <|> do { x <- anyChar; notFollowedBy alphaNum ; return [x] }
  return $ Character $ case value of
                         "space"   -> ' '
                         "newline" -> '\n'
                         otherwise -> (value !! 0)

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

parseUnQuoted :: Parser LispVal
parseUnQuoted = do
  char ','
  x <- parseExpr
  return $ List [Atom "unquoted", x]

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

parseBracketed :: Parser LispVal -> Parser LispVal
parseBracketed p = try $ do
  char '('
  many space
  x <- p
  char ')'
  return x

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

parseList :: Parser LispVal
parseList = liftM List $ sepEndBy parseExpr spaces

parseVector :: Parser LispVal
parseVector = do
  string "#("
  many space
  x <- liftM Vector $ sepEndBy parseExpr spaces
  char ')'
  return x

parseExpr :: Parser LispVal
parseExpr = parseAtom
  <|> parseString
  <|> try parseFloat
  <|> try parseRatio
  <|> try parseComplex
  <|> try parseNumber
  <|> try parseBool
  <|> try parseCharacter
  <|> try parseQuoted
  <|> try parseUnQuoted
  <|> try parseQuasiQuoted
  <|> parseBracketed parseList
  <|> parseBracketed parseDottedList
  <|> try parseVector


{-# LANGUAGE GADTs,  TypeFamilies,
 StandaloneDeriving, UndecidableInstances,
ScopedTypeVariables, FlexibleInstances, DataKinds,
FunctionalDependencies, PolyKinds,
TypeOperators, RankNTypes,  MultiParamTypeClasses,
 TypeApplications, FlexibleContexts, AllowAmbiguousTypes #-}
-- AllowAmbiguousTypes, ImpredicativeTypes, InstanceSigs, NoImplicitPrelude,

module Fib where

import Vec
import Data.Complex
import Control.Monad  ((<=<))
import GHC.TypeNats
import Data.Proxy

-- data FibAnyon = Id | Tau
-- I started using DataKinds and it ended up being a problem.

data Tau
--data Id
type Id = ()
data FibTree root leaves where
   TTT :: FibTree Tau l -> FibTree Tau r -> FibTree Tau (l,r)
   ITT :: FibTree Tau l -> FibTree Tau r -> FibTree Id (l,r)
   TIT :: FibTree Id l -> FibTree Tau r -> FibTree Tau (l,r)
   TTI :: FibTree Tau l -> FibTree Id r -> FibTree Tau (l,r)
   III :: FibTree Id l -> FibTree Id r -> FibTree Id (l,r)
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

-- pretty printing would be hella nice
deriving instance Show (FibTree a b)
deriving instance Eq (FibTree a b)

instance Ord (FibTree a b) where
    compare (ITT l r) (ITT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (ITT _ _) _ = LT
    compare _ (ITT _ _) = GT
    compare (TTI l r) (TTI l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (TIT l r) (TIT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (TTT l r) (TTT l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'
    compare (III l r) (III l' r') | l < l' = LT
                                  | l > l' = GT
                                  | otherwise = compare r r'

    compare (TTI _ _) _ = LT
    compare _ (TTI _ _) = GT
    compare (TIT _ _) _ = LT
    compare _ (TIT _ _) = GT
    compare (TTT _ _) _ = LT
    compare _ (TTT _ _) = GT
    compare (III _ _) _ = LT
    compare _ (III _ _) = GT
    compare TLeaf TLeaf = EQ
    compare ILeaf ILeaf = EQ

lmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (b,d) -> Q (FibTree e (c,d)))
lmap f (ITT l r) = fmap (\l' -> ITT l' r) (f l)
lmap f (TTI l r) = fmap (\l' -> TTI l' r) (f l)
lmap f (TIT l r) = fmap (\l' -> TIT l' r) (f l)
lmap f (TTT l r) = fmap (\l' -> TTT l' r) (f l)
lmap f (III l r) = fmap (\l' -> III l' r) (f l)

rmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (d,b) -> Q (FibTree e (d,c)))
rmap f (ITT l r) = fmap (\r' -> ITT l r') (f r)
rmap f (TTI l r) = fmap (\r' -> TTI l r') (f r)
rmap f (TIT l r) = fmap (\r' -> TIT l r') (f r)
rmap f (TTT l r) = fmap (\r' -> TTT l r') (f r)
rmap f (III l r) = fmap (\r' -> III l r') (f r)

braid :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid (ITT l r) = W [(ITT r l,  cis $ 4 * pi / 5)]  -- different scalar factors for trivial and non trivial fusion
braid (TTT l r) = W [(TTT r l,  (cis $ - 3 * pi / 5))]
braid (TTI l r) = pure $ TIT r l-- exchange with trivial means nothing
braid (TIT l r) = pure $ TTI r l
braid (III l r) = pure $ III r l

-- The inverse of braid
braid' :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid' = star . braid

tau :: Complex Double
tau =  ((sqrt 5) - 1) / 2 :+ 0 -- 0.618 :+ 0

fmove :: FibTree a (c,(d,e)) -> Q (FibTree a ((c,d),e))
-- fmove (ITT  a  (TTI b c)) = pure $ ITI ( TTT  a b) c -- pure (auto (auto a b) c) -- no maybe not. The internal one isn't auto
fmove (ITT  a  (TIT b c)) = pure $ ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = pure $ ITT ( TTT  a b) c
fmove (ITT  a  (TTI b c)) = pure $ III ( ITT  a b) c

fmove (TIT  a  (TTT b c)) = pure $ TTT ( TIT  a b) c
fmove (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
fmove (TIT  a  (TIT b c)) = pure $ TIT ( III  a b) c

-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root
-- internal I
fmove (TTI  a  (III b c)) = pure $ TTI ( TTI  a b) c
fmove (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, tau)         , (TTT ( TTT  a b) c, sqrt tau)]
-- internal T
fmove (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, sqrt tau)  ,   (TTT ( TTT  a b) c, - tau   )]
fmove (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c

fmove (III  a  (ITT b c)) = pure $ ITT ( TIT  a b) c
fmove (III  a  (III b c)) = pure $ III ( III  a b) c

-- largely just a tranpose of the above case.
fmove' :: FibTree a ((c,d),e) -> Q (FibTree a (c,(d,e)))
fmove' (ITT ( TTI  a b) c) = pure $ (ITT  a  (TIT b c))
fmove' (ITT ( TTT  a b) c) = pure $  (ITT  a  (TTT b c))
fmove' (ITT ( TIT  a b) c) = pure $  (III  a  (ITT b c))

--fmoveq (ITT  a  (TTT b c)) = pure $

fmove' (TTI ( TTT  a b) c) = pure $ (TTT  a  (TTI b c))
fmove' (TTI ( TTI  a b) c) = pure $ (TTI  a  (III b c))
fmove' (TTI ( TIT  a b) c) = pure $ TIT  a  (TTI b c)
--fmoveq (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c

--fmoveq (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c
fmove' (TIT ( ITT  a b) c) = W [(TTI  a  (ITT b c), tau)         , (TTT  a  (TTT b c) , sqrt tau)]
fmove' (TIT ( III  a b) c ) = pure $ TIT  a  (TIT b c)

fmove' (TTT ( TTI  a b) c ) = pure $ TTT  a  (TIT b c)
fmove' (TTT ( TIT  a b) c ) = pure $ TIT  a  (TTT b c)
fmove' (TTT ( TTT  a b) c) = W [(TTI  a  (ITT b c), sqrt tau)  , (TTT  a  (TTT b c),   - tau  )]

fmove' (III ( III  a b) c ) = pure $ III  a  (III b c)
fmove' (III ( ITT  a b) c ) = pure $ ITT  a  (TTI b c)

rightUnit :: FibTree e (a,Id) -> Q (FibTree e a)
rightUnit (TTI t _) = pure t
rightUnit (III t _) = pure t

rightUnit' :: FibTree e a -> Q (FibTree e (a,Id))
rightUnit' t@(TTT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TTI _ _) = pure (TTI  t ILeaf)
rightUnit' t@(TIT _ _) = pure (TTI  t ILeaf)
rightUnit' t@(III _ _) = pure (III  t ILeaf)
rightUnit' t@(ITT _ _) = pure (III  t ILeaf)
rightUnit' t@(ILeaf) = pure (III t ILeaf)
rightUnit' t@(TLeaf) = pure (TTI t ILeaf)

leftUnit :: FibTree e (Id,a) -> Q (FibTree e a)
leftUnit = rightUnit <=< braid

-- braid vs braid' doesn't matter, but it has a nice symmettry.
leftUnit' :: FibTree e a -> Q (FibTree e (Id,a))
leftUnit' = braid' <=< rightUnit'

dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
dot x@(TTI _ _) y@(TTI _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TIT _ _) y@(TIT _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TTT _ _) y@(TTT _ _) | x == y = pure TLeaf
                            | otherwise = mempty
dot x@(III _ _) y@(III _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot x@(ITT _ _) y@(ITT _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot _ _ = mempty

class AllTrees a b where
    allTrees :: [FibTree a b]
instance AllTrees Tau Tau where
    allTrees = [TLeaf]
instance AllTrees Id Tau where
        allTrees = []
instance AllTrees Tau Id where
        allTrees = []
instance AllTrees Id Id where
        allTrees = [ILeaf]
instance (AllTrees Tau a,
         AllTrees Id a,
         AllTrees Tau b,
         AllTrees Id b) => AllTrees Id (a,b) where
        allTrees = (III <$> ia <*> ib) <> (ITT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
instance (AllTrees Tau a,
        AllTrees Id a,
        AllTrees Tau b,
        AllTrees Id b) => AllTrees Tau (a,b) where
       allTrees = (TIT <$> ia <*> tb) <> (TTI <$> ta <*> ib) <> (TTT <$> ta <*> tb) where
            ta = allTrees @Tau @a
            ia = allTrees @Id @a
            tb = allTrees @Tau @b
            ib = allTrees @Id @b
t9 = allTrees @Tau @((Tau,Tau),Tau)

-- Resulting type depends on input
-- I think a typefamily type computation might be necessary?
-- pullLeft and pullRight might not need a type class.
-- pullLeft (TTT l r) = fmove (TTT pl r)
--                         where pl = pullLeft l
{-
type family Assoc a where
   Assoc ((a,b),c) = (a,(b,c))

type family PullLeft a where
   PullLeft ((a,b),c) = Assoc (PullLeft (a, b), c)
   PullLeft a = a

type Test1 = PullLeft (((Int,Int),Int),Int)
pullLeft :: forall a b c. (c ~ PullLeft b) => FibTree a b -> Q (FibTree a c)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
            t' <- lmap pullLeft t
            fmove' t'

pullLeft t@(ITT l r) = helper t
pullLeft t@(TTT l r) = helper t
pullLeft t@(TTI l r) = helper t
pullLeft t@(TIT l r) = helper t
pullLeft t@(III l r) = helper t where
                       helper :: (d ~ PullLeft (b,c)) => FibTree a (b,c) -> Q (FibTree a d)
                       helper t = do
                                t' <- lmap pullLeft t
                                fmove' t'

-}

{-
Spiritually this function is this. Can't write it this way unforturnately
-- There are typing problems
-- this isn't even right
-- we need to pattern match on ((),)
pullLeft TLeaf = pure Tleaf
pullLeft ILeaf = pure ILeaf
pullLeft t =  do
                t' <- lmap pullLeft t
                fmove' t'

-- but actually a two level tree desctruct
--treeDestruct()

-}
{-
type family Assoc a where
    Assoc ((a,b),c) = (a,(b,c))

 type family PullLeft' a where
    PullLeft' ((a,b),c) = Assoc (PullLeft' (a, b), c)
    PullLeft' a = a
    -}
-- could seperate the typeclass so that I'm using type families instead.

class PullLeftLeaf a b | a -> b where
  pullLeftLeaf :: FibTree c a -> Q (FibTree c b)
instance PullLeftLeaf (Tau,c) (Tau,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf (Id,c) (Id,c) where
  pullLeftLeaf = pure
instance PullLeftLeaf Tau Tau where
  pullLeftLeaf = pure
instance PullLeftLeaf Id Id where
  pullLeftLeaf = pure
instance (PullLeftLeaf (a,b) (a',b'),
          r ~ (a',(b',c))) => PullLeftLeaf ((a, b),c) r where
  pullLeftLeaf t = do
           t' <- lmap pullLeftLeaf t
           fmove' t'

class PullRightLeaf a b | a -> b where -- | a -> b functional dependency causes errors?
    pullRightLeaf :: FibTree c a -> Q (FibTree c b)
{-
instance (PullRightLeaf a a', r ~ ((b,c),a')) => PullRightLeaf ((b,c),a) r where
pullRightLeaf t = rmap pullRightLeaf t
-}
instance PullRightLeaf Tau Tau where
    pullRightLeaf = pure

instance PullRightLeaf Id Id where
    pullRightLeaf = pure

instance PullRightLeaf (c,Tau) (c,Tau) where
    pullRightLeaf = pure

instance PullRightLeaf (c,Id) (c,Id) where
    pullRightLeaf = pure

instance (PullRightLeaf (a,b) (a',b'), r ~ ((c,a'),b')) => PullRightLeaf (c,(a, b)) r where
    pullRightLeaf t = do
        t' <- rmap pullRightLeaf t
        fmove t'

class RightAssoc a b | a -> b where
  rightAssoc :: FibTree c a -> Q (FibTree c b)
instance RightAssoc Tau Tau where
  rightAssoc = pure
instance RightAssoc Id Id where
  rightAssoc = pure
instance (PullLeftLeaf (a,b) (a',b'),
          RightAssoc b' b'',
          r ~ (a', b'')) => RightAssoc (a,b) r where
  rightAssoc t = do
           t' <- pullLeftLeaf t
           rmap rightAssoc t'
-- usually you'll want to force not the root, but the leaves to be some type
-- hence the ordering b c d a for type application
bmove :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove t = do
           t'  :: FibTree a ((b,c),d) <- fmove t
           t'' :: FibTree a ((c,b),d) <-  lmap braid t'
           fmove' t''
bmove' :: forall b c d a. FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))
bmove' = fmove' <=< (lmap braid') <=< fmove

-- Therei s a general pattern for digging into
-- n and a tell us b, the subpiece of a
-- c and a and n tell use what a looks like with b replaced = d
--data Nat = Z | S Nat

rmapN :: forall n gte s t a b e. (RMapN n gte s t a b, gte ~ (CmpNat n 0)) => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
rmapN f t = rmapN' @n @gte f t

class RMapN n gte s t a b | n gte s b -> a t where
    rmapN' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
instance (a ~ s, b ~ t) => RMapN 0 'EQ s t a b where
    rmapN' f t = f t -- rmapN = id
instance (RMapN (n-1) gte r r' a b,
               gte ~ (CmpNat (n-1) 0),
               t ~ (l,r')) => RMapN n 'GT (l,r) t a b where
    rmapN' f t = rmap (rmapN @(n-1) f) t
{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => FibTree f a -> Q (FibTree f e)
bmoveN t = rmapN @n (bmove @b @c @d) t

-- dotN :: forall n a b c d e f. dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
fuseN :: forall n a b c d e f g. RMapN n f (b,(c,d)) (a,d) g => FibTree a (b, c) -> FibTree e f -> Q (FibTree e g)
fuseN q t = rmapN @n f t where
    f :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (a,d))
    f t' = do
            t'' <- fmove t'
            lmap (dot q) t''
-}

type family Count a where
    Count Tau = 1
    Count Id = 1
    Count (a,b) = (Count a) + (Count b)

type family LeftCount a where
    LeftCount (a,b) = Count a

lcamap ::  forall n s t a b e gte .
           (gte ~ CmpNat (LeftCount s) n,
           LCAMap n gte s t a b)
           => (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)
lcamap f t = lcamap' @n @gte f t

class LCAMap n gte s t a b | n gte s b -> t a where
    lcamap' :: (forall r. FibTree r a -> Q (FibTree r b)) -> (FibTree e s) -> Q (FibTree e t)

instance (n' ~ (n - Count l), -- We're searching in the right subtree. Subtract the leaf number in the left subtree
            lc ~ (LeftCount r), -- dip one level down to order which way we have to go next
            gte ~ (CmpNat lc n'), -- Do we go left, right or have we arrived in the next layer?
            LCAMap n' gte r r' a b,  -- recursive call
            t ~ (l,r') -- reconstruct total return type from recursive return type. Left tree is unaffected by lcamapping
            ) => LCAMap n 'LT (l,r) t a b where
        lcamap' f x = rmap (lcamap @n' f) x

instance (lc ~ (LeftCount l),
            gte ~ (CmpNat lc n),
            LCAMap n gte l l' a b,
            t ~ (l',r)
            ) => LCAMap n 'GT (l,r) t a b where
        lcamap' f x = lmap (lcamap @n f) x

instance (t ~ b, a ~ s) => LCAMap n 'EQ s t a b where -- Base case
    lcamap' f x = f x

class Twiddle s t a b | s b -> t a where
    twiddle :: (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)

instance Twiddle ((l,x),(y,r)) ((l,c),r) (x,y) c where
    twiddle f x = do
            x'  <- fmove x -- (((l',x),y),r')
            x'' <- lmap fmove' x' -- ((l',(x,y)),r')
            lmap (rmap f) x''
instance Twiddle (Tau, (y,r)) (c,r) (Tau, y) c where
    twiddle f x = fmove x >>= lmap f
instance Twiddle (Id, (y,r)) (c,r)  (Id, y) c where
    twiddle f x = fmove x >>= lmap f
instance Twiddle ((l,x), Tau) (l,c) (x,Tau) c where
    twiddle f x = fmove' x >>= rmap f
instance Twiddle ((l,x), Id) (l,c) (x,Id) c where
    twiddle f x = fmove' x >>= rmap f
instance Twiddle (Tau, Tau) c (Tau,Tau) c where
    twiddle f x = f x
instance Twiddle (Id, Id) c (Id,Id)  c where
    twiddle f x = f x
instance Twiddle (Tau, Id) c (Tau,Id)  c where
    twiddle f x = f x
instance Twiddle (Id, Tau) c (Id,Tau) c where
    twiddle f x = f x

nmap :: forall (n :: Nat) s t a b a' b' l l' r r' e gte.
    (gte ~ CmpNat (LeftCount s) n,
    LCAMap n gte s t a' b',
    a' ~ (l,r),
    PullRightLeaf l l',
    PullLeftLeaf r r',
    Twiddle (l',r') b' a b) =>
    (forall r. FibTree r a -> Q (FibTree r b)) -> FibTree e s -> Q (FibTree e t)
nmap f z = lcamap @n @s @t @a' @b' (\x -> do
                                           x'  <- lmap pullRightLeaf x
                                           x'' <- rmap pullLeftLeaf x'
                                           twiddle f x'') z

t1 = nmap @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t5 = nmap @2 pure (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf)) >>= nmap @3 pure
t2 = nmap @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t4 = nmap @1 braid (TTT TLeaf (TTT TLeaf TLeaf))
t3 = nmap @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t6 = rightAssoc (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
t7 = t6 >>= bmove
t8 = t6 >>= rmapN @0 bmove
-- For the category, we probably don't want the association structure.
-- '[Tau, Id, Tau, Id] typelevel list of particles
--

ttt = TTT TLeaf TLeaf
example = starttree >>=
        nmap @1 braid >>=
        nmap @2 braid >>=
        nmap @1 (dot ttt) >>=
        nmap @1 braid' >>=
        nmap @2 (dot ttt) >>=
        nmap @1 (dot ttt) where
        starttree = pure (TTT (TTT TLeaf
                              (TTT TLeaf
                                   TLeaf))
                          TLeaf
                         )
-- would be nice to use an unknot
-- example2 =
-- I should use neighbormap for this.
-- maybe I should call f fuse and make it global?

-- bmoveN t = rmapN @n (bmove :: forall r. FibTree r (b,(c,d)) -> Q (FibTree r (c,(b,d)))) t
-- bmoveN t = rmapN @n @a @(b, (c, d)) @(c, (b, d)) @e bmove t

{-
bmoveN :: forall n a b c d e f. RMapN n a (b, (c, d)) (c, (b, d)) e => Proxy n -> FibTree f a -> Q (FibTree f e)
bmoveN p t = rmapN p (bmove :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(b,d)))) t
-}
   {- do
            t'  <- fmove t
            t'' <-  lmap braid' t'
            fmove' t''
-}

           --
{-
pullLeftLeaf
pullRightLeaf

-}
{-
class Standardize a b | a -> b where
    standardize :: FibTree c a -> Q (FibTree c b)

instance Standardize a b where

type family Append a b where
    Append (a,b) c = (a, Append b c)
    Append a c = (a, c)

type family Leftize where
    Leftize (a,b) = Append (Leftize a) (Leftize b)
    Leftize a = a

fullrightassoc = standardize
completerightassoc
bmove = fmove braid fmove'

standardize ::  FibTree c a -> Q (FibTree c (Leftize a))
standardize (TLeaf) = TLeaf
standardize (ILeaf) = ILeaf
standardize t = do
                x <- pullLeft t
                rmap standardize x

lcamap :: FibTree a b -> Proxy (a :: Int) -> FibTree a ?
lcamap f n t@(TTT l r) | count l == n  = f t
                       | count l < n   = lcamap f (n - count l) r
                       | otherwise     = lcamap f n l

pullRight ()

treerecursor :: FibTree a b -> (l -> r -> c) -> (leaf -> c) -> c
treerecursor (TTT l r) f g = f l r
treerecursor (TLeaf) = g Tleaf

-}


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
module FokkerPlanck.FourierSeriesGPU where

import           Control.Monad                      as M
import           Control.Monad.Parallel             as MP
import qualified Data.Array.Accelerate              as A
import qualified Data.Array.Accelerate.Data.Complex as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Complex                       as C
import           Data.DList                         as DL
import           Data.List                          as L
import           Data.Vector.Unboxed                as VU
import           FokkerPlanck.Analytic
import           FokkerPlanck.BrownianMotion        (Particle (..),
                                                     generateRandomNumber,
                                                     moveParticle, scalePlus,
                                                     thetaPlus)
import           FokkerPlanck.FourierSeries
import           FokkerPlanck.GPUKernel
import           FokkerPlanck.Histogram
import           GHC.Float
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import           System.Random.MWC
import           Utils.Parallel
import Debug.Trace
import Utils.List

{-# INLINE computeFourierCoefficientsGPU #-}
computeFourierCoefficientsGPU ::
     [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Acc (A.Vector (Float, Float, Float, Float))
  -> A.Exp Float
  -> A.Exp Float
  -> PTX
  -> [DList Particle]
  -> Histogram (C.Complex Double)
computeFourierCoefficientsGPU !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !freqArr !sigmaExp !logPeriodExp !ptx !xs =
  let particles =
        L.map
          (\(Particle phi rho theta r w) ->
             ( double2Float phi
             , double2Float (log rho)
             , double2Float theta
             , double2Float (log r)
             , double2Float w)) .
        DL.toList . DL.concat $
        xs
      !len = L.length particles
   in Histogram
        [ L.length phiFreqs
        , L.length rhoFreqs
        , L.length thetaFreqs
        , L.length rFreqs
        ]
        len .
      VU.map (\(a :+ b) -> float2Double a :+ float2Double b) .
      VU.fromList .
      A.toList .
      runNWith ptx (gpuKernel'' sigmaExp logPeriodExp freqArr) .
      A.fromList (A.Z A.:. len) $
      particles

{-# INLINE computeFrequencyArray #-}
computeFrequencyArray ::
     (A.Elt a) => [a] -> [a] -> [a] -> [a] -> A.Vector (a, a, a, a)
computeFrequencyArray !phiFreqs !rhoFreqs !thetaFreqs !rFreqs =
  A.fromList
    (A.Z A.:.
     ((L.length rFreqs) * (L.length thetaFreqs) * (L.length rhoFreqs) *
      L.length phiFreqs))
    [ (rFreq', thetaFreq', rhoFreq', phiFreq')
    | rFreq' <- rFreqs
    , thetaFreq' <- thetaFreqs
    , rhoFreq' <- rhoFreqs
    , phiFreq' <- phiFreqs
    ]

-- {-# INLINE computeFourierCoefficientsGPU #-}
-- computeFourierCoefficientsGPU ::
--      Acc (A.Vector (Double, Double, Double, Double))
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> Exp Double
--   -> Exp Double
--   -> Exp (AC.Complex Double)
--   -> PTX
--   -> [DList Particle]
--   -> Histogram (C.Complex Double)
-- computeFourierCoefficientsGPU !freqArr !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !deltaLogRho !maxScaleExp !halfLogPeriodExp !deltaLogRhoComplexExp !ptx !xs =
--   let !particles = DL.toList . DL.concat $ xs
--       !len = L.length particles
--       !scaleSampledParticles =
--         L.concat .
--         parMap
--           rdeepseq
--           (\particle@(Particle phi rho theta r) ->
--              let !n = Prelude.floor $ (log r + halfLogPeriod) / deltaLogRho
--              in (phi, rho, theta, r) :
--                 [ ( phi
--                   , rho
--                   , theta
--                   , (exp $
--                      (Prelude.fromIntegral i * deltaLogRho - halfLogPeriod)))
--                 | i <- [0 .. n - 1]
--                 ]) $
--         particles
--       !movedParticleArr =
--         compute .
--         afst .
--         A.filter
--           (\particle ->
--              let (_, rho, _, _) =
--                    unlift particle :: ( Exp Double
--                                       , Exp Double
--                                       , Exp Double
--                                       , Exp Double)
--              in rho A./= 0) .
--         A.map (moveParticle maxScaleExp) .
--         use . A.fromList (Z :. (L.length scaleSampledParticles)) $
--         scaleSampledParticles
--   in Histogram
--        [ L.length phiFreqs
--        , L.length rhoFreqs
--        , L.length thetaFreqs
--        , L.length rFreqs
--        ]
--        len .
--      VU.fromList . A.toList $
--      runWith ptx .
--      A.map
--        (\(unlift -> (rFreq, thetaFreq, rhoFreq, phiFreq)) ->
--           sfoldl
--             (\s particle ->
--                s +
--                deltaLogRhoComplexExp *
--                (coefficient
--                   halfLogPeriodExp
--                   rFreq
--                   thetaFreq
--                   rhoFreq
--                   phiFreq
--                   particle))
--             0
--             (constant Z)
--             movedParticleArr) $
--      freqArr

computeFourierCoefficientsGPU' ::
     Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> PTX
  -> [(Double, Double, Double, Double, Double)]
  -> Histogram (Complex Double)
computeFourierCoefficientsGPU' !sigma !period !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !ptx !xs =
  let -- !freqArr =
      --   A.use $
      --   computeFrequencyArray
      --     (L.map double2Float phiFreqs)
      --     (L.map double2Float rhoFreqs)
      --     (L.map double2Float thetaFreqs)
      --     (L.map double2Float rFreqs)
      !freqArr =
        A.use $
        computeFrequencyArray
          ( phiFreqs)
          ( rhoFreqs)
          ( thetaFreqs)
          ( rFreqs)
   in Histogram
        [ L.length phiFreqs
        , L.length rhoFreqs
        , L.length thetaFreqs
        , L.length rFreqs
        ]
        1 .
      VU.fromList .
      -- L.map (\(a :+ b) -> (float2Double a) :+ (float2Double b)) .
      A.toList .
      runNWith
        ptx
        (gpuKernel''
           (A.constant ( sigma))
           (A.constant ( (log period)))
           freqArr) .
      A.fromList (A.Z A.:. (L.length xs)) -- .
      -- L.map
      --   (\(a, b, c, d, e) ->
      --      ( double2Float a
      --      , double2Float b
      --      , double2Float c
      --      , double2Float d
      --      , double2Float e))
        $
      xs

{-# INLINE computePinwheelCoefficients #-}
computePinwheelCoefficients ::
     Int
  -> Int
  -> Int
  -> Int
  -> PTX
  -> [(Double, Double, Double, Double, Complex Double)]
  -> Histogram (Complex Double)
computePinwheelCoefficients !maxPhiFreq !maxRhoFreq !maxThetaFreq !maxRFreq !ptx !xs =
  let freqFunc maxFreq = [-(fromIntegral maxFreq) .. (fromIntegral maxFreq)]
      phiFreqs = freqFunc maxPhiFreq
      rhoFreqs = freqFunc maxRhoFreq
      thetaFreqs = freqFunc maxThetaFreq
      rFreqs = freqFunc maxRFreq
      !freqArr =
        A.use $ computeFrequencyArray phiFreqs rhoFreqs thetaFreqs rFreqs
  in Histogram
       [ L.length phiFreqs
       , L.length rhoFreqs
       , L.length thetaFreqs
       , L.length rFreqs
       ]
       1 .
     VU.fromList .
     A.toList .
     runNWith ptx (pinwheelCoefficientsAcc freqArr) .
     A.fromList (A.Z A.:. (L.length xs)) $
     xs


{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UndecidableInstances #-}

module Network.Network
( Network(..) ) where

import           Network.Layer
import           Numeric.LinearAlgebra
import           System.Random

-- | A network is
class Network a where
  type Parameters :: *

  predict :: Vector Double -> a -> Vector Double
  createNetwork :: (RandomGen g) => RandomTransform -> g -> Parameters -> a


{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE UndecidableInstances   #-}

module Numeric.Optimization.TwoPhase.Tableau
    ( -- * Tableaus
      Tableau
    , IsTableau
    , mkPhaseI
    , mkPhaseII
    , maximizeEntering
    , maximizeOptimal
    , minimizeEntering
    , minimizeOptimal
    , phaseIEntering
    , phaseIOptimal
    , tableauInfeasible
    , tableauLeaving
    , tableauResult
    , tableauStep
    ) where

import           Data.Finite                                    (Finite)
import qualified Data.List as List
import           Data.Maybe                                     (fromJust)
import           Data.Ord                                       (comparing)
import qualified Data.Vector.Sized as Vec
import qualified Data.Vector.Storable.Sized as SVec
import           GHC.TypeLits
import           Numeric.LinearAlgebra                          (Extractor(..), (??))
import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static                   (L)
import qualified Numeric.LinearAlgebra.Static as LS
import qualified Numeric.LinearAlgebra.Static.Vector as LS
import           Text.Printf                                    (printf)
import           Unsafe.Coerce                                  (unsafeCoerce)

import           Numeric.Optimization.Problem
import           Numeric.Optimization.TwoPhase.Tableau.Build
import           Numeric.Optimization.TwoPhase.Tableau.Pivot
import           Numeric.Optimization.TwoPhase.Tableau.VarMap
import           Numeric.Optimization.TwoPhase.Types

type IsTableau p v s a c =
    ( KnownNat v
    , KnownNat s
    , KnownNat a
    , KnownNat c
    , KnownNat (Rows p c)
    , KnownNat (Cols p v s a)
    )

data Tableau (p :: Phase) (d :: Direction) v s a c =
    Tableau
        { varMap :: VarMap (Rows p c) (Cols p v s a)
        , table  :: L (Rows p c) (Cols p v s a)
        }

instance (IsTableau p v s a c) => Show (Tableau p d v s a c) where
    show (Tableau vs x) = unlines (cols : rows)
      where
        cols        = concatMap (printf "%8s" . show)
            . Vec.toList $ columnVars vs

        rows        = zipWith (printf "%s   %s") rowVals rowNames
        rowVals     = lines . LA.format "" (printf "%8.2f") $ LS.unwrap x
        rowNames    = Vec.toList . Vec.map show $ rowVars vs

mkPhaseI
    :: (IsTableau 'PhaseI v s a c)
    => Problem d v s a c
    -> Tableau 'PhaseI d v s a c
mkPhaseI x =
    Tableau (mkVarMap x) (mkTable x)

mkPhaseII
    :: forall d v s a c.
        ( IsTableau 'PhaseI  v s a c
        , IsTableau 'PhaseII v s a c
        )
    => Tableau 'PhaseI  d v s a c
    -> Tableau 'PhaseII d v s a c
mkPhaseII (Tableau vs x) =
    Tableau (VarMap newRows newColumns) newMatrix
  where
    newRows     = Vec.tail . unsafeCoerce $ rowVars vs

    newColumns  =
        let is = fromJust $ Vec.fromList indices
        in  Vec.backpermute (columnVars vs) is

    newMatrix   =
        let xs = LS.unwrap x ?? (Drop 1, Pos (LA.idxs indices))
        in  fromJust $ LS.create xs

    indices     = fromIntegral <$> findIndicesColumns isPhaseII vs
    isPhaseII n = n /= Objective PhaseI && not (isArtificial n)

-- Read the value associated with a variable in the tableau. If the variable is
-- in the basis (has a row in the tableau), it's value is the value in the RHS
-- of that row divided by the value in the variable's column in that row. If
-- the variable is not in the basis it has a value of 0 (non-basic).
--
readValue
    :: (IsTableau p v s a c)
    => VarName
    -> Tableau p d v s a c
    -> Double
readValue n (Tableau vs x) =
    case (valRowIx, valColIx, rhsColIx) of
        (Just i,  Just j, Just r) -> index (i, r) x / index (i, j) x
        (Nothing, Just _, Just _) -> 0
        _                         -> error "readValue: Variable not in tableau"
  where
    valRowIx    = elemIndexRows n vs
    valColIx    = elemIndexColumns n vs
    rhsColIx    = elemIndexColumns RHS vs

-- Read the value associated with the current objective and each decision
-- variable, and store them in the Vars type for this solver. This is what is
-- returned when optimization succeeds.
--
tableauResult
    :: (IsTableau p v s a c)
    => Tableau p d v s a c
    -> TwoPhaseResult d v
tableauResult x =
    TwoPhaseResult . Vec.zip names $ Vec.map toValue names
  where
    toValue i   = readValue i x
    names       = fromJust . Vec.fromList . findColumns isVar $ varMap x

    isVar n     =
        let objective = indexColumns 0 (varMap x)
        in  n == objective || isDecision n

-- A phase I tableau is infeasible if the value of the objective function is
-- greater than 0. This means that at least one artificial variable has to be
-- non-zero, so there can be no feasible BFS to start phase II from.
--
tableauInfeasible
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Bool
tableauInfeasible x
    | phaseIOptimal x   = readValue (Objective PhaseI) x /= 0
    | otherwise         = False

-- A tableau is optimal if the given test passes for each cell corresponding to
-- a variable (everything apart from columns for objectives and the RHS column)
-- in the current objective.
--
tableauOptimal
    :: (IsTableau p v s a c)
    => (Double -> Bool)
    -> Tableau p d v s a c
    -> Bool
tableauOptimal f (Tableau vs x) =
    SVec.and . SVec.imap optimal . LS.rVec . head . Vec.toList $ LS.lRows x
  where
    coeffs                  = findIndicesColumns isCoeff vs
    isCoeff                 = not . isSpecial

    optimal i a
        | i `elem` coeffs   = f a
        | otherwise         = True

-- A phase I problem is always a minimization. This is a separate function
-- from 'minimizeOptimal' as the problem being solved in phase II may not be
-- a minimization problem.
--
phaseIOptimal
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Bool
phaseIOptimal =
    tableauOptimal (<= 0)

-- A phase II problem which is a maximization is optimal when all coefficients
-- in the objective function are non-negative.
--
maximizeOptimal
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Max v s a c
    -> Bool
maximizeOptimal =
    tableauOptimal (>= 0)

-- A phase I problem which is a minimization is optimal when all coefficients
-- in the objective function are non-negative.
--
minimizeOptimal
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Min v s a c
    -> Bool
minimizeOptimal =
    tableauOptimal (<= 0)

-- The entering variable is the lowest indexed variable where the value in the
-- objective column passes a test, and the name of the variable passes another
-- test. The tests change depending on the phase, and whether the optimization
-- is a maximization or a minimization.
--
enteringVar
    :: (IsTableau p v s a c)
    => (Finite (Cols p v s a) -> Bool)
    -> (VarName -> Bool)
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Finite (Cols p v s a))
enteringVar validCol validVar (Tableau vs _) =
    case filter canEnter columns of
        [] -> Left NoEntering
        xs -> Right . fst $ List.minimumBy (comparing snd) xs
  where
    canEnter (j, n) = validCol j && validVar n
    columns         = Vec.toList . Vec.indexed $ columnVars vs

-- In phase I, the entering variable must be a positive value, as the problem
-- is a minimization. Artificial variables are still in the table at this
-- point, so the valid variable test must make sure artificial variables cannot
-- enter the basis (as they should be non-basic).
--
phaseIEntering
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseI v s a))
phaseIEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) > 0
    validVar n = isDecision n || isSlack n

-- When maximizing, the variable to enter the basis is the lowest indexed
-- variable with a negative coefficient in the objective function row.
--
maximizeEntering
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Max v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseII v s a))
maximizeEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) < 0
    validVar   = not . isSpecial

-- When minimizing, the variable to enter the basis is the lowest indexed
-- variable with a positive coefficient in the objective function row.
--
minimizeEntering
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Min v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseII v s a))
minimizeEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) > 0
    validVar   = not . isSpecial

-- The variable to leave the basis is the lowest indexed variable which has
-- the minimum ratio between it's coefficient and the RHS value. A variable
-- with a coefficient of 0 is not allowed to leave (as the ratio is infinite).
-- This function is the same regardless of the direction of optimization.
--
tableauLeaving
    :: (IsTableau p v s a c)
    => Finite (Cols p v s a)
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Finite (Rows p c))
tableauLeaving enter (Tableau vs x) =
    case toRatio <$> filter canLeave rows of
        [] -> Left Unbounded
        xs -> Right . fst $ List.minimumBy (comparing snd) xs
  where
    rows            = Vec.toList . Vec.indexed $ rowVars vs
    rhsColumn       = fromJust $ elemIndexColumns RHS vs

    toRatio (i, _)  =
        let val = index (i, enter) x
            rhs = index (i, rhsColumn) x
        in  (i, rhs / val)

    canLeave (i, n) =
        let val = index (i, enter) x
            rhs = index (i, rhsColumn) x
        in  val > 0 && rhs /= 0 && not (isSpecial n)

tableauStep
    :: (IsTableau p v s a c)
    => (Tableau p d v s a c
        -> Either TwoPhaseError (Finite (Cols p v s a)))
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Tableau p d v s a c)
tableauStep tableauEntering t@(Tableau vs x) = do
    enter <- tableauEntering t
    leave <- tableauLeaving enter t

    let cell = (leave, enter)
    pure $ Tableau (updateRow cell vs) (pivot cell x)


module Calibrate where

import Data.WAVE
import Data.Complex
import Data.List
import Numeric.FFT

import Util
import ExtProg

-- Calibrated or uncalibrated voice: calibrated means we know the fundamental frequency
-- for pitch 50, uncalibrated means we use the given voice name and text to get it.

type Calibration = Either (String, String, String) Double

-- Determine the fundamental frequency for a voice at the given pitch. Get the spectrum,
-- smoothen it and get the first peak.

voiceFundamental :: String -> String -> Int -> String -> IO (Maybe Double)

voiceFundamental exec text pitch voice = do
  p <- voiceSpectrumPitchText exec text pitch voice
  let smpp = filter ((> 0.05) . snd) $ smoothen 3 p
      lmpp = locMax smpp
      fund = case lmpp of
        [] -> Nothing
        _ -> Just $ fst $ head lmpp
  return fund

-- The most general case: set text, calibration pitch, and voice name.

voiceSpectrumPitchText :: String -> String -> Int -> String -> IO [(Double, Double)]

voiceSpectrumPitchText exec text pitch voice = do
  spd <- findSpeed exec 120 voice pitch 0.25 text (1, 300) >>= return . fst
  wav <- runEspeak exec 120 voice pitch spd text >>= getSample
  let fft = filter ((< 2000) . fst) $ waveFFT wav
  return fft

-- Run a FFT on a signal given WAVE, returning pairs of frequency-normalized magnitude
-- sorted by decreasing magnitude

waveFFT :: WAVE -> [(Double, Double)]

waveFFT wav =
  let sampldbl = map (sampleToDouble . head) $ waveSamples wav
      sampcmplx = map (:+ 0) sampldbl
      freqdom = map magnitude $ fft sampcmplx
      rate = waveFrameRate $ waveHeader wav
      freqs = map (\i -> fromIntegral (i * fromIntegral rate) / fromIntegral (length freqdom))
                  [0 .. length freqdom - 1]
      spectrum = zip freqs (normalize freqdom)
      sortDescSnd = sortBy $ \(x1, y1) (x2, y2) -> compare y2 y1
  in  sortDescSnd spectrum

-- Find a minimal range of speed that approximates the given time in seconds

findSpeed :: String -> Int -> String -> Int -> Rational -> String -> (Int, Int) -> IO (Int, Int)

findSpeed exec ampl voice pitch tsec utter (lo, hi) =
  rangeSearch 1 (lo, hi) $ \mid -> do
    wav <- runEspeak exec ampl voice pitch mid utter >>= getSample
    return $ compare (soundLength wav) tsec


module Game.Physics.Types ( Position
                          , Velocity
                          , Force
                          , PhysicsObject(..)
                          ) where

import Prelewd

import Control.DeepSeq
import Numeric.LinearAlgebra hiding ((<>))

import Text.Show

type Position     = Vector Double
type Velocity     = Vector Double
type Force        = Vector Double

-- | Represents a simulatable object's current state. Pass a Data.Vector to of
--   these to the integrator to step the simulation.
data PhysicsObject = PObject { pos      :: Position
                             , vel      :: Velocity
                             , invMass  :: Double -- ^ 1/mass
                             -- | Given a function which retrieves a physics
                             --   object from its index in the vector and the
                             --   Δt from the start of the frame, return the
                             --   net force acting on the object.
                             --
                             --   The retrieved physics objects will have
                             --   updated positions and velocities.
                             , netForce :: (Int -> PhysicsObject)
                                        -> Double
                                        -> Force
                             }

instance Show PhysicsObject where
    show (PObject p v im _) = "pos = " <> show p
                            <> " | vel = " <> show v
                            <> " | invMass = " <> show im

instance NFData PhysicsObject


{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
module Lib where

import Numeric.Vector.Sized
import Dependent.Size
import qualified Data.Array.Accelerate as Acc
import Data.Array.Accelerate(Acc,Array,Double)
import qualified Data.Array.Accelerate.Interpreter as AI

run :: ShapeSize size => SizedArray size -> Array (ShapeOf size) Double
run = runUsing AI.run

runUsing
  :: (ShapeSize size, Acc.Arrays (Array (ShapeOf size) Double))
  => (Acc (Array (ShapeOf size) Double) -> Array (ShapeOf size) Double)
  -> SizedArray size
  -> Array (ShapeOf size) Double
runUsing runner (Sized arr) = runner arr


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs     #-}
{-|
Module      : Grenade.Utils.PascalVoc
Description : Post-processing for TinyYoloV2
-}

module Grenade.Utils.PascalVoc (DetectedObject, labels, processOutput) where

import           Data.List                    (elemIndices)
import qualified Numeric.LinearAlgebra.Data   as NLA
import qualified Numeric.LinearAlgebra.Static as H

import           Grenade.Core.Shape
import           Grenade.Types

-- | List of possible labels of TinyYoloV2
labels :: [String]
labels
  = [ "aeroplane"
    , "bicycle"
    , "bird"
    , "boat"
    , "bottle"
    , "bus"
    , "car"
    , "cat"
    , "chair"
    , "cow"
    , "diningtable"
    , "dog"
    , "horse"
    , "motorbike"
    , "person"
    , "pottedplant"
    , "sheep"
    , "sofa"
    , "train"
    , "tvmonitor"
    ]

-- | Type alias for output from Yolo neural network.
--
--   Stores (x position, y position, width height, confidence score, object label).
type DetectedObject = (Int, Int, Int, Int, RealNum, String)

-- | Given the output of TinyYoloV2, finds all bounding boxes
--   with confidence higher than a threshold, and their labels
processOutput :: S ('D3 13 13 125) -- ^ Output tensor from TinyYoloV2
              -> RealNum           -- ^ Minimum confidence for bounding boxes, between 0 and 1. A good default is 0.3.
              -> RealNum           -- ^ Maximum IoU ratio to consider boxes not overlapping, between 0 and 1. A good default is 0.5.
              -> [DetectedObject]
processOutput (S3D mat) confThreshold iouThreshold = collapseBoxes iouThreshold $ map toDetectedObject filtered
  where
    filtered = filter (\(_, _, _, _, c, probs) -> snd (argMax probs) * c > confThreshold ) boxDescs
    out      = H.extract mat
    boxDescs = [ getBoundingBoxDesc out i j box | i <- [0..12], j <- [0..12], box <- [0..4]]

    toDetectedObject (x, y, w, h, c, probs) = (l, r, t, b, c, labels !! (fst $ argMax probs))
      where
        l = floor $ x - w / 2
        r = floor $ x + w / 2
        t = floor $ y - h / 2
        b = floor $ y + h / 2

    -- Returns (index of maximum element, maximum element)
    argMax probs = (i, m)
      where
        m = maximum probs
        i = head $ elemIndices m probs

-- | Given an IoU threshold and a list of boxes, find overlapping boxes and collapse them into one.
--
--   Boxes are determined to be overlapping if their intersection over union value passes the threshold.
--   When boxes are merged we pick the box with higher confidence score.
--
--   IoU threshold should be between 0 and 1, we've found 0.5 to work quite well.
collapseBoxes :: RealNum -> [DetectedObject] -> [DetectedObject]
collapseBoxes threshold xs = filter (\y -> all (doesNotOverlapWith y) xs) xs
  where
    -- Consider objects as overlapping if they have same label and their intersection over union area is too big.
    -- Only consider overlapping the object which has lower confidence.
    doesNotOverlapWith :: DetectedObject -> DetectedObject -> Bool
    doesNotOverlapWith y@(left, right, top, bottom, conf, label) x@(left', right', top', bottom', conf', label')
      = label /= label' || y == x || iou < threshold || conf > conf'
        where
          ix1   = max left left'
          ix2   = min right right'
          iy1   = max top top'
          iy2   = min bottom bottom'
          area  = (right  - left)  * (bottom  - top)
          area' = (right' - left') * (bottom' - top')
          inter = if ix1 < ix2 && iy1 < iy2 then (ix2 - ix1) * (iy2 - iy1) else 0
          iou   = fromIntegral inter / fromIntegral (area + area' - inter)

-- | Given the output matrix from YOLO, a cell position cy, cx and a bounding box b,
--   outputs (x, y, width, height, confidence score, probability distribution over 20 classes)
--   for said box.
--
--   Must have that 0 <= cy, cx <= 12 and 0 <= b <= 4.
getBoundingBoxDesc :: NLA.Matrix RealNum -> Int -> Int -> Int -> (RealNum, RealNum, RealNum, RealNum, RealNum, [RealNum])
getBoundingBoxDesc out cy cx b = (x, y, w, h, confidence * (maximum classes), classes)
  where
    anchors    = [1.08, 1.19, 3.42, 4.41, 6.63, 11.38, 9.42, 5.11, 16.62, 10.52]
    numClasses = 20
    channel    = b * 25

    tx = NLA.atIndex out ((channel + 0) * 13 + cy, cx)
    ty = NLA.atIndex out ((channel + 1) * 13 + cy, cx)
    tw = NLA.atIndex out ((channel + 2) * 13 + cy, cx)
    th = NLA.atIndex out ((channel + 3) * 13 + cy, cx)
    tc = NLA.atIndex out ((channel + 4) * 13 + cy, cx)

    x = (fromIntegral cx + sigmoid tx) * 32
    y = (fromIntegral cy + sigmoid ty) * 32

    w = exp tw * (anchors !! (2 * b    )) * 32
    h = exp th * (anchors !! (2 * b + 1)) * 32

    confidence = sigmoid tc

    classes = softmax [ NLA.atIndex out ((channel + 5 + c) * 13 + cy, cx) | c <- [0..numClasses - 1] ]

    sigmoid :: RealNum -> RealNum
    sigmoid = (1 /) . (1 +) . exp . negate

    softmax :: [RealNum] -> [RealNum]
    softmax xs = map (/sumxs) expxs
      where
        expxs = map exp xs
        sumxs = sum expxs


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      :  Spiral.Util.Pretty.LaTeX
-- Copyright   :  (c) 2017-2020 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.Util.Pretty.LaTeX (
    Pretty(..)
  ) where

import Data.Complex (Complex(..))
import Data.Complex.Cyclotomic
import Data.List (intersperse)
import qualified Data.Map as Map
import qualified Data.Matrix as M
import Data.Modular (Mod, unMod)
import Data.Monoid ((<>))
import Data.Ratio (Ratio, numerator, denominator)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Symbol (unintern)
import qualified Data.Text as T
import qualified Data.Vector as V
import GHC.TypeLits (KnownNat)
import Text.LaTeX
import Text.LaTeX.Base.Class (comm0,
                              comm1)
import Text.LaTeX.Packages.AMSFonts
import Text.LaTeX.Packages.AMSMath

import Spiral.Array (IArray,
                     Matrix,
                     (!))
import qualified Spiral.Array as A
import Spiral.Array.Shape
import Spiral.Exp
import Spiral.Permutation hiding (Inv)
import qualified Spiral.Permutation as P
import Spiral.SPL
import Spiral.Util.Name (Name(..))
import Spiral.Util.Pretty (Assoc(..),
                           Fixity(..),
                           HasFixity(..),
                           addPrec,
                           infixl_,
                           precOf)
import Spiral.Util.Uniq (Uniq(..))

class Pretty a where
  {-# MINIMAL pprPrec | ppr #-}
  ppr     :: a -> LaTeX
  pprPrec :: Int -> a -> LaTeX

  pprList     :: [a] -> LaTeX
  pprPrecList :: Int -> [a] -> LaTeX

  ppr       = pprPrec 0
  pprPrec _ = ppr

  pprPrecList _ = pprList
  pprList       = autoSquareBrackets . commasep . map ppr

parensIf :: Bool -> LaTeX -> LaTeX
parensIf False l = l
parensIf True l  = autoParens l

hsep :: LaTeX -> [LaTeX] -> LaTeX
hsep sep = mconcat . intersperse sep

commasep :: [LaTeX] -> LaTeX
commasep = hsep ","

mskip :: LaTeX -> LaTeX
mskip = comm1 "mskip"

nicefrac :: LaTeX -> LaTeX -> LaTeX
nicefrac e1 e2 = mempty ^: e1 <> (mskip "-2mu" <> "/" <> mskip "-1mu") !: e2
--nicefrac e1 e2 = e1 <> "/" <> e2

instance Pretty a => Pretty [a] where
  ppr     = pprList
  pprPrec = pprPrecList

instance Pretty a => Pretty (Set a) where
  ppr = autoBraces . commasep . map ppr . Set.toList

instance Pretty Bool where
  ppr True  = mathfrak "T"
  ppr False = mathfrak "F"

pprSignedIntegral :: (Integral a, Show a) => a -> LaTeX
pprSignedIntegral x
  | x < 0     = "-" <> pprSignedIntegral (-x)
  | otherwise = (fromString . show) x

pprRealFloat :: (RealFloat a, Show a) => a -> LaTeX
pprRealFloat x
    | isIntegral x = ppr (ceiling x :: Integer)
    | x < 0        = fromString $ showFloat (-x)
    | otherwise    = fromString $ showFloat x
  where
    isIntegral :: RealFrac a => a -> Bool
    isIntegral x = fromIntegral (ceiling x :: Integer) == x

instance Pretty Int where
  ppr = pprSignedIntegral

instance Pretty Integer where
  ppr = pprSignedIntegral

instance Pretty Float where
  ppr = pprRealFloat

instance Pretty Double where
  ppr = pprRealFloat

instance (Eq a, Num a, Pretty a) => Pretty (Ratio a) where
  ppr x
      | d == 1    = ppr n
      | otherwise = nicefrac (ppr n) (ppr d)
    where
      n = numerator x
      d = denominator x

-- | Pretty-print an imaginary number
pprIm :: (Eq a, Num a, Pretty a) => a -> LaTeX
pprIm 0    = mempty
pprIm 1    = "i"
pprIm (-1) = "-i"
pprIm i    = ppr i <> "i"

instance (Eq a, Num a, Pretty a) => Pretty (Complex a) where
  ppr (r :+ 0)    = ppr r
  ppr (0 :+ 1)    = "i"
  ppr (0 :+ (-1)) = "-i"
  ppr (0 :+ i)    = pprIm i
  ppr (r :+ i)    = case (T.unpack . render) (pprIm i) of
                      '-' : _ -> ppr r <> pprIm i
                      _ -> ppr r + pprIm i

instance Pretty Cyclotomic where
  pprPrec p (Cyclotomic n mp)
    | Map.null mp = "0"
    | null xs     = leadingTerm rat n ex
    | otherwise   = parensIf (p > addPrec) $
                    leadingTerm rat n ex <> mconcat (map (followingTerm n) xs)
    where
      ((ex,rat):xs) = Map.toList mp

      pprBaseExp :: Integer -> Integer -> LaTeX
      pprBaseExp n 1  = omega !: ppr n
      pprBaseExp n ex = omega !^ (ppr n, ppr ex)

      leadingTerm :: Rational -> Integer -> Integer -> LaTeX
      leadingTerm r _ 0 = ppr r
      leadingTerm r n ex
        | r == 1     = t
        | r == (-1)  = "-" <> t
        | r > 0      = ppr r <> t
        | r < 0      = "-" <> ppr (-r) <> t
        | otherwise  = mempty
        where
          t = pprBaseExp n ex

      followingTerm :: Integer -> (Integer, Rational) -> LaTeX
      followingTerm n (ex, r)
        | r == 1     = "+" <> t
        | r == (-1)  = "-" <> t
        | r > 0      = "+" <> ppr r <> t
        | r < 0      = "-" <> ppr (-r) <> t
        | otherwise  = mempty
        where
          t = pprBaseExp n ex

instance (Pretty i, KnownNat p) => Pretty (Mod i p) where
  ppr z = ppr (unMod z)

instance Pretty (Const a) where
  ppr (BoolC x)      = ppr x
  ppr (IntC x)       = ppr x
  ppr (IntegerC x)   = ppr x
  ppr (FloatC x)     = ppr x
  ppr (DoubleC x)    = ppr x
  ppr (RationalC x)  = ppr x
  ppr (ComplexC r i) = ppr (r :+ i)
  ppr (W _ 0 _)      = "1"
  ppr (W n 1 _)      = omega !: ppr n
  ppr (W n k _)      = omega !^ (ppr n, ppr k)
  ppr (CycC x)       = ppr x
  ppr (PiC k)        = ppr k <> pi_
  ppr (ModularC x)   = ppr x

instance Pretty Uniq where
  ppr (Uniq u) = ppr u

instance Pretty Name where
  ppr (Name sym Nothing)  = fromString (unintern sym)
  ppr (Name sym (Just u)) = fromString (unintern sym) !: ppr u

instance Pretty Var where
  ppr (Var n) = ppr n

instance (Pretty e, IArray r DIM2 e) => Pretty (Matrix r e) where
  ppr mat = bmatrix Nothing $ M.matrix m n f
    where
      Z :. m :. n = A.extent mat

      f :: (Int, Int) -> LaTeX
      f (i, j) = ppr (mat ! (i-1, j-1))

instance Pretty (Exp a) where
  pprPrec p (ConstE c) = pprPrec p c
  pprPrec p (VarE v)   = pprPrec p v

  pprPrec _ (UnopE Abs e) =
    autoBrackets "|" "|" (ppr e)

  pprPrec _ (UnopE Sqrt e) =
    tsqrt Nothing (ppr e)

  pprPrec p (UnopE op e) =
    unop p op e

  pprPrec p (BinopE op e1 e2) =
    infixop p op e1 e2

  pprPrec _ (IdxE ev eis) =
    ppr ev <> autoSquareBrackets (commasep (map ppr eis))

  pprPrec p (ComplexE er ei) =
    parensIf (p > addPrec) $
    ppr (er :+ ei)

  pprPrec _ (ReE e) =
    comm0 "Re" <> pprPrec 10 e

  pprPrec _ (ImE e) =
    comm0 "Im" <> pprPrec 10 e

  pprPrec p (BBinopE op e1 e2) =
    infixop p op e1 e2

  pprPrec p (IfE e1 e2 e3) =
    parensIf (p > 10) $
    operatorname (mathsf "if") <> ppr e1 <>
    operatorname (mathsf "then") <> ppr e2 <>
    operatorname (mathsf "else") <> ppr e3

instance Pretty Unop where
  ppr Neg    = "-"
  ppr Abs    = operatorname "abs"
  ppr Signum = operatorname "sgn"
  ppr Exp    = texp
  ppr Log    = tlog
  ppr Sqrt   = operatorname "sqrt"
  ppr Sin    = tsin
  ppr Cos    = tcos
  ppr Asin   = arcsin
  ppr Acos   = arccos
  ppr Atan   = arctan
  ppr Sinh   = tsinh
  ppr Cosh   = tcosh
  ppr Asinh  = operatorname "arsinh"
  ppr Acosh  = operatorname "arcosh"
  ppr Atanh  = operatorname "artanh"

instance Pretty Binop where
  ppr Add  = "+"
  ppr Sub  = "-"
  ppr Mul  = comm0 "cdot"
  ppr Quot = comm0 "quot"
  ppr Rem  = operatorname "rem"
  ppr Div  = operatorname "div"
  ppr Mod  = operatorname "mod"
  ppr FDiv = "/"

instance Pretty BBinop where
  ppr Eq = "="
  ppr Ne = comm0 "neq"
  ppr Lt = "<"
  ppr Le = comm0 "le"
  ppr Ge = comm0 "ge"
  ppr Gt = ">"

instance Pretty Permutation where
  ppr (L mn n)        =  "L" !^ (ppr n, ppr mn)
  ppr (J n)           = operatorname "J" !: ppr n
  ppr (CS n k)        = operatorname "S" !: ppr n <> autoParens (ppr k)
  ppr (CRT m n _ _)   = operatorname "CRT" !^ (ppr n, ppr m)
  ppr (CRT' m n _ _)  = operatorname "CRT'" !^ (ppr n, ppr m)
  ppr (Good m n _ _)  = operatorname "Good" !^ (ppr n, ppr m)
  ppr (Good' m n _ _) = operatorname "Good'" !^ (ppr n, ppr m)
  ppr (R p a)         = operatorname "R" !^ (ppr p, ppr a)
  ppr (P.Inv p _)     = ppr p ^: "-1"

data MatrixBinop = KOp
                 | DSOp
  deriving (Eq, Ord, Show)

instance HasFixity MatrixBinop where
  fixity KOp  = infixl_ 7
  fixity DSOp = infixl_ 6

instance Pretty MatrixBinop where
  ppr KOp  = comm0 "otimes"
  ppr DSOp = comm0 "oplus"

pprArgs :: Pretty a => [a] -> LaTeX
pprArgs = autoParens . commasep . map ppr

instance (Num e, Pretty e) => Pretty (SPL e) where
  pprPrec p m@E{}       = pprPrec p (toMatrix m)
  pprPrec _ (I n)       = operatorname "I" !: ppr n
  pprPrec _ (T e)       = pprPrec 10 e ^: comm0 "intercal"
  pprPrec _ (Inv e)     = pprPrec 10 e ^: "-1"
  pprPrec p (Pi pi)     = pprPrec p pi
  pprPrec _ (Rot alpha) = operatorname "R" !: ppr alpha
  pprPrec _ (Diag xs)   = operatorname "diag" <> pprArgs (V.toList xs)
  pprPrec p (KDiag _ e) = pprPrec p e

  pprPrec _ ((tl `Beside` tr) `Above` (bl `Beside` br)) =
    bmatrix Nothing $ M.matrix 2 2 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr tl
        f (1, 2) = ppr tr
        f (2, 1) = ppr bl
        f (2, 2) = ppr br
        f _      = error "can't happen"

  pprPrec _ (Above a b) =
    bmatrix Nothing $ M.matrix 2 1 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr a
        f (2, 1) = ppr b
        f _      = error "can't happen"

  pprPrec _ (Beside a b) =
    bmatrix Nothing $ M.matrix 1 2 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr a
        f (1, 2) = ppr b
        f _      = error "can't happen"

  pprPrec p (Kron a b)  = infixop p KOp a b
  pprPrec p (DSum a b)  = infixop p DSOp a b
  pprPrec _ (Prod a b)  = pprPrec 10 a <> pprPrec 10 b
  pprPrec _ (Circ xs)   = operatorname "circ" <> pprArgs (V.toList xs)
  pprPrec _ (Skew xs)   = operatorname "skew" <> pprArgs (V.toList xs)
  pprPrec _ (Toep xs)   = operatorname "toep" <> pprArgs (V.toList xs)
  pprPrec _ (Re a)      = comm0 "Re" <> autoParens (ppr a)
  pprPrec _ F2          = operatorname "F" !: "2"
  pprPrec _ (DFT n)     = operatorname "DFT" !: ppr n
  pprPrec _ (DFT' n)    = operatorname "DFT" !^ (ppr n, "-1")
  pprPrec _ (F n w)     = operatorname "DFT" !: ppr n <> autoParens (ppr w)
  pprPrec _ (F' n w)    = operatorname "DFT" !^ (ppr n, "-1") <> autoParens (ppr w)

unop :: (Pretty a, Pretty op, HasFixity op)
     => Int -- ^ precedence of context
     -> op  -- ^ operator
     -> a
     -> LaTeX
unop prec op x =
    parensIf (prec > precOf op) $
    ppr op <> pprPrec (precOf op) x

infixop :: (Pretty a, Pretty b, Pretty op, HasFixity op)
        => Int -- ^ precedence of context
        -> op  -- ^ operator
        -> a   -- ^ left argument
        -> b   -- ^ right argument
        -> LaTeX
infixop prec op l r =
    parensIf (prec > opPrec) $
    pprPrec leftPrec l <> ppr op <> pprPrec rightPrec r
  where
    leftPrec | opAssoc == RightAssoc = opPrec + 1
             | otherwise             = opPrec

    rightPrec | opAssoc == LeftAssoc = opPrec + 1
              | otherwise            = opPrec

    Fixity opAssoc opPrec = fixity op


{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, StandaloneDeriving #-}

import           Prelude hiding (mapM)

import           Options.Applicative
import           Data.Monoid ((<>))
import           Control.Monad.Trans.Class

import           Data.Vector (Vector)
import qualified Data.Vector.Generic as V
import           Statistics.Sample (mean)

import           Data.Traversable (mapM)
import qualified Data.Set as S
import           Data.Set (Set)
import qualified Data.Map.Strict as M

import           ReadData
import           SerializeText
import qualified RunSampler as Sampler
import           BayesStack.DirMulti
import           BayesStack.Models.Topic.LDA
import           BayesStack.UniqueKey

import qualified Data.Text as T
import qualified Data.Text.IO as TIO

import           System.FilePath.Posix ((</>))
import           Data.Binary
import qualified Data.ByteString as BS
import           Text.Printf

import           Data.Random
import           System.Random.MWC

data RunOpts = RunOpts { nodesFile       :: FilePath
                       , stopwords       :: Maybe FilePath
                       , nTopics         :: Int
                       , samplerOpts     :: Sampler.SamplerOpts
                       , hyperParams     :: HyperParams
                       }

runOpts :: Parser RunOpts
runOpts = RunOpts
    <$> strOption  ( long "nodes"
                  <> short 'n'
                  <> metavar "FILE"
                  <> help "File containing nodes and their associated items"
                   )
    <*> nullOption ( long "stopwords"
                  <> short 's'
                  <> metavar "FILE"
                  <> reader (pure . Just)
                  <> value Nothing
                  <> help "Stop word list"
                   )
    <*> option     ( long "topics"
                  <> short 't'
                  <> metavar "N"
                  <> value 20
                  <> help "Number of topics"
                   )
    <*> Sampler.samplerOpts
    <*> hyperOpts

hyperOpts = HyperParams
    <$> option     ( long "prior-theta"
                  <> value 1
                  <> help "Dirichlet parameter for prior on theta"
                   )
    <*> option     ( long "prior-phi"
                  <> value 0.1
                  <> help "Dirichlet parameter for prior on phi"
                   )

mapMKeys :: (Ord k, Ord k', Monad m, Applicative m)
         => (a -> m a') -> (k -> m k') -> M.Map k a -> m (M.Map k' a')
mapMKeys f g x = M.fromList <$> (mapM (\(k,v)->(,) <$> g k <*> f v) $ M.assocs x)

termsToItems :: M.Map NodeName [Term]
             -> (M.Map Node [Item], (M.Map Item Term, M.Map Node NodeName))
termsToItems nodes =
    let ((d', nodeMap), itemMap) =
            runUniqueKey' [Item i | i <- [0..]] $
            runUniqueKeyT' [Node i | i <- [0..]] $ do
                mapMKeys (mapM (lift . getUniqueKey)) getUniqueKey nodes
    in (d', (itemMap, nodeMap))

netData :: HyperParams -> M.Map Node [Item] -> Int -> NetData
netData hp nodeItems nTopics =
    NetData { dHypers           = hp
            , dItems            = S.unions $ map S.fromList $ M.elems nodeItems
            , dTopics           = S.fromList [Topic i | i <- [1..nTopics]]
            , dNodeItems        = M.fromList
                                  $ zip [NodeItem i | i <- [0..]]
                                  $ do (n,items) <- M.assocs nodeItems
                                       item <- items
                                       return (n, item)
            , dNodes            = M.keysSet nodeItems
            }

opts :: ParserInfo RunOpts
opts = info runOpts (  fullDesc
                    <> progDesc "Learn LDA model"
                    <> header "run-lda - learn LDA model"
                    )

instance Sampler.SamplerModel MState where
    estimateHypers = reestimate
    modelLikelihood = modelLikelihood
    summarizeHypers ms =
        "  phi  : "++show (dmAlpha $ snd $ M.findMin $ stPhis ms)++"\n"++
        "  theta: "++show (dmAlpha $ snd $ M.findMin $ stThetas ms)++"\n"

main :: IO ()
main = do
    args <- execParser opts
    stopWords <- case stopwords args of
                     Just f  -> S.fromList . T.words <$> TIO.readFile f
                     Nothing -> return S.empty
    printf "Read %d stopwords\n" (S.size stopWords)

    (nodeItems, (itemMap, nodeMap)) <- termsToItems
                            <$> readNodeItems stopWords (nodesFile args)

    Sampler.createSweeps $ samplerOpts args
    let sweepsDir = Sampler.sweepsDir $ samplerOpts args
    encodeFile (sweepsDir </> "item-map") itemMap
    encodeFile (sweepsDir </> "node-map") nodeMap

    let termCounts = V.fromListN (M.size nodeItems)
                     $ map length $ M.elems nodeItems :: Vector Int
    printf "Read %d nodes\n" (M.size nodeItems)
    printf "Mean items per node:  %1.2f\n" (mean $ V.map realToFrac termCounts)

    withSystemRandom $ \mwc->do
    let nd = netData (hyperParams args) nodeItems (nTopics args)
    encodeFile (sweepsDir </> "data") nd
    mInit <- runRVar (randomInitialize nd) mwc
    let m = model nd mInit
    Sampler.runSampler (samplerOpts args) m (updateUnits nd)
    return ()


-- |
-- Module    : Statistics.Matrix.Mutable
-- Copyright : (c) 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic mutable matrix operations.

module Statistics.Matrix.Mutable
    (
      MMatrix(..)
    , MVector
    , replicate
    , thaw
    , bounds
    , unsafeFreeze
    , unsafeRead
    , unsafeWrite
    , unsafeModify
    , immutably
    , unsafeBounds
    ) where

import Control.Applicative ((<$>))
import Control.DeepSeq (NFData(..))
import Control.Monad.ST (ST)
import Statistics.Matrix.Types (Matrix(..), MMatrix(..), MVector)
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M
import Prelude hiding (replicate)

replicate :: Int -> Int -> Double -> ST s (MMatrix s)
replicate r c k = MMatrix r c 0 <$> M.replicate (r*c) k

thaw :: Matrix -> ST s (MMatrix s)
thaw (Matrix r c e v) = MMatrix r c e <$> U.thaw v

unsafeFreeze :: MMatrix s -> ST s Matrix
unsafeFreeze (MMatrix r c e mv) = Matrix r c e <$> U.unsafeFreeze mv

unsafeRead :: MMatrix s -> Int -> Int -> ST s Double
unsafeRead mat r c = unsafeBounds mat r c M.unsafeRead
{-# INLINE unsafeRead #-}

unsafeWrite :: MMatrix s -> Int -> Int -> Double -> ST s ()
unsafeWrite mat row col k = unsafeBounds mat row col $ \v i ->
  M.unsafeWrite v i k
{-# INLINE unsafeWrite #-}

unsafeModify :: MMatrix s -> Int -> Int -> (Double -> Double) -> ST s ()
unsafeModify mat row col f = unsafeBounds mat row col $ \v i -> do
  k <- M.unsafeRead v i
  M.unsafeWrite v i (f k)
{-# INLINE unsafeModify #-}

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector.
bounds :: MMatrix s -> Int -> Int -> (MVector s -> Int -> r) -> r
bounds (MMatrix rs cs _ mv) r c k
  | r < 0 || r >= rs = error "row out of bounds"
  | c < 0 || c >= cs = error "column out of bounds"
  | otherwise        = k mv $! r * cs + c
{-# INLINE bounds #-}

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: MMatrix s -> Int -> Int -> (MVector s -> Int -> r) -> r
unsafeBounds (MMatrix _ cs _ mv) r c k = k mv $! r * cs + c
{-# INLINE unsafeBounds #-}

immutably :: NFData a => MMatrix s -> (Matrix -> a) -> ST s a
immutably mmat f = do
  k <- f <$> unsafeFreeze mmat
  rnf k `seq` return k
{-# INLINE immutably #-}