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

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

import Data.Complex

import Test.Hspec

import Numerical.HBLAS.MatrixTypes as Matrix
import Numerical.HBLAS.BLAS.Level2 as BLAS
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Storable.Mutable as SM

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  gbmvSpec
  gemvSpec
  gerSpec
  gercSpec
  geruSpec
  hbmvSpec
  hemvSpec
  herSpec
  her2Spec
  hpmvSpec
  hprSpec
  hpr2Spec
  sbmvSpec
  spmvSpec
  sprSpec
  spr2Spec
  symvSpec
  syrSpec
  syr2Spec
  tbmvSpec
  tbsvSpec
  tpmvSpec
  tpsvSpec
  trmvSpec
  trsvSpec
  gemvBugJune2017

gemvBugJune2017 :: Spec
gemvBugJune2017  =
  --context "?gemv bug"

 {-

 -}
    describe "dgemv abstraction" $ do
      it "3*2 matrix vector product sadness" $
             let m = S.fromList [1,2,3, 4,5,6]
                 x = S.fromList [10,20,30]
             in
             do

              m' <- S.unsafeThaw m
              x' <- S.unsafeThaw x
              y' <- SM.new 2

              let mat   = MutableDenseMatrix SRow 3 2 3 m'
                  xvec  = MutableDenseVector SDirect 3 1 x'
                  yvec  = MutableDenseVector SDirect 2 1 y'

              dgemv NoTranspose 1 0 mat xvec yvec
              resList <-mutableVectorToList $ _bufferMutDenseVector yvec
              resList `shouldBe` [140, 320]

gbmvSpec :: Spec
gbmvSpec =
  context "?GBMV" $ do
    describe "SGBMV" $ do
      it "3x5 a(5x5 matrix) all 1's" $ do
        matvecTest1SGBMV
    describe "DBGMV" $ do
      it "3x5 a(10x5 matrix) all 1's with beta 1.0" $ do
        matvecTest1DGBMV
    describe "CGBMV" $ do
      it "2x4 a(4x4 matrix) all 1+i's" $ do
        matvecTest1CGBMV
     -- it "gbmv on 2x4 a(4x4 matrix) all 1+i s with conjnotranspose" $ do
     --   matvecTest2CGBMV -- conjnotranspose is invalid
    describe "ZGBMV" $ do
      it "2x10 a(5x10 matrix) all 1+i s with transpose and alpha 1.0" $ do
        matvecTest1ZGBMV
      it "2x10 a(5x10 matrix) all 1+i s with conjtranspose and alpha 1.0" $ do
        matvecTest2ZGBMV

matvecTest1SGBMV :: IO ()
matvecTest1SGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3, 5) (\_ -> (1.0))
    x <- Matrix.generateMutableDenseVector 5 (\_ -> (1.0))
    res <- Matrix.generateMutableDenseVector 5 (\_ -> (1.0))
    BLAS.sgbmv Matrix.NoTranspose 5 5 1 1 1.0 a x 0.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2, 3, 3, 3, 2]

matvecTest1DGBMV :: IO ()
matvecTest1DGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3, 5) (\_ -> (1.0))
    x <- Matrix.generateMutableDenseVector 5 (\_ -> (1.0))
    res <- Matrix.generateMutableDenseVector 10 (\_ -> (1.0))
    BLAS.dgbmv Matrix.NoTranspose 10 5 1 1 1.0 a x 1.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [3, 4, 4, 4, 3, 1, 1, 1, 1, 1]
    -- not [3, 4, 4, 4, 3, 2, 1, 1, 1, 1]

matvecTest1CGBMV :: IO ()
matvecTest1CGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2, 4) (\(x, y) -> [1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 0] !! (x * 4 + y))
    x <- Matrix.generateMutableDenseVector 4 (\_ -> (1.0:+1.0))
    res <- Matrix.generateMutableDenseVector 4 (\_ -> (1.0:+1.0))
    BLAS.cgbmv Matrix.NoTranspose 4 4 0 1 1.0 a x 0.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [0:+4.0, 0:+4.0, 0:+4.0, 0:+2.0]

matvecTest2CGBMV :: IO ()
matvecTest2CGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2, 4) (\(x, y) -> [1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 0] !! (x * 4 + y))
    x <- Matrix.generateMutableDenseVector 4 (\_ -> (1.0:+1.0))
    res <- Matrix.generateMutableDenseVector 4 (\_ -> (1.0:+1.0))
    BLAS.cgbmv Matrix.ConjNoTranspose 4 4 0 1 1.0 a x 0.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [4.0:+0, 4.0:+0, 4.0:+0, 2.0:+0]

matvecTest1ZGBMV :: IO ()
matvecTest1ZGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2, 10) (\_ -> (1.0:+1.0))
    x <- Matrix.generateMutableDenseVector 10 (\_ -> (1.0:+1.0))
    res <- Matrix.generateMutableDenseVector 5 (\_ -> (1.0:+1.0))
    BLAS.cgbmv Matrix.Transpose 10 5 1 0 2.0 a x 0.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [0:+8.0, 0:+8.0, 0:+8.0, 0:+8.0, 0:+8.0]
    -- not [0:+8.0, 0:+8.0, 0:+8.0, 0:+8.0, 0:+4.0]

matvecTest2ZGBMV :: IO ()
matvecTest2ZGBMV = do
    a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2, 10) (\_ -> (1.0:+1.0))
    x <- Matrix.generateMutableDenseVector 10 (\_ -> (1.0:+1.0))
    res <- Matrix.generateMutableDenseVector 5 (\_ -> (1.0:+1.0))
    BLAS.cgbmv Matrix.ConjTranspose 10 5 1 0 2.0 a x 0.0 res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [8.0:+0, 8.0:+0, 8.0:+0, 8.0:+0, 8.0:+0]
    -- not [0:+8.0, 0:+8.0, 0:+8.0, 0:+8.0, 0:+4.0]

-- notes for gbmv
-- column > rows: the tail of the supper diagonals is considered.
-- rows > column: the tail of the sub diagonals is not considered.

gemvSpec :: Spec
gemvSpec =
  context "?GEMV" $ do
    describe "SGEMV" $ do
      it "2x2 all 1's" $ do
        matmatTest1SGEMV
    describe "DGEMV" $ do
      it "2x2 all 1's" $ do
        matmatTest1DGEMV
    describe "CGEMV" $ do
      it "2x2 all 1's" $ do
        matmatTest1CGEMV
    describe "ZGEMV" $ do
      it "2x2 all 1's" $ do
        matmatTest1ZGEMV

matmatTest1SGEMV :: IO ()
matmatTest1SGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0::Float))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 :: Float))
    res  <- Matrix.generateMutableDenseVector  2 (\_ -> (0.0 :: Float))
    BLAS.sgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2,2]

matmatTest1DGEMV :: IO ()
matmatTest1DGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector 2  (\_ -> (0.0 ))
    BLAS.dgemv Matrix.NoTranspose 1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2.0,2.0]

matmatTest1CGEMV :: IO ()
matmatTest1CGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector  2 (\_ -> (0.0 ))
    BLAS.cgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2.0,2.0]

matmatTest1ZGEMV :: IO ()
matmatTest1ZGEMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2) (\_ -> (1.0))
    right <- Matrix.generateMutableDenseVector 2 (\_ -> (1.0 ))
    res  <- Matrix.generateMutableDenseVector 2 (\_ -> (0.0 ))
    BLAS.zgemv Matrix.NoTranspose  1.0 1.0 left right res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2.0,2.0]

