{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Commons where

import Control.Applicative (liftA2)
import Control.Monad (foldM, forM)
import Data.Array
import Data.Complex
import Data.Function.HT (nest)
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate, sort)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (catMaybes, fromJust, mapMaybe)
import Data.Set (Set, fromList, toList)
import Data.Time (diffUTCTime, getCurrentTime)
import Data.Typeable (Typeable)
import Debug.Trace (traceShowId)
import GHC.IO.Unsafe (unsafePerformIO)
import GHC.TypeLits (KnownNat, Nat)
import HashedExpression.Internal.Expression
import HashedExpression.Internal.Normalize
import HashedExpression.Internal.Utils
import HashedExpression.Interp
import HashedExpression.Operation
import qualified HashedExpression.Operation
import HashedExpression.Prettify
import HashedExpression.Value
import Test.HUnit
import Test.Hspec
import Test.QuickCheck
import Var
import Prelude hiding ((^))

-- |
--
-- | Remove duplicate but also sort
removeDuplicate :: (Ord a) => [a] -> [a]
removeDuplicate = toList . fromList

-- |
vectorOfDifferent :: Eq a => Int -> Gen a -> Gen [a]
vectorOfDifferent sz gen = foldM f [] [1 .. sz]
  where
    f acc _ = (: acc) <$> gen `suchThat` (not . flip elem acc)

-- | Format
format :: [(String, String)] -> String
format = intercalate "\n" . map oneLine
  where
    oneLine (f, s) = f ++ ": " ++ s

-- |
inspect :: (Typeable d, Typeable rc) => Expression d rc -> Expression d rc
inspect x =
  unsafePerformIO $ do
    showExp x
    return x

-- | Vars list
type Vars = [[String]] -- Vars 0D, 1D, 2D, 3D, ..

mergeVars :: [Vars] -> Vars
mergeVars = foldl f [[], [], [], []]
  where
    f x y = map removeDuplicate $ zipWith (++) x y

genDouble :: Gen Double
genDouble = arbitrary `suchThat` inSmallRange
  where
    inSmallRange x = x >= 0 && x <= 10

-- |
genValMap ::
  forall size1D size2D1 size2D2.
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Vars ->
  Gen ValMaps
genValMap vars = do
  let sz1D = valueFromNat @size1D
      sz2D1 = valueFromNat @size2D1
      sz2D2 = valueFromNat @size2D2
  let [names0d, names1d, names2d, names3d] = vars
  list0d <- vectorOf (length names0d) genDouble
  let vm0 = Map.fromList . zip names0d $ map VScalar list0d
  list1d <- vectorOf (length names1d) . vectorOf sz1D $ genDouble
  let vm1 =
        Map.fromList . zip names1d . map (V1D . listArray (0, sz1D - 1)) $
          list1d
  list2d <- vectorOf (length names2d) . vectorOf (sz2D1 * sz2D2) $ genDouble
  let vm2 =
        Map.fromList
          . zip names2d
          . map (V2D . listArray ((0, 0), (sz2D1 - 1, sz2D2 - 1)))
          $ list2d
  let vm3 = Map.empty -- TODO: not testing with 3D yet.
  return $ Map.unions [vm0, vm1, vm2, vm3]

shouldApprox :: (HasCallStack, Approximable a) => a -> a -> Expectation
shouldApprox x y = assertBool msg (x ~= y)
  where
    msg = "Expected: " ++ prettifyShow y ++ "\nGot: " ++ prettifyShow x

infix 1 `shouldApprox`

liftE1 ::
  (Expression d1 et1 -> Expression d2 et2) ->
  (Expression d1 et1, Vars) ->
  (Expression d2 et2, Vars)
liftE1 op (e, v) = (op e, v)

liftE2 ::
  (Expression d1 et1 -> Expression d2 et2 -> Expression d3 et3) ->
  (Expression d1 et1, Vars) ->
  (Expression d2 et2, Vars) ->
  (Expression d3 et3, Vars)
liftE2 op (e1, v1) (e2, v2) = ((op e1 e2), (mergeVars [v1, v2]))

-------------------------------------------------------------------------------
primitiveScalarR :: Gen (Expression Scalar R, Vars)
primitiveScalarR = do
  name <- elements . map pure $ ['a' .. 'z']
  dbl <- genDouble
  r <- (`mod` 10) <$> arbitrary
  if r < 6
    then return ((variable name), [[(name)], [], [], []])
    else return ((constant dbl), [[], [], [], []])

primitiveScalarC :: Gen (Expression Scalar C, Vars)
primitiveScalarC = liftE2 (+:) <$> primitiveScalarR <*> primitiveScalarR

-------------------------------------------------------------------------------
primitive1DR ::
  forall n.
  KnownNat n =>
  Gen (Expression n R, Vars)
primitive1DR = do
  name <- elements . map ((++ "1") . pure) $ ['a' .. 'z']
  dbl <- genDouble
  r <- (`mod` 10) <$> arbitrary
  if r < 6
    then return ((variable1D @n name), [[], [(name)], [], []])
    else return ((constant1D @n dbl), [[], [], [], []])

primitive1DC ::
  forall n.
  KnownNat n =>
  Gen (Expression n C, Vars)
primitive1DC = liftE2 (+:) <$> primitive1DR @n <*> primitive1DR @n

