{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif

---------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Simple matrix operation for low-dimensional primitives.
---------------------------------------------------------------------------
module Linear.Matrix
  ( (!*!), (!+!), (!-!), (!*), (*!), (!!*), (*!!), (!!/)
  , column
  , adjoint
  , M22, M23, M24, M32, M33, M34, M42, M43, M44
  , m33_to_m44, m43_to_m44
  , det22, det33, det44, inv22, inv33, inv44
  , identity
  , Trace(..)
  , translation
  , transpose
  , fromQuaternion
  , mkTransformation
  , mkTransformationMat
  , _m22, _m23, _m24
  , _m32, _m33, _m34
  , _m42, _m43, _m44
#if MIN_VERSION_base(4,8,0)
  , lu
  , luFinite
  , forwardSub
  , forwardSubFinite
  , backwardSub
  , backwardSubFinite
  , luSolve
  , luSolveFinite
  , luInv
  , luInvFinite
  , luDet
  , luDetFinite
#endif
  ) where

#if __GLASGOW_HASKELL__ < 710
import Control.Applicative
#endif
import Control.Lens hiding (index)
import Control.Lens.Internal.Context
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Rep
import Linear.Quaternion
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector
import Linear.Conjugate
import Linear.Trace

#if MIN_VERSION_base(4,8,0)
import GHC.TypeLits
import Linear.V
#endif

-- $setup
-- >>> import Control.Lens hiding (index)
-- >>> import Data.Complex (Complex (..))
-- >>> import Linear.V2
-- >>> import Linear.V3
-- >>> import Linear.V
-- >>> import Data.IntMap
-- >>> import Debug.SimpleReflect.Vars

-- | This is a generalization of 'Control.Lens.inside' to work over any corepresentable 'Functor'.
--
-- @
-- 'column' :: 'Representable' f => 'Lens' s t a b -> 'Lens' (f s) (f t) (f a) (f b)
-- @
--
-- In practice it is used to access a column of a matrix.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) ^._x
-- V3 1 2 3
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) ^.column _x
-- V2 1 4
column :: Representable f => LensLike (Context a b) s t a b -> Lens (f s) (f t) (f a) (f b)
column l f es = o <$> f i where
   go = l (Context id)
   i = tabulate $ \ e -> ipos $ go (index es e)
   o eb = tabulate $ \ e -> ipeek (index eb e) (go (index es e))