----
----
gerSpec :: Spec
gerSpec =
  context "?GER" $ do
    describe "SGER" $ do
      it "2x2 all 1's" $ do
        matmatTest1SGER
    describe "DGER" $ do
      it "2x2 all 1's" $ do
        matmatTest1DGER

matmatTest1SGER :: IO ()
matmatTest1SGER = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2,2) (\_ -> 1.0)
  x <- Matrix.generateMutableDenseVector 2 (\_ -> 2.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0)
  BLAS.sger 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [13.0,13.0,13.0,13.0]

matmatTest1DGER :: IO ()
matmatTest1DGER = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (2,2) (\_ -> 1.0)
  x <- Matrix.generateMutableDenseVector 2 (\_ -> 2.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0)
  BLAS.sger 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [13.0,13.0,13.0,13.0]

gercSpec :: Spec
gercSpec =
  context "?GERC" $ do
    describe "CGERC" $ do
      it "2x3 all 1+i's" $ do
        matmatTest1CGERC
    describe "ZGERC" $ do
      it "2x3 all 1+i's" $ do
        matmatTest1ZGERC

matmatTest1CGERC :: IO ()
matmatTest1CGERC = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3,2) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 2.0:+3.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0:+2.0)
  BLAS.cgerc 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [25.0:+11.0, 25.0:+11.0, 1.0:+1.0, 25.0:+11.0, 25.0:+11.0, 1.0:+1.0]
  -- why the following is not correct...
  -- resList `shouldBe` [25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0]

matmatTest1ZGERC :: IO ()
matmatTest1ZGERC = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3,2) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 2.0:+3.0)
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0:+2.0)
  BLAS.zgerc 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [25.0:+11.0, 25.0:+11.0, 1.0:+1.0, 25.0:+11.0, 25.0:+11.0, 1.0:+1.0]
  -- why the following is not correct...
  -- resList `shouldBe` [25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0, 25.0:+11.0]

geruSpec :: Spec
geruSpec =
  context "?GERU" $ do
    describe "CGERU" $ do
      it "2x3 all 1+i's" $ do
        matmatTest1CGERU
    describe "ZGERU" $ do
      it "2x3 all 1+i's" $ do
        matmatTest1ZGERU

matmatTest1CGERU :: IO ()
matmatTest1CGERU = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3,2) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 2.0:+(-3.0))
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0:+(-2.0))
  BLAS.cgeru 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [1.0:+(-25.0), 1.0:+(-25.0), 1.0:+1.0, 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+1.0]
  -- why the following is not correct...
  -- resList `shouldBe` [1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0)]

matmatTest1ZGERU :: IO ()
matmatTest1ZGERU = do
  res <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3,2) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 2.0:+(-3.0))
  y <- Matrix.generateMutableDenseVector 2 (\_ -> 3.0:+(-2.0))
  BLAS.zgeru 2.0 x y res
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat res
  resList `shouldBe` [1.0:+(-25.0), 1.0:+(-25.0), 1.0:+1.0, 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+1.0]
  -- why the following is not correct...
  -- resList `shouldBe` [1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0), 1.0:+(-25.0)]

-- [1:+0    1:+1    1:+1    0:+0]
-- [1:+(-1) 1:+0    1:+1    1:+1]
-- [1:+(-1) 1:+(-1) 1:+0    1:+1]
-- [0:+0    1:+(-1) 1:+(-1) 1:+0]
--
-- [1:+1]
-- [1:+1]
-- [1:+1]
-- [1:+1]

hbmvSpec :: Spec
hbmvSpec =
  context "?HBMV" $ do
    describe "CHBMV" $ do
      it "4*3 a(4x4 matrix) upper all 1+i's" $ do
        matvecTest1CHBMV
      it "4*3 a(4x4 matrix) upper all 1+i's (column oriented)" $ do
        matvecTest2CHBMV
    describe "ZHBMV" $ do
      it "4*3 a(4x4 matrix) lower all 1+i's" $ do
        matvecTest1ZHBMV

matvecTest1CHBMV :: IO ()
matvecTest1CHBMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 3) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.chbmv Matrix.MatUpper 2 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+5.0, 3.0:+5.0, 5.0:+3.0, 5.0:+1.0]

matvecTest2CHBMV :: IO ()
matvecTest2CHBMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (4, 3) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.chbmv Matrix.MatUpper 2 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+5.0, 3.0:+5.0, 5.0:+3.0, 5.0:+1.0]

-- [1:+0    1:+1    1:+1    0:+0]
-- [1:+(-1) 1:+0    1:+1    1:+1]
-- [1:+(-1) 1:+(-1) 1:+0    1:+1]
-- [0:+0    1:+(-1) 1:+(-1) 1:+0]
--
-- [1:+1]
-- [1:+1]
-- [1:+1]
-- [1:+1]
matvecTest1ZHBMV :: IO ()
matvecTest1ZHBMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 3) (\_ -> 1.0:+(-1.0))
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.zhbmv Matrix.MatLower 2 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+5.0, 3.0:+5.0, 5.0:+3.0, 5.0:+1.0]

-- [1:+0    1:+1    1:+1    1:+1]
-- [1:+(-1) 1:+0    1:+1    1:+1]
-- [1:+(-1) 1:+(-1) 1:+0    1:+1]
-- [1:+(-1) 1:+(-1) 1:+(-1) 1:+0]
--
-- [1:+1]
-- [1:+1]
-- [1:+1]
-- [1:+1]
--
hemvSpec :: Spec
hemvSpec =
  context "?HEMV" $ do
    describe "CHEMV" $ do
      it "4*3 a(4x4 matrix) upper all 1+i's" $ do
        matvecTest1CHEMV
      it "4*3 a(4x4 matrix) upper all 1+i s (column oriented)" $ do
        matvecTest2CHEMV
    describe "ZHEMV" $ do
      it "4*3 a(4x4 matrix) lower all 1+i's" $ do
        matvecTest1ZHEMV

matvecTest1CHEMV :: IO ()
matvecTest1CHEMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.chemv Matrix.MatUpper 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+7.0, 3.0:+5.0, 5.0:+3.0, 7.0:+1.0]

matvecTest2CHEMV :: IO ()
matvecTest2CHEMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.chemv Matrix.MatUpper 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+7.0, 3.0:+5.0, 5.0:+3.0, 7.0:+1.0]

matvecTest1ZHEMV :: IO ()
matvecTest1ZHEMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+(-1.0))
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1.0:+1.0)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 0.0:+0.0)
  BLAS.zhemv Matrix.MatLower 1.0 a x 0.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [1.0:+7.0, 3.0:+5.0, 5.0:+3.0, 7.0:+1.0]

herSpec :: Spec
herSpec =
  context "?HER" $ do
    describe "CHER" $ do
      it "4*4 a upper all 1+i's" $ do
        matvecTest1CHER
    describe "ZHER" $ do
      it "4*4 a upper all 1+i's" $ do
        matvecTest1ZHER

matvecTest1CHER :: IO ()
matvecTest1CHER = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.cher Matrix.MatUpper 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [3.0:+0.0, 5.0:+1.0,  7.0:+1.0,  9.0:+1.0,
               1.0:+1.0, 9.0:+0.0, 13.0:+1.0, 17.0:+1.0,
               1.0:+1.0, 1.0:+1.0, 19.0:+0.0, 25.0:+1.0,
               1.0:+1.0, 1.0:+1.0,  1.0:+1.0, 33.0:+0.0]