-------------------------------------------------------------------------------
primitive2DR ::
  forall m n.
  (KnownNat m, KnownNat n) =>
  Gen (Expression '(m, n) R, Vars)
primitive2DR = do
  name <- elements . map ((++ "2") . pure) $ ['a' .. 'z']
  dbl <- genDouble
  r <- (`mod` 10) <$> arbitrary
  if r < 6
    then return (variable2D @m @n name, [[], [], [(name)], []])
    else return (constant2D @m @n dbl, [[], [], [], []])

primitive2DC ::
  forall m n.
  (KnownNat m, KnownNat n) =>
  Gen (Expression '(m, n) C, Vars)
primitive2DC = liftE2 (+:) <$> primitive2DR @m @n <*> primitive2DR @m @n

--type  = 10
-------------------------------------------------------------------------------
genScalarR ::
  forall default1D default2D1 default2D2.
  (KnownNat default1D, KnownNat default2D1, KnownNat default2D2) =>
  Int ->
  Gen (Expression Scalar R, Vars)
genScalarR size
  | size == 0 = primitiveScalarR
  | otherwise =
    let sub = genScalarR @default1D @default2D1 @default2D2 (size `div` 10)
        subC = genScalarC @default1D @default2D1 @default2D2 (size `div` 10)
        sub1D = gen1DR @default1D @default2D1 @default2D2 (size `div` 10)
        sub2D = gen2DR @default1D @default2D1 @default2D2 (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- sub
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (*.),
            binary (-),
            binary (<.>),
            unary negate,
            unary (^ 2),
            liftE1 xRe <$> subC,
            liftE1 xIm <$> subC,
            liftE2 (<.>) <$> sub1D <*> sub1D,
            liftE2 (<.>) <$> sub2D <*> sub2D
          ]

-------------------------------------------------------------------------------
genScalarC ::
  forall default1D default2D1 default2D2.
  (KnownNat default1D, KnownNat default2D1, KnownNat default2D2) =>
  Int ->
  Gen (Expression Scalar C, Vars)
genScalarC size
  | size == 0 = primitiveScalarC
  | otherwise =
    let sub = genScalarC @default1D @default2D1 @default2D2 (size `div` 10)
        subR = genScalarR @default1D @default2D1 @default2D2 (size `div` 10)
        sub1D = gen1DC @default1D @default2D1 @default2D2 (size `div` 10)
        sub2D = gen2DC @default1D @default2D1 @default2D2 (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- subR
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (*.),
            binary (-),
            binary (<.>),
            unary negate,
            unary (^ 2),
            liftE2 (+:) <$> subR <*> subR,
            liftE2 (<.>) <$> sub1D <*> sub1D,
            liftE2 (<.>) <$> sub2D <*> sub2D
          ]

-------------------------------------------------------------------------------
gen1DR ::
  forall n default2D1 default2D2.
  (KnownNat n, KnownNat default2D1, KnownNat default2D2) =>
  Int ->
  Gen (Expression n R, Vars)
gen1DR size
  | size == 0 = primitive1DR
  | otherwise =
    let sub = gen1DR @n @default2D1 @default2D2 (size `div` 10)
        subC = gen1DC @n @default2D1 @default2D2 (size `div` 10)
        subScalar = genScalarR @n @default2D1 @default2D2 (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- sub
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        fromRotate = do
          amount <- elements [- (valueFromNat @n) .. valueFromNat @n]
          liftE1 (rotate amount) <$> sub
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (-),
            unary negate,
            unary (^ 2),
            liftE1 xRe <$> subC,
            liftE1 xIm <$> subC,
            liftE2 (*.) <$> subScalar <*> sub,
            fromRotate,
            liftE1 (xRe . ft) <$> sub,
            liftE1 (xIm . ft) <$> sub
          ]

-------------------------------------------------------------------------------
gen1DC ::
  forall n default2D1 default2D2.
  (KnownNat n, KnownNat default2D1, KnownNat default2D2) =>
  Int ->
  Gen (Expression n C, Vars)
gen1DC size
  | size == 0 = primitive1DC
  | otherwise =
    let sub = gen1DC @n @default2D1 @default2D2 (size `div` 10)
        subR = gen1DR @n @default2D1 @default2D2 (size `div` 10)
        subScalar = genScalarC @n @default2D1 @default2D2 (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- subR
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        fromRotate = do
          amount <- elements [- (valueFromNat @n) .. valueFromNat @n]
          liftE1 (rotate amount) <$> sub
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (-),
            unary negate,
            unary (^ 2),
            liftE2 (+:) <$> subR <*> subR,
            liftE2 (*.) <$> subScalar <*> sub,
            liftE1 ft <$> sub,
            fromRotate
          ]

-------------------------------------------------------------------------------
gen2DR ::
  forall default1D m n.
  (KnownNat default1D, KnownNat m, KnownNat n) =>
  Int ->
  Gen (Expression '(m, n) R, Vars)
gen2DR size
  | size == 0 = primitive2DR
  | otherwise =
    let sub = gen2DR @default1D @m @n (size `div` 10)
        subC = gen2DC @default1D @m @n (size `div` 10)
        subScalar = genScalarR @default1D @m @n (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- sub
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        fromRotate = do
          amount1 <- elements [- (valueFromNat @m) .. valueFromNat @m]
          amount2 <- elements [- (valueFromNat @n) .. valueFromNat @n]
          liftE1 (rotate (amount1, amount2)) <$> sub
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (-),
            unary negate,
            unary (^ 2),
            liftE1 xRe <$> subC,
            liftE1 xIm <$> subC,
            liftE2 (*.) <$> subScalar <*> sub,
            fromRotate,
            liftE1 (xRe . ft) <$> sub,
            liftE1 (xIm . ft) <$> sub
          ]

-------------------------------------------------------------------------------
gen2DC ::
  forall default1D m n.
  (KnownNat default1D, KnownNat m, KnownNat n) =>
  Int ->
  Gen (Expression '(m, n) C, Vars)
gen2DC size
  | size == 0 = primitive2DC
  | otherwise =
    let sub = gen2DC @default1D @m @n (size `div` 10)
        subR = gen2DR @default1D @m @n (size `div` 10)
        subScalar = genScalarC @default1D @m @n (size `div` 10)
        fromPiecewise = do
          numBranches <- elements [2 .. 4]
          branches <- vectorOf numBranches sub
          condition <- subR
          marks <- sort <$> vectorOfDifferent (numBranches - 1) arbitrary
          let vars = mergeVars $ map snd branches ++ [snd condition]
              exp = piecewise marks (fst condition) $ map fst branches
          return (exp, vars)
        fromRotate = do
          amount1 <- elements [- (valueFromNat @m) .. valueFromNat @m]
          amount2 <- elements [- (valueFromNat @n) .. valueFromNat @n]
          liftE1 (rotate (amount1, amount2)) <$> sub
        binary op = liftE2 op <$> sub <*> sub
        unary op = liftE1 op <$> sub
     in oneof
          [ fromPiecewise,
            binary (+),
            binary (*),
            binary (-),
            unary negate,
            unary (^ 2),
            liftE2 (+:) <$> subR <*> subR,
            liftE2 (*.) <$> subScalar <*> sub,
            liftE1 ft <$> sub,
            fromRotate
          ]

-------------------------------------------------------------------------------
data Suite (size1D :: Nat) (size2D1 :: Nat) (size2D2 :: Nat) d et
  = Suite (Expression d et) ValMaps
  deriving (Show)

-------------------------------------------------------------------------------
type TestSuite = Suite Default1D Default2D1 Default2D2

type SuiteScalarR = TestSuite Scalar R

type SuiteScalarC = TestSuite Scalar C

type SuiteOneR = TestSuite Default1D R

type SuiteOneC = TestSuite Default1D C

type SuiteTwoR = TestSuite '(Default2D1, Default2D2) R

type SuiteTwoC = TestSuite '(Default2D1, Default2D2) C

-------------------------------------------------------------------------------
instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 Scalar R)
  where
  arbitrary = do
    (exp, vars) <- sized $ genScalarR @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 Scalar C)
  where
  arbitrary = do
    (exp, vars) <- sized $ genScalarC @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 size1D R)
  where
  arbitrary = do
    (exp, vars) <- sized $ gen1DR @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 size1D C)
  where
  arbitrary = do
    (exp, vars) <- sized $ gen1DC @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 '(size2D1, size2D2) R)
  where
  arbitrary = do
    (exp, vars) <- sized $ gen2DR @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