infixl 7 !*!
-- | Matrix product. This can compute any combination of sparse and dense multiplication.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !*! V3 (V2 1 2) (V2 3 4) (V2 4 5)
-- V2 (V2 19 25) (V2 43 58)
--
-- >>> V2 (fromList [(1,2)]) (fromList [(2,3)]) !*! fromList [(1,V3 0 0 1), (2, V3 0 0 5)]
-- V2 (V3 0 0 2) (V3 0 0 15)
(!*!) :: (Functor m, Foldable t, Additive t, Additive n, Num a) => m (t a) -> t (n a) -> m (n a)
f !*! g = fmap (\ f' -> Foldable.foldl' (^+^) zero $ liftI2 (*^) f' g) f

infixl 6 !+!
-- | Entry-wise matrix addition.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !+! V2 (V3 7 8 9) (V3 1 2 3)
-- V2 (V3 8 10 12) (V3 5 7 9)
(!+!) :: (Additive m, Additive n, Num a) => m (n a) -> m (n a) -> m (n a)
as !+! bs = liftU2 (^+^) as bs

infixl 6 !-!
-- | Entry-wise matrix subtraction.
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !-! V2 (V3 7 8 9) (V3 1 2 3)
-- V2 (V3 (-6) (-6) (-6)) (V3 3 3 3)
(!-!) :: (Additive m, Additive n, Num a) => m (n a) -> m (n a) -> m (n a)
as !-! bs = liftU2 (^-^) as bs

infixl 7 !*
-- | Matrix * column vector
--
-- >>> V2 (V3 1 2 3) (V3 4 5 6) !* V3 7 8 9
-- V2 50 122
(!*) :: (Functor m, Foldable r, Additive r, Num a) => m (r a) -> r a -> m a
m !* v = fmap (\r -> Foldable.sum $ liftI2 (*) r v) m

infixl 7 *!
-- | Row vector * matrix
--
-- >>> V2 1 2 *! V2 (V3 3 4 5) (V3 6 7 8)
-- V3 15 18 21

-- (*!) :: (Metric r, Additive n, Num a) => r a -> r (n a) -> n a
-- f *! g = dot f <$> distribute g

(*!) :: (Num a, Foldable t, Additive f, Additive t) => t a -> t (f a) -> f a
f *! g = sumV $ liftI2 (*^) f g

infixl 7 *!!
-- | Scalar-matrix product
--
-- >>> 5 *!! V2 (V2 1 2) (V2 3 4)
-- V2 (V2 5 10) (V2 15 20)
(*!!) :: (Functor m, Functor r, Num a) => a -> m (r a) -> m (r a)
s *!! m = fmap (s *^) m
{-# INLINE (*!!) #-}

infixl 7 !!*
-- | Matrix-scalar product
--
-- >>> V2 (V2 1 2) (V2 3 4) !!* 5
-- V2 (V2 5 10) (V2 15 20)
(!!*) :: (Functor m, Functor r, Num a) => m (r a) -> a -> m (r a)
(!!*) = flip (*!!)
{-# INLINE (!!*) #-}

infixl 7 !!/
-- | Matrix-scalar division
(!!/) :: (Functor m, Functor r, Fractional a) => m (r a) -> a -> m (r a)
m !!/ s = fmap (^/ s) m
{-# INLINE (!!/) #-}

-- | Hermitian conjugate or conjugate transpose
--
-- >>> adjoint (V2 (V2 (1 :+ 2) (3 :+ 4)) (V2 (5 :+ 6) (7 :+ 8)))
-- V2 (V2 (1.0 :+ (-2.0)) (5.0 :+ (-6.0))) (V2 (3.0 :+ (-4.0)) (7.0 :+ (-8.0)))
adjoint :: (Functor m, Distributive n, Conjugate a) => m (n a) -> n (m a)
adjoint = collect (fmap conjugate)
{-# INLINE adjoint #-}

-- * Matrices
--
-- Matrices use a row-major representation.

-- | A 2x2 matrix with row-major representation
type M22 a = V2 (V2 a)
-- | A 2x3 matrix with row-major representation
type M23 a = V2 (V3 a)
-- | A 2x4 matrix with row-major representation
type M24 a = V2 (V4 a)
-- | A 3x2 matrix with row-major representation
type M32 a = V3 (V2 a)
-- | A 3x3 matrix with row-major representation
type M33 a = V3 (V3 a)
-- | A 3x4 matrix with row-major representation
type M34 a = V3 (V4 a)
-- | A 4x2 matrix with row-major representation
type M42 a = V4 (V2 a)
-- | A 4x3 matrix with row-major representation
type M43 a = V4 (V3 a)
-- | A 4x4 matrix with row-major representation
type M44 a = V4 (V4 a)

-- | Build a rotation matrix from a unit 'Quaternion'.
fromQuaternion :: Num a => Quaternion a -> M33 a
fromQuaternion (Quaternion w (V3 x y z)) =
  V3 (V3 (1-2*(y2+z2)) (2*(xy-zw)) (2*(xz+yw)))
     (V3 (2*(xy+zw)) (1-2*(x2+z2)) (2*(yz-xw)))
     (V3 (2*(xz-yw)) (2*(yz+xw)) (1-2*(x2+y2)))
  where x2 = x*x
        y2 = y*y
        z2 = z*z
        xy = x*y
        xz = x*z
        xw = x*w
        yz = y*z
        yw = y*w
        zw = z*w
{-# INLINE fromQuaternion #-}

-- | Build a transformation matrix from a rotation matrix and a
-- translation vector.
mkTransformationMat :: Num a => M33 a -> V3 a -> M44 a
mkTransformationMat (V3 r1 r2 r3) (V3 tx ty tz) =
  V4 (snoc3 r1 tx) (snoc3 r2 ty) (snoc3 r3 tz) (V4 0 0 0 1)
  where snoc3 (V3 x y z) = V4 x y z
{-# INLINE mkTransformationMat #-}

-- |Build a transformation matrix from a rotation expressed as a
-- 'Quaternion' and a translation vector.
mkTransformation :: Num a => Quaternion a -> V3 a -> M44 a
mkTransformation = mkTransformationMat . fromQuaternion
{-# INLINE mkTransformation #-}

-- | Convert from a 4x3 matrix to a 4x4 matrix, extending it with the @[ 0 0 0 1 ]@ column vector
m43_to_m44 :: Num a => M43 a -> M44 a
m43_to_m44
  (V4 (V3 a b c)
      (V3 d e f)
      (V3 g h i)
      (V3 j k l)) =
  V4 (V4 a b c 0)
     (V4 d e f 0)
     (V4 g h i 0)
     (V4 j k l 1)

-- | Convert a 3x3 matrix to a 4x4 matrix extending it with 0's in the new row and column.
m33_to_m44 :: Num a => M33 a -> M44 a
m33_to_m44 (V3 r1 r2 r3) = V4 (vector r1) (vector r2) (vector r3) (point 0)

-- |The identity matrix for any dimension vector.
--
-- >>> identity :: M44 Int
-- V4 (V4 1 0 0 0) (V4 0 1 0 0) (V4 0 0 1 0) (V4 0 0 0 1)
-- >>> identity :: V3 (V3 Int)
-- V3 (V3 1 0 0) (V3 0 1 0) (V3 0 0 1)
identity :: (Num a, Traversable t, Applicative t) => t (t a)
identity = scaled (pure 1)

-- |Extract the translation vector (first three entries of the last
-- column) from a 3x4 or 4x4 matrix.
translation :: (Representable t, R3 t, R4 v) => Lens' (t (v a)) (V3 a)
translation = column _w._xyz
{-
translation f rs = aux <$> f (view _w <$> view _xyz rs)
 where aux (V3 x y z) = (_x._w .~ x) . (_y._w .~ y) . (_z._w .~ z) $ rs

-- translation :: (R3 t, R4 v, Functor f, Functor t) => (V3 a -> f (V3 a)) -> t (v a) -> f (t a)
-- translation = (. fmap (^._w)) . _xyz where
--   x ^. l = getConst (l Const x)
-}

-- |Extract a 2x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m22 :: (Representable t, R2 t, R2 v) => Lens' (t (v a)) (M22 a)
_m22 = column _xy._xy

-- |Extract a 2x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m23 :: (Representable t, R2 t, R3 v) => Lens' (t (v a)) (M23 a)
_m23 = column _xyz._xy

-- |Extract a 2x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m24 :: (Representable t, R2 t, R4 v) => Lens' (t (v a)) (M24 a)
_m24 = column _xyzw._xy

-- |Extract a 3x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m32 :: (Representable t, R3 t, R2 v) => Lens' (t (v a)) (M32 a)
_m32 = column _xy._xyz

-- |Extract a 3x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m33 :: (Representable t, R3 t, R3 v) => Lens' (t (v a)) (M33 a)
_m33 = column _xyz._xyz

-- |Extract a 3x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m34 :: (Representable t, R3 t, R4 v) => Lens' (t (v a)) (M34 a)
_m34 = column _xyzw._xyz

-- |Extract a 4x2 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m42 :: (Representable t, R4 t, R2 v) => Lens' (t (v a)) (M42 a)
_m42 = column _xy._xyzw

-- |Extract a 4x3 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m43 :: (Representable t, R4 t, R3 v) => Lens' (t (v a)) (M43 a)
_m43 = column _xyz._xyzw

-- |Extract a 4x4 matrix from a matrix of higher dimensions by dropping excess
-- rows and columns.
_m44 :: (Representable t, R4 t, R4 v) => Lens' (t (v a)) (M44 a)
_m44 = column _xyzw._xyzw

-- |2x2 matrix determinant.
--
-- >>> det22 (V2 (V2 a b) (V2 c d))
-- a * d - b * c
det22 :: Num a => M22 a -> a
det22 (V2 (V2 a b) (V2 c d)) = a * d - b * c
{-# INLINE det22 #-}

-- |3x3 matrix determinant.
--
-- >>> det33 (V3 (V3 a b c) (V3 d e f) (V3 g h i))
-- a * (e * i - f * h) - d * (b * i - c * h) + g * (b * f - c * e)
det33 :: Num a => M33 a -> a
det33 (V3 (V3 a b c)
          (V3 d e f)
          (V3 g h i)) = a * (e*i-f*h) - d * (b*i-c*h) + g * (b*f-c*e)
{-# INLINE det33 #-}

-- |4x4 matrix determinant.
det44 :: Num a => M44 a -> a
det44 (V4 (V4 i00 i01 i02 i03)
          (V4 i10 i11 i12 i13)
          (V4 i20 i21 i22 i23)
          (V4 i30 i31 i32 i33)) =
  let
    s0 = i00 * i11 - i10 * i01
    s1 = i00 * i12 - i10 * i02
    s2 = i00 * i13 - i10 * i03
    s3 = i01 * i12 - i11 * i02
    s4 = i01 * i13 - i11 * i03
    s5 = i02 * i13 - i12 * i03

    c5 = i22 * i33 - i32 * i23
    c4 = i21 * i33 - i31 * i23
    c3 = i21 * i32 - i31 * i22
    c2 = i20 * i33 - i30 * i23
    c1 = i20 * i32 - i30 * i22
    c0 = i20 * i31 - i30 * i21
  in s0 * c5 - s1 * c4 + s2 * c3 + s3 * c2 - s4 * c1 + s5 * c0
{-# INLINE det44 #-}

-- |2x2 matrix inverse.
--
-- >>> inv22 $ V2 (V2 1 2) (V2 3 4)
-- V2 (V2 (-2.0) 1.0) (V2 1.5 (-0.5))
inv22 :: Fractional a => M22 a -> M22 a
inv22 m@(V2 (V2 a b) (V2 c d)) = (1 / det) *!! V2 (V2 d (-b)) (V2 (-c) a)
  where det = det22 m
{-# INLINE inv22 #-}

-- |3x3 matrix inverse.
--
-- >>> inv33 $ V3 (V3 1 2 4) (V3 4 2 2) (V3 1 1 1)
-- V3 (V3 0.0 0.5 (-1.0)) (V3 (-0.5) (-0.75) 3.5) (V3 0.5 0.25 (-1.5))
inv33 :: Fractional a => M33 a -> M33 a
inv33 m@(V3 (V3 a b c)
            (V3 d e f)
            (V3 g h i))
  = (1 / det) *!! V3 (V3 a' b' c')
                     (V3 d' e' f')
                     (V3 g' h' i')
  where a' = cofactor (e,f,h,i)
        b' = cofactor (c,b,i,h)
        c' = cofactor (b,c,e,f)
        d' = cofactor (f,d,i,g)
        e' = cofactor (a,c,g,i)
        f' = cofactor (c,a,f,d)
        g' = cofactor (d,e,g,h)
        h' = cofactor (b,a,h,g)
        i' = cofactor (a,b,d,e)
        cofactor (q,r,s,t) = det22 (V2 (V2 q r) (V2 s t))
        det = det33 m
{-# INLINE inv33 #-}

-- | 'transpose' is just an alias for 'distribute'
--
-- > transpose (V3 (V2 1 2) (V2 3 4) (V2 5 6))
-- V2 (V3 1 3 5) (V3 2 4 6)
transpose :: (Distributive g, Functor f) => f (g a) -> g (f a)
transpose = distribute
{-# INLINE transpose #-}

-- |4x4 matrix inverse.
inv44 :: Fractional a => M44 a -> M44 a
inv44 (V4 (V4 i00 i01 i02 i03)
          (V4 i10 i11 i12 i13)
          (V4 i20 i21 i22 i23)
          (V4 i30 i31 i32 i33)) =
  let s0 = i00 * i11 - i10 * i01
      s1 = i00 * i12 - i10 * i02
      s2 = i00 * i13 - i10 * i03
      s3 = i01 * i12 - i11 * i02
      s4 = i01 * i13 - i11 * i03
      s5 = i02 * i13 - i12 * i03
      c5 = i22 * i33 - i32 * i23
      c4 = i21 * i33 - i31 * i23
      c3 = i21 * i32 - i31 * i22
      c2 = i20 * i33 - i30 * i23
      c1 = i20 * i32 - i30 * i22
      c0 = i20 * i31 - i30 * i21
      det = s0 * c5 - s1 * c4 + s2 * c3 + s3 * c2 - s4 * c1 + s5 * c0
      invDet = recip det
  in invDet *!! V4 (V4 (i11 * c5 - i12 * c4 + i13 * c3)
                       (-i01 * c5 + i02 * c4 - i03 * c3)
                       (i31 * s5 - i32 * s4 + i33 * s3)
                       (-i21 * s5 + i22 * s4 - i23 * s3))
                   (V4 (-i10 * c5 + i12 * c2 - i13 * c1)
                       (i00 * c5 - i02 * c2 + i03 * c1)
                       (-i30 * s5 + i32 * s2 - i33 * s1)
                       (i20 * s5 - i22 * s2 + i23 * s1))
                   (V4 (i10 * c4 - i11 * c2 + i13 * c0)
                       (-i00 * c4 + i01 * c2 - i03 * c0)
                       (i30 * s4 - i31 * s2 + i33 * s0)
                       (-i20 * s4 + i21 * s2 - i23 * s0))
                   (V4 (-i10 * c3 + i11 * c1 - i12 * c0)
                       (i00 * c3 - i01 * c1 + i02 * c0)
                       (-i30 * s3 + i31 * s1 - i32 * s0)
                       (i20 * s3 - i21 * s1 + i22 * s0))
{-# INLINE inv44 #-}

#if MIN_VERSION_base(4,8,0)
-- | Compute the (L, U) decomposition of a square matrix using Crout's
--   algorithm. The 'Index' of the vectors must be 'Integral'.
lu :: ( Num a
      , Fractional a
      , Foldable m
      , Traversable m
      , Applicative m
      , Additive m
      , Ixed (m a)
      , Ixed (m (m a))
      , i ~ Index (m a)
      , i ~ Index (m (m a))
      , Eq i
      , Integral i
      , a ~ IxValue (m a)
      , m a ~ IxValue (m (m a))
      , Num (m a)
      )
   => m (m a)
   -> (m (m a), m (m a))
lu a =
    let n = fromIntegral (length a)
        initU = identity
        initL = zero
        buildLVal !i !j !l !u =
            let go !k !s
                    | k == j = s
                    | otherwise = go (k+1)
                                     ( s
                                      + ( (l ^?! ix i ^?! ix k)
                                        * (u ^?! ix k ^?! ix j)
                                        )
                                      )
                s' = go 0 0
            in l & (ix i . ix j) .~ ((a ^?! ix i ^?! ix j) - s')
        buildL !i !j !l !u
            | i == n = l
            | otherwise = buildL (i+1) j (buildLVal i j l u) u
        buildUVal !i !j !l !u =
            let go !k !s
                    | k == j = s
                    | otherwise = go (k+1)
                                     ( s
                                     + ( (l ^?! ix j ^?! ix k)
                                       * (u ^?! ix k ^?! ix i)
                                       )
                                     )
                s' = go 0 0
            in u & (ix j . ix i) .~ ( ((a ^?! ix j ^?! ix i) - s')
                                    / (l ^?! ix j ^?! ix j)
                                    )
        buildU !i !j !l !u
            | i == n = u
            | otherwise = buildU (i+1) j l (buildUVal i j l u)
        buildLU !j !l !u
            | j == n = (l, u)
            | otherwise =
                let l' = buildL j j l u
                    u' = buildU j j l' u
                in buildLU (j+1) l' u'
    in buildLU 0 initL initU

-- | Compute the (L, U) decomposition of a square matrix using Crout's
--   algorithm, using the vector's 'Finite' instance to provide an index.
luFinite :: ( Num a
            , Fractional a
            , Functor m
            , Finite m
            , n ~ Size m
            , KnownNat n
            , Num (m a)
            )
         => m (m a)
         -> (m (m a), m (m a))
luFinite a =
    bimap (fmap fromV . fromV)
          (fmap fromV . fromV)
          (lu (fmap toV (toV a)))

-- | Solve a linear system with a lower-triangular matrix of coefficients with
--   forwards substitution.
forwardSub :: ( Num a
              , Fractional a
              , Foldable m
              , Additive m
              , Ixed (m a)
              , Ixed (m (m a))
              , i ~ Index (m a)
              , i ~ Index (m (m a))
              , Eq i
              , Ord i
              , Integral i
              , a ~ IxValue (m a)
              , m a ~ IxValue (m (m a))
              )
           => m (m a)
           -> m a
           -> m a
forwardSub a b =
    let n = fromIntegral (length b)
        initX = zero
        coeff !i !j !s !x
            | j == i = s
            | otherwise = coeff i (j+1) (s + ((a ^?! ix i ^?! ix j) * (x ^?! ix j))) x
        go !i !x
            | i == n = x
            | otherwise = go (i + 1) (x & ix i .~ ( ((b ^?! ix i) - coeff i 0 0 x)
                                                  / (a ^?! ix i ^?! ix i)
                                                  ))
    in go 0 initX

-- | Solve a linear system with a lower-triangular matrix of coefficients with
--   forwards substitution, using the vector's 'Finite' instance to provide an
--   index.
forwardSubFinite :: ( Num a
                    , Fractional a
                    , Foldable m
                    , n ~ Size m
                    , KnownNat n
                    , Additive m
                    , Finite m
                    )
                 => m (m a)
                 -> m a
                 -> m a
forwardSubFinite a b = fromV (forwardSub (fmap toV (toV a)) (toV b))

-- | Solve a linear system with an upper-triangular matrix of coefficients with
--   backwards substitution.
backwardSub :: ( Num a
               , Fractional a
               , Foldable m
               , Additive m
               , Ixed (m a)
               , Ixed (m (m a))
               , i ~ Index (m a)
               , i ~ Index (m (m a))
               , Eq i
               , Ord i
               , Integral i
               , a ~ IxValue (m a)
               , m a ~ IxValue (m (m a))
               )
            => m (m a)
            -> m a
            -> m a
backwardSub a b =
    let n = fromIntegral (length b)
        initX = zero
        coeff !i !j !s !x
            | j == n = s
            | otherwise = coeff i
                                (j+1)
                                (s + ((a ^?! ix i ^?! ix j) * (x ^?! ix j)))
                                x
        go !i !x
            | i < 0 = x
            | otherwise = go (i-1)
                             (x & ix i .~ ( ((b ^?! ix i) - coeff i (i+1) 0 x)
                                          / (a ^?! ix i ^?! ix i)
                                          ))
    in go (n-1) initX

-- | Solve a linear system with an upper-triangular matrix of coefficients with
--   backwards substitution, using the vector's 'Finite' instance to provide an
--   index.
backwardSubFinite :: ( Num a
                     , Fractional a
                     , Foldable m
                     , n ~ Size m
                     , KnownNat n
                     , Additive m
                     , Finite m
                     )
                  => m (m a)
                  -> m a
                  -> m a
backwardSubFinite a b = fromV (backwardSub (fmap toV (toV a)) (toV b))

-- | Solve a linear system with LU decomposition.
luSolve :: ( Num a
           , Fractional a
           , Foldable m
           , Traversable m
           , Applicative m
           , Additive m
           , Ixed (m a)
           , Ixed (m (m a))
           , i ~ Index (m a)
           , i ~ Index (m (m a))
           , Eq i
           , Integral i
           , a ~ IxValue (m a)
           , m a ~ IxValue (m (m a))
           , Num (m a)
           )
        => m (m a)
        -> m a
        -> m a
luSolve a b =
    let (l, u) = lu a
    in backwardSub u (forwardSub l b)

-- | Solve a linear system with LU decomposition, using the vector's 'Finite'
--   instance to provide an index.
luSolveFinite :: ( Num a
                 , Fractional a
                 , Functor m
                 , Finite m
                 , n ~ Size m
                 , KnownNat n
                 , Num (m a)
                 )
              => m (m a)
              -> m a
              -> m a
luSolveFinite a b = fromV (luSolve (fmap toV (toV a)) (toV b))

-- | Invert a matrix with LU decomposition.
luInv :: ( Num a
         , Fractional a
         , Foldable m
         , Traversable m
         , Applicative m
         , Additive m
         , Distributive m
         , Ixed (m a)
         , Ixed (m (m a))
         , i ~ Index (m a)
         , i ~ Index (m (m a))
         , Eq i
         , Integral i
         , a ~ IxValue (m a)
         , m a ~ IxValue (m (m a))
         , Num (m a)
         )
      => m (m a)
      -> m (m a)
luInv a =
    let n = fromIntegral (length a)
        initA' = zero
        (l, u) = lu a
        go !i !a'
            | i == n = a'
            | otherwise = let e   = zero & ix i .~ 1
                              a'r = backwardSub u (forwardSub l e)
                          in go (i+1) (a' & ix i .~ a'r)
    in transpose (go 0 initA')

-- | Invert a matrix with LU decomposition, using the vector's 'Finite' instance
--   to provide an index.
luInvFinite :: ( Num a
               , Fractional a
               , Functor m
               , Finite m
               , n ~ Size m
               , KnownNat n
               , Num (m a)
               )
            => m (m a)
            -> m (m a)
luInvFinite a = fmap fromV (fromV (luInv (fmap toV (toV a))))

-- | Compute the determinant of a matrix using LU decomposition.
luDet :: ( Num a
         , Fractional a
         , Foldable m
         , Traversable m
         , Applicative m
         , Additive m
         , Trace m
         , Ixed (m a)
         , Ixed (m (m a))
         , i ~ Index (m a)
         , i ~ Index (m (m a))
         , Eq i
         , Integral i
         , a ~ IxValue (m a)
         , m a ~ IxValue (m (m a))
         , Num (m a)
         )
      => m (m a)
      -> a
luDet a =
    let (l, u) = lu a
        p      = Foldable.foldl (*) 1
    in (p (diagonal l)) * (p (diagonal u))

-- | Compute the determinant of a matrix using LU decomposition, using the
--   vector's 'Finite' instance to provide an index.
luDetFinite :: ( Num a
               , Fractional a
               , Functor m
               , Finite m
               , n ~ Size m
               , KnownNat n
               , Num (m a)
               )
            => m (m a)
            -> a
luDetFinite = luDet . fmap toV . toV
#endif


{-# LANGUAGE BangPatterns , RankNTypes, GADTs, DataKinds #-}

{- | The 'Numerical.HBLAS.BLAS.Level2' module provides a fully general
yet type safe Level2 BLAS API.

When in doubt about the semantics of an operation,
consult your system's BLAS api documentation, or just read the documentation
for
<https://software.intel.com/sites/products/documentation/hpc/mkl/mklman/index.htm the Intel MKL BLAS distribution>

A few basic notes about how to invoke BLAS routines.

Many BLAS operations take one or more arguments of type 'Transpose'.
'Tranpose' has the following different constructors, which tell BLAS
routines what transformation to implicitly apply to an input matrix @mat@ with dimension @n x m@.

*  'NoTranspose' leaves the matrix @mat@ as is.

* 'Transpose' treats the @mat@ as being implicitly transposed, with dimension
    @m x n@. Entry @mat(i,j)@ being treated as actually being the entry
    @mat(j,i)@. For Real matrices this is also the matrix adjoint operation.
    ie @Tranpose(mat)(i,j)=mat(j,i)@

*  'ConjNoTranspose' will implicitly conjugate @mat@, which is a no op for Real ('Float' or 'Double') matrices, but for
'Complex Float' and 'Complex Double' matrices, a given matrix entry @mat(i,j)==x':+'y@
will be treated as actually being  @conjugate(mat)(i,j)=y':+'x@.

* 'ConjTranpose' will implicitly transpose and conjugate the input matrix.
ConjugateTranpose acts as matrix adjoint for both real and complex matrices.

The *gemm operations  work as follows (using 'sgemm' as an example):

* @'sgemm trLeft trRight alpha beta left right result'@, where @trLeft@ and @trRight@
are values of type 'Transpose' that respectively act on the matrices @left@ and @right@.

* the generalized matrix computation thusly formed can be viewed as being
@result = alpha * trLeft(left) * trRight(right) + beta * result@

the *gemv operations are akin to the *gemm operations, but with @right@ and @result@
being vectors rather than matrices.

the *trsv operations solve for @x@ in the equation @A x = y@ given @A@ and @y@.
The 'MatUpLo' argument determines if the matrix should be treated as upper or
lower triangular and 'MatDiag' determines if the triangular solver should treat
the diagonal of the matrix as being all 1's or not.  A general pattern of invocation
would be @'strsv' matuplo  tranposeMatA  matdiag  matrixA  xVector@.
A key detail to note is that the input vector is ALSO the result vector,
ie 'strsv' and friends updates the vector place.

-}

module Numerical.HBLAS.BLAS.Level2(
        sgbmv
        ,dgbmv
        ,cgbmv
        ,zgbmv

        ,sgemv
        ,dgemv
        ,cgemv
        ,zgemv

        ,sger
        ,dger
        ,cgerc
        ,zgerc
        ,cgeru
        ,zgeru

        ,chbmv
        ,zhbmv

        ,chemv
        ,zhemv

        ,cher
        ,zher
        ,cher2
        ,zher2

        ,chpmv
        ,zhpmv

        ,chpr
        ,zhpr
        ,chpr2
        ,zhpr2

        ,ssbmv
        ,dsbmv

        ,sspmv
        ,dspmv

        ,sspr
        ,dspr
        ,sspr2
        ,dspr2

        ,ssymv
        ,dsymv

        ,ssyr
        ,dsyr
        ,ssyr2
        ,dsyr2

        ,stbmv
        ,dtbmv
        ,ctbmv
        ,ztbmv

        ,stbsv
        ,dtbsv
        ,ctbsv
        ,ztbsv

        ,stpmv
        ,dtpmv
        ,ctpmv
        ,ztpmv

        ,stpsv
        ,dtpsv
        ,ctpsv
        ,ztpsv

        ,strmv
        ,dtrmv
        ,ctrmv
        ,ztrmv

        ,strsv
        ,dtrsv
        ,ctrsv
        ,ztrsv
            ) where

import Numerical.HBLAS.UtilsFFI
import Numerical.HBLAS.BLAS.FFI.Level2
import Numerical.HBLAS.BLAS.Internal.Level2
import Control.Monad.Primitive
import Data.Complex

sgbmv :: PrimMonad m => GbmvFun Float orient (PrimState m) m
sgbmv = gbmvAbstraction "sgbmv" cblas_sgbmv_safe cblas_sgbmv_unsafe (\x f -> f x)

dgbmv :: PrimMonad m => GbmvFun Double orient (PrimState m) m
dgbmv = gbmvAbstraction "dgbmv" cblas_dgbmv_safe cblas_dgbmv_unsafe (\x f -> f x)

cgbmv :: PrimMonad m => GbmvFun (Complex Float) orient (PrimState m) m
cgbmv = gbmvAbstraction "cgbmv" cblas_cgbmv_safe cblas_cgbmv_unsafe withRStorable_

zgbmv :: PrimMonad m => GbmvFun (Complex Double) orient (PrimState m) m
zgbmv = gbmvAbstraction "zgbmv" cblas_zgbmv_safe cblas_zgbmv_unsafe withRStorable_

sgemv :: PrimMonad m => GemvFun Float orient (PrimState m) m
sgemv = gemvAbstraction "sgemv" cblas_sgemv_safe cblas_sgemv_unsafe (flip ($))

dgemv :: PrimMonad m => GemvFun Double orient (PrimState m) m
dgemv = gemvAbstraction "dgemv" cblas_dgemv_safe cblas_dgemv_unsafe (flip ($))

cgemv :: PrimMonad m => GemvFun (Complex Float) orient (PrimState m) m
cgemv = gemvAbstraction "cgemv" cblas_cgemv_safe cblas_cgemv_unsafe withRStorable_

zgemv :: PrimMonad m => GemvFun (Complex Double) orient (PrimState m) m
zgemv = gemvAbstraction "zgemv" cblas_zgemv_safe cblas_zgemv_unsafe withRStorable_

sger :: PrimMonad m => GerFun Float orient (PrimState m) m
sger = gerAbstraction "sger" cblas_sger_safe cblas_sger_unsafe (\x f -> f x)

dger :: PrimMonad m => GerFun Double orient (PrimState m) m
dger = gerAbstraction "dger" cblas_dger_safe cblas_dger_unsafe (\x f -> f x)

cgerc :: PrimMonad m => GerFun (Complex Float) orient (PrimState m) m
cgerc = gerAbstraction "cgerc" cblas_cgerc_safe cblas_cgerc_unsafe withRStorable_

zgerc :: PrimMonad m => GerFun (Complex Double) orient (PrimState m) m
zgerc = gerAbstraction "zgerc" cblas_zgerc_safe cblas_zgerc_unsafe withRStorable_

cgeru :: PrimMonad m => GerFun (Complex Float) orient (PrimState m) m
cgeru = gerAbstraction "cgeru" cblas_cgeru_safe cblas_cgeru_unsafe withRStorable_

zgeru :: PrimMonad m => GerFun (Complex Double) orient (PrimState m) m
zgeru = gerAbstraction "zgeru" cblas_zgeru_safe cblas_zgeru_unsafe withRStorable_

chbmv :: PrimMonad m => HbmvFun (Complex Float) orient (PrimState m) m
chbmv = hbmvAbstraction "chbmv" cblas_chbmv_safe cblas_chbmv_unsafe withRStorable_

zhbmv :: PrimMonad m => HbmvFun (Complex Double) orient (PrimState m) m
zhbmv = hbmvAbstraction "zhbmv" cblas_zhbmv_safe cblas_zhbmv_unsafe withRStorable_

chemv :: PrimMonad m => HemvFun (Complex Float) orient (PrimState m) m
chemv = hemvAbstraction "chemv" cblas_chemv_safe cblas_chemv_unsafe withRStorable_

zhemv :: PrimMonad m => HemvFun (Complex Double) orient (PrimState m) m
zhemv = hemvAbstraction "zhemv" cblas_zhemv_safe cblas_zhemv_unsafe withRStorable_

cher :: PrimMonad m => HerFun Float (Complex Float) orient (PrimState m) m
cher = herAbstraction "cher" cblas_cher_safe cblas_cher_unsafe (\x f -> f x)

zher :: PrimMonad m => HerFun Double (Complex Double) orient (PrimState m) m
zher = herAbstraction "zher" cblas_zher_safe cblas_zher_unsafe (\x f -> f x)

cher2 :: PrimMonad m => Her2Fun (Complex Float) orient (PrimState m) m
cher2 = her2Abstraction "cher2" cblas_cher2_safe cblas_cher2_unsafe withRStorable_

zher2 :: PrimMonad m => Her2Fun (Complex Double) orient (PrimState m) m
zher2 = her2Abstraction "zher2" cblas_zher2_safe cblas_zher2_unsafe withRStorable_

chpmv :: PrimMonad m => HpmvFun (Complex Float) orient (PrimState m) m
chpmv = hpmvAbstraction "chpmv" cblas_chpmv_safe cblas_chpmv_unsafe withRStorable_

zhpmv :: PrimMonad m => HpmvFun (Complex Double) orient (PrimState m) m
zhpmv = hpmvAbstraction "zhpmv" cblas_zhpmv_safe cblas_zhpmv_unsafe withRStorable_

chpr :: PrimMonad m => HprFun Float (Complex Float) orient (PrimState m) m
chpr = hprAbstraction "chpr" cblas_chpr_safe cblas_chpr_unsafe (\x f -> f x)

zhpr :: PrimMonad m => HprFun Double (Complex Double) orient (PrimState m) m
zhpr = hprAbstraction "zhpr" cblas_zhpr_safe cblas_zhpr_unsafe (\x f -> f x)

chpr2 :: PrimMonad m => Hpr2Fun (Complex Float) orient (PrimState m) m
chpr2 = hpr2Abstraction "chpr2" cblas_chpr2_safe cblas_chpr2_unsafe withRStorable_

zhpr2 :: PrimMonad m => Hpr2Fun (Complex Double) orient (PrimState m) m
zhpr2 = hpr2Abstraction "zhpr2" cblas_zhpr2_safe cblas_zhpr2_unsafe withRStorable_

ssbmv :: PrimMonad m => SbmvFun Float orient (PrimState m) m
ssbmv = sbmvAbstraction "ssbmv" cblas_ssbmv_safe cblas_ssbmv_unsafe (\x f -> f x)

dsbmv :: PrimMonad m => SbmvFun Double orient (PrimState m) m
dsbmv = sbmvAbstraction "dsbmv" cblas_dsbmv_safe cblas_dsbmv_unsafe (\x f -> f x)

sspmv :: PrimMonad m => SpmvFun Float orient (PrimState m) m
sspmv = spmvAbstraction "sspmv" cblas_sspmv_safe cblas_sspmv_unsafe (\x f -> f x)

dspmv :: PrimMonad m => SpmvFun Double orient (PrimState m) m
dspmv = spmvAbstraction "dspmv" cblas_dspmv_safe cblas_dspmv_unsafe (\x f -> f x)

sspr :: PrimMonad m => SprFun Float orient (PrimState m) m
sspr = sprAbstraction "sspr" cblas_sspr_safe cblas_sspr_unsafe (\x f -> f x)

dspr :: PrimMonad m => SprFun Double orient (PrimState m) m
dspr = sprAbstraction "dspr" cblas_dspr_safe cblas_dspr_unsafe (\x f -> f x)

sspr2 :: PrimMonad m => Spr2Fun Float orient (PrimState m) m
sspr2 = spr2Abstraction "sspr2" cblas_sspr2_safe cblas_sspr2_unsafe (\x f -> f x)

dspr2 :: PrimMonad m => Spr2Fun Double orient (PrimState m) m
dspr2 = spr2Abstraction "dspr2" cblas_dspr2_safe cblas_dspr2_unsafe (\x f -> f x)

ssymv :: PrimMonad m => SymvFun Float orient (PrimState m) m
ssymv = symvAbstraction "ssymv" cblas_ssymv_safe cblas_ssymv_unsafe (\x f -> f x)

dsymv :: PrimMonad m => SymvFun Double orient (PrimState m) m
dsymv = symvAbstraction "dsymv" cblas_dsymv_safe cblas_dsymv_unsafe (\x f -> f x)

ssyr :: PrimMonad m => SyrFun Float orient (PrimState m) m
ssyr = syrAbstraction "ssyr" cblas_ssyr_safe cblas_ssyr_unsafe (\x f -> f x)

dsyr :: PrimMonad m => SyrFun Double orient (PrimState m) m
dsyr = syrAbstraction "dsyr" cblas_dsyr_safe cblas_dsyr_unsafe (\x f -> f x)

ssyr2 :: PrimMonad m => Syr2Fun Float orient (PrimState m) m
ssyr2 = syr2Abstraction "ssyr2" cblas_ssyr2_safe cblas_ssyr2_unsafe (\x f -> f x)

dsyr2 :: PrimMonad m => Syr2Fun Double orient (PrimState m) m
dsyr2 = syr2Abstraction "dsyr2" cblas_dsyr2_safe cblas_dsyr2_unsafe (\x f -> f x)

stbmv :: PrimMonad m => TbmvFun Float orient (PrimState m) m
stbmv = tbmvAbstraction "stbmv" cblas_stbmv_safe cblas_stbmv_unsafe

dtbmv :: PrimMonad m => TbmvFun Double orient (PrimState m) m
dtbmv = tbmvAbstraction "dtbmv" cblas_dtbmv_safe cblas_dtbmv_unsafe

ctbmv :: PrimMonad m => TbmvFun (Complex Float) orient (PrimState m) m
ctbmv = tbmvAbstraction "ctbmv" cblas_ctbmv_safe cblas_ctbmv_unsafe

ztbmv :: PrimMonad m => TbmvFun (Complex Double) orient (PrimState m) m
ztbmv = tbmvAbstraction "ztbmv" cblas_ztbmv_safe cblas_ztbmv_unsafe

stbsv :: PrimMonad m => TbsvFun Float orient (PrimState m) m
stbsv = tbsvAbstraction "stbsv" cblas_stbsv_safe cblas_stbsv_unsafe

dtbsv :: PrimMonad m => TbsvFun Double orient (PrimState m) m
dtbsv = tbsvAbstraction "dtbsv" cblas_dtbsv_safe cblas_dtbsv_unsafe

ctbsv :: PrimMonad m => TbsvFun (Complex Float) orient (PrimState m) m
ctbsv = tbsvAbstraction "ctbsv" cblas_ctbsv_safe cblas_ctbsv_unsafe

ztbsv :: PrimMonad m => TbsvFun (Complex Double) orient (PrimState m) m
ztbsv = tbsvAbstraction "ztbsv" cblas_ztbsv_safe cblas_ztbsv_unsafe

stpmv :: PrimMonad m => TpmvFun Float orient (PrimState m) m
stpmv = tpmvAbstraction "stpmv" cblas_stpmv_safe cblas_stpmv_unsafe

dtpmv :: PrimMonad m => TpmvFun Double orient (PrimState m) m
dtpmv = tpmvAbstraction "dtpmv" cblas_dtpmv_safe cblas_dtpmv_unsafe

ctpmv :: PrimMonad m => TpmvFun (Complex Float) orient (PrimState m) m
ctpmv = tpmvAbstraction "ctpmv" cblas_ctpmv_safe cblas_ctpmv_unsafe

ztpmv :: PrimMonad m => TpmvFun (Complex Double) orient (PrimState m) m
ztpmv = tpmvAbstraction "ztpmv" cblas_ztpmv_safe cblas_ztpmv_unsafe

stpsv :: PrimMonad m => TpsvFun Float orient (PrimState m) m
stpsv = tpsvAbstraction "stpsv" cblas_stpsv_safe cblas_stpsv_unsafe

dtpsv :: PrimMonad m => TpsvFun Double orient (PrimState m) m
dtpsv = tpsvAbstraction "dtpsv" cblas_dtpsv_safe cblas_dtpsv_unsafe

ctpsv :: PrimMonad m => TpsvFun (Complex Float) orient (PrimState m) m
ctpsv = tpsvAbstraction "ctpsv" cblas_ctpsv_safe cblas_ctpsv_unsafe

ztpsv :: PrimMonad m => TpsvFun (Complex Double) orient (PrimState m) m
ztpsv = tpsvAbstraction "ztpsv" cblas_ztpsv_safe cblas_ztpsv_unsafe

strmv :: PrimMonad m => TrmvFun Float orient (PrimState m) m
strmv = trmvAbstraction "strmv" cblas_strmv_safe cblas_strmv_unsafe

dtrmv :: PrimMonad m => TrmvFun Double orient (PrimState m) m
dtrmv = trmvAbstraction "dtrmv" cblas_dtrmv_safe cblas_dtrmv_unsafe

ctrmv :: PrimMonad m => TrmvFun (Complex Float) orient (PrimState m) m
ctrmv = trmvAbstraction "ctrmv" cblas_ctrmv_safe cblas_ctrmv_unsafe

ztrmv :: PrimMonad m => TrmvFun (Complex Double) orient (PrimState m) m
ztrmv = trmvAbstraction "ztrmv" cblas_ztrmv_safe cblas_ztrmv_unsafe

strsv :: PrimMonad m => TrsvFun Float orient (PrimState m) m
strsv = trsvAbstraction "strsv" cblas_strsv_safe cblas_strsv_unsafe

dtrsv :: PrimMonad m => TrsvFun Double orient (PrimState m) m
dtrsv = trsvAbstraction "dtrsv" cblas_dtrsv_safe cblas_dtrsv_unsafe

ctrsv :: PrimMonad m => TrsvFun (Complex Float) orient (PrimState m) m
ctrsv = trsvAbstraction "ctrsv" cblas_ctrsv_safe cblas_ctrsv_unsafe

ztrsv :: PrimMonad m => TrsvFun (Complex Double) orient (PrimState m) m
ztrsv = trsvAbstraction "ztrsv" cblas_ztrsv_safe cblas_ztrsv_unsafe


{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Test.QuickCheck.LinearAlgebra
-- Copyright  : Copyright (c) , Patrick Perry <patperry@stanford.edu>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability  : experimental
--
--
-- Test generators for linear algebra types.
--

module Test.QuickCheck.LinearAlgebra (
    -- Testable element types
    TestElem(..),

    -- * Generating random objects
    -- ** Dimensions
    dim,
    dim2,
    Dim(..),
    Dim2(..),

    -- ** Indices
    index,
    index2,
    Index(..),
    Index2(..),

    -- ** Elements
    elem,
    elems,
    -- realElem,
    -- realElems,

    -- ** Association lists
    assocs,
    assocs2,
    Assocs(..),
    Assocs2(..),
    -- bandedAssocs,

    -- ** Vectors
    vector,
    VectorPair(..),
    VectorTriple(..),
    VectorList(..),
    WeightedVectorList(..),
    NonEmptyVectorList(..),
    NonEmptyWeightedVectorList(..),

    --  ** Matrices
    matrix,
    MatrixPair(..),
    MatrixTriple(..),
    -- hermMatrix,
    -- triMatrix,

    --  Banded matrices
    -- bandwidths,
    -- banded,
    -- bandedWith,
    -- hermBanded,
    -- triBanded,

    -- BandedAssocs(..),

    ) where

import Prelude hiding ( elem )

-- import BLAS.Types( UpLoEnum(..), DiagEnum(..) )
import Control.Monad
import Data.Complex( magnitude )
import Data.Maybe( fromJust )

import Test.QuickCheck hiding ( vector )
import qualified Test.QuickCheck as QC

import Numeric.LinearAlgebra.Types
import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Matrix( Matrix )
import qualified Numeric.LinearAlgebra.Matrix as M
-- import Data.Matrix.Banded( Banded, maybeBandedFromMatrixStorage )
-- import Data.Matrix.Banded.ST( runSTBanded, unsafeThawBanded,
--     diagViewBanded )

class TestElem e where
    maybeToReal :: e -> Maybe Double
    toReal :: e -> Double
    norm1 :: e -> Double
    norm :: e -> Double
    conjugate :: e -> e

    toReal = fromJust . maybeToReal

instance (TestElem Double) where
    maybeToReal = Just
    toReal = id
    norm1 = abs
    norm = abs
    conjugate = id

instance (TestElem (Complex Double)) where
    maybeToReal (x :+ y) | y == 0    = Just x
                         | otherwise = Nothing
    norm1 (x :+ y) = abs x + abs y
    norm z = magnitude z
    conjugate (x :+ y) = x :+ (-y)

{-
-- | Element types that can be tested with QuickCheck properties.
class (Storable e, Arbitrary e, CoArbitrary e) => TestElem e where
    -- | Inicates whether or not the value should be used in tests.  For
    -- 'Double's, @isTestElem e@ is defined as
    -- @not (isNaN e || isInfinite e || isDenormalized e)@.
    isTestElem :: e -> Bool

instance TestElem Double where
    isTestElem e = not (isNaN e || isInfinite e || isDenormalized e)
    {-# INLINE isTestElem #-}

instance TestElem (Complex Double) where
    isTestElem (x :+ y) = isTestElem x && isTestElem y
    {-# INLINE isTestElem #-}
-}

minit:: Int ->Int
minit  n=  n `mod` 20

-- | Generate a random element.
elem :: (Arbitrary e) => Gen e
elem = arbitrary

-- | Generate a list of elements suitable for testing with.
elems :: (Arbitrary e) => Int -> Gen [e]
elems n = replicateM n elem

-- Generate a random element that has no imaginary part.
-- realElem :: (TestElem e) => Gen e
-- realElem = liftM fromReal elem

-- Generate a list of elements for testing that have no imaginary part.
-- realElems :: (TestElem e) => Int -> Gen [e]
-- realElems n = replicateM n realElem

-- | Get an appropriate dimension for a random vector
dim :: Gen Int
dim = sized $ \s -> do
    (NonNegative n) <- resize (s `div` 4) $ arbitrary
    return (minit n)

-- | Get an appropriate dimension for a random matrix
dim2 :: Gen (Int,Int)
dim2 = sized $ \s -> do
    m <- resize (s `div` 2) $ dim
    n <- resize (s `div` 2) $ dim
    return (m,n)

-- | A vector dimension.
newtype Dim = Dim Int
  deriving (Eq, Ord, Num, Integral, Real, Enum, Show, Read)

instance Arbitrary Dim where
    arbitrary = sized $ \s -> do
        (NonNegative n) <- resize (s `div` 4) $ arbitrary
        return $ Dim n

    shrink (Dim a) =
        [ Dim a' | (NonNegative a') <- shrink (NonNegative a) ]

-- | A matrix dimension.
newtype Dim2 = Dim2 (Int,Int)
  deriving (Eq, Show, Read)

instance Arbitrary Dim2 where
    arbitrary = do
        mn <- dim2
        return $ Dim2 mn

-- | Given a dimension generate a valid index.  The dimension must be positive.
index :: Int -> Gen Int
index n | n <= 0 =
            error $ "index " ++ (show n) ++ ":"
                  ++ " dimension must be positive (QuickCheck error)"
        | otherwise =
            choose (0,n-1)

-- | Given a matrix dimension generate a valid index.
index2 :: (Int,Int) -> Gen (Int,Int)
index2 (m,n) = do
    i <- index m
    j <- index n
    return (i,j)

-- | A dimension and a valid index for it.
data Index = Index Int Int deriving (Eq,Show)
instance Arbitrary Index where
    arbitrary = do
        n <- (1+) `fmap` dim
        i <- index n
        return $ Index n i

-- | A matrix dimension and a valid index for it.
data Index2 = Index2 (Int,Int) (Int,Int) deriving (Eq,Show)
instance Arbitrary Index2 where
    arbitrary = do
        (m',n') <- dim2
        let mn = (m' + 1, n' + 1)
        ij <- index2 mn
        return $ Index2 mn ij

-- | Generate an associations list for a vector of the given dimension.
assocs :: (Arbitrary e) => Int -> Gen [(Int,e)]
assocs n | n == 0    = return []
         | otherwise = do
    l <- choose(0, 2*n)
    is <- replicateM l $ index n
    es <- elems l
    return $ zip is es

-- | Generate an associations list for a matrix of the given shape.
assocs2 :: (Arbitrary e) => (Int,Int) -> Gen [((Int,Int),e)]
assocs2 (m,n) | m*n == 0  = return []
              | otherwise = do
    l <- choose(0, 2*m*n)
    is <- replicateM l $ index2 (m,n)
    es <- elems l
    return $ zip is es

-- | A dimension and an associations list.
data Assocs e = Assocs Int [(Int,e)] deriving (Eq,Show)

instance (Arbitrary e) => Arbitrary (Assocs e) where
    arbitrary = do
        n   <- dim
        ies <- assocs n
        return $ Assocs n ies

    shrink (Assocs n ies) =
        [ Assocs n' $ filter ((< n') . fst) ies
        | n' <- shrink n
        ] ++
        [ Assocs n ies'
        | ies' <- shrink ies
        ]

-- | A shape and an associations list.
data Assocs2 e = Assocs2 (Int,Int) [((Int,Int),e)] deriving (Eq,Show)
instance (Arbitrary e) => Arbitrary (Assocs2 e) where
    arbitrary = do
        mn   <- dim2
        ies  <- assocs2 mn
        return $ Assocs2 mn ies

-- | Generate a random vector of the given size.
vector :: (Arbitrary e, Storable e) => Int -> Gen (Vector e)
vector n = do
    es <- elems n
    return $ V.fromList n es

instance (Arbitrary e, Storable e) => Arbitrary (Vector e) where
    arbitrary = dim >>= vector

    shrink x =
        [ V.slice 0 n x
        | (NonNegative n) <- shrink (NonNegative $ V.dim x)
        ]

-- | Two vectors with the same dimension.
data VectorPair e f =
    VectorPair (Vector e) (Vector f) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f) =>
    Arbitrary (VectorPair e f) where
        arbitrary = do
            x <- arbitrary
            y <- vector (V.dim x)
            return $ VectorPair x y

        shrink (VectorPair x y) =
            [ VectorPair (V.slice 0 n' x) (V.slice 0 n' y)
            | n' <- shrink (V.dim x)
            ]

-- | Three vectors with the same dimension.
data VectorTriple e f g =
    VectorTriple (Vector e) (Vector f) (Vector g) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f,
          Arbitrary g, Storable g) =>
    Arbitrary (VectorTriple e f g) where
        arbitrary = do
            x <- arbitrary
            y <- vector (V.dim x)
            z <- vector (V.dim x)
            return $ VectorTriple x y z

-- | A nonempty list of vectors with the same dimension.
data NonEmptyVectorList e = NonEmptyVectorList Int [Vector e] deriving (Eq, Show)
instance (Arbitrary e, Storable e) => Arbitrary (NonEmptyVectorList e) where
    arbitrary = do
        x <- arbitrary
        n <- choose (0,20)
        let p = V.dim x
        xs <- replicateM n $ vector p
        return $ NonEmptyVectorList p $ x:xs

-- | A nonempty list of (weight, vector) pairs, with the weights all non-negative
-- and the vectors all having the same dimension.
data NonEmptyWeightedVectorList e = NonEmptyWeightedVectorList Int [(e, Vector e)]
    deriving (Eq, Show)
instance (Arbitrary e, Storable e, Num e) => Arbitrary (NonEmptyWeightedVectorList e) where
    arbitrary = do
        (NonEmptyVectorList p xs) <- arbitrary
        ws <- replicateM (length xs) $ fmap abs arbitrary
        return $ NonEmptyWeightedVectorList p $ zip ws xs

-- | A list of vectors with the same dimension.
data VectorList e = VectorList Int [Vector e] deriving (Eq, Show)
instance (Arbitrary e, Storable e) => Arbitrary (VectorList e) where
    arbitrary = do
        (NonEmptyVectorList p (_:xs)) <- arbitrary
        return $ VectorList p xs

-- | A list of (weight, vector) pairs, with the weights all non-negative
-- and the vectors all having the same dimension.
data WeightedVectorList e = WeightedVectorList Int [(e, Vector e)]
    deriving (Eq, Show)
instance (Arbitrary e, Storable e, Num e) => Arbitrary (WeightedVectorList e) where
    arbitrary = do
        (NonEmptyWeightedVectorList p (_:wxs)) <- arbitrary
        return $ WeightedVectorList p wxs

-- | Generate a random matrix of the given size.
matrix :: (Arbitrary e, Storable e) => (Int,Int) -> Gen (Matrix e)
matrix (m,n) =
    oneof [ raw, sub ]
  where
    raw = do
        es <- elems $ m * n
        return $ M.fromList (m,n) es
    sub = do
        m' <- choose (m, 2*m)
        es <- elems $ m' * n
        return $ M.slice (0,0) (m,n) (M.fromList (m',n) es)

instance (Arbitrary e, Storable e) => Arbitrary (Matrix e) where
    arbitrary = dim2 >>= matrix

-- | Two matrices with the same dimension.
data MatrixPair e f =
    MatrixPair (Matrix e) (Matrix f) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f) =>
    Arbitrary (MatrixPair e f) where
        arbitrary = do
            x <- arbitrary
            y <- matrix (M.dim x)
            return $ MatrixPair x y

-- | Three matrices with the same dimension.
data MatrixTriple e f g =
    MatrixTriple (Matrix e) (Matrix f) (Matrix g) deriving (Eq, Show)
instance (Arbitrary e, Storable e, Arbitrary f, Storable f,
          Arbitrary g, Storable g) =>
    Arbitrary (MatrixTriple e f g) where
        arbitrary = do
            x <- arbitrary
            y <- matrix (M.dim x)
            z <- matrix (M.dim x)
            return $ MatrixTriple x y z

instance Arbitrary Trans where
    arbitrary = elements [ NoTrans, Trans, ConjTrans ]
{-

-- | Generate a triangular dense matrix.
triMatrix :: (TestElem e) => (Int,Int) -> Gen (Tri Matrix e)
triMatrix (m,n) = do
    a <- matrix (m,n)
    u <- QC.elements [ Lower, Upper  ]
    d <- QC.elements [ Unit, NonUnit ]
    return $ Tri u d a

-- | Generate a Hermitian dense matrix.
hermMatrix :: (TestElem e) => Int -> Gen (Herm Matrix e)
hermMatrix n = do
    a <- matrix (n,n)
    d <- realElems n
    let a' = runSTMatrix $ do
                 ma <- unsafeThawMatrix a
                 setElems (diagView ma 0) d
                 return ma
    u <- QC.elements [ Lower, Upper ]
    return $ Herm u a'

rawMatrix :: (TestElem e) => (Int,Int) -> Gen (Matrix e)
rawMatrix (m,n) = do
    es <- elems (m*n)
    return $ M.fromList (m,n) es

subMatrix :: (TestElem e) => (Int,Int) -> Gen (SubMatrix e)
subMatrix (m,n) =
    oneof [ rawSubMatrix (m,n)
          , rawSubMatrix (n,m) >>= \(SubMatrix a (i,j) (m',n')) ->
                return $ SubMatrix (herm a) (j,i) (n',m')
          ]

rawSubMatrix :: (TestElem e) => (Int,Int) -> Gen (SubMatrix e)
rawSubMatrix (m,n) = do
    i <- choose (0,5)
    j <- choose (0,5)
    e <- choose (0,5)
    f <- choose (0,5)
    x <- rawMatrix (i+m+e, j+n+f)
    return $ SubMatrix x (i,j) (m,n)

instance (TestElem e) => Arbitrary (SubMatrix e) where
    arbitrary = do
        (m,n) <- shape
        (SubMatrix a ij mn) <- subMatrix (m,n)
        return $ SubMatrix a ij mn

-- | Generate valid bandwidth for a given matrix dimension size
bandwidth :: Int -> Gen Int
bandwidth n = if n == 0 then return 0 else choose (0,n-1)

-- | Generate valid bandwidths for the given matrix shape.
bandwidths :: (Int,Int) -> Gen (Int,Int)
bandwidths (m,n) = liftM2 (,) (bandwidth m) (bandwidth n)

-- | Generate a random banded matrix of the given shape.
banded :: (TestElem e) => (Int,Int) -> Gen (Banded e)
banded mn = do
    lu <- bandwidths mn
    bandedWith lu mn

-- | Generate a random banded matrix with the given bandwidths.
bandedWith :: (TestElem e)
           => (Int,Int) -> (Int,Int) -> Gen (Banded e)
bandedWith lu mn = frequency [ (3, rawBanded mn lu)
                             , (2, hermedBanded mn lu)
                             ]

-- | Generate a triangular banded matrix.
triBanded :: (TestElem e) => Int -> Gen (Tri Banded e)
triBanded n = do
    a <- banded (n,n)
    u <- QC.elements [ Lower, Upper  ]
    d <- QC.elements [ Unit, NonUnit ]
    return $ Tri u d a

-- | Generate a Hermitian banded matrix.
hermBanded :: (TestElem e) => Int -> Gen (Herm Banded e)
hermBanded n = do
    a <- banded (n,n)
    d <- realElems n
    let a' = runSTBanded $ do
                 ma <- unsafeThawBanded a
                 setElems (diagViewBanded ma 0) d
                 return ma
    u <- QC.elements [ Lower, Upper ]
    return $ Herm u a'

rawBanded :: (TestElem e) =>
    (Int,Int) -> (Int,Int) -> Gen (Banded e)
rawBanded (m,n) (kl,ku) =
    let bw = kl+ku+1
    in do
        a <- frequency [ (2, rawMatrix (bw,n))
                       , (1, rawSubMatrix (bw,n) >>= \(SubMatrix b ij _) ->
                                 return $ submatrix b ij (bw,n))
                       ]
        return $ fromJust (maybeBandedFromMatrixStorage (m,n) (kl,ku) a)

hermedBanded :: (TestElem e) =>
    (Int,Int) -> (Int,Int) -> Gen (Banded e)
hermedBanded (m,n) (kl,ku) = do
    x <- rawBanded (n,m) (ku,kl)
    return $ herm x

-- | Generate an associations list for a banded matrix of the given shape
-- and bandwidths.
bandedAssocs :: (TestElem e) => (Int,Int) -> (Int,Int) -> Gen [((Int,Int),e)]
bandedAssocs (m,n) (kl,ku) | m*n == 0  = return []
                           | otherwise = do
    (Nat l) <- arbitrary
    ijs     <- replicateM l $ index2 (kl+1+ku,n)
    let ijs' = mapMaybe (\(i,j) -> let i' = i - j - ku in
                                   if 0 <= i' && i' < m then Just (i',j)
                                                        else Nothing    ) ijs
    es      <- replicateM l elem
    return $ zip ijs' es

-- | A shape, bandwidths, and an associations list.
data BandedAssocs e = BandedAssocs (Int,Int) (Int,Int) [((Int,Int),e)] deriving (Eq,Show)
instance (TestElem e) => Arbitrary (BandedAssocs e) where
    arbitrary = do
        mn  <- shape
        bw  <- bandwidths mn
        ies <- bandedAssocs mn bw
        return $ BandedAssocs mn bw ies
-}


{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE DeriveGeneric        #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ViewPatterns         #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

import           Control.DeepSeq
import           Criterion.Main
import           Criterion.Types
import           Data.Char
import           Data.Functor.Identity
import           Data.Time
import           GHC.Generics                 (Generic)
import           GHC.TypeLits
import           Lens.Micro
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.Backprop.Class
import           Numeric.LinearAlgebra.Static
import           System.Directory
import qualified Data.Vector                  as V
import qualified Numeric.LinearAlgebra        as HM
import qualified System.Random.MWC            as MWC

type family HKD f a where
    HKD Identity a = a
    HKD f        a = f a

data Layer' i o f =
    Layer { _lWeights :: !(HKD f (L o i))
          , _lBiases  :: !(HKD f (R o))
          }
  deriving (Generic)

type Layer i o = Layer' i o Identity

deriving instance (KnownNat i, KnownNat o) => Show (Layer i o)
instance NFData (Layer i o)

makeLenses ''Layer'

data Network' i h1 h2 o f =
    Net { _nLayer1 :: !(HKD f (Layer i  h1))
        , _nLayer2 :: !(HKD f (Layer h1 h2))
        , _nLayer3 :: !(HKD f (Layer h2 o ))
        }
  deriving (Generic)

type Network i h1 h2 o = Network' i h1 h2 o Identity

deriving instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Show (Network i h1 h2 o)
instance NFData (Network i h1 h2 o)

makeLenses ''Network'

main :: IO ()
main = do
    g     <- MWC.initialize
           . V.fromList
           . map (fromIntegral . ord)
           $ "hello world"
    test0 <- MWC.uniformR @(R 784, R 10) ((0,0),(1,1)) g
    net0  <- MWC.uniformR @(Network 784 300 100 10) (-0.5, 0.5) g
    t     <- getZonedTime
    let tstr = formatTime defaultTimeLocale "%Y%m%d-%H%M%S" t
    createDirectoryIfMissing True "bench-results"
    defaultMainWith defaultConfig
          { reportFile = Just $ "bench-results/mnist-bench_" ++ tstr ++ ".html"
          , timeLimit  = 10
          } [
        bgroup "gradient"
          [ let runTest x y     = gradNetManual x y net0
            in  bench "manual"  $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (netErr x y) net0
            in  bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (netErrHKD x y) net0
            in  bench "bp-hkd"  $ nf (uncurry runTest) test0
          , let runTest x y     = gradBP (\n' -> netErrHybrid n' y x) net0
            in  bench "hybrid"  $ nf (uncurry runTest) test0
          ]
      , bgroup "descent"
          [ let runTest x y     = trainStepManual 0.02 x y net0
            in  bench "manual"  $ nf (uncurry runTest) test0
          , let runTest x y     = trainStep 0.02 x y net0
            in  bench "bp-lens" $ nf (uncurry runTest) test0
          , let runTest x y     = trainStepHKD 0.02 x y net0
            in  bench "bp-hkd"  $ nf (uncurry runTest) test0
          , let runTest x y     = trainStepHybrid 0.02 x y net0
            in  bench "hybrid"  $ nf (uncurry runTest) test0
          ]
      , bgroup "run"
          [ let runTest         = runNetManual net0
            in  bench "manual"  $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetwork` x) net0
            in  bench "bp-lens" $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetworkHKD` x) net0
            in  bench "bp-hkd"  $ nf runTest (fst test0)
          , let runTest x       = evalBP (`runNetHybrid` x) net0
            in  bench "hybrid"  $ nf runTest (fst test0)
          ]
      ]

-- ------------------------------
-- - "Backprop" Lens Mode       -
-- ------------------------------

runLayer
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (Layer i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLayer l x = (l ^^. lWeights) #>! x + (l ^^. lBiases)
{-# INLINE runLayer #-}

softMax :: (KnownNat n, Reifies s W) => BVar s (R n) -> BVar s (R n)
softMax x = konst' (1 / sumElements' expx) * expx
  where
    expx = exp x
{-# INLINE softMax #-}

runNetwork
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetwork n = softMax
             . runLayer (n ^^. nLayer3)
             . logistic
             . runLayer (n ^^. nLayer2)
             . logistic
             . runLayer (n ^^. nLayer1)
             . auto
{-# INLINE runNetwork #-}

crossEntropy
    :: (KnownNat n, Reifies s W)
    => R n
    -> BVar s (R n)
    -> BVar s Double
crossEntropy t r = negate $ log r <.>! auto t
{-# INLINE crossEntropy #-}

netErr
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => R i
    -> R o
    -> BVar s (Network i h1 h2 o)
    -> BVar s Double
netErr x t n = crossEntropy t (runNetwork n x)
{-# INLINE netErr #-}

trainStep
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStep r !x !t !n = n - realToFrac r * gradBP (netErr x t) n
{-# INLINE trainStep #-}

-- ------------------------------
-- - "Backprop" HKD Mode        -
-- ------------------------------

runLayerHKD
    :: (KnownNat i, KnownNat o, Reifies s W)
    => BVar s (Layer i o)
    -> BVar s (R i)
    -> BVar s (R o)
runLayerHKD (splitBV->Layer w b) x = w #>! x + b
{-# INLINE runLayerHKD #-}

runNetworkHKD
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetworkHKD (splitBV->Net l1 l2 l3) = softMax
                                      . runLayerHKD l3
                                      . logistic
                                      . runLayerHKD l2
                                      . logistic
                                      . runLayerHKD l1
                                      . auto
{-# INLINE runNetworkHKD #-}

netErrHKD
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => R i
    -> R o
    -> BVar s (Network i h1 h2 o)
    -> BVar s Double
netErrHKD x t n = crossEntropy t (runNetworkHKD n x)
{-# INLINE netErrHKD #-}

trainStepHKD
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepHKD r !x !t !n = n - realToFrac r * gradBP (netErrHKD x t) n
{-# INLINE trainStepHKD #-}

-- ------------------------------
-- - "Manual" Mode              -
-- ------------------------------

runLayerManual
    :: (KnownNat i, KnownNat o)
    => Layer i o
    -> R i
    -> R o
runLayerManual l x = (l ^. lWeights) #> x + (l ^. lBiases)
{-# INLINE runLayerManual #-}

softMaxManual :: KnownNat n => R n -> R n
softMaxManual x = konst (1 / sumElements expx) * expx
  where
    expx = exp x
{-# INLINE softMaxManual #-}

runNetManual
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Network i h1 h2 o
    -> R i
    -> R o
runNetManual n = softMaxManual
               . runLayerManual (n ^. nLayer3)
               . logistic
               . runLayerManual (n ^. nLayer2)
               . logistic
               . runLayerManual (n ^. nLayer1)
{-# INLINE runNetManual #-}

gradNetManual
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
gradNetManual x t (Net (Layer w1 b1) (Layer w2 b2) (Layer w3 b3)) =
    let y1 = w1 #> x
        z1 = y1 + b1
        x2 = logistic z1
        y2 = w2 #> x2
        z2 = y2 + b2
        x3 = logistic z2
        y3 = w3 #> x3
        z3 = y3 + b3
        o0 = exp z3
        o1 = HM.sumElements (extract o0)
        o2 = o0 / konst o1
        -- o3 = - (log o2 <.> t)
        dEdO3 = 1
        dEdO2 = dEdO3 * (- t / o2)
        dEdO1 = - (dEdO2 <.> o0) / (o1 ** 2)
        dEdO0 = konst dEdO1 + dEdO2 / konst o1
        dEdZ3 = dEdO0 * o0
        dEdY3 = dEdZ3
        dEdX3 = tr w3 #> dEdY3
        dEdZ2 = dEdX3 * (x3 * (1 - x3))
        dEdY2 = dEdZ2
        dEdX2 = tr w2 #> dEdY2
        dEdZ1 = dEdX2 * (x2 * (1 - x2))
        dEdY1 = dEdZ1
        dEdB3 = dEdZ3
        dEdW3 = dEdY3 `outer` x3
        dEdB2 = dEdZ2
        dEdW2 = dEdY2 `outer` x2
        dEdB1 = dEdZ1
        dEdW1 = dEdY1 `outer` x
    in  Net (Layer dEdW1 dEdB1) (Layer dEdW2 dEdB2) (Layer dEdW3 dEdB3)
{-# INLINE gradNetManual #-}

trainStepManual
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepManual r !x !t !n =
    let gN = gradNetManual x t n
    in  n - (realToFrac r * gN)

-- ------------------------------
-- - "Hybrid" Mode              -
-- ------------------------------

layerOp :: (KnownNat i, KnownNat o) => Op '[Layer i o, R i] (R o)
layerOp = op2 $ \(Layer w b) x ->
    ( w #> x + b
    , \g -> (Layer (g `outer` x) g, tr w #> g)
    )
{-# INLINE layerOp #-}

logisticOp
    :: Floating a
    => Op '[a] a
logisticOp = op1 $ \x ->
    let lx = logistic x
    in  (lx, \g -> lx * (1 - lx) * g)
{-# INLINE logisticOp #-}

softMaxOp
    :: KnownNat n
    => Op '[R n] (R n)
softMaxOp = op1 $ \x ->
    let expx   = exp x
        tot    = sumElements expx
        invtot = 1 / tot
        res    = konst invtot * expx
    in  ( res
        , \g -> res - konst (invtot ** 2) * exp (2 * x) * g
        )
{-# INLINE softMaxOp #-}

softMaxCrossEntropyOp
    :: KnownNat n
    => R n
    -> Op '[R n] Double
softMaxCrossEntropyOp targ = op1 $ \x ->
    let expx   = exp x
        sm     = konst (1 / sumElements expx) * expx
        ce     = negate $ log sm <.> targ
    in  ( ce
        , \g -> (sm - targ) * konst g
        )
{-# INLINE softMaxCrossEntropyOp #-}

runNetHybrid
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R i
    -> BVar s (R o)
runNetHybrid n = liftOp1 softMaxOp
               . liftOp2 layerOp (n ^^. nLayer3)
               . liftOp1 logisticOp
               . liftOp2 layerOp (n ^^. nLayer2)
               . liftOp1 logisticOp
               . liftOp2 layerOp (n ^^. nLayer1)
               . auto
{-# INLINE runNetHybrid #-}

netErrHybrid
    :: (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o, Reifies s W)
    => BVar s (Network i h1 h2 o)
    -> R o
    -> R i
    -> BVar s Double
netErrHybrid n t = liftOp1 (softMaxCrossEntropyOp t)
                 . liftOp2 layerOp (n ^^. nLayer3)
                 . liftOp1 logisticOp
                 . liftOp2 layerOp (n ^^. nLayer2)
                 . liftOp1 logisticOp
                 . liftOp2 layerOp (n ^^. nLayer1)
                 . auto
{-# INLINE netErrHybrid #-}

trainStepHybrid
    :: forall i h1 h2 o. (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o)
    => Double
    -> R i
    -> R o
    -> Network i h1 h2 o
    -> Network i h1 h2 o
trainStepHybrid r !x !t !n =
    let gN = gradBP (\n' -> netErrHybrid n' t x) n
    in  n - (realToFrac r * gN)
{-# INLINE trainStepHybrid #-}

-- ------------------------------
-- - Operations                 -
-- ------------------------------

infixr 8 #>!
(#>!)
    :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (L m n)
    -> BVar s (R n)
    -> BVar s (R m)
(#>!) = liftOp2 . op2 $ \m v ->
  ( m #> v, \g -> (g `outer` v, tr m #> g) )
{-# INLINE (#>!) #-}

infixr 8 <.>!
(<.>!)
    :: (KnownNat n, Reifies s W)
    => BVar s (R n)
    -> BVar s (R n)
    -> BVar s Double
(<.>!) = liftOp2 . op2 $ \x y ->
  ( x <.> y, \g -> (konst g * y, x * konst g)
  )
{-# INLINE (<.>!) #-}

konst'
    :: (KnownNat n, Reifies s W)
    => BVar s Double
    -> BVar s (R n)
konst' = liftOp1 . op1 $ \c -> (konst c, HM.sumElements . extract)
{-# INLINE konst' #-}

sumElements :: KnownNat n => R n -> Double
sumElements = HM.sumElements . extract
{-# INLINE sumElements #-}

sumElements'
    :: (KnownNat n, Reifies s W)
    => BVar s (R n)
    -> BVar s Double
sumElements' = liftOp1 . op1 $ \x -> (sumElements x, konst)
{-# INLINE sumElements' #-}

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))
{-# INLINE logistic #-}

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

instance (KnownNat i, KnownNat o) => Num (Layer i o) where
    Layer w1 b1 + Layer w2 b2 = Layer (w1 + w2) (b1 + b2)
    Layer w1 b1 - Layer w2 b2 = Layer (w1 - w2) (b1 - b2)
    Layer w1 b1 * Layer w2 b2 = Layer (w1 * w2) (b1 * b2)
    abs    (Layer w b)        = Layer (abs    w) (abs    b)
    signum (Layer w b)        = Layer (signum w) (signum b)
    negate (Layer w b)        = Layer (negate w) (negate b)
    fromInteger x             = Layer (fromInteger x) (fromInteger x)

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Num (Network i h1 h2 o) where
    Net a b c + Net d e f = Net (a + d) (b + e) (c + f)
    Net a b c - Net d e f = Net (a - d) (b - e) (c - f)
    Net a b c * Net d e f = Net (a * d) (b * e) (c * f)
    abs    (Net a b c)    = Net (abs    a) (abs    b) (abs    c)
    signum (Net a b c)    = Net (signum a) (signum b) (signum c)
    negate (Net a b c)    = Net (negate a) (negate b) (negate c)
    fromInteger x         = Net (fromInteger x) (fromInteger x) (fromInteger x)

instance (KnownNat i, KnownNat o) => Fractional (Layer i o) where
    Layer w1 b1 / Layer w2 b2 = Layer (w1 / w2) (b1 / b2)
    recip (Layer w b)         = Layer (recip w) (recip b)
    fromRational x            = Layer (fromRational x) (fromRational x)

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Fractional (Network i h1 h2 o) where
    Net a b c / Net d e f = Net (a / d) (b / e) (c / f)
    recip (Net a b c)     = Net (recip a) (recip b) (recip c)
    fromRational x        = Net (fromRational x) (fromRational x) (fromRational x)

instance KnownNat n => MWC.Variate (R n) where
    uniform g = randomVector <$> MWC.uniform g <*> pure Uniform
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where
    uniform g = uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat i, KnownNat o) => MWC.Variate (Layer i o) where
    uniform g = Layer <$> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => MWC.Variate (Network i h1 h2 o) where
    uniform g = Net <$> MWC.uniform g <*> MWC.uniform g <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance Backprop (R n) where
    zero = zeroNum
    add  = addNum
    one  = oneNum

instance (KnownNat n, KnownNat m) => Backprop (L m n) where
    zero = zeroNum
    add  = addNum
    one  = oneNum

instance (KnownNat i, KnownNat o) => Backprop (Layer i o)
instance (KnownNat i, KnownNat h1, KnownNat h2, KnownNat o) => Backprop (Network i h1 h2 o)


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

module Language.Hakaru.Distribution where

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

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

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

bern :: Double -> Dist Bool
bern p = Dist {logDensity = (\ (Discrete x) -> log (if x then p else 1 - p)),
               distSample = (\ g -> case randomR (0, 1) g of
                                  (t, g') -> (Discrete $ t <= p, g'))}

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

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

marsaglia :: (RandomGen g, Random a, Ord a, Floating a) => g -> ((a, a), g)
marsaglia g0 = -- "Marsaglia polar method"
  let (x, g1) = randomR (-1,1) g0
      (y, g ) = randomR (-1,1) g1
      s       = x * x + y * y
      q       = sqrt ((-2) * log s / s)
  in if 1 >= s && s > 0 then ((x * q, y * q), g) else marsaglia g

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

chooseIndex :: (RandomGen g) => [Double] -> g -> (Int, g)
chooseIndex probs g0 =
  let (p, g) = random g0
      k = fromMaybe (error ("chooseIndex: failure p=" ++ show p))
                    (findIndex (p <=) (scanl1 (+) probs))
  in (k, g)

normal_rng :: (Real a, Floating a, Random a, RandomGen g) =>
              a -> a -> g -> (a, g)
normal_rng mu sd g | sd > 0 = case marsaglia g of
                                ((x, _), g1) -> (mu + sd * x, g1)
normal_rng _ _ _ = error "normal: invalid parameters"

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

normal :: Double -> Double -> Dist Double
normal mu sd = Dist {logDensity = normalLogDensity mu sd . fromLebesgue,
                     distSample = mapFst Lebesgue . normal_rng mu sd}

categoricalLogDensity :: (Eq b, Floating a) => [(b, a)] -> b -> a
categoricalLogDensity list x = log $ fromMaybe 0 (lookup x list)
categoricalSample :: (Num b, Ord b, RandomGen g, Random b) =>
    [(t,b)] -> g -> (t, g)
categoricalSample list g = (elem', g1)
    where
      (p, g1) = randomR (0, total) g
      elem' = fst $ head $ filter (\(_,p0) -> p <= p0) sumList
      sumList = scanl1 (\acc (a, b) -> (a, b + snd(acc))) list
      total = sum $ map snd list

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

lnFact :: Integer -> Double
lnFact = logFactorial

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

         make_poisson :: (RandomGen g) => g -> (Integer,g)
         make_poisson g = let (u, g1) = randomR (-0.5,0.5) g
                              (v, g2) = randomR (0,1) g1
                              us = 0.5 - abs u
                              k = floor $ (2*a / us + b)*u + lambda + 0.43 in
                          case () of
                            () | us >= 0.07 && v <= vr -> (k, g2)
                            () | k < 0 -> make_poisson g2
                            () | us <= 0.013 && v > us -> make_poisson g2
                            () | accept_region us v k -> (k, g2)
                            _  -> make_poisson g2

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

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

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

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

gamma :: Double -> Double -> Dist Double
gamma shape scl = Dist {logDensity = gammaLogDensity shape scl . fromLebesgue,
                          distSample = mapFst Lebesgue . gamma_rng shape scl}

beta_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
beta_rng a b g | a <= 1.0 && b <= 1.0 =
                 let (u, g1) = randomR (0.0, 1.0) g
                     (v, g2) = randomR (0.0, 1.0) g1
                     x = u ** (recip a)
                     y = v ** (recip b)
                 in  case (x+y) <= 1.0 of
                       True -> (x / (x + y), g2)
                       False -> beta_rng a b g2
beta_rng a b g = let (ga, g1) = gamma_rng a 1 g
                     (gb, g2) = gamma_rng b 1 g1
                 in (ga / (ga + gb), g2)

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

beta :: Double -> Double -> Dist Double
beta a b = Dist {logDensity = betaLogDensity a b . fromLebesgue,
                 distSample = mapFst Lebesgue . beta_rng a b}

laplace_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
laplace_rng mu sd g = sample (randomR (0.0, 1.0) g)
   where sample (u, g1) = case u < 0.5 of
                            True  -> (mu + sd * log (u + u), g1)
                            False -> (mu - sd * log (2.0 - u - u), g1)

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

laplace :: Double -> Double -> Dist Double
laplace mu sd = Dist {logDensity = laplaceLogDensity mu sd . fromLebesgue,
                      distSample = mapFst Lebesgue . laplace_rng mu sd}

-- Consider having dirichlet return Vector
-- Note: This is acutally symmetric dirichlet
dirichlet_rng :: (RandomGen g) => Int ->  Double -> g -> ([Double], g)
dirichlet_rng n' a g' = normalize (gammas g' n')
  where gammas g 0 = ([], 0, g)
        gammas g n = let (xs, total, g1) = gammas g (n-1)
                         ( x, g2) = gamma_rng a 1 g1
                     in ((x : xs), x+total, g2)
        normalize (b, total, h) = (map (/ total) b, h)

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


module Main where

import Numeric.GSL.Integration
import Complex
import System.IO
import Data.Fixed
import Data.Array
import Data.Maybe
import Control.Monad
import Data.Packed.Matrix
import Data.Packed.Vector
import Numeric.LinearAlgebra.Algorithms
import qualified Data.IntMap as M

data Ctree a b = End (a,b)
               | Branche (a,b) [Ctree a b]
    deriving (Show)

-- Given N boxes, return all possible combinations of P particles in these boxes.

c n p = g 1 0
    where g i a | i == n     = [ End (n, p - a) ]
                | a == p - 1 = [ End (i, 1) ] ++ g (i + 1) a
                | otherwise  = [ if r == p - a
                                 then End (i,r)
                                 else Branche (i,r) $ g (i+1) (a + r) |
                                                         i <- [i..n-1],
                                                         r <- [1..p-a] ]
                                                            ++ [ End (n, p-a) ]

-- Convert combinations of Ctree do list format

cl [] = []
cl ((End p):l) = [p]:cl l
cl ((Branche p b):l) = (map (p:) $ cl b) ++ cl l

-- Create newstates aplying A+(i)A+(j)A(k)A(l) operator on state 'a'

newState i j k l n a = do cp i =<< cp j =<< dp k =<< dp l a
    where
        cp i a = if a!i < n
                 then Just $ a//[(i,a!i + 1)]
                 else Nothing

        dp i a = if a!i > 0
                 then Just $ a//[(i,a!i - 1)]
                 else Nothing

-- List all states that can have operator A+(i)A+(j)A(k)A(l) applied

possibleStates i j k l nb np = filter (isJust . newState i j k l np) . aStates nb

-- Autovalue of operator A+(i)A+(j)A(k)A(l)

value i j k l n a = sqrt $ ((a!i + 1)*(a!j + 1)*a!k*a!l) / ((prod $ elems a')*(prod $ elems a))
    where a' = fromJust $ newState i j k l n a
          prod x = prod' x 1
          prod' [] a = a
          prod' (x:xs) a | x == 0 = prod' xs a
                         | otherwise = prod' xs a*x

sumValues i j k l n = sum . map (value i j k l n)

-- Convert state from List to Array type

aStates nb = map (listArray (1,nb) (repeat 0) //)

-- Total quasi-momentum of base state 'a'

qt a = qt' a 0 (min + 1)
    where qt' a t i | i == max  = mod (t + (max - 1)*a!max) max
                    | otherwise = qt' a (t + (i-1)*a!i) (i + 1)
          (min,max) = bounds a

-- Group states by quasi-momentum

groupQt states = M.fromListWith (++) $ map (\x -> (qt x, [x])) states

-- Print grouped states

printGroups states = do mapM_ printG states
    where printG (n,s) = do putStrLn $ (show n) ++ ":"
                            mapM_ (putStrLn.show.elems) s
                            putStrLn ""

na = array (0, 4) [ (0, 1.81772),
                    (1, 1.75048),
                    (2, 0.998368),
                    (3, 0.111417),
                    (4, 1.19828) ]

a = array (0,4) [ (0, 1.94699),
                  (1, 1.98301),
                  (2, 2.29932),
                  (3, 2.49687),
                  (4, 2.22725) ]

n = 5 :: Int
nd = 5 :: Double

bloch phi w a b = exp(0 :+ a*phi)*exp(0 :+ (-wa*phy))*(cos(b*phy)*sin(b*pi/nd)*cos(wa*pi/nd) :+ sin(b*phy)*cos(b*pi/nd)*sin(wa*pi/nd))
    where phy = -pi/nd + mod' (phi + pi/nd) (2*pi/nd)
          wa = a - w

cbloch phi w a b = conjugate $ bloch phi w a b

toInt phi w i j k l = 1/((na!i)*(na!j)*(na!k)*(na!l))*
                            cbloch phi w (fi i) (a!i) * cbloch phi w (fi j) (a!j) *
                            bloch phi w (fi k) (a!k) * bloch phi w (fi l) (a!l)
      where fi = fromIntegral

-- Calculate matrix for each quasi-momentum

matFrom :: Int -> Int -> Array Int (Array Int Int) -> Double -> Array (Int,Int) Double
matFrom nb np gs c = matrix
    where matrix = array ((1,1),(size,size)) [((i,j), f i j) | i <- [1..size], j <- [1..size]]
          f i j | i > j =  matrix!(j,i)
                | otherwise = energy i j + c*(sum $ map (\x -> sq x (gs!j)*(im!x)) [ (oi,oj,ok,ol) | oi <- [1..nb],
                                                                                oj <- [1..nb],
                                                                                ok <- [1..nb], ol <- [1..nb],
                                                                                newState oi oj ok ol np (gs!j) == Just (gs!i) ])
          sq (i,j,k,l) s = sqrt $ (fi $ (s!i) + 1)*(fi $ (s!j) + 1)*(fi $ s!k)*(fi $ s!l)
          energy i j | i == j = e (gs!j) 1 0
                     | otherwise = 0
          e s i t | i == nb = t*(fi $ s!i)
                  | otherwise = e s (i+1) (t + (a!(i-1))*(fi $ s!i))
          size = snd $ bounds gs
          im = intMatrix nb
          fi = fromIntegral

intMatrix n = array ((1,1,1,1),(n,n,n,n)) [ ((i,j,k,l), integrate (i,j,k,l) n) | i <- [1..n], j <- [1..n], k <- [1..n], l <- [1..n] ]

integrate (i,j,k,l) n | mod (i+j-k-l) n == 0 = fst $ integrateQAGS 1E-6 1000 (\x -> realPart $ toInt x 0.4 (i-1) (j-1) (k-1) (l-1)) 0 (2*pi)
                      | otherwise = 0

-- Create Matrix

genMatrix nb np k cst = matFrom nb np (listArray (1,length $ qmGroup M.! k) $ qmGroup M.! k) cst
    where qmGroup = groupQt $ aStates nb $ cl $ c nb np
          qmKeys = M.keys qmGroup

-- Print eigenvalues and eigenstates to all states, group by quasi-momentum

printStates nb np cst file = do
    let qmGroup = groupQt $ aStates nb $ cl $ c nb np
        qmKeys = M.keys qmGroup

        calEig n x = eig $ (n><n) $ elems x
        printEig (s,v) = do forM_ (zip (toList s) (toLists v)) $ \(e,v) -> do
                                 hPutStrLn file $ "Eigenvalue: " ++ (show e) ++ "\n"
                                 mapM_ (\n -> hPutStrLn file $ (show n) ++ " ") v
                                 hPutStrLn file "\n"

    forM_ qmKeys $ \k -> do
        let kSize = length $ qmGroup M.! k
        hPutStrLn file $ "quasi-momentum: " ++ (show k) ++ "\n"
        printEig $ calEig (kSize) $ matFrom nb np (listArray (1, kSize) $ qmGroup M.! k) cst

main = do
    file <- openFile "output" WriteMode
    printStates 5 10 1 file
    hClose file


{-# LANGUAGE QuantifiedConstraints,
             TupleSections,
             DefaultSignatures,
             FlexibleInstances,
             FlexibleContexts,
             UndecidableInstances,
             ConstraintKinds,
             MultiParamTypeClasses,
             TypeFamilies #-}

module Tensor where
import Data.Complex
import Helper
import EitherTrans
import Control.Monad.Trans.State
import Data.Monoid
import Data.List
import Data.Function
import System.Random
import System.Random.Stateful

-- orphan instance
-- sort complex numbers by their magnitudes
instance RealFloat a => Ord (Complex a) where
    compare a b = compare (magnitude a) (magnitude b)

instance (UniformRange a) => UniformRange (Complex a) where
    uniformRM (lr :+ li, hr :+ hi) g = do
        mr <- uniformRM (lr, hr) g
        mi <- uniformRM (li, hi) g
        return $ mr :+ mi

instance (Uniform a) => Uniform (Complex a) where
    uniformM g = do
        r <- uniformM g
        i <- uniformM g
        return $ r :+ i

{-
    The definition of tensor product should follow conventional tensor product properties.
    Tensor product should form a monoid upto isomorphism
-}

-- higer-kinded definition
-- Polymorphic constraints

class Cmplx c where
    conj :: c -> c
    conj = error "please define conjugate function"

instance (Num a) => Cmplx (Complex a) where
    conj = conjugate

class (Monoid (t a), Cmplx a) => Tensor t a where
    (|*|) :: t a -> t a -> t a
    (|*|) = (<>)
    tconcat :: [t a] -> t a
    tconcat = mconcat
    norm :: Floating a => t a -> t a
    fromList :: Integral b => [a] -> b -> t a
    rank :: Integral b => t a -> b
    isSquare :: t a -> Bool
    conjTr :: t a -> t a
    inner :: t a -> t a -> t a

class (Tensor q a, Num (q a)) => QuantumRegister q a where

    densityOp :: q a
        -> q a

    densityOp x = x * conjTr x

    -- Builds a transformation tensor to apply input gate to selected qubit
    build :: (Integral b, Integral d) => b -- ^ Number of qubits this tensor applies to
        -> [(d, q a)]-- ^ Index / Gate pair (Pre-sorted)
        -> q a -- ^ Default Gate
        -> q a
    build i xs d = tconcat $ go i (sortOnFst xs) d 0
        where
            go i2 [] d2 n = replicate (fromIntegral i2 - n) d2
            go i2 (x:xs) d2 n = replicate (fromIntegral (fst x) - n) d2 ++ (snd x : go i2 xs d2 (fromIntegral $ fst x + 1))

    getProb :: Int -- ^ Index of the qubit
        -> q a -- ^ Qubits
        -> Either String (a, a) -- ^ Resulting Probability
    default getProb :: (Foldable q, Floating a) => Int -> q a -> Either String (a, a) -- ^ Resulting Probability
    getProb i q
        | i < 0 || i >= rank q = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid Quantum Register"
        | otherwise = do
            a <- applyGate mask0 i q
            b <- applyGate mask1 i q
            return (f a, f b)
        where
            f = getSum . foldMap (Sum . (\x -> conj x * x))

    initQubit :: a -- ^ Constant for |0>
        -> a -- ^ Constant for |1>
        -> Either String (q a) -- ^ Resulting Qubit

    initQubit0 :: q a -- ^ Initialize |0> qubit at z-direction

    initNumQubit0 :: (Integral b) => b -- ^ Number of qubits to initialize
        -> Either String (q a) -- ^ Initialized quantum register

    toQuantumRegister :: [a] -- ^ State vector array to convert
        -> Either String (q a) -- ^ Converted Quantum Register

    collapse :: Int -- ^ Index of the qubit
        -> Bool -- ^ Collapsed State
        -> q a -- ^ Qubits to collapse
        -> Either String (q a) -- ^ Resulting Tensor

    default collapse :: (Floating a) => Int -> Bool -> q a -> Either String (q a)
    collapse i s q
        | i < 0 || i >= rank q = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid state vector"
        | s = norm <$> applyGate mask1 i q
        | otherwise = norm <$> applyGate mask0 i q

    measure :: (Monad m, StatefulGen g m, UniformRange a) => Int -- ^ Index of the qubit
        -> g -- ^ generator
        -> StateT (q a) (EitherT String m) Bool -- ^ Resulting State Transformer

    default measure :: (Monad m, Num a, Ord a, StatefulGen g m, UniformRange a) => Int -- ^ Index of the qubit
        -> g -- ^ generator
        -> StateT (q a) (EitherT String m) Bool -- ^ Resulting State Transformer
    measure i g = StateT go
        where
            go q
                | i < 0 || i >= rank q = EitherT $ return $ Left "Index out of bounds"
                | not (isSquare q) = EitherT $ return $ Left "Invalid state vector"
                | otherwise = EitherT $ do
                    p <- uniformRM (0, 1) g
                    return $ do
                        t <- getProb i q
                        if p < fst t
                        then (False, ) <$> collapse i False q
                        else (True, ) <$> collapse i True q

    applyGate :: q a -- ^ Gate tensor
        -> Int -- ^ Index
        -> q a -- ^ Quantum Registers
        -> Either String (q a)
    applyGate a i b
        | not (isSquare a) = Left "Invalid Gate Tensor"
        | not (isSquare b) = Left "Invalid Statevector"
        | i < 0 || i >= rank b = Left "Index out of bounds"
        | otherwise = Right $ build (rank b) [(i, a)] pauliId * b

    applyGateAll :: q a -- ^ Gate tensor
        -> q a -- ^ Quantum Registers
        -> Either String (q a)

    applyGateAll a b
        | not (isSquare a) = Left "Invalid Gate Tensor"
        | not (isSquare b) = Left "Invalid Statevector"
        | otherwise = Right $ build (rank b) [] a * b

    applyControl :: q a -- ^ Gate tensor
        -> Int -- ^ Control qubit index
        -> Int -- ^ Apply qubit index
        -> q a -- ^ Quantum Registers
        -> Either String (q a)

    applyControl a ctl i b
        | not (isSquare a) = Left "Invalid gate tensor"
        | not (isSquare b) = Left "Invalid state vector"
        | i < 0 || i >= rank b = Left "Application qubit index out of bounds"
        | ctl < 0 || ctl >= rank b = Left "Control qubit index out of bounds"
        | ctl == i = Left "Control and application qubit cannot be the same"
        | otherwise = Right $ (build (rank b) [(ctl, mask0)] pauliId
            + build (rank b) (sortBy (compare `on` fst) [(ctl, mask1), (i, a)]) pauliId) * b

    illegalPeek :: q a
        -> String

    subSystem :: [Int] -> q a -> Either String (q a)

    isEntangled :: [Int] -> q a -> Either String Bool

    mask0 :: q a
    mask1 :: q a
    pauliId :: q a

class (QuantumRegister q a) => Gates q a where
    pauliX :: q a
    pauliY :: q a
    pauliZ :: q a
    hadamard :: q a

-- Num-only declaration

-- class (forall a. Num a => Monoid (t a)) => Tensor t where
--     (|*|) :: (Num a) => t a -> t a -> t a
--     (|*|) = (<>)
--     norm :: Floating a => t a -> t a
--     fromList :: Integral b => [a] -> b -> t a
--     rank :: Integral b => t a -> b
--     isSquare :: t a -> Bool

-- class (Tensor q, forall a. Num a => Num (q a)) => QuantumRegister q where
--     -- Builds a transformation tensor to apply input gate to selected qubit
--     build :: (RealFloat a, Integral b, Integral c) => b -- ^ Number of qubits this tensor applies to
--         -> [(c, q (Complex a))]-- ^ Index / Gate pair (Pre-sorted)
--         -> q (Complex a) -- ^ Default Gate
--         -> q (Complex a)
--     build i xs d = mconcat $ go i (sortOnFst xs) d 0
--         where
--             go i2 [] d2 n = replicate (fromIntegral i2 - n) d2
--             go i2 (x:xs) d2 n = replicate (fromIntegral (fst x) - n) d2 ++ (snd x : go i2 xs d2 (fromIntegral $ fst x + 1))

--     getProb :: (RealFloat a) => Int -- ^ Index of the qubit
--         -> q (Complex a) -- ^ Qubits
--         -> Either String (a, a) -- ^ Resulting Probability

--     initQubit :: (RealFloat a) => Complex a -- ^ Constant for |0>
--         -> Complex a -- ^ Constant for |1>
--         -> Either String (q (Complex a)) -- ^ Resulting Qubit

--     initQubit0 :: (RealFloat a) => q (Complex a) -- ^ Initialize |0> qubit at z-direction

--     initNumQubit0 :: (RealFloat a, Integral b) => b -- ^ Number of qubits to initialize
--         -> Either String (q (Complex a)) -- ^ Initialized quantum register

--     toQuantumRegister :: (RealFloat a) => [Complex a] -- ^ State vector array to convert
--         -> Either String (q (Complex a)) -- ^ Converted Quantum Register

--     collapse :: (RealFloat a) => Int -- ^ Index of the qubit
--         -> Complex a -- ^ Collapsed State
--         -> q (Complex a) -- ^ Qubits to collapse
--         -> EitherT String IO (q (Complex a)) -- ^ Resulting Tensor

--     measure :: (RealFloat a, Monad m) => Int -- ^ Index of the qubit
--         -> StateT (q (Complex a)) (EitherT String m) (Complex a)
--     -- -> q (Complex a) -- ^ Qubits
--     -- -> Either String ( m (Complex a, q (Complex a))) -- ^ Resulting measurement and the collapsed qubits

--     pauliX :: (RealFloat a) => q (Complex a)
--     pauliY :: (RealFloat a) => q (Complex a)
--     pauliZ :: (RealFloat a) => q (Complex a)
--     pauliId :: (RealFloat a) => q (Complex a)
--     hadamard :: (RealFloat a) => q (Complex a)
--     mask0 :: (RealFloat a) => q (Complex a)
--     mask1 :: (RealFloat a) => q (Complex a)


{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, GeneralizedNewtypeDeriving #-}

-- | This module defines the Syntax of the Quantum IO Monad, which is an embedded
-- language for writing quantum computations.
module QIO.QioSyn where

import Data.Monoid as Monoid
import Data.Complex
import Control.Applicative (Applicative(..))
import Control.Monad       (liftM, ap)

-- | For Real numbers, we simply use the built in Double type
type RR = Double

-- | For Complex numbers, we use the built in Complex numbers, over our Real
-- number type (i.e. Double)
type CC = Complex RR

-- | The amplitude of a complex number is the magnitude squared.
amp :: CC -> RR
amp k = (magnitude k)*(magnitude k)

-- | The type of Qubits in QIO are simply integer references.
newtype Qbit = Qbit Int deriving (Num, Enum, Eq, Ord)

-- | A rotation is in essence a two-by-two complex valued matrix
type Rotation = ((Bool,Bool) -> CC)

-- | The underlying data type of a U unitary operation
data U = UReturn | Rot Qbit Rotation U
       | Swap Qbit Qbit U | Cond Qbit (Bool -> U) U | Ulet Bool (Qbit -> U) U

-- | The underlying data type of a QIO Computation
data QIO a = QReturn a | MkQbit Bool (Qbit -> QIO a) | ApplyU U (QIO a)
           | Meas Qbit (Bool -> QIO a)

-- | The type "U" forms a Monoid
instance Semigroup U where
    --mempty = UReturn
    UReturn <> u = u
    (Rot x a u) <> u' = Rot x a (u <> u')
    (Swap x y u) <> u' = Swap x y (u <> u')
    (Cond x br u') <> u'' = Cond x br (u' <> u'')
    (Ulet b f u) <> u' = Ulet b f (u <> u')

instance Monoid  U where
  mempty = UReturn

-- | Apply the given rotation to the given qubit
rot :: Qbit -> Rotation -> U
rot x r = Rot x r UReturn

-- | Swap the state of the two given qubits
swap :: Qbit -> Qbit -> U
swap x y = Swap x y UReturn

-- | Apply the conditional unitary, depending on the value of the given qubit
cond :: Qbit -> (Bool -> U) -> U
cond x br = Cond x br UReturn

-- | Introduce an Ancilla qubit in the given state, for use in the sub-unitary
ulet :: Bool -> (Qbit -> U) -> U
ulet b ux = Ulet b ux UReturn

-- | Returns the inverse (or reverse) of the given unitary operation
urev :: U -> U
urev UReturn = UReturn
urev (Rot x r u) = urev u <> rot x (rrev r)
urev (Swap x y u) = urev u <> swap x y
urev (Cond x br u) = urev u <> cond x (urev.br)
urev (Ulet b xu u) = urev u <> ulet b (urev.xu)

-- | Apply a not rotation to the given qubit
unot :: Qbit -> U
unot x = rot x rnot

-- | Apply a hadamard rotation to the given qubit
uhad :: Qbit -> U
uhad x = rot x rhad

-- | Apply a phase rotation (of the given angle) to the given qubit
uphase :: Qbit -> RR -> U
uphase x r = rot x (rphase r)

instance Functor QIO where
    fmap = liftM

instance Applicative QIO  where
    pure  = QReturn
    (<*>) = ap

-- | The "QIO" type forms a Monad
instance Monad QIO where
    return = pure
    (QReturn a) >>= f = f a
    (MkQbit b g) >>= f = MkQbit b (\ x -> g x >>= f)
    (ApplyU u q) >>= f = ApplyU u (q >>= f)
    (Meas x g) >>= f = Meas x (\ b -> g b >>= f)

-- | Initialise a qubit in the given state (adding it to the overall quantum state)
mkQbit :: Bool -> QIO Qbit
mkQbit b = MkQbit b return

-- | Apply the given unitary operation to the current quantum state
applyU :: U -> QIO ()
applyU u =  ApplyU u (return ())

-- | Measure the given qubit, and return the measurement outcome (note that this
-- operation may affect the overall quantum state, as a measurement is destructive)
measQbit :: Qbit -> QIO Bool
measQbit x = Meas x return

-- | The identity rotation
rid :: Rotation
rid (x,y) = if x==y then 1 else 0

-- | The not rotation
rnot :: Rotation
rnot (x,y) = if x==y then 0 else 1

-- | The hadamard rotation
rhad :: Rotation
rhad (x,y) = if x && y then -h else h where h = (1/sqrt 2)

-- | The phase rotation
rphase :: RR -> Rotation
rphase _ (False,False)  = 1
rphase r (True,True)    = exp(0:+r)
rphase _ (_,_)          = 0

-- | Returns the inverse (or reverse) of the given rotation
rrev :: Rotation -> Rotation
rrev r (False,True)   = conjugate (r (True,False))
rrev r (True,False)   = conjugate (r (False,True))
rrev r xy             = conjugate (r xy)

-- | Rotations can be compared for equality.
-- They are equal if the define the same matrix.
instance Eq Rotation where
    f == g =    (f (False,False)  == g (False,False))
             && (f (False,True)   == g (False,True))
             && (f (True,False)   == g (True,False))
             && (f (True,True)    == g (True,True))
    f /= g =    (f (False,False)  /= g (False,False))
             || (f (False,True)   /= g (False,True))
             || (f (True,False)   /= g (True,False))
             || (f (True,True)    /= g (True,True))

-- | We can display a qubit reference
instance Show Qbit where
    show (Qbit q) = "(Qbit:" ++ show q ++ ")"

-- | We can display the matrix representation of a rotation
instance Show Rotation where
    show f = "(" ++ (show (f (False,False))) ++ "," ++ (show (f (False,True))) ++ "," ++ (show (f (True,False))) ++ "," ++ (show (f (True,True))) ++ ")"

-- | We can display a representation of a unitary
instance Show U where
    show u = show' u 0 (-1)

-- | A helper function for the show instance of U
show' :: U -> Int -> Int -> String
show' (UReturn) x fv = ""
show' (Rot q a u) x fv = spaces x ++ "Rotate " ++ show q ++ " by " ++ show a ++ ".\n" ++ show' u x fv
show' (Swap q1 q2 u) x fv = spaces x ++ "Swap " ++ show q1 ++ " and " ++ show q2 ++ ".\n" ++ show' u x fv
show' (Cond q f u) x fv = spaces x ++ "Cond (if " ++ show q ++ " then \n" ++ spaces (x+1) ++ "(\n" ++ show' (f True) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ spaces x ++ "else \n" ++ spaces (x+1) ++ "(\n" ++ show' (f False) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ show' u x fv
show' (Ulet b f u) x fv = spaces x ++ "Ulet " ++ show b ++ " (\\" ++ show (Qbit fv) ++ "->\n " ++ show' (f (Qbit fv)) x (fv-1) ++ ")\n" ++ show' u x fv

-- | A helper function that returns a string of 4\x\ spaces.
spaces :: Int -> String
spaces 0 = ""
spaces n = if (n < 0) then error "spaces: negative argument"
                      else "    " ++ spaces (n-1)