matvecTest1ZHER :: IO ()
matvecTest1ZHER = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.zher Matrix.MatLower 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [3.0:+0.0,  1.0:+1.0,  1.0:+1.0,  1.0:+1.0,
               5.0:+1.0,  9.0:+0.0,  1.0:+1.0,  1.0:+1.0,
               7.0:+1.0, 13.0:+1.0, 19.0:+0.0,  1.0:+1.0,
               9.0:+1.0, 17.0:+1.0, 25.0:+1.0, 33.0:+0.0]

her2Spec :: Spec
her2Spec =
  context "?HER2" $ do
    describe "CHER2" $ do
      it "4*4 a upper all 1+i's" $ do
        matvecTest1CHER2
    describe "ZHER2" $ do
      it "4*4 a upper all 1+i's" $ do
        matvecTest1ZHER2

matvecTest1CHER2 :: IO ()
matvecTest1CHER2 = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.cher2 Matrix.MatUpper 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [5.0:+0.0,  9.0:+1.0, 13.0:+1.0, 17.0:+1.0,
               1.0:+1.0, 17.0:+0.0, 25.0:+1.0, 33.0:+1.0,
               1.0:+1.0,  1.0:+1.0, 37.0:+0.0, 49.0:+1.0,
               1.0:+1.0,  1.0:+1.0,  1.0:+1.0, 65.0:+0.0]

matvecTest1ZHER2 :: IO ()
matvecTest1ZHER2 = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 4) (\_ -> 1.0:+1.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.zher2 Matrix.MatLower 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [ 5.0:+0.0,  1.0:+1.0,  1.0:+1.0,  1.0:+1.0,
                9.0:+1.0, 17.0:+0.0,  1.0:+1.0,  1.0:+1.0,
               13.0:+1.0, 25.0:+1.0, 37.0:+0.0,  1.0:+1.0,
               17.0:+1.0, 33.0:+1.0, 49.0:+1.0, 65.0:+0.0]

hpmvSpec :: Spec
hpmvSpec =
  context "?HPMV" $ do
    describe "CHPMV" $ do
      it "4*4 a upper (row oriented)" $ do
        matvecTest1CHPMV
    describe "ZHPMV" $ do
      it "4*4 a upper (column oriented)" $ do
        matvecTest1ZHPMV

-- [0:+0    1:+1    2:+2    3:+3]
-- [1:+(-1) 4:+0    5:+5    6:+6]
-- [2:+(-2) 5:+(-5) 7:+0    8:+8]
-- [3:+(-3) 6:+(-6) 8:+(-8) 9:+0]
--
-- [2:+2]
-- [2:+2]
-- [2:+2]
-- [2:+2]
--
-- [ 3:+27]
-- [15:+55]
-- [45:+49]
-- [89:+21]
matvecTest1CHPMV :: IO ()
matvecTest1CHPMV = do
  a <- Matrix.generateMutableDenseVector 10 (\idx -> [0.0:+0.0, 1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0, 5.0:+5.0, 6.0:+6.0, 7.0:+7.0, 8.0:+8.0, 9.0:+9.0] !! idx)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [2.0:+2.0, 2.0:+2.0, 2.0:+2.0, 2.0:+2.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [3.0:+3.0, 3.0:+3.0, 3.0:+3.0, 3.0:+3.0] !! idx)
  BLAS.chpmv Matrix.SRow Matrix.MatUpper 4 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [ 3.0:+27.0, 15.0:+55.0, 45.0:+49.0, 89.0:+21.0]

-- [0:+0    1:+1    3:+3    6:+6]
-- [1:+(-1) 2:+0    4:+4    7:+7]
-- [3:+(-3) 4:+(-4) 5:+0    8:+8]
-- [6:+(-6) 7:+(-7) 8:+(-8) 9:+0]
--
-- [2:+2]
-- [2:+2]
-- [2:+2]
-- [2:+2]
--
-- [  6:+46]
-- [ 14:+54]
-- [ 44:+48]
-- [108:+24]
matvecTest1ZHPMV :: IO ()
matvecTest1ZHPMV = do
  a <- Matrix.generateMutableDenseVector 10 (\idx -> [0.0:+0.0, 1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0, 5.0:+5.0, 6.0:+6.0, 7.0:+7.0, 8.0:+8.0, 9.0:+9.0] !! idx)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [2.0:+2.0, 2.0:+2.0, 2.0:+2.0, 2.0:+2.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [3.0:+3.0, 3.0:+3.0, 3.0:+3.0, 3.0:+3.0] !! idx)
  BLAS.zhpmv Matrix.SColumn Matrix.MatUpper 4 1.0 a x 2.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [ 6.0:+46.0, 14.0:+54.0, 44.0:+48.0, 108.0:+24.0]

-- [1:+1]
-- [2:+2]
-- [3:+3]
-- [4:+4]
--
-- [2:+0    4:+0    6:+0    8:+0 ]
-- [4:+0    8:+0    12:+0   16:+0]
-- [6:+0    12:+0   18:+0   24:+0]
-- [8:+0    16:+0   24:+0   32:+0]
--
-- [0:+0    1:+1    3:+3    6:+6]
-- [1:+(-1) 2:+0    4:+4    7:+7]
-- [3:+(-3) 4:+(-4) 5:+0    8:+8]
-- [6:+(-6) 7:+(-7) 8:+(-8) 9:+0]
--
-- [2:+0     5:+1     9:+3     14:+6]
-- [5:+(-1)  10:+0    16:+4    23:+7]
-- [9:+(-3)  16:+(-4) 23:+0    32:+8]
-- [14:+(-6) 23:+(-7) 32:+(-8) 41:+0]

hprSpec :: Spec
hprSpec =
  context "?HPR" $ do
    describe "CHPR" $ do
      it "4*4 a upper (column oriented)" $ do
        matvecTest1CHPR
    describe "ZHPR" $ do
      it "4*4 a lower (row oriented)" $ do
        matvecTest1ZHPR

matvecTest1CHPR :: IO ()
matvecTest1CHPR = do
  a <- Matrix.generateMutableDenseVector 10 (\idx -> [0.0:+0.0, 1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0, 5.0:+5.0, 6.0:+6.0, 7.0:+7.0, 8.0:+8.0, 9.0:+9.0] !! idx)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.chpr Matrix.SColumn Matrix.MatUpper 4 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [2.0:+0.0, 5.0:+1.0, 10.0:+0.0, 9.0:+3.0, 16.0:+4.0, 23.0:+0.0, 14.0:+6.0, 23.0:+7.0,  32.0:+8.0, 41.0:+0.0]

-- [1:+1]
-- [2:+2]
-- [3:+3]
-- [4:+4]
--
-- [2:+0    4:+0    6:+0    8:+0 ]
-- [4:+0    8:+0    12:+0   16:+0]
-- [6:+0    12:+0   18:+0   24:+0]
-- [8:+0    16:+0   24:+0   32:+0]
--
-- [0:+0    1:+1    3:+3    6:+6]
-- [1:+(-1) 2:+0    4:+4    7:+7]
-- [3:+(-3) 4:+(-4) 5:+0    8:+8]
-- [6:+(-6) 7:+(-7) 8:+(-8) 9:+0]
--
-- [2:+0     5:+1     9:+3     14:+6]
-- [5:+(-1)  10:+0    16:+4    23:+7]
-- [9:+(-3)  16:+(-4) 23:+0    32:+8]
-- [14:+(-6) 23:+(-7) 32:+(-8) 41:+0]
matvecTest1ZHPR :: IO ()
matvecTest1ZHPR = do
  a <- Matrix.generateMutableDenseVector 10 (\idx -> [0.0:+(-0.0), 1.0:+(-1.0), 2.0:+(-2.0), 3.0:+(-3.0), 4.0:+(-4.0), 5.0:+(-5.0), 6.0:+(-6.0), 7.0:+(-7.0), 8.0:+(-8.0), 9.0:+(-9.0)] !! idx)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 2.0:+2.0, 3.0:+3.0, 4.0:+4.0] !! idx)
  BLAS.zhpr Matrix.SRow Matrix.MatLower 4 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [2.0:+0.0, 5.0:+(-1.0), 10.0:+0.0, 9.0:+(-3.0), 16.0:+(-4.0), 23.0:+0.0, 14.0:+(-6.0), 23.0:+(-7.0), 32.0:+(-8.0), 41.0:+0.0]