instance
  (KnownNat size1D, KnownNat size2D1, KnownNat size2D2) =>
  Arbitrary (Suite size1D size2D1 size2D2 '(size2D1, size2D2) C)
  where
  arbitrary = do
    (exp, vars) <- sized $ gen2DC @size1D @size2D1 @size2D2
    valMaps <- genValMap @size1D @size2D1 @size2D2 vars
    return $ Suite exp valMaps

-------------------------------------------------------------------------------
instance Arbitrary (Expression Scalar R) where
  arbitrary = fst <$> sized (genScalarR @Default1D @Default2D1 @Default2D2)

instance Arbitrary (Expression Scalar C) where
  arbitrary = fst <$> sized (genScalarC @Default1D @Default2D1 @Default2D2)

instance KnownNat n => Arbitrary (Expression n R) where
  arbitrary = fst <$> sized (gen1DR @n @Default2D1 @Default2D2)

instance KnownNat n => Arbitrary (Expression n C) where
  arbitrary = fst <$> sized (gen1DC @n @Default2D1 @Default2D2)

instance (KnownNat m, KnownNat n) => Arbitrary (Expression '(m, n) R) where
  arbitrary = fst <$> sized (gen2DR @Default1D @m @n)

instance (KnownNat m, KnownNat n) => Arbitrary (Expression '(m, n) C) where
  arbitrary = fst <$> sized (gen2DC @Default1D @m @n)

-------------------------------------------------------------------------------
data ArbitraryExpresion = forall d et. (DimensionType d, ElementType et) => ArbitraryExpresion (Expression d et)

instance Show ArbitraryExpresion where
  show (ArbitraryExpresion exp) = show exp

instance Arbitrary ArbitraryExpresion where
  arbitrary =
    let option1 =
          fmap ArbitraryExpresion (arbitrary :: Gen (Expression Scalar R))
        option2 =
          fmap ArbitraryExpresion (arbitrary :: Gen (Expression Scalar C))
        option3 =
          fmap
            ArbitraryExpresion
            (arbitrary :: Gen (Expression Default1D R))
        option4 =
          fmap
            ArbitraryExpresion
            (arbitrary :: Gen (Expression Default1D C))
        option5 =
          fmap
            ArbitraryExpresion
            (arbitrary :: Gen (Expression '(Default2D1, Default2D2) R))
        option6 =
          fmap
            ArbitraryExpresion
            (arbitrary :: Gen (Expression '(Default2D1, Default2D2) C))
     in oneof [option1, option2, option3, option4, option5, option6]

-- |
getWrappedExp :: ArbitraryExpresion -> (ExpressionMap, NodeID)
getWrappedExp (ArbitraryExpresion (Expression n mp)) = (mp, n)

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

-- |
sz :: Expression d et -> Int
sz = IM.size . exMap

infix 1 `shouldNormalizeTo`

shouldNormalizeTo ::
  (HasCallStack, DimensionType d, ElementType et, Typeable et, Typeable d) =>
  Expression d et ->
  Expression d et ->
  IO ()
shouldNormalizeTo exp1 exp2 = do
  prettify (normalize exp1) `shouldBe` prettify (normalize exp2)
  normalize exp1 `shouldBe` normalize exp2


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
#if __GLASGOW_HASKELL__ >= 805
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE NoStarIsType #-}
#endif
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Eigen.Matrix
  (
    -- * Types
    Matrix(..)
  , Vec(..)
  , MatrixXf
  , MatrixXd
  , MatrixXcf
  , MatrixXcd

    -- * Common API
  , Elem
  , C
  , natToInt
  , Row(..)
  , Col(..)

    -- * Encode/Decode a Matrix
  , encode
  , decode

    -- * Querying a Matrix
  , null
  , square
  , rows
  , cols
  , dims

    -- * Constructing a Matrix
  , empty
  , constant
  , zero
  , ones
  , identity
  , random
  , diagonal

  , (!)
  , coeff
  , generate
  , sum
  , prod
  , mean
  , trace
  , all
  , any
  , count
  , norm
  , squaredNorm
  , blueNorm
  , hypotNorm
  , determinant
  , add
  , sub
  , mul
  , map
  , imap
  , TriangularMode(..)
  , triangularView
  , filter
  , ifilter
  , length
  , foldl
  , foldl'
  , inverse
  , adjoint
  , transpose
  , conjugate
  , normalize
  , modify
  , block
  , unsafeFreeze
  , unsafeWith
  , fromList
  , toList
  ) where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Prelude hiding
  (map, null, filter, length, foldl, any, all, sum)
import Control.Monad (forM_)
import Control.Monad.Primitive (PrimMonad(..))
import Data.Binary (Binary(..))
import qualified Data.Binary as Binary
import qualified Data.ByteString.Lazy as BSL
import Data.Complex (Complex)
import Data.Constraint.Nat
import Eigen.Internal
  ( Elem
  , Cast(..)
  , natToInt
  , Row(..)
  , Col(..)
  )
import qualified Eigen.Internal as Internal
import qualified Eigen.Matrix.Mutable as M
import qualified Data.List as List
import Data.Kind (Type)
import GHC.TypeLits (Nat, type (*), type (<=), KnownNat)
import Foreign.C.Types (CInt)
import Foreign.C.String (CString)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr)
import Foreign.Storable (peek)

import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM

-- | Matrix to be used in pure computations.
--
--   * Uses column majour memory layout.
--
--   * Has a copy-free FFI using the <http://eigen.tuxfamily.org Eigen> library.
--
newtype Matrix :: Nat -> Nat -> Type -> Type where
  Matrix :: Vec (n * m) a -> Matrix n m a

-- | Used internally to track the size and corresponding C type of the matrix.
newtype Vec :: Nat -> Type -> Type where
  Vec :: VS.Vector (C a) -> Vec n a

instance forall n m a. (Elem a, Show a, KnownNat n, KnownNat m) => Show (Matrix n m a) where
  show m = List.concat
    [ "Matrix ", show (rows m), "x", show (cols m)
    , "\n", List.intercalate "\n" $ List.map (List.intercalate "\t" . List.map show) $ toList m, "\n"
    ]

instance forall n m a. (KnownNat n, KnownNat m, Elem a) => Binary (Matrix n m a) where
  put (Matrix (Vec vals)) = do
    put $ Internal.magicCode (undefined :: C a)
    put $ natToInt @n
    put $ natToInt @m
    put vals

  get = do
    get >>= (`when` fail "wrong matrix type") . (/= Internal.magicCode (undefined :: C a))
    Matrix . Vec <$> get

-- | Encode the sparse matrix as a lazy bytestring
encode :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> BSL.ByteString
encode = Binary.encode

-- | Decode the sparse matrix from a lazy bytestring
decode :: (Elem a, KnownNat n, KnownNat m) => BSL.ByteString -> Matrix n m a
decode = Binary.decode

-- | Alias for single precision matrix
type MatrixXf n m = Matrix n m Float
-- | Alias for double precision matrix
type MatrixXd n m = Matrix n m Double
-- | Alias for single precision matrix of complex numbers
type MatrixXcf n m = Matrix n m (Complex Float)
-- | Alias for double precision matrix of complex numbers
type MatrixXcd n m = Matrix n m (Complex Double)

-- | Construct an empty 0x0 matrix
empty :: Elem a => Matrix 0 0 a
{-# INLINE empty #-}
empty = Matrix (Vec (VS.empty))

-- | Is matrix empty?
null :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Bool
{-# INLINE null #-}
null m = cols m == 0 && rows m == 0

-- | Is matrix square?
--
square :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Bool
{-# INLINE square #-}
square _ = natToInt @n == natToInt @m

-- | Matrix where all coeffs are filled with the given value
constant :: forall n m a. (Elem a, KnownNat n, KnownNat m) => a -> Matrix n m a
{-# INLINE constant #-}
constant !val =
  let !cval = toC val
  in withDims $ \rs cs -> VS.replicate (rs * cs) cval

-- | Matrix where all coeffs are filled with 0
zero :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a
{-# INLINE zero #-}
zero = constant 0

-- | Matrix where all coeffs are filled with 1
ones :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a
{-# INLINE ones #-}
ones = constant 1

-- | The identity matrix (not necessarily square)
identity :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a
identity =
  Internal.performIO $ do
     m :: M.IOMatrix n m a <- M.new
     Internal.call $ M.unsafeWith m Internal.identity
     unsafeFreeze m

-- | The random matrix of a given size
random :: forall n m a. (Elem a, KnownNat n, KnownNat m) => IO (Matrix n m a)
random = do
  m :: M.IOMatrix n m a <- M.new
  Internal.call $ M.unsafeWith m Internal.random
  unsafeFreeze m

withDims :: forall n m a. (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> VS.Vector (C a)) -> Matrix n m a
{-# INLINE withDims #-}
withDims f =
  let !r = natToInt @n
      !c = natToInt @m
  in Matrix $ Vec $ f r c

-- | The number of rows in the matrix
rows :: forall n m a. KnownNat n => Matrix n m a -> Int
{-# INLINE rows #-}
rows _ = natToInt @n

-- | The number of colums in the matrix
cols :: forall n m a. KnownNat m => Matrix n m a -> Int
{-# INLINE cols #-}
cols _ = natToInt @m

-- | Return Matrix size as a pair of (rows, cols)
dims :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> (Int, Int)
{-# INLINE dims #-}
dims _ = (natToInt @n, natToInt @m)

-- | Return the value at the given position.
(!) :: forall n m a r c. (Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m) => Row r -> Col c -> Matrix n m a -> a
{-# INLINE (!) #-}
(!) = coeff

-- | Return the value at the given position.
coeff :: forall n m a r c. (Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m) => Row r -> Col c -> Matrix n m a -> a
{-# INLINE coeff #-}
coeff _ _ m@(Matrix (Vec vals)) =
  let !row  = natToInt @r
      !col  = natToInt @c
  in fromC $! VS.unsafeIndex vals $! col * rows m + row

unsafeCoeff :: (Elem a, KnownNat n) => Int -> Int -> Matrix n m a -> a
{-# INLINE unsafeCoeff #-}
unsafeCoeff row col m@(Matrix (Vec vals)) = fromC $! VS.unsafeIndex vals $! col * rows m + row

-- | Given a generation function `f :: Int -> Int -> a`, construct a Matrix of known size
--   using points in the matrix as inputs.
generate :: forall n m a. (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a) -> Matrix n m a
generate f = withDims $ \rs cs -> VS.create $ do
  vals :: VSM.MVector s (C a) <- VSM.new (rs * cs)
  forM_ [0 .. pred rs] $ \r ->
    forM_ [0 .. pred cs] $ \c ->
      VSM.write vals (c * rs + r) (toC $! f r c)
  pure vals

-- | The sum of all coefficients in the matrix
sum :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
sum = _prop Internal.sum

-- | The product of all coefficients in the matrix
prod :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
prod = _prop Internal.prod

-- | The arithmetic mean of all coefficients in the matrix
mean :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
mean = _prop Internal.mean

-- | The trace of a matrix is the sum of the diagonal coefficients.
--
--   'trace' m == 'sum' ('diagonal' m)
trace :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
trace = _prop Internal.trace

-- | Given a predicate p, determine if all values in the Matrix satisfy p.
all :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Bool
all f (Matrix (Vec vals)) = VS.all (f . fromC) vals

-- | Given a predicate p, determine if any values in the Matrix satisfy p.
any :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Bool
any f (Matrix (Vec vals)) = VS.any (f . fromC) vals

-- | Given a predicate p, determine how many values in the Matrix satisfy p.
count :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Int
count f (Matrix (Vec vals)) = VS.foldl' (\n x-> if f (fromC x) then (n + 1) else n) 0 vals

norm, squaredNorm, blueNorm, hypotNorm :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a

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

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

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

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

-- | The determinant of the matrix
determinant :: forall n a. (Elem a, KnownNat n) => Matrix n n a -> a
determinant m = _prop Internal.determinant m

-- | Add two matrices.
add :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a -> Matrix n m a
add m1 m2 = _binop Internal.add m1 m2

-- | Subtract two matrices.
sub :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a -> Matrix n m a
sub m1 m2 = _binop Internal.sub m1 m2

-- | Multiply two matrices.
mul :: (Elem a, KnownNat p, KnownNat q, KnownNat r) => Matrix p q a -> Matrix q r a -> Matrix p r a
mul m1 m2 = _binop Internal.mul m1 m2

{- | Apply a given function to each element of the matrix.
Here is an example how to implement scalar matrix multiplication:

>>> let a = fromList [[1,2],[3,4]] :: MatrixXf 2 2
>>> a
Matrix 2x2
1.0 2.0
3.0 4.0
>>> map (*10) a
Matrix 2x2
10.0    20.0
30.0    40.0
-}
map :: Elem a => (a -> a) -> Matrix n m a -> Matrix n m a
map f (Matrix (Vec vals)) = Matrix $ Vec $ VS.map (toC . f . fromC) vals

{- | Apply a given function to each element of the matrix.
Here is an example how upper triangular matrix can be implemented:
>>> let a = fromList [[1,2,3],[4,5,6],[7,8,9]] :: MatrixXf
>>> a
Matrix 3x3
1.0 2.0 3.0
4.0 5.0 6.0
7.0 8.0 9.0
>>> imap (\row col val -> if row <= col then val else 0) a
Matrix 3x3
1.0 2.0 3.0
0.0 5.0 6.0
0.0 0.0 9.0
-}
imap :: (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a -> a) -> Matrix n m a -> Matrix n m a
imap f (Matrix (Vec vals)) =
  withDims $ \rs _ ->
    VS.imap (\n ->
      let (c,r) = divMod n rs
      in toC . f r c . fromC) vals

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

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

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

-- | Filter elements in the matrix with an indexed predicate. Filtered elements will be replaces by 0.
ifilter :: (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a -> Bool) -> Matrix n m a -> Matrix n m a
ifilter f = imap (\r c x -> if f r c x then x else 0)

-- | The length of the matrix.
length :: forall n m a r. (Elem a, KnownNat n, KnownNat m, r ~ (n * m), KnownNat r) => Matrix n m a -> Int
length _ = natToInt @r

-- | Left fold of a matrix, where accumulation is lazy.
foldl :: (Elem a, KnownNat n, KnownNat m) => (b -> a -> b) -> b -> Matrix n m a -> b
foldl f b (Matrix (Vec vals)) = VS.foldl (\a x -> f a (fromC x)) b vals

-- | Right fold of a matrix, where accumulation is strict.
foldl' :: Elem a => (b -> a -> b) -> b -> Matrix n m a -> b
foldl' f b (Matrix (Vec vals)) = VS.foldl' (\ !a x -> f a (fromC x)) b vals

-- | Return the diagonal of a matrix.
diagonal :: (Elem a, KnownNat n, KnownNat m, r ~ Min n m, KnownNat r) => Matrix n m a -> Matrix r 1 a
diagonal = _unop Internal.diagonal

{- | Inverse of the matrix
For small fixed sizes up to 4x4, this method uses cofactors. In the general case, this method uses PartialPivLU decomposition
-}
inverse :: forall n a. (Elem a, KnownNat n) => Matrix n n a -> Matrix n n a
inverse = _unop Internal.inverse

-- | Adjoint of the matrix
adjoint :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix m n a
adjoint = _unop Internal.adjoint

-- | Transpose of the matrix
transpose :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix m n a
transpose = _unop Internal.transpose

-- | Conjugate of the matrix
conjugate :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a
conjugate = _unop Internal.conjugate

-- | Normalise the matrix by dividing it on its 'norm'
normalize :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a
normalize (Matrix (Vec vals)) = Internal.performIO $ do
  vals' <- VS.thaw vals
  VSM.unsafeWith vals' $ \p ->
    let !rs = natToInt @n
        !cs = natToInt @m
    in Internal.call $ Internal.normalize p (toC rs) (toC cs)
  Matrix . Vec <$> VS.unsafeFreeze vals'

-- | Apply a destructive operation to a matrix. The operation will be performed in-place, if it is safe
--   to do so - otherwise, it will create a copy of the matrix.
modify :: (Elem a, KnownNat n, KnownNat m) => (forall s. M.MMatrix n m s a -> ST s ()) -> Matrix n m a -> Matrix n m a
modify f (Matrix (Vec vals)) = Matrix $ Vec $ VS.modify (f . M.fromVector ) vals

-- | Extract rectangular block from matrix defined by startRow startCol blockRows blockCols
block :: forall sr sc br bc n m a.
     (Elem a, KnownNat sr, KnownNat sc, KnownNat br, KnownNat bc, KnownNat n, KnownNat m)
  => (sr <= n, sc <= m, br <= n, bc <= m)
  => Row sr -- ^ starting row
  -> Col sc -- ^ starting col
  -> Row br -- ^ block of rows
  -> Col bc -- ^ block of cols
  -> Matrix n m a -- ^ extract from this
  -> Matrix br bc a -- ^ extraction
block _ _ _ _ m =
  let !startRow = natToInt @sr
      !startCol = natToInt @sc
  in generate $ \row col -> unsafeCoeff (startRow + row) (startCol + col) m

-- | Turn a mutable matrix into an immutable matrix without copying.
--   The mutable matrix should not be modified after this conversion.
unsafeFreeze :: (Elem a, KnownNat n, KnownNat m, PrimMonad p) => M.MMatrix n m (PrimState p) a -> p (Matrix n m a)
unsafeFreeze m = VS.unsafeFreeze (M.vals m) >>= pure . Matrix . Vec

-- | Pass a pointer to the matrix's data to the IO action. The data may not be modified through the pointer.
unsafeWith  :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> (Ptr (C a) -> CInt -> CInt -> IO b) -> IO b
unsafeWith m@(Matrix (Vec (vals))) f =
  VS.unsafeWith vals $ \p ->
    let !rs = toC $! rows m
        !cs = toC $! cols m
    in f p rs cs

_prop :: (Elem a, KnownNat n, KnownNat m) => (Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString) -> Matrix n m a -> a
{-# INLINE _prop #-}
_prop f m = fromC $ Internal.performIO $ alloca $ \p -> do
   Internal.call $ unsafeWith m (f p)
   peek p

_binop :: forall n m n1 m1 n2 m2 a. (Elem a, KnownNat n, KnownNat m, KnownNat n1, KnownNat m1, KnownNat n2, KnownNat m2)
  => (Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString)
  -> Matrix n m a
  -> Matrix n1 m1 a
  -> Matrix n2 m2 a
{-# INLINE _binop #-}
_binop g m1 m2 = Internal.performIO $ do
  m0 :: M.IOMatrix n2 m2 a <- M.new
  M.unsafeWith m0 $ \vals0 rows0 cols0 ->
      unsafeWith m1 $ \vals1 rows1 cols1 ->
          unsafeWith m2 $ \vals2 rows2 cols2 ->
              Internal.call $ g
                vals0 rows0 cols0
                vals1 rows1 cols1
                vals2 rows2 cols2
  unsafeFreeze m0

_unop :: forall n m n1 m1 a. (Elem a, KnownNat n, KnownNat m, KnownNat n1, KnownNat m1)
  => (Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString)
  -> Matrix n m a
  -> Matrix n1 m1 a
{-# INLINE _unop #-}
_unop g m1 = Internal.performIO $ do
  m0 :: M.IOMatrix n1 m1 a <- M.new
  M.unsafeWith m0 $ \vals0 rows0 cols0 ->
      unsafeWith m1 $ \vals1 rows1 cols1 ->
          Internal.call $ g
              vals0 rows0 cols0
              vals1 rows1 cols1
  unsafeFreeze m0

-- | Convert a matrix to a list.
toList :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> [[a]]
{-# INLINE toList #-}
toList m@(Matrix (Vec vals))
  | null m = []
  | otherwise = [[fromC $ vals `VS.unsafeIndex` (col * _rows + row) | col <- [0..pred _cols]] | row <- [0..pred _rows]]
  where
    !_rows = rows m
    !_cols = cols m

-- | Convert a list to a matrix. Returns 'Nothing' if the dimensions of the list do not match that
--   of the matrix.
fromList :: forall n m a. (Elem a, KnownNat n, KnownNat m) => [[a]] -> Maybe (Matrix n m a)
fromList list = do
  let myRows = natToInt @n
  let myCols = natToInt @m
  let _rows  = List.length list
  let _cols  = List.foldl' max 0 (List.map List.length list)
  if ((myRows /= _rows) || (myCols /= _cols))
    then Nothing
    else (Just . Matrix . Vec) $ VS.create $ do
      vm <- VSM.replicate (_rows * _cols) (toC (0 :: a))
      forM_ (zip [0..] list) $ \(row,vals) ->
        forM_ (zip [0..] vals) $ \(col, val) ->
          VSM.write vm (col * _rows + row) (toC val)
      pure vm


{-# LANGUAGE Rank2Types #-}
import Control.Monad.ST (RealWorld)
import Data.Complex (Complex)
import Data.Vector.Storable (MVector(MVector))
import Foreign (Ptr, Storable, withForeignPtr)
import Text.Printf (printf)
import Blas.Generic.Unsafe (Numeric)
import qualified Data.Vector.Storable as Vector
import qualified Blas.Primitive.Types as Blas
import qualified Blas.Generic.Unsafe as Blas
import qualified TestUtils as T

main :: IO ()
main = T.runTest $ do
  sequence_ (mapNumericTypes testGemm)

mapNumericTypes :: (forall a. (Eq a, Numeric a, Show a) => a -> b) -> [b]
mapNumericTypes f =
  [ f (dummy :: Float)
  , f (dummy :: Double)
  , f (dummy :: Complex Float)
  , f (dummy :: Complex Double)
  ]

dummy :: a
dummy = error "dummy value that shouldn't be used"

withMVector :: MVector s a -> (Ptr a -> IO b) -> IO b
withMVector (MVector _ foreignPtr) = withForeignPtr foreignPtr

mVectorFromList :: Storable a => [a] -> IO (MVector RealWorld a)
mVectorFromList = Vector.thaw . Vector.fromList

testGemm :: (Eq a, Numeric a, Show a) => a -> T.Test ()
testGemm numType = do
  let order    = Blas.RowMajor
      transa   = Blas.NoTrans
      transb   = Blas.Trans
      n        = 2
      size     = n * n
      alpha    = 1.0 `asTypeOf` numType
      beta     = 0.0
      a        = Vector.fromList [1, 2, 3, 4]
      b        = Vector.fromList [1, 2, 3, 5]
      expected = Vector.fromList [5, 13, 11, 29]

  c' <- T.liftIO $ do
    c <- mVectorFromList $ take size (repeat 0)
    Vector.unsafeWith a $ \ pa ->
      Vector.unsafeWith b $ \ pb ->
        withMVector c $ \ pc ->
          Blas.gemm order transa transb n n n alpha pa n pb n beta pc n
    Vector.freeze c

  if c' == expected
    then T.passTest "testGemm"
    else T.failTest $ printf "testGemm: c' does not match: %s != %s"
                      (show c') (show expected)


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE PartialTypeSignatures #-}
module Main where

import Test.Tasty
import           Test.Tasty.HUnit
import Data.List
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as VS
import Bio.Data.Bed.Types
import Data.Maybe
import Conduit
import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Generic as D
import qualified Data.Matrix.Static.Sparse as S
import Data.Matrix.Dynamic (Dynamic(..), matrix)
import Data.Matrix.Static.LinearAlgebra
import Statistics.Sample hiding (covariance)
import System.Random.MWC (create)

import Bio.ChromVAR
import Bio.ChromVAR.Utils

main :: IO ()
main = defaultMain $ testGroup "Main"
    [ testCase "sortBed" test
    , testCase "sortBed" bgTest]
    --, testCase "sortBed" peakTest ]

test :: Assertion
test = (map round' $ concat dev', map round' $ concat z') @=?
    (map round' $ concat deviations, map round' $ concat z)
  where
    [(dev', z')] = runIdentity $ runConduit $ yieldMany [(4, cellByPeak)] .|
        computeDeviation 3 3 expectation bg peakBymotif .| sinkList
    expectation = let e = D.replicate 1 @@ (S.fromTriplet (U.fromList cellByPeak) :: SparseMatrix 4 3 Double) :: Matrix 1 3 Double
                      s = VS.sum $ D.flatten e
                  in D.toList $ D.map (/s) e
    round' :: Double -> Double
    round' x = fromIntegral (round (x * 1e4)) / 1e4

cellByPeak :: [(Int, Int, Double)]
cellByPeak = [(0,1,1), (0,2,2), (1,0,1), (1,1,1), (2,0,2), (2,2,1), (3,1,1), (3,2,1)]

peakBymotif :: [(Int, Int, Double)]
peakBymotif = [(0,0,1), (0,1,1), (1,0,1), (1,1,1), (1,2,1), (2,0,1), (2,2,1)]

peaks :: [BED3]
peaks = map (\(chr, i) -> BED3 chr i (i + 500)) dat
  where
    dat = [ ("chr1", 76585873)
          , ("chr2", 42772928)
          , ("chr2", 100183786) ]

bg :: [[Int]]
bg = [ [1,0,1]
     , [1,0,1]
     , [1,2,0] ]

deviations = transpose [ [0.1728395, -0.4444444, 0.1728395, -0.07407407]
             , [-0.2910053, 0.3174603, 0.2116402, -0.19841270]
             , [0.7407407, -0.6349206, -0.7407407,  0.63492063] ]

z = transpose [ [1.1547005, -1.1547005, 1.1547005, -1.1547005]
    , [-0.5773503, 0.5773503, 0.5773503, -0.5773503]
    , [3.2331615, -1.1547005, -4.0414519, 9.2376043] ]

bgTest :: Assertion
bgTest = do
    input <- readData "tests/data/coordinates.tsv"
    expected <- readData "tests/data/transformed.tsv"
    let actual = case matrix (map (\(a,b) -> [a,b]) $ U.toList input) of
            Dynamic mat@(D.Matrix _) -> U.fromList $ map (\[a,b] -> (a,b)) $
                map VS.toList $ D.toRows $ whiten Cholesky mat
    U.map (\(x,y) -> (round' x, round' y)) actual @=?
        U.map (\(x,y) -> (round' x, round' y)) expected
  where
    readData fl = U.fromList . map (f . words) . lines <$> readFile fl
      where
        f [x,y] = (read x, read y)
    round' :: Double -> Double
    round' x = fromIntegral (round (x * 1e4)) / 1e4

peakTest :: Assertion
peakTest = do
    coord <- readData "tests/data/coordinates.tsv"

    expected <- map (map (\x -> coord U.! (read x - 1)) . words) .
        lines <$> readFile "tests/data/background.tsv"
    actual <- fmap (transpose . map (map (coord U.!) . U.toList)) $
        create >>= getBackgroundPeaks 50 coord

    print $ zipWith f expected actual

    --print $ flip map input' $ \x -> meanVarianceUnb $ U.fromList $ map dist $ comb x
    return undefined
  where
    f xs ys = meanVarianceUnb $ U.fromList $ map dist $ map (\[a,b] -> (a,b)) $ sequence [xs, ys]
    readData fl = U.fromList . map (f . words) . lines <$> readFile fl
      where
        f [x,y] = (read x, read y)
    comb (x:xs) = zip (repeat x) xs ++ comb xs
    comb _ = []
    dist ((x1,y1), (x2,y2)) = sqrt $ (x1 - x2)**2 + (y1-y2)**2


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Monad
import           Control.Monad.Random
import           Data.Constraint                 (Dict (..))
import           Data.List                       (foldl')
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.Prelude.List
import qualified Data.Vector.Storable            as V
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static    as SA
import           Options.Applicative
import           System.Random.MWC
import           System.Random.MWC.Distributions
import           Unsafe.Coerce                   (unsafeCoerce)

import           Grenade

-- | Choose a backend: 0 BLAS, 1 HMATRIX
#define HMATRIX 0

netSpec :: SpecNet
netSpec = specFullyConnected 2 400 |=> specRelu1D 400 |=> specFullyConnected 400 1 |=> specLogit1D 1 |=> specNil1D 1

-- netSpec :: SpecNet
-- netSpec = specFullyConnected 2 4000 |=> specRelu1D 4000 |=> specFullyConnected 4000 1 |=> specLogit1D 1 |=> specNil1D 1

-- netSpec :: SpecNet
-- netSpec =
--       specFullyConnected 2 4500
--   |=> specTanh1D 4500
--   |=> specFullyConnected 4500 750
--   |=> specRelu1D 750
--   |=> specFullyConnected 750 150
--   |=> specRelu1D 150
--   |=> specFullyConnected 150 150
--   |=> specRelu1D 150
--   |=> specFullyConnected 150 30
--   |=> specRelu1D 30
--   |=> specFullyConnected 30 20
--   |=> specRelu1D 20
--   |=> specFullyConnected 20 10
--   |=> specRelu1D 10
--   |=> specFullyConnected 10 1
--   |=> specLogit1D 1
--   |=> specNil1D 1

-- -- -- | The definition for a feed forward network using the dynamic module. Note the nested networks. This network clearly is over-engeneered for this example!
-- netSpec :: SpecNet
-- netSpec =
--  specFullyConnected 2 40 |=> specTanh1D 40 |=>
--  specDropout 400 0.95 Nothing |=>
--  netSpecInner |=>
--  specFullyConnected 20 30 |=> specRelu1D 30 |=>
--  specFullyConnected 30 20 |=> specRelu1D 20 |=>
--  specFullyConnected 20 10 |=> specRelu1D 10 |=>
--  specFullyConnected 10 1 |=> specLogit1D 1 |=> specNil1D 1
--  where netSpecInner = specFullyConnected 40 30 |=> specRelu1D 30 |=> specFullyConnected 30 20 |=> specReshape1D2D 20 (2, 10) |=> specReshape2D1D (2, 10) 20 |=> specNil1D 20

goals :: [((RealNum, RealNum), RealNum)]
goals = [((0.50, 0.50), 0.25), ((-0.50, -0.50), 0.25)]

-- | Specifications can be built using the following interface also. Here one does not have to carefully pay attention to match the dimension inputs and outputs as when using specification directly.
buildNetViaInterface :: IO SpecConcreteNetwork
buildNetViaInterface =
  buildModelWith (NetworkInitSettings UniformInit HMatrix (Just 10)) (DynamicBuildSetup { printResultingSpecification = False })  $
  inputLayer1D 2 >>                            -- 1. Every model has to start with an input dimension
  fullyConnected 10 >> dropout 0.89 >> relu >> -- 2. Layers are simply added as desired. The input of each layer is determined automatically,
  fullyConnected 20 >> relu >>                 --    whereas the output is specified as in `fullyConnected 20`. Thus, the layer empits 20 signals
  fullyConnected 4 >> tanhLayer >>             --    which arethe input of the next layer.
  networkLayer (
    inputLayer1D 4 >> fullyConnected 10 >> relu >> fullyConnected 4 >> sinusoid
    ) >>
  reshape (4, 1, 1) >>                         -- 3. Reshape is ignored if the input and output is the same
  fullyConnected 1                             -- 4. The output is simply the number of signals the alst layer emits.

netTrain :: forall shapes len1 len2 layers o .
     (SingI (Last shapes), KnownNat len1, KnownNat len2, Head shapes ~ 'D1 len1, Last shapes ~ 'D1 len2)
  => Network layers shapes
  -> Optimizer o
  -> Int
  -> IO (Network layers shapes)
netTrain net0 op n = do
#if HMATRIX
    inps <- replicateM n $ do
      s  <- getRandom
      return $ S1D (SA.randomVector s SA.Uniform * 2 - 1 :: SA.R len1)
#else
    inps <- replicateM n $ do
        s  <- getRandom
        return $ S1DV $ SA.extract (SA.randomVector s SA.Uniform * 2 - 1 :: SA.R len1)
#endif
    let mkVal :: S ('D1 len1) -> Rational
        mkVal (S1D v) = mkVal (S1DV $ SA.extract v)
        mkVal (S1DV v) = if any (v `inCircleV`) goals
                         then 1
                         else 0
#if HMATRIX
        constr = S1D . fromRational
#else
        constr = S1DV . fromRational
#endif
    let outs = map (constr . mkVal) inps
    let trained = foldl' trainEach net0 (zip inps outs)
    return trained

  where trainEach !network (i,o) = train op network i o

netScore :: (KnownNat len, Head shapes ~ 'D1 len, Last shapes ~ 'D1 1) => Network layers shapes -> IO ()
netScore network = do
    let testIns = [ [ (x,y)  | x <- [0..50] ]
                             | y <- [0..20] ]
#if HMATRIX
        outMat  = fmap (fmap (\(x,y) -> (render (x/25-1) (y/10-1) . normx) $ runNet network (S1D $ SA.vector [x / 25 - 1,y / 10 - 1]))) testIns
#else
        outMat  = fmap (fmap (\(x,y) -> (render (x/25-1) (y/10-1) . normx) $ runNet network (S1DV $ V.fromList [x / 25 - 1,y / 10 - 1]))) testIns
#endif
    putStrLn $ unlines outMat

  where
    render x y n'  | x == 0 && y == 0 = '+'
                   | y == 0 = '-'
                   | x == 0 = '|'
                   | n' <= 0.2  = ' '
                   | n' <= 0.4  = '.'
                   | n' <= 0.6  = '-'
                   | n' <= 0.8  = '='
                   | otherwise = '#'

normx :: S ('D1 1) -> RealNum
normx (S1D r)  = SA.mean r
normx (S1DV v) = V.sum v / fromIntegral (V.length v)

testValues :: forall len shapes layers . (KnownNat len, Head shapes ~ 'D1 len, Last shapes ~ 'D1 1) => Network layers shapes -> IO ()
testValues network = do
  let n = 1000
#if HMATRIX
  inps <- replicateM n $ do
      s  <- getRandom
      return $ S1D (SA.randomVector s SA.Uniform * 2 - 1 :: SA.R len)
#else
  inps <- replicateM n $ do
       s  <- getRandom
       return $ S1DV $ SA.extract (SA.randomVector s SA.Uniform * 2 - 1 :: SA.R len)
#endif

  let mkVal :: S ('D1 len1) -> Integer
      mkVal (S1D v) = mkVal (S1DV $ SA.extract v :: S ('D1 len))
      mkVal (S1DV v) = if any (v `inCircleV`) goals
                         then 1
                         else 0
  let outs :: [Integer]
      outs = map mkVal inps
  let ress = zip outs (map (round . normx . runNet network) inps)
      correct = length $ filter id $ map (uncurry (==)) ress
      incorrect = length $ filter id $ map (uncurry (/=)) ress
      falsePositives = length $ filter id $ map (uncurry (\shd nn -> shd == 0 && nn == 1)) ress
      falseNegatives = length $ filter id $ map (uncurry (\shd nn -> shd == 1 && nn == 0)) ress
  putStr $ show correct  ++ " | "
  putStr $ show incorrect ++ " | "
  putStr $ show falsePositives ++ " | "
  putStrLn $ show falseNegatives ++ " | "

inCircle :: KnownNat n => SA.R n -> (SA.R n, RealNum) -> Bool
v `inCircle` (o, r) = SA.norm_2 (v - o) <= r

inCircleV :: V.Vector RealNum -> ((RealNum, RealNum), RealNum) -> Bool
v `inCircleV` ((x,y), r) = norm (V.zipWith (-) v circle) <= r
  where
    circle = V.fromList [x,y]
    norm x = sqrt (V.sum $ V.map (^ (2 :: Int)) x)

data FeedForwardOpts = FeedForwardOpts Int (Optimizer 'Adam)

feedForward' :: Parser FeedForwardOpts
feedForward' =
  FeedForwardOpts <$> option auto (long "examples" <> short 'e' <> value 10000)
                  <*> (OptAdam
                       <$> option auto (long "alpha" <> short 'r' <> value 0.001)
                       <*> option auto (long "beta1" <> value 0.9)
                       <*> option auto (long "beta2" <> value 0.999)
                       <*> option auto (long "epsilon" <> value 1e-4)
                       <*> option auto (long "lambda" <> value 1e-3))

main :: IO ()
main = do

  FeedForwardOpts examples rate <- execParser (info (feedForward' <**> helper) idm)
  putStrLn "| Nr | Correct | Incorrect | FalsePositives | FalseNegatives |"
  putStrLn "--------------------------------------------------------------"
  let nr = 10 :: Int
  mapM_
    (\n -> do
       putStr $ "| " ++ show n ++ " | "
#if HMATRIX
       SpecConcreteNetwork1D1D (net0 :: Network layers shapes) <- networkFromSpecificationWith (NetworkInitSettings HeEtAl HMatrix Nothing) netSpec
#else
       SpecConcreteNetwork1D1D (net0 :: Network layers shapes) <- networkFromSpecificationWith (NetworkInitSettings HeEtAl BLAS Nothing) netSpec
#endif
       -- We need to specify the actual number of output nodes, as our functions requiere that!
       case (unsafeCoerce (Dict :: Dict ()) :: Dict (('D1 1) ~ Last shapes)) of
         Dict -> do
           net <- netTrain net0 rate examples
           unsafeCoerce $ -- only needed as GADTs are enabled, which disallowes the type to escape and thus prevents the type inference to work. The result is not needed anyways.
             testValues net)
    [1 .. nr]

  -- Features of dynamic networks:
  SpecConcreteNetwork1D1D (net' :: Network layers shapes) <- networkFromSpecificationWith (NetworkInitSettings HeEtAl BLAS (Just 50000)) netSpec
  net2 <- netTrain net' rate 0
  let spec' = networkToSpecification net2
  putStrLn "String represenation of the network specification: "
  print spec'
  let serializedSpec = encode spec'   -- only the specification (not the weights) are serialized here! The weights can be serialized using the networks serialize instance!
  let _ = encode net2                 -- E.g. like this.
  case decode serializedSpec of
    Left err -> print err
    Right spec2 -> do
      print spec2
      SpecConcreteNetwork1D1D (net3 :: Network layers3 shapes3) <- networkFromSpecificationWith (NetworkInitSettings HeEtAl BLAS (Just 50000)) spec2
      net4 <- foldM (\n _ -> netTrain n rate examples) net3 [(1 :: Int)..30]
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (('D1 1) ~ Last shapes3)) of
        Dict -> netScore net4

  -- -- There is a also nice interface available also
  -- SpecConcreteNetwork1D1D newNet <- buildNetViaInterface
  -- print (networkToSpecification newNet)


-- | Helpers for testing
module Tests.Helpers (
    -- * helpers
    T(..)
  , typeName
  , eq
  , eqC
    -- * Generic QC tests
  , monotonicallyIncreases
    -- * HUnit helpers
  , testAssertion
  , testEquality
  ) where

import Data.Complex
import Data.Typeable

import qualified Test.HUnit      as HU
import Test.Framework
import Test.Framework.Providers.HUnit

import Numeric.MathFunctions.Comparison

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

-- | Phantom typed value used to select right instance in QC tests
data T a = T

-- | String representation of type name
typeName :: Typeable a => T a -> String
typeName = show . typeOf . typeParam
  where
    typeParam :: T a -> a
    typeParam _ = undefined

-- | Approximate equality for 'Double'. Doesn't work well for numbers
--   which are almost zero.
eq :: Double                    -- ^ Relative error
   -> Double -> Double -> Bool
eq = eqRelErr

-- | Approximate equality for 'Complex Double'
eqC :: Double                   -- ^ Relative error
    -> Complex Double
    -> Complex Double
    -> Bool
eqC eps a@(ar :+ ai) b@(br :+ bi)
  | a == 0 && b == 0 = True
  | otherwise        = abs (ar - br) <= eps * d
                    && abs (ai - bi) <= eps * d
  where
    d = max (realPart $ abs a) (realPart $ abs b)

----------------------------------------------------------------
-- Generic QC
----------------------------------------------------------------

-- Check that function is nondecreasing
monotonicallyIncreases :: (Ord a, Ord b) => (a -> b) -> a -> a -> Bool
monotonicallyIncreases f x1 x2 = f (min x1 x2) <= f (max x1 x2)

----------------------------------------------------------------
-- HUnit helpers
----------------------------------------------------------------

testAssertion :: String -> Bool -> Test
testAssertion str cont = testCase str $ HU.assertBool str cont

testEquality :: (Show a, Eq a) => String -> a -> a -> Test
testEquality msg a b = testCase msg $ HU.assertEqual msg a b


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Utils where

import Data.Complex
import Data.IntMap.Strict
import Expression
import Hash
import Prelude hiding ((*), (+), lookup)

-- | Auxiliary functions for operations
--
expressionNumType :: (NumType rc) => Expression d rc -> RC
expressionNumType (Expression n mp) =
    case lookup n mp of
        Just (_, node) -> nodeNumType node
        _ -> error "expression not in map"

expressionShape :: (DimensionType d) => Expression d rc -> Shape
expressionShape (Expression n mp) =
    case lookup n mp of
        Just (dim, _) -> dim
        _ -> error "expression not in map"

retrieveNode :: ExpressionMap -> Int -> Node
retrieveNode mp n =
    case lookup n mp of
        Just (_, node) -> node
        _ -> error "node not in map"

ensureSameShape :: (Ring d rc) => Expression d rc -> Expression d rc -> a -> a
ensureSameShape e1 e2 after =
    if expressionShape e1 == expressionShape e2
        then after
        else error "Ensure same shape failed"

fromReal :: Double -> Complex Double
fromReal x = x :+ 0


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module ImageLoader
(
  loadExamples
, loadNonExamples
)

where

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

import Control.Exception          ( Exception
                                  , try )
import Control.Monad              ( sequence )
import System.FilePath            ( FilePath
                                  , (</>) )
import Vision.Image.Type          ( manifestVector )
import Vision.Image.RGB.Type      ( RGB )
import Vision.Image.Grey.Type     ( Grey )
import Vision.Image.Storage.DevIL ( Autodetect (..)
                                  , load )
import Paths_img_neural           ( getDataFileName )

import qualified Data.Vector.Storable       as VS
import qualified Numeric.LinearAlgebra.Data as VD

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

type Samples = [(VD.Vector Double, VD.Vector Double)]

-- Load the vectors which represent an example of the letter A
loadExamples :: Exception e => IO (Either e Samples)
loadExamples = try $ do
    images <- loadImages exampleImageNames
    return $ map (\x -> (x, output)) images
  where output = VD.fromList [1 :: Double]

-- Load the vectors which do not represent an example of the letter A
loadNonExamples :: Exception e => IO (Either e Samples)
loadNonExamples = try $ do
    images <- loadImages nonExampleImageNames
    return $ map (\x -> (x, output)) images
  where output = VD.fromList [0 :: Double]

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

-- Lists all images that are examples of the letter A
exampleImageNames :: [FilePath]
exampleImageNames = map ("A" </>) [
                      "A1.png"
                    , "A2.png"
                    , "A3.png"
                    , "A4.png"
                    , "A5.png"
                    , "A6.png" ]

-- Lists all images that are _not_ examples of the letter A
nonExampleImageNames :: [FilePath]
nonExampleImageNames = map ("not-A" </>) [
                         "B1.png"
                       , "B2.png"
                       , "B3.png"
                       , "B4.png"
                       , "B5.png"
                       , "B6.png"
                       , "B7.png" ]

-- Loads images from the given list
loadImages :: [FilePath] -> IO [VD.Vector Double]
loadImages fp = do
    fnames  <- map' getFileName fp
    vectors <- map' loadFile fnames
    return $ map convertToLinearVector vectors
  where map' f xs   = sequence $ map f xs
        getFileName = getDataFileName . ("res" </>)

-- Loads the given file to an image, converted to a Vector
loadFile :: FilePath -> IO (VS.Vector Double)
loadFile f = do
    dat <- load Autodetect f
    case dat
      of Left  e           -> error  $ show e
         Right (i :: Grey) -> return $ greyManifestToVector i

-- Converts a Grey into a Vector Int
greyManifestToVector :: Grey -> VS.Vector Double
greyManifestToVector m = VS.map (\i -> 1 - (fromIntegral $ round $ (fromIntegral i)/256)) (manifestVector m)

-- Converts a Data/Vector to a LinearAlgebra/Vector
convertToLinearVector :: VS.Storable a => VS.Vector a -> VD.Vector a
convertToLinearVector = VD.fromList . VS.toList