hpr2Spec :: Spec
hpr2Spec =
  context "?HPR2" $ do
    describe "CHRP2" $ do
      it "4*4 a upper (column oriented)" $ do
        matvecTest1CHPR2
    describe "ZHRP2" $ do
      it "4*4 a upper (row oriented)" $ do
        matvecTest1ZHPR2

-- [12:+0 ...]
-- [.       ...]
-- [.       ...]
-- [.       ...]
--
matvecTest1CHPR2 :: IO ()
matvecTest1CHPR2 = do
  a <- Matrix.generateMutableDenseVector 10 (\_ -> 0.0:+0.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+2.0, 1.0:+2.0, 1.0:+2.0, 1.0:+2.0] !! idx)
  BLAS.chpr2 Matrix.SColumn Matrix.MatUpper 4 2.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0]

matvecTest1ZHPR2 :: IO ()
matvecTest1ZHPR2 = do
  a <- Matrix.generateMutableDenseVector 10 (\_ -> 0.0:+0.0)
  x <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+1.0, 1.0:+1.0, 1.0:+1.0, 1.0:+1.0] !! idx)
  y <- Matrix.generateMutableDenseVector 4 (\idx -> [1.0:+2.0, 1.0:+2.0, 1.0:+2.0, 1.0:+2.0] !! idx)
  BLAS.zhpr2 Matrix.SColumn Matrix.MatUpper 4 2.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0, 12.0:+0.0]

-- [1 2 0 0]
-- [2 3 4 0]
-- [0 4 5 6]
-- [0 0 6 7]
--
-- [1 2]
-- [3 4]
-- [5 6]
-- [7 0]
--
-- [5]
-- [11]
-- [17]
-- [15]

sbmvSpec :: Spec
sbmvSpec =
  context "?SBMV" $ do
    describe "SSBMV" $ do
      it "4*4 a upper (row oriented)" $ do
        matvecTest1SSBMV
    describe "DSBMV" $ do
      it "4*4 a lower (column oriented)" $ do
        matvecTest1DSBMV

matvecTest1SSBMV :: IO ()
matvecTest1SSBMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (4, 2) (\(x, y) -> [1, 2, 3, 4, 5, 6, 7, 0] !! (y * 4 + x))
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 2)
  BLAS.ssbmv Matrix.MatUpper 1 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [5, 11, 17, 15]

-- [1 3 5 7]
-- [2 4 6 0]

matvecTest1DSBMV :: IO ()
matvecTest1DSBMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (4, 2) (\(x, y) -> [1, 3, 5, 7, 2, 4, 6, 0] !! (y * 4 + x))
  x <- Matrix.generateMutableDenseVector 4 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 4 (\_ -> 2)
  BLAS.dsbmv Matrix.MatLower 1 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [5, 11, 17, 15]

spmvSpec :: Spec
spmvSpec =
  context "?SPMV" $ do
    describe "SSPMV" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSPMV
    describe "DSPMV" $ do
      it "3*3 a lower (column oriented)" $ do
        matvecTest1DSPMV

matvecTest1SSPMV :: IO ()
matvecTest1SSPMV = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\_ -> 2)
  BLAS.sspmv Matrix.SRow Matrix.MatUpper 3 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [8, 13, 16]

matvecTest1DSPMV :: IO ()
matvecTest1DSPMV = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\_ -> 2)
  BLAS.dspmv Matrix.SColumn Matrix.MatLower 3 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [8, 13, 16]

sprSpec :: Spec
sprSpec =
  context "?SPR" $ do
    describe "SSPR" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSPR
    describe "DSPR" $ do
      it "3*3 a upper (column oriented)" $ do
        matvecTest1DSPR

matvecTest1SSPR :: IO ()
matvecTest1SSPR = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  BLAS.sspr Matrix.SRow Matrix.MatUpper 3 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [2, 3, 4, 5, 6, 7]

matvecTest1DSPR :: IO ()
matvecTest1DSPR = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  BLAS.dspr Matrix.SRow Matrix.MatUpper 3 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [2, 3, 4, 5, 6, 7]

spr2Spec :: Spec
spr2Spec =
  context "?SPR2" $ do
    describe "SSPR2" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSPR2
    describe "DSPR2" $ do
      it "3*3 a upper (column oriented)" $ do
        matvecTest1DSPR2

matvecTest1SSPR2 :: IO ()
matvecTest1SSPR2 = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 2, 3] !! idx)
  BLAS.sspr2 Matrix.SRow Matrix.MatUpper 3 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [3, 5, 7, 8, 10, 12]

matvecTest1DSPR2 :: IO ()
matvecTest1DSPR2 = do
  a <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 3, 4, 5, 6] !! idx)
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 2, 3] !! idx)
  BLAS.dspr2 Matrix.SRow Matrix.MatUpper 3 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector a
  resList `shouldBe` [3, 5, 7, 8, 10, 12]

symvSpec :: Spec
symvSpec =
  context "?SYMV" $ do
    describe "SSYMV" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSYMV
    describe "DSYMV" $ do
      it "3*3 a lower (column oriented)" $ do
        matvecTest1DSYMV

-- [1 2 3]
-- [0 5 6]
-- [0 0 7]
matvecTest1SSYMV :: IO ()
matvecTest1SSYMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3, 3) (\(x, y) -> [1, 2, 3, 0, 5, 6, 0, 0, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\_ -> 2)
  BLAS.ssymv Matrix.MatUpper 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [8, 15, 18]

matvecTest1DSYMV :: IO ()
matvecTest1DSYMV = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (3, 3) (\(x, y) -> [1, 0, 0, 2, 5, 0, 3, 6, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\_ -> 2)
  BLAS.dsymv Matrix.MatLower 1.0 a x 1.0 y
  resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector y
  resList `shouldBe` [8, 15, 18]

syrSpec :: Spec
syrSpec =
  context "?SYR" $ do
    describe "SSYR" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSYR
    describe "DSYR" $ do
      it "3*3 a upper (column oriented)"$ do
        matvecTest1DSYR
-- [1 2 3]
-- [0 5 6]
-- [0 0 7]
matvecTest1SSYR :: IO ()
matvecTest1SSYR = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3, 3) (\(x, y) -> [1, 2, 3, 0, 5, 6, 0, 0, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  BLAS.ssyr Matrix.MatUpper 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [2, 3, 4, 0, 6, 7, 0, 0, 8]

matvecTest1DSYR :: IO ()
matvecTest1DSYR = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (3, 3) (\(x, y) -> [1, 0, 0, 2, 5, 0, 3, 6, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  BLAS.dsyr Matrix.MatLower 1.0 x a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [2, 3, 4, 0, 6, 7, 0, 0, 8]

syr2Spec :: Spec
syr2Spec =
  context "?SYR2" $ do
    describe "SSYR2" $ do
      it "3*3 a upper (row oriented)" $ do
        matvecTest1SSYR2
    describe "DSYR2" $ do
      it "3*3 a upper (column oriented)" $ do
        matvecTest1DSYR2

matvecTest1SSYR2 :: IO ()
matvecTest1SSYR2 = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SRow) (3, 3) (\(x, y) -> [1, 2, 3, 0, 5, 6, 0, 0, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 2, 3] !! idx)
  BLAS.ssyr2 Matrix.MatUpper 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [3, 5, 7, 0, 9, 11, 0, 0, 13]

matvecTest1DSYR2 :: IO ()
matvecTest1DSYR2 = do
  a <- Matrix.generateMutableDenseMatrix (Matrix.SColumn) (3, 3) (\(x, y) -> [1, 0, 0, 2, 5, 0, 3, 6, 7] !! (y * 3 + x))
  x <- Matrix.generateMutableDenseVector 3 (\_ -> 1)
  y <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 2, 3] !! idx)
  BLAS.dsyr2 Matrix.MatLower 1.0 x y a
  resList <- Matrix.mutableVectorToList $ _bufferDenMutMat a
  resList `shouldBe` [3, 5, 7, 0, 9, 11, 0, 0, 13]

tbmvSpec :: Spec
tbmvSpec =
  context "?TBMV" $ do
    describe "STBMV" $ do
      it "3x3 upper no trans (row oriented)" $ do
        matmatTest1STBMV
    describe "DTBMV" $ do
      it "3x3 lower trans (column oriented)" $ do
        matmatTest1DTBMV
    describe "CTBMV" $ do
      it "3x3 upper conj trans (row oriented)" $ do
        matmatTest1CTBMV
    describe "ZTBMV" $ do
      it "3x3 lower no trans (row oriented)" $ do
        matmatTest1ZTBMV

-- [ 1 2 0]   [-2]   [4]
-- [ 0 1 1] * [ 3] = [5]
-- [ 0 0 1]   [ 2]   [2]
matmatTest1STBMV:: IO ()
matmatTest1STBMV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (3, 2) (\(x, y) -> [0, 2, 1, 1, 1, 1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.stbmv MatUpper NoTranspose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [4, 5, 2]

-- [1 1 0]   [-2]   [1]
-- [0 1 1] * [ 3] = [5]
-- [0 0 1]   [ 2]   [2]
matmatTest1DTBMV:: IO ()
matmatTest1DTBMV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SColumn)  (3, 2) (\(x, y) -> [0, 1, 1, 1, 1, 1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.dtbmv MatLower Transpose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [1, 5, 2]

-- [ 1:+0    0:+0    0:+0]   [-2]   [-2:+0   ]
-- [ 2:+(-2) 1:+0    0:+0] * [ 3] = [-1:+4   ]
-- [ 0:+0    1:+(-1) 1:+0]   [ 2]   [ 5:+(-3)]
matmatTest1CTBMV:: IO ()
matmatTest1CTBMV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (3, 2) (\(x, y) -> [0:+0, 2:+2, 1:+1, 1:+1, 1:+1, 1:+1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.ctbmv MatUpper Matrix.ConjTranspose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [(-2):+0, (-1):+4, 5:+(-3)]

-- [1 1 0]   [2]   [5]
-- [0 1 1] * [3] = [5]
-- [0 0 1]   [2]   [2]
matmatTest1ZTBMV:: IO ()
matmatTest1ZTBMV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SColumn)  (3, 2) (\(x, y) -> [0:+0, 1:+0, 1:+0, 1:+0, 1:+0, 1:+0] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [2, 3, 2] !! idx)
    BLAS.ztbmv MatUpper NoTranspose MatNonUnit 1 a x -- TODO: NAN error when using Lower Transpose
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [5, 5, 2]

tbsvSpec :: Spec
tbsvSpec =
  context "?TBSV" $ do
    describe "STBSV" $ do
      it "3x3 upper no trans (row oriented)" $ do
        matmatTest1STBSV
    describe "DTBSV" $ do
      it "3x3 lower trans (column oriented)" $ do
        matmatTest1DTBSV
    describe "CTBSV" $ do
      it "3x3 upper conj trans (row oriented)" $ do
        matmatTest1CTBSV
    describe "ZTBSV" $ do
      it "3x3 lower no trans (row oriented)" $ do
        matmatTest1ZTBSV

-- [ 1 2 0]   [-2]   [4]
-- [ 0 1 1] * [ 3] = [5]
-- [ 0 0 1]   [ 2]   [2]
matmatTest1STBSV:: IO ()
matmatTest1STBSV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (3, 2) (\(x, y) -> [0, 2, 1, 1, 1, 1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [4, 5, 2] !! idx)
    BLAS.stbsv MatUpper NoTranspose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [1 1 0]   [-2]   [1]
-- [0 1 1] * [ 3] = [5]
-- [0 0 1]   [ 2]   [2]
matmatTest1DTBSV:: IO ()
matmatTest1DTBSV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SColumn)  (3, 2) (\(x, y) -> [0, 1, 1, 1, 1, 1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 5, 2] !! idx)
    BLAS.dtbsv MatLower Transpose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [ 1:+0    0:+0    0:+0]   [-2]   [-2:+0   ]
-- [ 2:+(-2) 1:+0    0:+0] * [ 3] = [-1:+4   ]
-- [ 0:+0    1:+(-1) 1:+0]   [ 2]   [ 5:+(-3)]
matmatTest1CTBSV:: IO ()
matmatTest1CTBSV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (3, 2) (\(x, y) -> [0:+0, 2:+2, 1:+1, 1:+1, 1:+1, 1:+1] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [(-2):+0, (-1):+4, 5:+(-3)] !! idx)
    BLAS.ctbsv MatUpper Matrix.ConjTranspose MatUnit 1 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [1 1 0]   [2]   [5]
-- [0 1 1] * [3] = [5]
-- [0 0 1]   [2]   [2]
matmatTest1ZTBSV:: IO ()
matmatTest1ZTBSV = do
    a  <- Matrix.generateMutableDenseMatrix (Matrix.SColumn)  (3, 2) (\(x, y) -> [0:+0, 1:+0, 1:+0, 1:+0, 1:+0, 1:+0] !! (y * 3 + x))
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [5, 5, 2] !! idx)
    BLAS.ztbsv MatUpper NoTranspose MatNonUnit 1 a x -- TODO: NAN error when using Lower Transpose
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [2, 3, 2]

tpmvSpec :: Spec
tpmvSpec =
  context "?TPMV" $ do
    describe "STPMV" $ do
      it "3x3 upper no trans (row oriented)" $ do
        matmatTest1STPMV
    describe "DTPMV" $ do
      it "3x3 lower trans (column oriented)" $ do
        matmatTest1DTPMV
    describe "CTPMV" $ do
      it "3x3 upper conj trans (row oriented)" $ do
        matmatTest1CTPMV
    describe "ZTPMV" $ do
      it "3x3 lower no trans (row oriented)" $ do
        matmatTest1ZTPMV

-- [ 1 2 0]   [-2]   [4]
-- [ 0 1 1] * [ 3] = [5]
-- [ 0 0 1]   [ 2]   [2]
matmatTest1STPMV:: IO ()
matmatTest1STPMV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 0, 1, 1, 1] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.stpmv Matrix.SRow MatUpper NoTranspose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [4, 5, 2]

-- [1 1 0]   [-2]   [1]
-- [0 1 1] * [ 3] = [5]
-- [0 0 1]   [ 2]   [2]
matmatTest1DTPMV:: IO ()
matmatTest1DTPMV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 1, 0, 1, 1, 1] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.dtpmv Matrix.SColumn MatLower Transpose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [1, 5, 2]

-- [ 1:+0    0:+0    0:+0]   [-2]   [-2:+0   ]
-- [ 2:+(-2) 1:+0    0:+0] * [ 3] = [-1:+4   ]
-- [ 0:+0    1:+(-1) 1:+0]   [ 2]   [ 5:+(-3)]
matmatTest1CTPMV:: IO ()
matmatTest1CTPMV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1:+0, 2:+2, 0:+0, 1:+0, 1:+1, 1:+0] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [-2, 3, 2] !! idx)
    BLAS.ctpmv Matrix.SRow MatUpper Matrix.ConjTranspose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [(-2):+0, (-1):+4, 5:+(-3)]

-- [1 1 0]   [2]   [5]
-- [0 1 1] * [3] = [5]
-- [0 0 1]   [2]   [2]
matmatTest1ZTPMV:: IO ()
matmatTest1ZTPMV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1:+0, 1:+0, 1:+0, 0:+0, 1:+0, 1:+0] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [2, 3, 2] !! idx)
    BLAS.ztpmv Matrix.SColumn MatUpper NoTranspose MatNonUnit 3 a x -- TODO: NAN error when using Lower Transpose
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [5, 5, 2]

tpsvSpec :: Spec
tpsvSpec =
  context "?TPSV" $ do
    describe "STPSV" $ do
      it "3x3 upper no trans (row oriented)" $ do
        matmatTest1STPSV
    describe "DTPSV" $ do
      it "3x3 lower trans (column oriented)" $ do
        matmatTest1DTPSV
    describe "CTPSV" $ do
      it "3x3 upper conj trans (row oriented)" $ do
        matmatTest1CTPSV
    describe "ZTPSV" $ do
      it "3x3 lower no trans (row oriented)" $ do
        matmatTest1ZTPSV

-- [ 1 2 0]   [-2]   [4]
-- [ 0 1 1] * [ 3] = [5]
-- [ 0 0 1]   [ 2]   [2]
matmatTest1STPSV:: IO ()
matmatTest1STPSV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 2, 0, 1, 1, 1] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [4, 5, 2] !! idx)
    BLAS.stpsv Matrix.SRow MatUpper NoTranspose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [1 1 0]   [-2]   [1]
-- [0 1 1] * [ 3] = [5]
-- [0 0 1]   [ 2]   [2]
matmatTest1DTPSV:: IO ()
matmatTest1DTPSV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1, 1, 0, 1, 1, 1] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [1, 5, 2] !! idx)
    BLAS.dtpsv Matrix.SColumn MatLower Transpose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [ 1:+0    0:+0    0:+0]   [-2]   [-2:+0   ]
-- [ 2:+(-2) 1:+0    0:+0] * [ 3] = [-1:+4   ]
-- [ 0:+0    1:+(-1) 1:+0]   [ 2]   [ 5:+(-3)]
matmatTest1CTPSV:: IO ()
matmatTest1CTPSV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1:+0, 2:+2, 0:+0, 1:+0, 1:+1, 1:+0] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [(-2):+0, (-1):+4, 5:+(-3)]!! idx)
    BLAS.ctpsv Matrix.SRow MatUpper Matrix.ConjTranspose MatUnit 3 a x
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [-2, 3, 2]

-- [1 1 0]   [2]   [5]
-- [0 1 1] * [3] = [5]
-- [0 0 1]   [2]   [2]
matmatTest1ZTPSV:: IO ()
matmatTest1ZTPSV = do
    a  <- Matrix.generateMutableDenseVector 6 (\idx -> [1:+0, 1:+0, 1:+0, 0:+0, 1:+0, 1:+0] !! idx)
    x  <- Matrix.generateMutableDenseVector 3 (\idx -> [5, 5, 2] !! idx)
    BLAS.ztpsv Matrix.SColumn MatUpper NoTranspose MatNonUnit 3 a x -- TODO: NAN error when using Lower Transpose
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector x
    resList `shouldBe` [2, 3, 2]

trmvSpec :: Spec
trmvSpec =
  context "?TRMV" $ do
    describe "STRMV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1STRMV
    describe "DTRMV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1DTRMV
    describe "CTRMV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1CTRMV
    describe "ZTRMV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1ZTRMV

matmatTest1STRMV:: IO ()
matmatTest1STRMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
            (\(i,j) -> if i >= j then (1.0::Float) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 2 else 1)
    BLAS.strmv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [3,1]

matmatTest1DTRMV:: IO ()
matmatTest1DTRMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::Double) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 2 else 1)
    BLAS.dtrmv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [3,1]

matmatTest1CTRMV:: IO ()
matmatTest1CTRMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Float)) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 2 else 1)
    BLAS.ctrmv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [3,1]

matmatTest1ZTRMV:: IO ()
matmatTest1ZTRMV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Double )) else 0 )
    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 2 else 1)
    BLAS.ztrmv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [3,1]

trsvSpec :: Spec
trsvSpec =
  context "?TRSV" $ do
    describe "STRSV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1STRSV
    describe "DTRSV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1DTRSV
    describe "CTRSV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1CTRSV
    describe "ZTRSV" $ do
      it "2x2 upper 1's" $ do
        matmatTest1ZTRSV

matmatTest1STRSV:: IO ()
matmatTest1STRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
            (\(i,j) -> if i >= j then (1.0::Float) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.strsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2,1]

matmatTest1DTRSV:: IO ()
matmatTest1DTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::Double) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.dtrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2,1]

matmatTest1CTRSV:: IO ()
matmatTest1CTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Float)) else 0 )

    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.ctrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2,1]

matmatTest1ZTRSV:: IO ()
matmatTest1ZTRSV = do
    left  <- Matrix.generateMutableDenseMatrix (Matrix.SRow)  (2,2)
                (\(i,j) -> if i >= j then (1.0::(Complex Double )) else 0 )
    res  <- Matrix.generateMutableDenseVector  2 (\i -> if i == 0 then 3 else 1)
    BLAS.ztrsv MatUpper NoTranspose MatUnit left res
    resList <- Matrix.mutableVectorToList $ _bufferMutDenseVector res
    resList `shouldBe` [2,1]


{-# LANGUAGE FlexibleContexts,
             MultiParamTypeClasses #-}
module Network.Trainer.BackpropTrainer
( BackpropTrainer(..)

, backprop
, inputs
, outputs
, deltas
, hiddenDeltas
, calculateNablas
, fit
, evaluate
) where

import Network.Network
import Network.Network.FeedForwardNetwork
import Network.Neuron
import Network.Layer
import Network.Trainer

import Numeric.LinearAlgebra

-- | A BackpropTrainer performs simple backpropagation on a neural network.
--   It can be used as the basis for more complex trainers.
data BackpropTrainer = BackpropTrainer { eta :: Double
                                       , cost :: CostFunction
                                       , cost' :: CostFunction'
                                       }

-- | Declare the BackpropTrainer to be an instance of Trainer.
instance Trainer BackpropTrainer FeedForwardNetwork where
  fit s t n examples = foldl (backprop t) n $ s examples
  -- | Use the cost function to determine the error of a network
  evaluate t n example = (cost t) (snd example) (predict (fst example) n)

-- | Perform backpropagation on a single training data instance.
backprop :: BackpropTrainer -> FeedForwardNetwork -> [TrainingData] -> FeedForwardNetwork
backprop t n es =
  updateNetwork (length es) t (foldl (calculateNablas t n) emptyFeedForwardNetwork es) n

-- | Given the size of the minibatch, the trainer, the nablas for each layer, given
--   as a network, and the network itself, return a network with updated wieghts.
updateNetwork :: Int -> BackpropTrainer -> FeedForwardNetwork -> FeedForwardNetwork -> FeedForwardNetwork
updateNetwork mag t nablas n = addFeedForwardNetworks n
  (FeedForwardNetwork $ map (scaleLayer $ -1 * (eta t) / (fromIntegral mag)) (layers nablas))

-- | Calculate the nablas for a minibatch and return them as a network (so each
--   weight and bias gets its own nabla).
calculateNablas :: BackpropTrainer -> FeedForwardNetwork -> FeedForwardNetwork -> TrainingData -> FeedForwardNetwork
calculateNablas t n nablas e = FeedForwardNetwork $ map (updateLayer t) (zip3 (layers n) ds os)
  where ds = deltas t n e
        os = outputs (fst e) n

-- | The mapped function to update the weight and biases in a single layer
updateLayer :: BackpropTrainer -> (Layer, Vector Double, Vector Double) -> Layer
updateLayer t (l, delta, output) = Layer newWeight newBias n
  where n = neuron l
        newWeight = ((reshape 1 delta) <> (reshape (dim output) output))
        newBias = delta

-- | The outputs function scans over each layer of the network and stores the
--   activated results
outputs :: Vector Double -> FeedForwardNetwork -> [Vector Double]
outputs input network = scanl apply input (layers network)

-- | The inputs function performs a similar task to outputs, but returns a list
--   of vectors of unactivated inputs
inputs :: Vector Double -> FeedForwardNetwork -> [Vector Double]
inputs input network = if null (layers network) then []
  else unactivated : inputs activated (FeedForwardNetwork (tail $ layers network))
  where unactivated = weightMatrix layer <> input + biasVector layer
        layer = head $ layers network
        activated = mapVector (activation (neuron layer)) unactivated

-- | The deltas function returns a list of layer deltas.
deltas :: BackpropTrainer -> FeedForwardNetwork -> TrainingData -> [Vector Double]
deltas t n example = hiddenDeltas
  (FeedForwardNetwork (reverse (layers n))) outputDelta (tail $ reverse is)
    ++ [outputDelta]
  where outputDelta = costd (snd example) output *
          mapVector activationd lastInput
        costd = cost' t
        activationd = activation' (neuron (last (layers n)))
        output = last os
        lastInput = last is
        is = inputs (fst example) n
        os = outputs (fst example) n

-- | Compute the hidden layer deltas
hiddenDeltas :: FeedForwardNetwork -> Vector Double -> [Vector Double] -> [Vector Double]
hiddenDeltas n prevDelta is = if length (layers n) <= 1 then []
  else delta : hiddenDeltas rest delta (tail is)
  where rest = FeedForwardNetwork (tail $ layers n)
        delta = (trans w) <> prevDelta * spv
        w = weightMatrix (head $ layers n)
        spv = mapVector (activation' (neuron (head $ layers n))) (head is)


{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE EmptyCase                                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE PartialTypeSignatures                    #-}
{-# LANGUAGE QuantifiedConstraints                    #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TupleSections                            #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Extra.Solver    #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

import           Backprop.Learn
import           Control.Applicative
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Primitive
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State
import           Data.Char
import           Data.Conduit
import           Data.Default
import           Data.Kind
import           Data.Proxy
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Time
import           Data.Type.Equality
import           Data.Type.Tuple
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Numeric.LinearAlgebra.Static.Backprop hiding ((<>))
import           Numeric.Natural
import           Numeric.Opto
import           Options.Applicative
import           Statistics.Distribution
import           Statistics.Distribution.Normal
import           Text.Printf
import qualified Data.Conduit.Combinators                     as C
import qualified Data.Vector.Sized                            as SV
import qualified System.Random.MWC                            as MWC

data Mode = CSimulate (Maybe Int)
          | forall s p. ( Floating s
                        , Floating p
                        , NFData s
                        , NFData p
                        , Metric Double s
                        , Metric Double p
                        , Show s
                        , Show p
                        , Initialize p
                        , Initialize s
                        , Backprop p
                        , Backprop s
                        , forall m. PrimMonad m => LinearInPlace m Double p
                        , forall m. PrimMonad m => LinearInPlace m Double s
                        )
                        => CLearn (ModelPick p s)

data ModelPick :: Type -> Type -> Type where
    MARIMA :: Sing p -> Sing d -> Sing q
           -> ModelPick (ARIMAp p q) (ARIMAs p d q)
    MFCRNN :: Sing h
           -> ModelPick (LRp h 1 :# LRp (1 + h) h) (R h)
    MLSTM  :: Sing h
           -> ModelPick (LRp h 1 :# LSTMp 1 h)     (R h :# R h)
    MGRU   :: Sing h
           -> ModelPick (LRp h 1 :# GRUp 1 h)      (R h)

data Process = PSin

modelLearn :: ModelPick p s -> Model ('Just p) ('Just s) Double Double
modelLearn = \case
    MARIMA (SNat :: Sing pp) (SNat :: Sing d) (SNat :: Sing q) ->
              arima @pp @d @q
    MFCRNN (SNat :: Sing h) ->
              funcD sumElements
           <~ fc
           <~ fcra logistic logistic
           <~ funcD konst
    MLSTM  (SNat :: Sing h) ->
              funcD sumElements
           <~ fc
           <~ lstm
           <~ funcD konst
    MGRU   (SNat :: Sing h) ->
              funcD sumElements
           <~ fc
           <~ gru
           <~ funcD konst

data Options = O { oMode     :: Mode
                 , oProcess  :: Process
                 , oNoise    :: Double
                 , oInterval :: Int
                 , oLookback :: Natural
                 }

processConduit
    :: PrimMonad m
    => Process
    -> MWC.Gen (PrimState m)
    -> ConduitT i Double m ()
processConduit = \case
    PSin -> \g -> void . flip (foldr (>=>) pure) (0, 1) . repeat $ \(t, v) -> do
      dv <- genContVar (normalDistr 0 0.025) g
      let v' = min 2 . max 0.5 $ v + dv
          t' = t + v'
      yield (sin (2 * pi * (1/25) * t))
      return (t', v')

noisyConduit
    :: PrimMonad m
    => Double
    -> MWC.Gen (PrimState m)
    -> ConduitT Double Double m ()
noisyConduit σ g = C.mapM $ \x ->
    (x + ) <$> genContVar (normalDistr 0 σ) g

-- type Context = ConduitT (Vector)

main :: IO ()
main = MWC.withSystemRandom @IO $ \g -> do
    O{..} <- execParser $ info (parseOpt <**> helper)
                            ( fullDesc
                           <> progDesc "Learning ARIMA"
                           <> header "backprop-learn-arima - backprop-learn demo"
                            )

    SomeNat (Proxy :: Proxy n) <- pure $ someNatVal oLookback
    Just Refl <- pure $ Proxy @1 `isLE` Proxy @n
    let generator = processConduit oProcess g
                 .| noisyConduit oNoise g

    case oMode of
      CSimulate lim ->
        runConduit $ generator
                  .| maybe (C.map id) C.take lim
                  .| C.mapM_ print
                  .| C.sinkNull

      CLearn (modelLearn->model) -> do
        let unrolled = trainState . unrollFinal @(SV.Vector n) $ model
        p0 <- initParamNormal unrolled 0.2 g

        let report n b = do
              liftIO $ printf "(Batch %d)\n" (b :: Int)
              t0 <- liftIO getCurrentTime
              C.drop (n - 1)
              mp <- mapM (liftIO . evaluate . force) =<< await
              t1 <- liftIO getCurrentTime
              case mp of
                Nothing -> liftIO $ putStrLn "Done!"
                Just p -> do
                  chnk <- lift . state $ (,[])
                  liftIO $ do
                    printf "Trained on %d points in %s.\n"
                      (length chnk)
                      (show (t1 `diffUTCTime` t0))
                    let trainScore = testModelAll absErrorTest unrolled (TJust p) chnk
                    printf "Training error:   %.8f\n" trainScore
                  report n (b + 1)

        flip evalStateT []
            . runConduit
            $ transPipe lift generator
           .| leadings
           .| skipSampling 0.05 g
           .| C.iterM (modify . (:))
           .| optoConduit def p0
                (adam def (modelGrad squaredError noReg unrolled))
           .| report oInterval 0
           .| C.sinkNull

parseOpt :: Parser Options
parseOpt = O <$> subparser ( command "simulate" (info (parseSim   <**> helper) (progDesc "Simulate ARIMA only"))
                          <> command "learn"    (info (parseLearn <**> helper) (progDesc "Simulate and learn model"))
                           )
             <*> option (maybeReader parseProcess)
                   ( long "process"
                  <> help "Process to learn"
                  <> showDefaultWith showProcess
                  <> value PSin
                  <> metavar "PROCESS"
                   )
             <*> option auto
                   ( short 'e'
                  <> help "Standard deviation of noise term"
                  <> showDefault
                  <> value 0.05
                  <> metavar "DOUBLE"
                   )
             <*> option auto
                   ( long "interval"
                  <> short 'i'
                  <> help "Report interval"
                  <> showDefault
                  <> value 5000
                  <> metavar "INT"
                   )
             <*> option auto
                   ( long "lookback"
                  <> short 'l'
                  <> help "Learn lookback"
                  <> showDefault
                  <> value 10
                  <> metavar "INT"
                   )
  where
    parseSim :: Parser Mode
    parseSim = CSimulate <$> optional (option auto
                                          ( short 'n'
                                         <> help "Number of items to generate (infinite, if none given)"
                                         <> metavar "INT"
                                          )
                                      )
    parseLearn :: Parser Mode
    parseLearn = subparser
        ( command "arima" (info parseARIMA (progDesc "Learn ARIMA(p,d,q) model"))
       <> command "fcrnn" (info parseFCRNN (progDesc "Learn Fully Connected RNN model"))
       <> command "lstm"  (info parseLSTM  (progDesc "Learn LSTM model"))
       <> command "gru"   (info parseGRU   (progDesc "Learn GRU model"))
        )
    parseARIMA :: Parser Mode
    parseARIMA = do
        p <- argument auto ( help "AR(p): Autoregressive lookback"
                          <> metavar "INT"
                           )
        d <- argument auto ( help "I(d): Differencing degree"
                          <> metavar "INT"
                           )
        q <- argument auto ( help "MA(q): Moving average lookback"
                          <> metavar "INT"
                           )
        pure $ withSomeSing @Nat p $ \sp@SNat ->
          withSomeSing @Nat d $ \sd@SNat ->
          withSomeSing @Nat q $ \sq@SNat ->
            CLearn (MARIMA sp sd sq)
    parseFCRNN :: Parser Mode
    parseFCRNN = do
        h <- argument auto ( help "Hidden layer size"
                          <> metavar "INT"
                          <> showDefault
                          <> value 10
                           )
        pure $ withSomeSing @Nat h $ \sh@SNat ->
          CLearn (MFCRNN sh)
    parseLSTM :: Parser Mode
    parseLSTM = do
        h <- argument auto ( help "Hidden layer size"
                          <> metavar "INT"
                          <> showDefault
                          <> value 10
                           )
        pure $ withSomeSing @Nat h $ \sh@SNat ->
          CLearn (MLSTM sh)
    parseGRU :: Parser Mode
    parseGRU = do
        h <- argument auto ( help "Hidden layer size"
                          <> metavar "INT"
                          <> showDefault
                          <> value 10
                           )
        pure $ withSomeSing @Nat h $ \sh@SNat ->
          CLearn (MGRU sh)

showProcess :: Process -> String
showProcess PSin = "sin"

parseProcess :: String -> Maybe Process
parseProcess s = case map toLower s of
    "sin" -> Just PSin
    _     -> Nothing


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables   #-}
module Bio.ChromVAR.Background (getBackgroundPeaks) where

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as S
import qualified Data.Vector as V
import qualified Data.Matrix.Unboxed as MU
import qualified Data.Vector.Unboxed.Mutable as UM
import Control.Monad
import Data.Ord
import Data.List
import Control.Arrow
import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Generic as D
import Data.Matrix.Dynamic (Dynamic(..), matrix)

import Statistics.Sample
import System.Random.MWC
import System.Random.MWC.Distributions (categorical)

import Bio.ChromVAR.Utils

type PeakGroup = ( MU.Matrix Double   -- ^ Group to group distance
                 , V.Vector (U.Vector Int)  -- ^ Peaks in each group
                 , U.Vector Int       -- ^ Peak-group membership
                 )

getBackgroundPeaks :: Int -> U.Vector (Double, Double) -> GenIO -> IO [U.Vector Int]
getBackgroundPeaks n xs gen =
    let pg = mkPeakGroup xs
    in replicateM n $ getBackgroundPeak gen pg

getBackgroundPeak :: GenIO -> PeakGroup -> IO (U.Vector Int)
getBackgroundPeak gen (weightMat, bins, membership) = U.generateM n $ \i -> do
    let ws = weightMat `MU.takeRow` (membership U.! i)
    grp <- categorical ws gen
    let peaks = bins V.! grp
    idx <- uniformR (0, U.length peaks - 1) gen
    return $ peaks U.! idx
  where
    n = U.length membership
{-# INLINE getBackgroundPeak #-}

-- | Make peak groups.
mkPeakGroup :: U.Vector (Double, Double) -> PeakGroup
mkPeakGroup raw = (weights, V.fromList groups, membership)
  where
    transformed = case matrix (map (\(x,y) -> [x,y]) $ U.toList raw) of
        Dynamic mat@(D.Matrix _) -> U.fromList $
           map ((\[x,y] -> (x,y)) . S.toList) $ D.toRows $ whiten Cholesky mat
    weights = MU.generate (U.length points, U.length points) $ \(i,j) ->
        weight (points U.! i) (points U.! j)
      where
        points = U.fromList $ flip map groups $ \is -> mean *** mean $
            U.unzip $ U.map (transformed U.!) is
    membership = U.create $ do
        v <- UM.new $ U.length transformed
        forM_ (zip [0..] groups) $ \(x, is) -> U.forM_ is $ \i ->
            UM.unsafeWrite v i x
        return v
    groups = go [] 0 0 $ sortBy (comparing fst) $ zip (U.toList transformed) [0..]
      where
        go acc i j (((x,y), idx) : rest)
            | null acc || i' == i || j' == j = go (idx : acc) i' j' rest
            | i' > i || j' > j = U.fromList acc : go [idx] i' j' rest
            | otherwise = error "Impossible"
          where
            i' = truncate $ (x - x_min) / x_step :: Int
            j' = truncate $ (y - y_min) / y_step :: Int
        go _ _ _ _ = []
    (xs, ys) = U.unzip transformed
    x_min = U.minimum xs
    x_max = U.maximum xs
    x_step = (x_max - x_min) / n
    y_min = U.minimum ys
    y_max = U.maximum ys
    y_step = (y_max - y_min) / n
    n = 50
{-# INLINE mkPeakGroup #-}