module Numeric.Tensile.Operations.Linear.Internal where

import Data.Vector.Sized (Vector)
import Numeric.Tensile.Dimensions
import Numeric.Tensile.Tensor.Internal
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Finite as F
import qualified Data.Vector.Sized as N
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Storable.Mutable as M

-- import qualified Numeric.LinearAlgebra.HMatrix as Ma

-- f :: (Dims d' -> Perm n -> Perm n) -> Perm n -> Tensor d e -> Tensor d e'
-- f dim2Idx perm t = Tensor $ reifySomeDims (permuteDims perm (dims @d)) $ \p ->
--   modifyIdx (reflect p) (modify (permuteIdxs (dim2Idx (reflect p) perm) _)) (reflect p)) t -- basically make user derive the Idxs d' -> Idxs d'

-- dim2Idx :: Rank d ~ n => Dims d -> Perm n -> Perm n
-- takes a perm on dimensions and derives a perm in indices, eg
-- dim2Idx d p = lowerPerm' ...
-- --
-- otherwise consider using the raw index fold and lowerPerm???
-- could also create :  Perm d d'

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

transpose
  :: Elt e
  => Permutable d d'
  => Dims d -> Perm (Rank d) -> Tensor d e -> Tensor d' e
transpose d p (Tensor v) = Tensor v'
  where v' = modifyIdxs d v $ \i m ->
               remapIdxs p d i $ \d' i' ->
                 M.modify m (const $ v S.! fromIdxs d' (unsafePermute p i)) (fromIdxs d' i')

{-
gen#
  :: forall s t d. PrimBytes t
  => Int# -- ^ number of elements, not checked!
             --   Avoid using this argument if possible.
  -> (s -> (# s, t #))
  -> s -> (# s, ArrayBase t d #)
gen# n f z0 = go (byteSize @e undefined *# n)
  where
    go bsize = case runRW#
     ( \s0 -> case newByteArray# bsize s0 of
         (# s1, mba #) -> case loop0 mba 0# z0 s1 of
           (# s2, z1 #) -> case unsafeFreezeByteArray# mba s2 of
             (# s3, ba #) -> (# s3, (# z1, ba #) #)
     ) of (# _, (# z1, ba #) #) -> (# z1, ArrayBase (# | (# 0# , n , ba #) #) #)
    {-# NOINLINE go #-}
    loop0 mba i z s
      | isTrue# (i ==# n) = (# s, z #)
      | otherwise = case f z of
          (# z', x #) -> loop0 mba (i +# 1#) z' (writeArray mba i x s)

3-D tensor `a` w shape=[2, 2, 3]
[[[ 1,  2,  3],
  [ 4,  5,  6]],
 [[ 7,  8,  9],
  [10, 11, 12]]]

3-D tensor `b` w shape=[2, 3, 2]
[[[13, 14],
  [15, 16],
  [17, 18]],
 [[19, 20],
  [21, 22],
  [23, 24]]]
                )

mul `a` `b` has shape=[2,2,2]
[[[ 94, 100],
  [229, 244]],
 [[508, 532],
  [697, 730]]]
-}

{-

-- <#
-- same as tf.matmul
productN
  :: forall a b c x. All KnownDim '[a, b, c]
  => KnownDims x
  => T (x +: a +: b)
  -> T (x +: b +: c)
  -> T (x +: a +: c)
productN = undefined
  where
    d = dims @x
    --s = unsafeCoerce a

product
  :: forall m x y. KnownDim m
  => KnownDims x
  => KnownDims y
  => T (x +: m) -> T (m :+ y) -> T (x ++ y)
product t u
    | I# m <- fromIntegral $ dimVal @m
    , I# n <- fromIntegral $ totalDim' @x
    , I# k <- fromIntegral $ totalDim' @y
    , nk <- n *# k
    = let loop1 i j l r | isTrue# (l ==# m) = r
                        | otherwise = loop1 i j (l +# 1#)
                          (r + ix# (i +# n *# l) t * ix# (l +# m *# j) u)

          loop2 (T# i j) | isTrue# (j ==# k) = (# T# i j, 0 #)
                         | isTrue# (i ==# n) = loop2 (T# 0# (j +# 1#))
                         | otherwise = (# T# (i +# 1#) j, loop1 i j 0# 0 #)
      in case gen# nk loop2 (T# 0# 0#) of
          (# _, r #) -> r

data T# = T# Int# Int#

-- mode-i tensor-matrix product
-- see http://citeseerx.ist.psu.edu/viewdoc/download;jsessionid=4A0663C7848627DADDBA6A243BC43E78?doi=10.1.1.130.782&rep=rep1&type=pdf
product
  :: forall m x y. KnownDim m
  => KnownDims x
  => KnownDims y
  => T (x ++ [m] ++ y)
  -> T '[m, n]
  -> T (x ++ n ++ y)

-- #>
productR
  :: All KnownDim '[a, b, c]
  => KnownDims x
  => T (a :+ b :+ x)
  -> T (b :+ c :+ x)
  -> T (a :+ c :+ x)
productR = undefined

-- <#
-- same as tf.matmul
productN
  :: All KnownDim '[a, b, c]
  => KnownDims x
  => T (x +: a +: b)
  -> T (x +: b +: c)
  -> T (x +: a +: c)
productN = undefined

transpose
  :: forall m n x. All KnownDim '[m, n]
  => T '[n, m] --(n :+ m :+ x)
  -> T '[m, n] --(m :+ n :+ x)
transpose t = case elemSize0 t of
  0# -> broadcast (ix# 0# t)
  nm | I# m <- fromIntegral $ dimVal @m
     , I# n <- fromIntegral $ dimVal @n
     -> let f ( I# i,  I# j )
              | isTrue# (i ==# m) = f ( 0 , I# (j +# 1#) ) -- skip to next col
              | otherwise         = (# ( I# (i +# 1#), I# j ), ix# (i *# n +# j) t #) --col-major indexing
        in case gen# nm f (0,0) of
          (# _, r #) -> r

-}


module Numeric.LinearAlgebra.Matrix.Mat44 where

import Numeric.LinearAlgebra.Matrix.Class
import Numeric.LinearAlgebra.Matrix.Mat33 ( det33 )

import Data.List ( foldl1' )
import Data.NumInstances() -- Add function instances to Num

data Mat44 a = Mat44
  { m00 :: !a, m01 :: !a, m02 :: !a, m03 :: !a
  , m10 :: !a, m11 :: !a, m12 :: !a, m13 :: !a
  , m20 :: !a, m21 :: !a, m22 :: !a, m23 :: !a
  , m30 :: !a, m31 :: !a, m32 :: !a, m33 :: !a
  }
  deriving (Read, Show, Eq, Ord)

instance Matrix Mat44 where
  {-# INLINE mDim #-}
  mDim _ = 4
  {-# INLINE mElement #-}
  mElement m 0 0 = m00 m
  mElement m 0 1 = m01 m
  mElement m 0 2 = m02 m
  mElement m 0 3 = m03 m
  mElement m 1 0 = m10 m
  mElement m 1 1 = m11 m
  mElement m 1 2 = m12 m
  mElement m 1 3 = m13 m
  mElement m 2 0 = m20 m
  mElement m 2 1 = m21 m
  mElement m 2 2 = m22 m
  mElement m 2 3 = m23 m
  mElement m 3 0 = m30 m
  mElement m 3 1 = m31 m
  mElement m 3 2 = m32 m
  mElement m 3 3 = m33 m
  mElement _ i j = error ("Index " ++ show i ++ ", " ++ show j ++ ": out of range, must be 0,0 to 3,3")
  {-# INLINE mZip #-}
  mZip f m n = Mat44 (f (m00 m) (m00 n)) (f (m01 m) (m01 n)) (f (m02 m) (m02 n)) (f (m03 m) (m03 n))
                     (f (m10 m) (m10 n)) (f (m11 m) (m11 n)) (f (m12 m) (m12 n)) (f (m13 m) (m13 n))
                     (f (m20 m) (m20 n)) (f (m21 m) (m21 n)) (f (m22 m) (m22 n)) (f (m23 m) (m23 n))
                     (f (m30 m) (m30 n)) (f (m31 m) (m31 n)) (f (m32 m) (m32 n)) (f (m33 m) (m33 n))
  {-# INLINE mFold #-}
  mFold f m = foldl1' f [ mElement m i j | i <- [ 0 .. 3 ], j <- [ 0 .. 3 ] ]
  {-# INLINE mIndexOf #-}
  mIndexOf p m = fst (foldl1' p' [ ((i,j), mElement m i j) | j <- [ 3, 2 .. 0 ], i <- [ 3, 2 .. 0 ] ])
    where
    p' x@(_, a) y@(_, a') | a `p` a'  = x
                          | otherwise = y
  {-# INLINE det #-}
  det = det44

instance Functor Mat44 where
  {-# INLINE fmap #-}
  fmap f m = Mat44 (f (m00 m)) (f (m01 m)) (f (m02 m)) (f (m03 m))
                   (f (m10 m)) (f (m11 m)) (f (m12 m)) (f (m13 m))
                   (f (m20 m)) (f (m21 m)) (f (m22 m)) (f (m23 m))
                   (f (m30 m)) (f (m31 m)) (f (m32 m)) (f (m33 m))

{-# SPECIALIZE det44 :: Mat44 Double -> Double #-}
{-# SPECIALIZE det44 :: Mat44 Float  -> Float #-}
det44 :: Num a => Mat44 a -> a
det44 = m00 * det33 m11 m12 m13
                    m21 m22 m23
                    m31 m32 m33
      - m01 * det33 m10 m12 m13
                    m20 m22 m23
                    m30 m32 m33
      + m02 * det33 m10 m11 m13
                    m20 m21 m23
                    m30 m31 m33
      - m03 * det33 m10 m11 m12
                    m20 m21 m22
                    m30 m31 m32

{-# SPECIALIZE inv44 :: Mat44 Double -> Mat44 Double #-}
{-# SPECIALIZE inv44 :: Mat44 Float  -> Mat44 Float  #-}
inv44 :: Fractional a => Mat44 a -> Mat44 a
inv44 m = mApply m' m
  where
  d  = const $ det44 m
  m' = Mat44 (det33 m11 m12 m13
                    m21 m22 m23
                    m31 m32 m33 / d)
           (-(det33 m01 m02 m03
                    m21 m22 m23
                    m31 m32 m33 / d))
             (det33 m01 m02 m03
                    m11 m12 m13
                    m31 m32 m33 / d)
           (-(det33 m01 m02 m03
                    m11 m12 m13
                    m21 m22 m23 / d))

           (-(det33 m10 m12 m13
                    m20 m22 m23
                    m30 m32 m33 / d))
             (det33 m00 m02 m03
                    m20 m22 m23
                    m30 m32 m33 / d)
           (-(det33 m00 m02 m03
                    m10 m12 m13
                    m30 m32 m33 / d))
             (det33 m00 m02 m03
                    m10 m12 m13
                    m20 m22 m23 / d)

             (det33 m10 m11 m13
                    m20 m21 m23
                    m30 m31 m33 / d)
           (-(det33 m00 m01 m03
                    m20 m21 m23
                    m30 m31 m33 / d))
             (det33 m00 m01 m03
                    m10 m11 m13
                    m30 m31 m33 / d)
           (-(det33 m00 m01 m03
                    m10 m11 m13
                    m20 m21 m23 / d))

           (-(det33 m10 m11 m12
                    m20 m21 m22
                    m30 m31 m32 / d))
             (det33 m00 m01 m02
                    m20 m21 m22
                    m30 m31 m32 / d)
           (-(det33 m00 m01 m02
                    m10 m11 m12
                    m30 m31 m32 / d))
             (det33 m00 m01 m02
                    m10 m11 m12
                    m20 m21 m22 / d)


module Numbers where

import Control.Monad.Except (throwError)
import Data.Complex
import Data.Ratio
import DataTypes
       (Arity(..), LispError(..), LispVal(..), PrimitiveFunc, ThrowsError)

numPrimitives :: [(String, PrimitiveFunc)]
numPrimitives =
  [ ("+", numAdd)
  , ("-", numSub)
  , ("*", numMul)
  , ("/", numDiv)
  , ("modulo", numMod)
  , ("Integer?", isNumber)
  , ("complex?", isComplex)
  , ("real?", isReal)
  , ("rational?", isRational)
  , ("integer?", isInteger)
  , ("=", numBoolBinopEq)
  , ("/=", numBoolBinopNeq)
  , (">", numBoolBinopGt)
  , ("<", numBoolBinopLt)
  , (">=", numBoolBinopGte)
  , ("<=", numBoolBinopLte)
  , ("quotient", numQuotient)
  , ("remainder", numRem)
  , ("sin", numSine)
  , ("cos", numCos)
  ]

foldlM :: Monad m => (a -> b -> m a) -> a -> [b] -> m a
foldlM f v (x:xs) = (f v x) >>= \a -> foldlM f a xs
foldlM _ v [] = return v

foldl1M :: Monad m => (a -> a -> m a) -> [a] -> m a
foldl1M f (x:xs) = foldlM f x xs
foldl1M _ _ = error "Unexpected error in foldl1M"

numAdd :: PrimitiveFunc
numAdd [] = return $ Integer 0
numAdd a = foldlM (\b c -> doAdd =<< (numCast [b, c])) (Integer 0) a
  where
    doAdd :: LispVal -> ThrowsError LispVal
    doAdd (List [Integer c, Integer d]) = return $ Integer (c + d)
    doAdd (List [Rational c, Rational d]) = return $ Rational (c + d)
    doAdd (List [Float c, Float d]) = return $ Float (c + d)
    doAdd (List [Complex c, Complex d]) = return $ Complex (c + d)
    doAdd _ = throwError $ Default "Unexpected error in +"

numSub :: PrimitiveFunc
numSub [] = throwError $ NumArgs (Min 1) 0 []
numSub a = foldl1M (\b c -> doSub =<< (numCast [b, c])) a -- TODO zero args check
  where
    doSub :: LispVal -> ThrowsError LispVal
    doSub (List [Integer c, Integer d]) = return $ Integer (c - d)
    doSub (List [Rational c, Rational d]) = return $ Rational (c - d)
    doSub (List [Float c, Float d]) = return $ Float (c - d)
    doSub (List [Complex c, Complex d]) = return $ Complex (c - d)
    doSub _ = throwError $ Default "Unexpected error in -"

numMul :: PrimitiveFunc
numMul [] = return $ Integer 1
numMul a = foldl1M (\b c -> doMul =<< (numCast [b, c])) a
  where
    doMul :: LispVal -> ThrowsError LispVal
    doMul (List [Integer c, Integer d]) = return $ Integer (c * d)
    doMul (List [Rational c, Rational d]) = return $ Rational (c * d)
    doMul (List [Float c, Float d]) = return $ Float (c * d)
    doMul (List [Complex c, Complex d]) = return $ Complex (c * d)
    doMul _ = throwError $ Default "Unexpected error in -"

numDiv :: PrimitiveFunc
numDiv [] = throwError $ NumArgs (Min 1) 0 []
numDiv a = foldl1M (\b c -> doDiv =<< (numCast [b, c])) a -- TODO: Zero division error
  where
    doDiv :: LispVal -> ThrowsError LispVal
    doDiv (List [Integer c, Integer d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Rational (c % d)
    doDiv (List [Rational c, Rational d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Rational (c / d)
    doDiv (List [Float c, Float d]) =
      if d == 0.0
        then throwError $ Default "Zero division error"
        else return $ Float (c / d)
    doDiv (List [Complex c, Complex d]) =
      if d == 0
        then throwError $ Default "Zero division error"
        else return $ Complex (c / d)
    doDiv _ = throwError $ Default "Unexpected error in /"

numMod :: PrimitiveFunc
numMod [] = throwError $ NumArgs (MinMax 2 2) 0 []
numMod (a:[]) = throwError $ NumArgs (MinMax 2 2) 1 [a]
numMod [a, b] = do
  c <- numCast $ [a, b]
  doMod c
  where
    doMod :: LispVal -> ThrowsError LispVal
    doMod (List [Integer c, Integer d]) = return $ Integer (c `mod` d)
    doMod (List [c@(Rational _), d@(Rational _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Rational ((c' `mod` d') % 1)
    doMod (List [c@(Float _), d@(Float _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Float (fromInteger (c' `mod` d'))
    doMod (List [c@(Complex _), d@(Complex _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Complex (fromInteger (c' `mod` d') :+ 0)
    doMod _ = throwError $ Default "Unexpected error in modulo"
numMod a = throwError $ NumArgs (MinMax 2 2) (length a) a

numRem :: PrimitiveFunc
numRem [] = throwError $ NumArgs (MinMax 2 2) 0 []
numRem (a:[]) = throwError $ NumArgs (MinMax 2 2) 1 [a]
numRem [a, b] = do
  c <- numCast $ [a, b]
  doRem c
  where
    doRem :: LispVal -> ThrowsError LispVal
    doRem (List [Integer c, Integer d]) = return $ Integer (c `rem` d)
    doRem (List [c@(Rational _), d@(Rational _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Rational ((c' `rem` d') % 1)
    doRem (List [c@(Float _), d@(Float _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Float (fromInteger (c' `rem` d'))
    doRem (List [c@(Complex _), d@(Complex _)]) = do
      Integer c' <- numToInt c
      Integer d' <- numToInt d
      return $ Complex (fromInteger (c' `rem` d') :+ 0)
    doRem _ = throwError $ Default "Unexpected error in remainder"
numRem a = throwError $ NumArgs (MinMax 2 2) (length a) a

isInteger :: PrimitiveFunc
isInteger [Integer _] = return $ Bool True
isInteger [_] = return $ Bool False
isInteger a = throwError $ NumArgs (MinMax 1 1) (length a) a

isRational :: PrimitiveFunc
isRational [Rational _] = return $ Bool True
isRational a = isInteger a

isReal :: PrimitiveFunc
isReal [Float _] = return $ Bool True
isReal a = isRational a

isComplex :: PrimitiveFunc
isComplex [Complex _] = return $ Bool True
isComplex a = isReal a

isNumber :: PrimitiveFunc
isNumber = isComplex

numToInt :: LispVal -> ThrowsError LispVal
numToInt a@(Integer _) = return a
numToInt (Rational a) =
  if denominator a == 1
    then return $ Integer (numerator a)
    else throwError $ Default "Could not convert rational to integer"
numToInt (Float a) =
  if a == fromInteger (round a)
    then return $ Integer (round a)
    else throwError $ Default "Could not convert float to integer"
numToInt (Complex a) =
  let rp = realPart a
  in if imagPart a == 0 && rp == fromInteger (round rp)
       then return $ Integer (round rp)
       else throwError $ Default "Could not convert complex to integer"
numToInt _ = throwError $ Default "Could not convert non-number to integer"

numCast :: [LispVal] -> ThrowsError LispVal
numCast [a@(Integer _), b@(Integer _)] = return $ List [a, b]
numCast [a@(Rational _), b@(Rational _)] = return $ List [a, b]
numCast [a@(Float _), b@(Float _)] = return $ List [a, b]
numCast [a@(Complex _), b@(Complex _)] = return $ List [a, b]
-- Integer
numCast [(Integer a), b@(Rational _)] = return $ List [Rational (a % 1), b]
numCast [(Integer a), b@(Float _)] = return $ List [Float (fromInteger a), b]
numCast [(Integer a), b@(Complex _)] =
  return $ List [Complex (fromInteger a :+ 0), b]
-- Rational
numCast [a@(Rational _), (Integer b)] = return $ List [a, Rational (b % 1)]
numCast [(Rational a), b@(Float _)] = return $ List [Float (fromRational a), b]
numCast [(Rational a), b@(Complex _)] =
  return $ List [Complex (fromRational a :+ 0), b]
-- Float
numCast [a@(Float _), (Rational b)] = return $ List [a, Float (fromRational b)]
numCast [a@(Float _), (Integer b)] = return $ List [a, Float (fromInteger b)]
numCast [(Float a), b@(Complex _)] = return $ List [Complex (a :+ 0), b]
-- Complex
numCast [a@(Complex _), (Rational b)] =
  return $ List [a, Complex (fromRational b :+ 0)]
numCast [a@(Complex _), (Float b)] = return $ List [a, Complex (b :+ 0)]
numCast [a@(Complex _), (Integer b)] =
  return $ List [a, Complex (fromInteger b :+ 0)]
numCast [a, b] =
  case a of
    Integer _ -> doThrowError b
    Float _ -> doThrowError b
    Rational _ -> doThrowError b
    Complex _ -> doThrowError b
    _ -> doThrowError a
  where
    doThrowError :: LispVal -> ThrowsError LispVal
    doThrowError num = throwError $ TypeMismatch "Integer" num
numCast _ = throwError $ Default "Unexpected error in numCast"

numBoolBinop ::
     String
  -> (LispVal -> LispVal -> ThrowsError LispVal)
  -> LispVal
  -> [LispVal]
  -> ThrowsError LispVal
numBoolBinop name' op b (c:d) = do
  List [b', c'] <- numCast [b, c]
  result <- op b' c'
  case result of
    Bool True -> numBoolBinop name' op c' d
    Bool False -> return $ Bool False
    _ -> throwError $ Default $ "Unexpected error in " ++ name'
numBoolBinop _ _ _ _ = return $ Bool True

numBoolBinopEq :: [LispVal] -> ThrowsError LispVal
numBoolBinopEq [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopEq (x:xs) = numBoolBinop "=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c == d)
    fn (Rational c) (Rational d) = return $ Bool (c == d)
    fn (Float c) (Float d) = return $ Bool (c == d)
    fn (Complex c) (Complex d) = return $ Bool (c == d)
    fn _ _ = throwError $ Default "Unexpected error in ="

numBoolBinopNeq :: [LispVal] -> ThrowsError LispVal
numBoolBinopNeq [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopNeq (x:xs) = numBoolBinop "/=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c /= d)
    fn (Rational c) (Rational d) = return $ Bool (c /= d)
    fn (Float c) (Float d) = return $ Bool (c /= d)
    fn (Complex c) (Complex d) = return $ Bool (c /= d)
    fn _ _ = throwError $ Default "Unexpected error in /="

numBoolBinopLt :: [LispVal] -> ThrowsError LispVal
numBoolBinopLt [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopLt (x:xs) = numBoolBinop "<" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c < d)
    fn (Rational c) (Rational d) = return $ Bool (c < d)
    fn (Float c) (Float d) = return $ Bool (c < d)
    fn (Complex _) (Complex _) =
      throwError $ Default "< not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in <"

numBoolBinopLte :: [LispVal] -> ThrowsError LispVal
numBoolBinopLte [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopLte (x:xs) = numBoolBinop "<=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c <= d)
    fn (Rational c) (Rational d) = return $ Bool (c <= d)
    fn (Float c) (Float d) = return $ Bool (c <= d)
    fn (Complex _) (Complex _) =
      throwError $ Default "<= not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in <="

numBoolBinopGt :: [LispVal] -> ThrowsError LispVal
numBoolBinopGt [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopGt (x:xs) = numBoolBinop ">" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c > d)
    fn (Rational c) (Rational d) = return $ Bool (c > d)
    fn (Float c) (Float d) = return $ Bool (c > d)
    fn (Complex _) (Complex _) =
      throwError $ Default "> not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in >"

numBoolBinopGte :: [LispVal] -> ThrowsError LispVal
numBoolBinopGte [] = throwError $ NumArgs (Min 1) 0 []
numBoolBinopGte (x:xs) = numBoolBinop ">=" fn x xs
  where
    fn :: LispVal -> LispVal -> ThrowsError LispVal
    fn (Integer c) (Integer d) = return $ Bool (c >= d)
    fn (Rational c) (Rational d) = return $ Bool (c >= d)
    fn (Float c) (Float d) = return $ Bool (c >= d)
    fn (Complex _) (Complex _) =
      throwError $ Default ">= not defined for complex numbers"
    fn _ _ = throwError $ Default "Unexpected error in >="

numQuotient :: PrimitiveFunc
numQuotient args =
  if length args /= 2
    then throwError $ NumArgs (MinMax 2 2) (length args) args
    else do
      nums <- numCast args
      case nums of
        List [Integer a, Integer b] -> return $ Integer (a `quot` b)
        _ -> throwError $ Default "Unexpected error in <=" -- TODO better errors

unaryTrig ::
     (Double -> Double) -> (Double -> Double -> Complex Double) -> PrimitiveFunc
unaryTrig op complexOp args =
  if length args /= 1
    then throwError $ NumArgs (MinMax 1 1) (length args) args
    else case args of
           [Integer a] -> return $ Float (op $ fromInteger a)
           [Rational a] -> return $ Float (op $ fromRational a)
           [Float a] -> return $ Float (op a)
           [Complex a] -> return $ Complex $ complexOp (realPart a) (imagPart a)
           _ -> throwError $ Default "Numerical input expected"

numSine :: PrimitiveFunc
numSine = unaryTrig sin (\r i -> ((sin r) * (cosh i)) :+ ((cos r) * (sinh i)))

numCos :: PrimitiveFunc
numCos =
  unaryTrig cos (\r i -> ((cos r) * (cosh i)) :+ (-1 * ((sin r) * (sinh i))))


module HollowPinwheel where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           FokkerPlanck.FourierSeries
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.List
import           Utils.Time
import Graphics.Rendering.Chart.Easy
import Utils.Parallel
import Graphics.Rendering.Chart.Backend.Cairo
import Utils.SimpsonRule

main = do
  args@(deviceIDsStr:numPointsStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:stdStr:radiusStr:sStr:idxStr:numBatchStr:aStr:bStr:deltaStr:radialFreqStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      std = read stdStr :: Double
      radius = read radiusStr :: Double
      s = read sStr :: Double
      idx = read idxStr :: (Int, Int)
      numBatch = read numBatchStr :: Int
      numThread = read numThreadStr :: Int
      folderPath = "output/test/HollowPinwheel"
  createDirectoryIfMissing True folderPath
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  let harmonicsArray' =
        pinwheelFourierCoefficientsAnatical
          numR2Freq
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          (-s)
          periodR2
          (periodR2 * sqrt 2)
      hpf = idealLowPassFilter radius periodR2 numR2Freq
      hpfMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        hpf
  plotImageRepaComplex (folderPath </> "Filter_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    hpf
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPoints
      numPoints
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
  harmonicsArray <- convolveFrequency plan numR2Freq hpf harmonicsArray'
  -- let harmonicsArray = centerHollow numR2Freq harmonicsArray'
  let pinwheelMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        (harmonicsArray' IA.! idx)
      pinwheelHPMat =
        A.use .
        A.fromList (A.Z A.:. (numR2Freq ^ 2) A.:. (1 :: Int)) . VS.toList $
        (harmonicsArray IA.! idx)
  plotImageRepaComplex (folderPath </> "FilteredPinhweel_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    (harmonicsArray IA.! idx)
  pinwheelR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      pinwheelMat
  pinwheelHPR2 <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      pinwheelHPMat
  hpfR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freq
      numPoints
      1
      periodR2
      1
      numBatch
      hpfMat
  plotImageRepaComplex (folderPath </> "Pinwheel.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    pinwheelR2
  printf
    "pinwheel center: %f\n"
    (magnitude $
     (pinwheelR2 R.!
      (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "Pinwheel_HighPass.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    pinwheelHPR2
  printf
    "lowpassed pinwheel center: %f\n"
    (magnitude $
     (pinwheelHPR2 R.!
      (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "Pinwheel_Coef.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) . VS.convert $
    (harmonicsArray' IA.! idx)
  printf
    "filter center: %f\n"
    (magnitude $
     (hpfR2' R.! (Z :. (0 :: Int) :. (div numPoints 2) :. (div numPoints 2 - 1))))
  plotImageRepaComplex (folderPath </> "HighPass.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
    hpfR2'
  plotImageRepaComplex (folderPath </> "PinwheelHollow.png") .
    ImageRepa 8 .
    computeS .
    extend (Z :. (1 :: Int) :. All :. All) .
    sumS . rotate3D . R.zipWith (-) pinwheelR2 $
    pinwheelHPR2


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE CPP #-}

{-|

  Module      : AutoBench.Internal.Regression
  Description : Generating and fitting models for regression analysis.
  Copyright   : (c) 2018 Martin Handley
  License     : BSD-style
  Maintainer  : martin.handley@nottingham.ac.uk
  Stability   : Experimental
  Portability : GHC

  This module is responsible for generating and fitting regression models.
  See 'LinearType', 'LinearCandidate' and 'LinearFit' for more details.

  The system uses ridge regression because it has a closed-form solution so
  is fast. Previously used LASSO implemented using the HVX library but
  it was too computationally expensive, (see bottom of this file).

-}

{-
   ----------------------------------------------------------------------------
   <TO-DO>:
   ----------------------------------------------------------------------------
   - Too much white space;
   -
-}

module AutoBench.Internal.Regression
  (

  -- * Candidate generation
    generateLinearCandidate    -- Generate a 'LinearCandidate' from a 'LinearType' by filling
                               -- in the details of the model necessary to fit it to a given dataset.
  -- * Regression analysis
  , fitOLS                     -- Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
                               -- fit the model to the data using the OLS method.
  , fitRidgeRegress            -- Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
                               -- fit the model to the data using ridge regression.

  ) where

import           Data.List            (minimumBy, nub)
import           Data.Ord             (comparing)
import qualified Data.Vector.Storable as V

import Numeric.LinearAlgebra
  ( Element
  , Matrix
  , Vector
  , (<>)
  , (<\>)
  , (#>)
  , asColumn
  , cmap
  , diag
  , dot
  , fromColumns
  , fromRows
  , konst
  , norm_2
  , norm_Inf
  , qr
  , rows
  , scale
  , subVector
  , sumElements
  , takeRows
  , toColumns
  , tr'
  , vjoin
  )

import qualified AutoBench.Internal.Expr as E
import AutoBench.Internal.Types
  ( Coord
  , Exp
  , LinearCandidate(..)
  , LinearType(..)
  , numPredictors
  )

#ifdef MIN_VERSION_base
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif
#endif

-- * 'LinearCandidate' generation

-- | Generate a 'LinearCandidate' from a 'LinearType' by filling in the
-- details of the model necessary to fit it to a given dataset.
-- Makes use of the generation helper functions 'generateLinearCandidateExpr'
-- and 'generateLinearCandidateYHat'.
generateLinearCandidate :: LinearType -> LinearCandidate
generateLinearCandidate t@(Poly _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = id
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(Log b _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap $ logBase (fromIntegral b)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(PolyLog b _) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap (\x -> x * logBase (fromIntegral b) x)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }
generateLinearCandidate t@(Exp n) =
  LinearCandidate
   {
     _lct   = t
   , _fxs   = cmap (flip (**) $ fromIntegral n)
   , _fex   = generateLinearCandidateExpr t
   , _fyhat = generateLinearCandidateYHat t
   }

-- * Linear fitting

-- | Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
-- fit the model to the data using the ordinary least squares (OLS) method. The
-- result is the model's coefficients.
fitOLS :: [Coord] -> LinearCandidate -> Vector Double
fitOLS coords c = head $ toColumns $ (_preds <\> _resps)
  where (_preds, _resps) = genFitMatrices coords c

-- | Given a data set (a list of xy-'Coord's) and a 'LinearCandidate',
-- fit the model to the data using ridge regression. The result is the model's
-- coefficients.
fitRidgeRegress :: [Coord] -> LinearCandidate -> Vector Double
fitRidgeRegress coords c =
  -- Select the lambda that minimises mean squared error.
  fst $ minimumBy (comparing snd) $ zip fits errs
  where
    n        = length coords
    (xs, ys) = unzip  coords
    vXs      = V.fromList xs
    vYs      = V.fromList ys
    -- Fitting error is mean squared error.
    calcMse m vec = ((norm_2 $ V.zipWith (-) vYs vec) ** 2.0) / fromIntegral m
    -- Generate the matrices to solve.
    (_preds, _resps) = genFitMatrices coords c
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- Using ridge regression with maximum tuning parameter:
    -- https://stats.stackexchange.com/questions/289075/what-is-the-
    -- smallest-lambda-that-gives-a-0-component-in-lasso?noredirect=1&lq=1
    ---------------------------------------------------------------------------
    maxLam   = (1.0 / fromIntegral n) * norm_Inf (tr' _preds <> _resps)
    -- Minimise 'mse' to select best tuning parameter.
    lamStep  = maxLam / 10
    lamRange = [lamStep, 2 * lamStep .. maxLam]
    fits     = nub $ fmap (head
                            . toColumns
                            . ridgeRegress _preds _resps
                          ) lamRange
    -- ** Do not use vFxs here, use vXs **
    errs     = fmap (calcMse n . flip (_fyhat c) vXs) fits
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- Closed form solution of ridge regression.
    -- https://github.com/chris-taylor/aima-haskell/blob/master/src/AI/
    -- Learning/LinearRegression.hs
    ---------------------------------------------------------------------------
    ridgeRegress
      :: Matrix Double  -- X
      -> Matrix Double  -- Y
      -> Double         -- λ
      -> Matrix Double  -- (X'X + λI)^−1 X'Y.
    ridgeRegress x y lambda =
      let m      = rows x
          (_, r) = qr x
          rr     = takeRows m r
          ww     = diag $ vjoin [0, konst 1 (m - 1)]
      in (tr' rr <> rr + (lambda `scale` ww)) <\> (tr' x <> y)

-- * Helpers

-- | Generate an expression ('Exp') for a 'LinearCandidate' that will allow
-- its model's equation to be pretty printed at a later time. Important note:
-- the vectors are assumed to be of the correct form, /no/  checking is
-- performed currently.
generateLinearCandidateExpr
   :: LinearType      -- Model.
   -> Vector Double   -- Coefficients of the model.
   -> Exp             -- Model's equation as an expression for pretty printing.
--
-- Given coefficients:
--
--   a_0, a_1, ..., a_n
--
-- Generate:
--
--   a_0 + a_1 * x^1 + a_2 * x^2 + ... + a_n * x^n
--
generateLinearCandidateExpr (Poly n) coeffs =
  foldr1 E.Add . take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                         -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k) (E.Var "x")    -- a_1 * x
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)         -- a_k * x^k, k > 1
                                          (E.Pow (E.Var "x")
                                                 (E.Num p)
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1, ..., a_n
--
-- Generate:
--
-- a_0 + a_1 * log_b^1(x) + a_2 * log_b^2(x) + ... + a_n * log_b^n(x)
--
generateLinearCandidateExpr (Log b n) coeffs =
  foldr1 E.Add $ take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                                -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k)                       -- a_1 * log_b(x)
                                   (E.Log (E.Num $ fromIntegral b)
                                          (E.Var "x")
                                   )
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)                -- a_k * log_b^k(x), k > 1
                                          (E.Pow (E.Log (E.Num $ fromIntegral b)
                                                        (E.Var "x")
                                                 )
                                                 (E.Num p)
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1, ..., a_n
--
-- Generate:
--
-- a_0 + a_1 * x^1 * log_b^1(x) + a_2 * x^2 * log_b^2(x) + ... + a_n * x^n * log_b^n(x)
--
generateLinearCandidateExpr (PolyLog b n) coeffs =
  foldr1 E.Add $ take (n + 1) $ zipWith ($) terms (V.toList coeffs)
  where terms = E.Num                                                -- a_0
                 : ( \a_k -> E.Mul (E.Num a_k)                       -- a_1 * x * log_b(x)
                                   (E.Mul (E.Var "x")
                                          (E.Log (E.Num $ fromIntegral b)
                                                 (E.Var "x")
                                          )
                                   )
                   )
                 : fmap ( \p a_k -> E.Mul (E.Num a_k)                -- a_k * x^k * log_b^k(x), k > 1
                                          (E.Mul (E.Pow (E.Var "x")
                                                        (E.Num p)
                                                 )
                                                 (E.Pow (E.Log (E.Num $ fromIntegral b)
                                                               (E.Var "x")
                                                        )
                                                        (E.Num p)
                                                 )
                                          )
                        ) [2.0..]
--
-- Given coefficients:
--
-- a_0, a_1
--
-- Generate:
--
-- a_0 + a_1 * n^x
--
generateLinearCandidateExpr (Exp n) coeffs =
  foldr1 E.Add $ zipWith ($) terms (V.toList coeffs)
  where terms = [ E.Num  -- a_0
                , \a_1 -> E.Mul (E.Num a_1)   -- a_1 *
                                (E.Pow (E.Num $ fromIntegral n) -- n^x
                                        (E.Var "x")
                                )
                ]

-- | For a given set of x-coordinates, generate the y-coordinates that are
-- predicted by the model of a 'LinearCandidate'. Important note: the vectors
-- are assumed to be of the correct form, /no/ checking is performed currently.
generateLinearCandidateYHat
  :: LinearType
  -> Vector Double   -- Coefficients of the model.
  -> Vector Double   -- x-coordinates.
  -> Vector Double   -- Predicted y-coordinates.
--
-- Given coefficients and x-coords:
--
--   a_0, a_1, ..., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
--   a_0 + a_1 * x_i^1 + a_2 * x_i^2 + ... + a_n * x_i^n
--
generateLinearCandidateYHat (Poly n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith vecPow [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1, ..., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * log_b^1(x_i) + a_2 * log_b^2(x_i) + ... + a_n * log_b^n(x_i)
--
generateLinearCandidateYHat (Log b n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecLogPow $ fromIntegral b) [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1, .., a_n     x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * x_i^1 * log_b^1(x_i) + a_2 * x_i^2 * log_b^2(x_i) + ... + a_n * x_i^n * log_b^n(x_i)
--
generateLinearCandidateYHat (PolyLog b n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecPolyLogPow $ fromIntegral b) [1.0..] (replicate n xs)
--
-- Given coefficients and x-coords:
--
-- a_0, a_1      x_0, x_1, ..., x_m
--
-- For each x_i, generate:
--
-- a_0 + a_1 * n^x_i
--
generateLinearCandidateYHat (Exp n) coeffs xs =
  (fromColumns $ V.fromList (replicate m 1.0) : xs') #> coeffs
  where
    m   = V.length xs
    xs' = zipWith (vecBasePow $ fromIntegral n) [1.0..] (replicate 1 xs)

-- | Constructs predictor (@preds@) and responder (@resps@) matrices so that
-- regression equations can be solved in their matrix form.
genFitMatrices :: [Coord] -> LinearCandidate -> (Matrix Double, Matrix Double)
genFitMatrices coords c = gen (numPredictors (_lct c) - 1)
  where

    gen :: Int -> (Matrix Double, Matrix Double)
    gen k = let
               _xPows    = xPows k
               _xSumPows = xSumPows _xPows
               _ySums    = ySums _xPows k
               _preds    = preds _xSumPows k
               _resps    = resps _ySums
            in (_preds, asColumn _resps)

    ---------------------------------------------------------------------------
    --
    --
    (xs, ys) = unzip coords
    -- Transform the x-coords using the candidate's 'fxs' function. This allows
    -- us to fit different models to the data in a generic way. For example, a
    -- 'Log b n' candidate will first take 'LogBase b' of the x-coords. We can
    -- then use (a variant of) linear least squares to fit a straight line
    -- through the resulting data points. Although this is a straight line on
    -- the /resulting data/, it corresponds to a logarithmic function on the
    -- /initial/ data. This technique is standard.
    vFxs = (_fxs c) (V.fromList xs)
    vYs  = V.fromList ys
    -- Number of data points.
    n    = length coords
    -- As the x-'Coords' have been transformed by the candidate's 'fxs'
    -- function, this method works generically for any model and 'k'
    -- coefficients.
    ---------------------------------------------------------------------------
    --
    -- Helper functions to construct the matrices to solve using linear
    -- regression. Note: these functions are to be used in a particular way so
    -- their names reflect their use.
    --
    -- x_i^j for j = 1, 2, .., 2 * k
    xPows :: Int -> [Vector Double]
    xPows k = zipWith vecPow [1.0..] (replicate (2 * k) vFxs)
    -- SUM x_i^j for j = 1, 2, .., 2 * k
    xSumPows :: [Vector Double] -> [Double]
    xSumPows  = fmap sumElements
    -- SUM y_i, SUM x_i * y_i, SUM x_i^2 * y_i, .., SUM x_i^j * y_i
    ySums :: [Vector Double] -> Int -> [Double]
    ySums vs k = zipWith dot (V.fromList (replicate n 1.0) : vs) (replicate (k + 1) vYs)
    --
    -- Convert xSumPows into a matrix by 'chunking' off rows from the initial vector.
    --
    -- Rows are length k + 1, the initial rows are:
    --
    --   n      , SUM x_i  , SUM x_i^2, .., SUM x_i^k
    --
    --   SUM x_i, SUM x_i^2, SUM x_i^3, .., SUM x_i^(k + 1)
    --
    -- The last row is:
    --
    --   SUM x_i^j, SUM x_i^(j + 1), .., SUM x_i^(2 * k)
    --
    -- These values are the /predictors/.
    --
    preds :: [Double] -> Int -> Matrix Double
    preds ds k = vecToMatrix (V.fromList $ fromIntegral n : ds) (k + 1)
    -- 'ySums' are the /responders/.
    resps :: [Double] -> Vector Double
    resps  = V.fromList

-- | Raise each double in a vector of doubles to a given power.
vecPow :: Double -> Vector Double -> Vector Double
vecPow pow = cmap (** pow)

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> base ** (x * pow)
--
-- for a given 'base' and 'pow'.
vecBasePow :: Double -> Double -> Vector Double -> Vector Double
vecBasePow base pow = cmap (\x -> base ** (x * pow))

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> (x ** pow) * (logBase b x ** pow)
--
-- for a given 'b' and 'pow'.
vecPolyLogPow :: Double -> Double -> Vector Double -> Vector Double
vecPolyLogPow b pow = cmap (\x -> (x ** pow) * ((logBase b x) ** pow))

-- | For each double 'x' in a vector of doubles, apply the following
-- calculation:
--
-- > \x -> logBase b x ** pow
--
-- for a given 'b' and 'pow'.
vecLogPow :: Double -> Double -> Vector Double -> Vector Double
vecLogPow b pow = cmap (\x -> (logBase b x) ** pow)

-- | Create a matrix by 'chunking' off rows of a given length from a vector.
-- Works similarly to 'tails'.
vecToMatrix :: (V.Storable a, Element a) => Vector a -> Int -> Matrix a
vecToMatrix v j = fromRows [ subVector i j v | i <- [0..j - 1] ]

{-

LASSO has been retired for the time being due to computational expense.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DataKinds        #-}

import           HVX.Internal.Constraints ((<=~))
import           HVX.Primitives           ((+~), (*~), neg, norm, quadform)
import           HVX.Internal.Solvers     (ellipsoidMinimize)
import qualified HVX.Internal.Primitives  as HVX

-- ** LASSO (Least Absolute Shrinkage and Selection Operator)

fitLASSO :: [Coord] -> LinearCandidate -> Vector Double
fitLASSO coords c =
  -- Select the lambda that minimises mean squared error.
  fst $ minimumBy (comparing snd) $ zip fits errs
  where
    n        = length coords
    (xs, ys) = unzip  coords
    vXs      = V.fromList xs
    vYs      = V.fromList ys
    -- Fitting error is mean squared error.
    calcMse m vec = ((norm_2 $ V.zipWith (-) vYs vec) ** 2.0) / fromIntegral m
    -- Generate the matrices to solve.
    (_preds, _resps) = genFitMatrices coords c
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    initCoeffs = _preds <\> _resps
    maxT = V.sum $ cmap abs $ flatten initCoeffs
    -- Minimise 'mse' to select best tuning parameter.
    tStep  = maxT / 10
    tRange = [0 + tStep, 0 + (2 * tStep) .. maxT]
    fits   = filter (all (\m -> not (isNaN m || isInfinite m)) . V.toList)
              . nub
              $ fmap (head
                       . toColumns
                       . lassoRegress _preds _resps
                     ) tRange
    -- ** Do not use vFxs here, use vXs **
    errs   = fmap (calcMse n . flip (fyhat c) vXs) fits
    ---------------------------------------------------------------------------
    ---------------------------------------------------------------------------
    -- LASSO regression (convex equation using the subgradient method in HVX).
    lassoRegress
      :: Matrix Double      -- ^ n >< n matrix X.
      -> Matrix Double      -- ^ n >< 1 matrix Y.
      -> Double             -- ^ Constraint on sum of coefficients.
      -> Matrix Double      -- ^ Optimal variable values.
    lassoRegress x y t = snd . head $ res
      where
        (res, _, _) = ellipsoidMinimize expr [constr] [("a", m)] 1e-20 (2 * maxT)
        m      = rows x
        constr = norm 1 (HVX.EVar "a") <=~ (HVX.EConst $ konst t (1, 1))
        expr   = quadform (HVX.EConst $ ident m) $ ((HVX.EConst y) +~ neg (HVX.EConst x *~ HVX.EVar "a"))

-}


module Main where

import RayLib
import Numeric.LinearAlgebra
import qualified Graphics.Image as G

main :: IO ()
main = do let screen = Screen 256 256
          let ka = MatParam (fromList [0,1,1])
          let ks = MatParam (fromList [1,1,0])
          let kd = MatParam (fromList [0.4,0.4,0])
          let ke = MatParam (fromList [0,0,0])
          let kr = MatParam (fromList [0,0,0])
          let kt = MatParam (fromList [0,0,0])
          let shininess = 64
          let material = Material ka ks kd ke kr kt shininess
          let sphere1 = Sphere (fromList [0,0,0]) 1.5 material
          let scene = Scene [] [sphere1] (fromList [1,0,0.1])
          putStrLn . show $ scene
          let camera = Camera (fromList [0,0,-4]) (fromList [0,0,1]) (fromList [0,1,0]) 1 90
          putStrLn . show $ camera
          putStrLn . show $ screen
          let image = traceImage scene screen camera
          let toWrite = render image
          G.writeImage "test.png" toWrite


{-# LANGUAGE OverloadedStrings, ExistentialQuantification, FlexibleContexts #-}

module BNN where

import System.Random.MWC (createSystemRandom)

--import Data.Aeson (ToJSON(toJSON), Value)
--import Data.Aeson (encode)
import qualified Data.ByteString.Lazy.Char8 as BL

import Data.List (sort, partition)
import qualified Data.Histogram as DH
import qualified Data.Histogram.Fill as DF
import qualified Data.Histogram.Bin.Bin2D as Bin2D
import Data.Vector.Unboxed.Base (Unbox)
import qualified Data.Vector.Unboxed as Vec

import Control.Monad (liftM2, replicateM, forM, forM_)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Sampler
import Control.Monad.Bayes.Traced
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Inference.SMC as SMC
import Control.Monad.Bayes.Inference.RMSMC as RMSMC
import Control.Monad.Bayes.Sequential
import Control.Monad.Bayes.Population
import Control.Monad.Bayes.Traced.Static (Traced)
import Numeric.LinearAlgebra (Matrix, Vector, toList, vector, matrix, dot, (#>), (!), cmap, scalar, size)

import Numeric.Log

-- Plotting

histo2D :: (Foldable f, Unbox val, Num val) =>
    (Double, Int, Double)
    -> (Double, Int, Double)
    -> f (Double, Double)
    -> DH.Histogram (Bin2D.Bin2D DH.BinD DH.BinD) val
histo2D (xmin, xn, xmax) (ymin, yn, ymax) = DF.fillBuilder buildr
  where
    binX = DH.binD xmin xn xmax
    binY = DH.binD ymin yn ymax
    bins = Bin2D.Bin2D binX binY
    buildr = DF.mkSimple bins

histo :: (Foldable v, Unbox a, Num a) =>
         (Double, Int, Double)
      -> v Double
      -> DH.Histogram DF.BinD a
histo (xmin, n, xmax) = DF.fillBuilder buildr
  where
    bins = DH.binD xmin n xmax
    buildr = DF.mkSimple bins

-- A SIMPLE NEURAL NETWORK
-- Model Setup

data NN
  = NN
      { biass :: Vector Double,
        weights :: Vector Double,
        sigma :: Double
      }
  deriving (Eq, Show)

data Data
  = Data
      { xValue :: Double,
        yValue :: Double
      }
  deriving (Eq, Show)

forwardNN :: NN -> Double -> Double
forwardNN (NN bs ws _) x =
    ws `dot` cmap activation (scalar x - bs)
  where activation x = if x < 0 then 0 else 1

errorModel :: Double -> Double -> Double -> Log Double
errorModel mean std = normalPdf mean std

likelihood :: NN -> Data -> Log Double
likelihood nn (Data xObs yObs) =
    errorModel yMean ySigma yObs
  where
    ySigma = sigma nn
    yMean = forwardNN nn xObs

-- Prior, Posterior and Predictive Distribution

postNN :: MonadInfer m => m NN -> [Data] ->  m NN
postNN pr obs = do
  nn <- pr
  forM_ obs (score . likelihood nn)
  return nn

uniformVec :: MonadSample m => (Double, Double) -> Int -> m (Vector Double)
uniformVec (wmin, wmax) nelements =
  vector <$> replicateM nelements (uniform wmin wmax)

priorNN :: MonadSample m => Int -> m NN
priorNN nnodes = do
  bias <- uniformVec (0, 10) nnodes
  weight <- uniformVec (-10, 10) nnodes
  sigma <- uniform 0.5 1.5
  return $ NN bias weight sigma

predDist :: MonadInfer m => m NN -> m (NN, Data)
predDist pr = do
  nn <- pr
  x <- uniform 0 10
  y <- uniform (-5) 10
  score $ likelihood nn (Data x y)
  return (nn, Data x y)


module AI.Learning.Bootstrap where

import Control.Monad.Random
import Foreign.Storable (Storable)
import Numeric.LinearAlgebra
import qualified Data.List as L

import AI.Util.Matrix
import AI.Util.Util

---------------
-- Bootstrap --
---------------

-- |Generate a bootstrap sample of size @sz@.
genBootstrapSample :: RandomGen g => Int -> Rand g [Int]
genBootstrapSample sz = go sz []
    where go 0 accum = return accum
          go n accum = do
            i <- getRandomR (0,sz-1)
            go (n - 1) (i:accum)

sampleVector :: (Storable a, RandomGen g) => Vector a -> Rand g (Vector a)
sampleVector v = do
    idx <- genBootstrapSample (dim v)
    return (v `subRefVec` idx)

sampleMatrixRows :: (Element a, RandomGen g) => Matrix a -> Rand g (Matrix a)
sampleMatrixRows m = do
    idx <- genBootstrapSample (rows m)
    return $ m `subRefRows` idx

sampleMatrixCols :: (Element a, RandomGen g) => Matrix a -> Rand g (Matrix a)
sampleMatrixCols m = do
    idx <- genBootstrapSample (cols m)
    return $ m `subRefCols` idx

-- |Generate a bootstrap sample of a statistic from a data set of type /a/.
bootStrapResample :: RandomGen g =>
                     (a -> Rand g a)    -- Sampling function
                  -> (a -> b)           -- Statistic to be resampled
                  -> Int                -- Number of resamples
                  -> a                  -- Data
                  -> Rand g [b]
bootStrapResample sample func nSamples x = go [] nSamples
    where
        go samples 0 = return samples
        go samples n = do
            x' <- sample x
            go (func x' : samples) (n-1)


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    , normalDistr
    , normalDistrE
    , normalDistrErr
    , standard
    ) where

import Control.Applicative
import Data.Aeson            (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)
import qualified System.Random.MWC.Distributions as MWC
import qualified Data.Vector.Generic as G

import qualified Statistics.Distribution as D
import qualified Statistics.Sample as S
import Statistics.Internal

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show NormalDistribution where
  showsPrec i (ND m s _ _) = defaultShow2 "normalDistr" m s i
instance Read NormalDistribution where
  readPrec = defaultReadPrecM2 "normalDistr" normalDistrE

instance ToJSON NormalDistribution
instance FromJSON NormalDistribution where
  parseJSON (Object v) = do
    m  <- v .: "mean"
    sd <- v .: "stdDev"
    either fail return $ normalDistrErr m sd
  parseJSON _ = empty

instance Binary NormalDistribution where
    put (ND m sd _ _) = put m >> put sd
    get = do
      m  <- get
      sd <- get
      either fail return $ normalDistrErr m sd

instance D.Distribution NormalDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.ContDistr NormalDistribution where
    logDensity    = logDensity
    quantile      = quantile
    complQuantile = complQuantile

instance D.MaybeMean NormalDistribution where
    maybeMean = Just . D.mean

instance D.Mean NormalDistribution where
    mean = mean

instance D.MaybeVariance NormalDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Variance NormalDistribution where
    stdDev = stdDev

instance D.Entropy NormalDistribution where
  entropy d = 0.5 * log (2 * pi * exp 1 * D.variance d)

instance D.MaybeEntropy NormalDistribution where
  maybeEntropy = Just . D.entropy

instance D.ContGen NormalDistribution where
    genContVar d = MWC.normal (mean d) (stdDev d)

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistr :: Double            -- ^ Mean of distribution
            -> Double            -- ^ Standard deviation of distribution
            -> NormalDistribution
normalDistr m sd = either error id $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistrE :: Double            -- ^ Mean of distribution
             -> Double            -- ^ Standard deviation of distribution
             -> Maybe NormalDistribution
normalDistrE m sd = either (const Nothing) Just $ normalDistrErr m sd

-- | Create normal distribution from parameters.
--
normalDistrErr :: Double            -- ^ Mean of distribution
               -> Double            -- ^ Standard deviation of distribution
               -> Either String NormalDistribution
normalDistrErr m sd
  | sd > 0    = Right $ ND { mean       = m
                           , stdDev     = sd
                           , ndPdfDenom = log $ m_sqrt_2_pi * sd
                           , ndCdfDenom = m_sqrt_2 * sd
                           }
  | otherwise = Left $ errMsg m sd

errMsg :: Double -> Double -> String
errMsg _ sd = "Statistics.Distribution.Normal.normalDistr: standard deviation must be positive. Got " ++ show sd

-- | Variance is estimated using maximum likelihood method
--   (biased estimation).
--
--   Returns @Nothing@ if sample contains less than one element or
--   variance is zero (all elements are equal)
instance D.FromSample NormalDistribution Double where
  fromSample xs
    | G.length xs <= 1 = Nothing
    | v == 0           = Nothing
    | otherwise        = Just $! normalDistr m (sqrt v)
    where
      (m,v) = S.meanVariance xs

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0

complQuantile :: NormalDistribution -> Double -> Double
complQuantile d p
  | p == 0         = inf
  | p == 1         = -inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.complQuantile: p must be in [0,1] range. Got: "++show p
  where x          = invErfc (2 * p)
        inf        = 1/0


module DeltaSigma.PowerSpectrumDensity where

import Numeric.FFT
import Data.Complex

psd :: [Double] -> [Double]
psd xs = take n_half . fmap (/ (fromIntegral n_half))
        . fmap magnitude . fft . fmap (:+ 0) $ take n xs
  where
    n = length xs
    n_half = n `div` 2


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

{- | The 'Numerical.HBLAS.BLAS' module provides a fully general
yet type safe 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(
        GemvFun
        ,GemmFun
        ,TrsvFun

        ,dgemm
        ,sgemm
        ,cgemm
        ,zgemm

        ,sgemv
        ,dgemv
        ,cgemv
        ,zgemv

        ,sger
        ,dger

        ,strsv
        ,dtrsv
        ,ctrsv
        ,ztrsv
            ) where

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

sgemm :: PrimMonad m=>  GemmFun Float  orient  (PrimState m) m
sgemm =  gemmAbstraction "sgemm"  cblas_sgemm_safe cblas_sgemm_unsafe (\x f -> f x )

dgemm :: PrimMonad m=>  GemmFun  Double orient  (PrimState m) m
dgemm = gemmAbstraction "dgemm"  cblas_dgemm_safe cblas_dgemm_unsafe  (\x f -> f x )

cgemm :: PrimMonad m=>  GemmFun (Complex Float) orient  (PrimState m) m
cgemm = gemmAbstraction "cgemm" cblas_cgemm_safe cblas_cgemm_unsafe  withRStorable_

zgemm :: PrimMonad m=>  GemmFun (Complex Double) orient  (PrimState m) m
zgemm = gemmAbstraction "zgemm"  cblas_zgemm_safe cblas_zgemm_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_
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

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

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


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RecordWildCards #-}
module Data.Matrix.Static.LinearAlgebra
    ( module Data.Matrix.Static.LinearAlgebra.Types
    , Arithmetic(..)
    , Factorization(..)
    , LinearAlgebra(..)
    , EigenOptions(..)
    , defaultEigenOptions
    , SortRule(..)

      -- * Dense matrix operation
    , zeros
    , ones
    , inverse
    , eig
    , svd
    , cond
    ) where

import qualified Data.Vector.Storable as VS
import Data.Vector.Storable (Vector)
import System.IO.Unsafe (unsafePerformIO)
import Data.Complex (Complex)
import Data.Singletons.Prelude hiding ((@@), type (==), type (-), type (<=))
import Data.Type.Equality
import GHC.TypeLits
import Data.Maybe (fromMaybe)

import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Sparse as S
import qualified Data.Matrix.Static.Generic.Mutable as CM
import qualified Data.Matrix.Static.Generic as C
import qualified Data.Matrix.Static.Internal as Internal
import Data.Matrix.Static.LinearAlgebra.Types
import Data.Matrix.Static.LinearAlgebra.Internal

class Arithmetic (mat1 :: C.MatrixKind) (mat2 :: C.MatrixKind) where
    -- | Matrix multiplication between different types of matrices.
    (@@) :: ( Numeric a, SingI n, SingI m
            , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
         => mat1 n p Vector a
         -> mat2 p m Vector a
         -> mat3 n m Vector a
    infixr 8 @@

    -- | Element-wise addition between different types of matrices.
    (%+%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %+%

    -- | Element-wise substraction between different types of matrices.
    (%-%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 D.Matrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %-%

    -- | Element-wise multiplication between different types of matrices.
    (%*%) :: ( Numeric a, SingI n, SingI m
             , If (mat1 == mat2) mat1 S.SparseMatrix ~ mat3 )
          => mat1 n m Vector a
          -> mat2 n m Vector a
          -> mat3 n m Vector a
    infixr 8 %*%

instance Arithmetic D.Matrix D.Matrix where
    (@@) = withFun2 Internal.c_dd_mul
    (%+%) = (+)
    (%-%) = (-)
    (%*%) = (*)

instance Arithmetic D.Matrix S.SparseMatrix where
    (@@) = withDS Internal.c_ds_mul
    (%+%) = flip (%+%)
    (%-%) a b = a %+% C.map negate b
    (%*%) = undefined

instance Arithmetic S.SparseMatrix D.Matrix where
    (@@) = withSD Internal.c_sd_mul
    (%+%) = withSD Internal.c_sd_plus
    (%-%) a b = a %+% C.map negate b
    (%*%) = undefined

instance Arithmetic S.SparseMatrix S.SparseMatrix where
    (@@) = withSS Internal.c_ss_mul
    (%+%) = withSS Internal.c_ss_plus
    (%-%) a b = a %+% C.map negate b
    (%*%) = withSS Internal.c_ss_cmul

class LinearAlgebra (mat :: C.MatrixKind) where
    ident :: (Numeric a, SingI n) => mat n n Vector a

    colSum :: (Numeric a, SingI n, C.Matrix mat Vector a)
           => mat m n Vector a
           -> Matrix 1 n a
    colSum mat = D.create $ do
        m <- CM.replicate 0
        flip C.imapM_ mat $ \(_,j) v -> CM.unsafeModify m (+v) (0, j)
        return m
    {-# INLINE colSum #-}

    rowSum :: (Numeric a, SingI m, C.Matrix mat Vector a)
           => mat m n Vector a
           -> Matrix m 1 a
    rowSum mat = D.create $ do
        m <- CM.replicate 0
        flip C.imapM_ mat $ \(i,_) x -> CM.unsafeModify m (+x) (i, 0)
        return m
    {-# INLINE rowSum #-}

instance LinearAlgebra D.Matrix where
    ident = D.diag 0 $ D.replicate 1

instance LinearAlgebra S.SparseMatrix where
    ident = S.diag $ D.replicate 1

data EigenOptions = EigenOptions
    { _sigma :: Maybe Double -- ^ Find eigenvalues near sigma using shift-invert mode.
    , _ncv :: Maybe Int      -- ^ Parameter that controls the convergence
                             -- speed of the algorithm. Typically a larger ncv
                             -- means faster convergence, but it may also result
                             -- in greater memory use and more matrix operations
                             -- in each iteration. This parameter must satisfy
                             -- k+2 <= ncv <= n, and is advised to take ncv > 2k+1.
                             -- Default: min(n, max(2*k + 1, 20))
    , _maxit :: Maybe Int          -- ^ Maximum number of iterations allowed in the algorithm.
    , _tol   :: Double       -- Precision parameter for the calculated eigenvalues.
    , _sort_rule :: SortRule
    }

defaultEigenOptions :: EigenOptions
defaultEigenOptions = EigenOptions
    { _sigma = Nothing
    , _ncv = Nothing
    , _maxit = Nothing
    , _tol = 2.2204460492503131e-16
    , _sort_rule = LM }

data SortRule = LM
              | LR
              | LI
              | LA
              | SM
              | SR
              | SI
              | SA
              | BE
              deriving (Enum)

class Factorization mat where
    -- | Eigenvalues (from largest to smallest) and
    -- eigenvectors (as columns) of a general square matrix.
    eigS :: (SingI k, SingI n, k <= n - 2)
         => Sing k
         -> mat n n Vector Double
         -> (Matrix k 1 (Complex Double), Matrix n k (Complex Double))
    eigS a b = eigS' a b defaultEigenOptions

    eigS' :: (SingI k, SingI n, k <= n - 2)
         => Sing k
         -> mat n n Vector Double
         -> EigenOptions
         -> (Matrix k 1 (Complex Double), Matrix n k (Complex Double))

    -- | Eigenvalues (from largest to smallest) and
    -- eigenvectors (as columns) of a symmetric square matrix.
    eigSH :: (SingI k, SingI n, k <= n - 1)
          => Sing k
          -> mat n n Vector Double
          -> (Matrix k 1 Double, Matrix n k Double)
    eigSH a b = eigSH' a b defaultEigenOptions

    eigSH' :: (SingI k, SingI n, k <= n - 1)
           => Sing k
           -> mat n n Vector Double
           -> EigenOptions
           -> (Matrix k 1 Double, Matrix n k Double)

    -- | Generalized eigen solver for real symmetric matrices
    geigSH :: (SingI k, SingI n, k <= n - 1)
           => Sing k
           -> mat n n Vector Double
           -> SparseMatrix n n Double
           -> (Matrix k 1 Double, Matrix n k Double)
    geigSH a b c = geigSH' a b c defaultEigenOptions

    geigSH' :: (SingI k, SingI n, k <= n - 1)
            => Sing k
            -> mat n n Vector Double
            -> SparseMatrix n n Double
            -> EigenOptions
            -> (Matrix k 1 Double, Matrix n k Double)

    -- | Cholesky decomposition
    cholesky :: (Numeric a, SingI n) => mat n n Vector a -> mat n n Vector a

instance Factorization D.Matrix where
    eigS' s mat EigenOptions{..}
        | D.all (==0) mat = ( D.replicate 0, D.replicate 1)
        | otherwise = unsafePerformIO $ do
            m1 <- CM.new
            m2 <- CM.new
            info <- fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
                unsafeWith' m2 $ \v2 _ _ -> unsafeWith mat $ \v n _ ->
                    Internal.c_eigs k v1 v2 v n ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
            case info of
                Nothing -> do
                    m1' <- C.unsafeFreeze m1
                    m2' <- C.unsafeFreeze m2
                    return (m1', m2')
                Just err -> error err
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigS' #-}

    eigSH' s mat EigenOptions{..}
        | D.all (==0) mat = (D.replicate 0, D.replicate 1)
        | otherwise = unsafePerformIO $ do
            m1 <- CM.new
            m2 <- CM.new
            ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
                unsafeWith' m2 $ \v2 _ _ -> unsafeWith mat $ \v n _ ->
                    Internal.c_eigsh k v1 v2 v n ncv' maxit _tol mode sigma $
                    fromIntegral $ fromEnum _sort_rule ) >>= maybe
                    ( do
                        m1' <- C.unsafeFreeze m1
                        m2' <- C.unsafeFreeze m2
                        return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigSH' #-}

    geigSH' s matA matB EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWith matA $ \v n _ ->
            unsafeWithS matB $ \pv pin po _ _ size ->
            Internal.c_geigsh k v1 v2 v n pv po pin size
            ncv' maxit _tol $ fromIntegral $ fromEnum _sort_rule ) >>= maybe
            ( do
                m1' <- C.unsafeFreeze m1
                m2' <- C.unsafeFreeze m2
                return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows matA
    {-# INLINE geigSH' #-}

    cholesky mat = flip withFun1 mat $
        \code p1 c1 _ p2 _ _ -> Internal.c_cholesky code p1 p2 c1
    {-# INLINE cholesky #-}

instance Factorization S.SparseMatrix where
    eigS' s mat EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWithS mat $ \pv pin po n _ size ->
                Internal.c_seigs k v1 v2 pv po pin n size
                ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
                ) >>= maybe
            ( do
                m1' <- C.unsafeFreeze m1
                m2' <- C.unsafeFreeze m2
                return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigS' #-}

    eigSH' s mat EigenOptions{..} = unsafePerformIO $ do
        m1 <- CM.new
        m2 <- CM.new
        ( fmap Internal.computationInfo $ unsafeWith' m1 $ \v1 _ _ ->
            unsafeWith' m2 $ \v2 _ _ -> unsafeWithS mat $ \pv pin po n _ size ->
                Internal.c_seigsh k v1 v2 pv po pin n size
                ncv' maxit _tol mode sigma $ fromIntegral $ fromEnum _sort_rule
                ) >>= maybe
                ( do
                    m1' <- C.unsafeFreeze m1
                    m2' <- C.unsafeFreeze m2
                    return (m1', m2') ) error
      where
        ncv' = case _ncv of
            Just x -> fromIntegral x
            Nothing -> min d $ max 20 $ 2 * k + 1
        maxit = case _maxit of
            Just x -> fromIntegral x
            Nothing -> d * 10
        mode = case _sigma of
            Nothing -> 0
            _ -> 1
        sigma = fromMaybe 0 _sigma
        k = fromIntegral $ fromSing s
        d = fromIntegral $ D.rows mat
    {-# INLINE eigSH' #-}

    cholesky = undefined

type family R a where
    R Float = Float
    R Double = Double
    R (Complex Double) = Double
    R (Complex Float) = Float

zeros :: (SingI m, SingI n) => Matrix m n Double
zeros = D.replicate 0
{-# INLINE zeros #-}

ones :: (SingI m, SingI n) => Matrix m n Double
ones = D.replicate 1
{-# INLINE ones #-}

-- | The inverse of a dense matrix.
inverse :: (SingI n, Numeric a) => Matrix n n a -> Matrix n n a
inverse = withFun1 Internal.c_inverse
{-# INLINE inverse #-}

-- | Compute the full eigendecomposition for dense matrix.
eig :: forall n . SingI n
    => Matrix n n Double
    -> (Matrix n 1 (Complex Double), Matrix n n (Complex Double))
eig mat = unsafePerformIO $ do
    m1 <- CM.new
    m2 <- CM.new
    _ <- unsafeWith' m1 $ \v1 _ _ -> unsafeWith' m2 $ \v2 _ _ ->
        unsafeWith mat $ \v n _ -> Internal.c_eig v1 v2 v n
    m1' <- C.unsafeFreeze m1
    m2' <- C.unsafeFreeze m2
    return (m1', m2')
{-# INLINE eig #-}

-- | Compute the full singular value decomposition for dense matrix.
svd :: forall n p a m. (Numeric (R a), Numeric a, SingI n, SingI p, SingI m, m ~ Min n p)
    => Matrix n p a
    -> (Matrix n m a, Matrix m 1 (R a), Matrix p m a)
svd mat = unsafePerformIO $ do
    mu <- CM.new
    ms <- CM.new
    mv <- CM.new
    checkResult $ unsafeWith' mu $ \pu _ _ -> unsafeWith' ms $ \ps _ _ ->
        unsafeWith' mv $ \pv _ _ -> unsafeWith mat $ \px r c ->
            Internal.c_bdcsvd (foreignType (undefined :: a))
                pu ps pv px r c
    u <- C.unsafeFreeze mu
    s <- C.unsafeFreeze ms
    v <- C.unsafeFreeze mv
    return (u, s, v)
{-# INLINE svd #-}

-- | Condition number.
cond :: ( Numeric a, Numeric (R a), Ord (R a), Fractional (R a)
        , SingI n, SingI m, SingI (Min n m))
     => Matrix n m a -> R a
cond mat = VS.maximum val / VS.minimum val
  where
    val = VS.filter (/=0) $ D.flatten s
    (_,s,_) = svd mat
{-# INLINE cond #-}


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Control.Monad.Random
import Data.Binary                  as B
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TypeLits
import GHC.Generics                 (Generic)
import Numeric.LinearAlgebra.Static

-- * Weights
--

data Weights i o = W { wBiases :: !(R o)
                     , wNodes  :: !(L o i)
                     }
  deriving (Show, Generic)

instance (KnownNat i, KnownNat o) => Binary (Weights i o)

randomWeights :: (MonadRandom m, KnownNat i, KnownNat o)
              => m (Weights i o)
randomWeights = do
    s1 :: Int <- getRandom
    s2 :: Int <- getRandom
    let wB = randomVector  s1 Uniform * 2 - 1
        wN = uniformSample s2 (-1) 1
    return $ W wB wN

runLayer :: (KnownNat i, KnownNat o)
         => Weights i o
         -> R i
         -> R o
runLayer (W wB wN) v = wB + wN #> v

-- * Network
--

data Network :: Nat -> [Nat] -> Nat -> * where
    O     :: !(Weights i o)
          -> Network i '[] o
    (:&~) :: KnownNat h
          => !(Weights i h)
          -> !(Network h hs o)
          -> Network i (h ': hs) o
infixr 5 :&~

deriving instance (KnownNat i, KnownNat o) => Show (Network i hs o)

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

runNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> R i
       -> R o
runNet (O w)      !v = logistic (runLayer w v)
runNet (w :&~ n') !v = let v' = logistic (runLayer w v)
                       in  runNet n' v'

hiddenStruct :: Network i hs o -> [Integer]
hiddenStruct = \case
    O _    -> []
    _ :&~ (n' :: Network h hs' o)
           -> natVal (Proxy @h)
            : hiddenStruct n'

randomNet' :: forall m i hs o. (MonadRandom m, KnownNat i, KnownNat o)
           => Sing hs -> m (Network i hs o)
randomNet' = \case
    SNil            ->     O <$> randomWeights
    SNat `SCons` ss -> (:&~) <$> randomWeights <*> randomNet' ss

randomNet :: forall m i hs o. (MonadRandom m, KnownNat i, SingI hs, KnownNat o)
          => m (Network i hs o)
randomNet = randomNet' sing

putNet :: (KnownNat i, KnownNat o)
       => Network i hs o
       -> Put
putNet = \case
    O w     -> put w
    w :&~ n -> put w *> putNet n

getNet :: forall i hs o. (KnownNat i, KnownNat o)
       => Sing hs
       -> Get (Network i hs o)
getNet = \case
    SNil            ->     O <$> get
    SNat `SCons` ss -> (:&~) <$> get <*> getNet ss

instance (KnownNat i, SingI hs, KnownNat o) => Binary (Network i hs o) where
    put = putNet
    get = getNet sing

-- * OpaqueNet
--

data OpaqueNet :: Nat -> Nat -> * where
    ONet :: Network i hs o -> OpaqueNet i o

runOpaqueNet :: (KnownNat i, KnownNat o)
             => OpaqueNet i o
             -> R i
             -> R o
runOpaqueNet (ONet n) x = runNet n x

numHiddens :: OpaqueNet i o -> Int
numHiddens (ONet n) = go n
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'

randomONet :: (MonadRandom m, KnownNat i, KnownNat o)
           => [Integer]
           -> m (OpaqueNet i o)
randomONet hs = case toSing hs of
                  SomeSing ss -> ONet <$> randomNet' ss

putONet :: (KnownNat i, KnownNat o)
        => OpaqueNet i o
        -> Put
putONet (ONet net) = do
    put (hiddenStruct net)
    putNet net

getONet :: (KnownNat i, KnownNat o)
        => Get (OpaqueNet i o)
getONet = do
    hs <- get
    withSomeSing hs $ \ss ->
      ONet <$> getNet ss

instance (KnownNat i, KnownNat o) => Binary (OpaqueNet i o) where
    put = putONet
    get = getONet

-- * OpaqueNet'
--

type OpaqueNet' i o r = (forall hs. Network i hs o -> r) -> r

oNet' :: Network i hs o -> OpaqueNet' i o r
oNet' n = \f -> f n

runOpaqueNet' :: (KnownNat i, KnownNat o)
              => OpaqueNet' i o (R o)
              -> R i
              -> R o
runOpaqueNet' oN x = oN (\n -> runNet n x)
--            :: ((forall hs. Network i hs o -> R o) -> R o)
--            -> R i
--            -> R o

numHiddens' :: OpaqueNet' i o Int -> Int
numHiddens' oN = oN go
  where
    go :: Network i hs o -> Int
    go = \case
        O _      -> 0
        _ :&~ n' -> 1 + go n'
--          :: ((forall hs. Network i hs o -> Int) -> Int)
--          -> Int

withRandomONet' :: (MonadRandom m, KnownNat i, KnownNat o)
                => [Integer]
                -> (forall hs. Network i hs o -> m r)
                -> m r
--         aka, => [Integer]
--              -> OpaqueNet' i o (m r)
withRandomONet' hs f = withSomeSing hs $ \ss -> do
                         net <- randomNet' ss
                         f net

putONet' :: (KnownNat i, KnownNat o)
         => OpaqueNet' i o Put
         -> Put
putONet' oN = oN $ \net -> do
                      put (hiddenStruct net)
                      putNet net

getONet' :: (KnownNat i, KnownNat o)
         => (forall hs. Network i hs o -> Get r)
         -> Get r
--  aka, => OpaqueNet' i o (Get r)
getONet' f = do
    hs <- get :: Get [Integer]
    withSomeSing hs $ \ss -> do
      n <- getNet ss
      f n

main :: IO ()
main = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    n  <- randomONet hs
    case n of
      ONet (net :: Network 10 hs 3) -> do
        print net
        -- blah blah stuff with our dynamically generated net

main' :: IO ()
main' = do
    putStrLn "What hidden layer structure do you want?"
    hs <- readLn
    withRandomONet' hs $ \(net :: Network 10 hs 3) -> do
      print net
      -- blah blah stuff with our dynamically generated net

-- * SomeNet
--

data SomeNet where
    SNet :: (KnownNat i, KnownNat o)
         => Network i hs o
         -> SomeNet

withONet :: SomeNet
         -> (forall i o. (KnownNat i, KnownNat o) => OpaqueNet i o -> r)
         -> r
withONet (SNet n) f = f (ONet n)

randomSNet :: forall m. MonadRandom m
           => Integer
           -> [Integer]
           -> Integer
           -> m SomeNet
randomSNet i hs o =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) ->
      SNet <$> (randomNet' ss :: m (Network i hs o))

withRandomSNet' :: forall m r. MonadRandom m
                => Integer
                -> [Integer]
                -> Integer
                -> (forall i hs o. (KnownNat i, KnownNat o) => Network i hs o -> m r)
                -> m r
withRandomSNet' i hs o f =
    withSomeSing i  $ \(SNat :: Sing (i  :: Nat  )) ->
    withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
    withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
      n <- randomNet' ss :: m (Network i hs o)
      f n

instance Binary SomeNet where
    put (SNet (net :: Network i hs o)) = do
        put $ natVal (Proxy @i)
        put $ hiddenStruct net
        put $ natVal (Proxy @o)
        putNet net
    get = do
        i  <- get :: Get Integer
        hs <- get :: Get [Integer]
        o  <- get :: Get Integer
        withSomeSing i    $ \(SNat :: Sing (i  :: Nat  )) ->
          withSomeSing hs $ \(ss   :: Sing (hs :: [Nat])) ->
          withSomeSing o  $ \(SNat :: Sing (o  :: Nat  )) -> do
            n <- getNet ss :: Get (Network i hs o)
            return (SNet n)


module Main where

import MnistLoader
import Network

import Numeric.LinearAlgebra as L

main :: IO ()
main = testMnist

testMnist = do
  net <- initNet [784, 30, 10]
  exemplars <- loadDataWrapper
  test <- loadTestWrapper

  -- let eta = 0.050
  -- let batchSize = 20
  -- let epochs = 30

  let eta = 0.5
  let batchSize = 16
  let epochs = 30

  putStrLn "Running Stochastic Gradient Descent "
  putStrLn $ "eta: " ++ (show eta)
  putStrLn $ "Batch Size: " ++ (show batchSize)
  putStrLn $ "Epochs: " ++ (show epochs)

  trainedNetwork <- sgd net exemplars test epochs 0 batchSize eta
  return ()

test = do
  net <- initNet [3, 5]

  putStrLn "The Net"
  putStrLn $ show net

  putStrLn "Feed Forward"
  putStrLn $ show $ feedForward net (vect 3 [0, 1, 0])

  putStrLn "A ForwardPass"
  putStrLn $ show $ forwardPass net (vect 3 [0, 1, 0])

  putStrLn "Entire backprop"
  let eb = backprop net (vect 3 [0, 1, 0]) (vect 5 [1, 0, 1, 0, 1])
  putStrLn $ show $ length eb
  putStrLn $ show $ eb

  putStrLn "A Minibatch"
  let x = vect 3 [0,1,0]
  let y = vect 5 [1, 0, 1, 0, 1]
  let batch = [(x, y), (x, y), (x, y)]
  let newn = updateMiniBatch net batch 3.0
  putStrLn $ show newn
  return newn

vect n = (n><1)


{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module Math.Regression.LeastSquares where

import qualified Math.HMatrixUtils             as HU
import qualified Math.Regression.Regression    as RE

import qualified Polysemy                      as P
import qualified Knit.Effect.Logger            as Log

import qualified Data.List                     as List

import           Numeric.LinearAlgebra          ( (#>)
                                                , (<.>)
                                                , (<\>)
                                                )
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )

-- matrix dimensions given in (row x column) form
-- if A is (m x n), its transpose, A' is (n x m)
-- ||A|| is the Frobenius norm of A, the sum of the squares of the elements.
-- if A (m x l) and B (m x k) have the same number of rows, [A B] is the (m x (l+k)) matrix formed by appending the columns of B to the columns of A
-- A: an (m x n) matrix, of m observations of n observables
-- B: an (m x d) matrix, of observations (often d=1 and B is just a length m column-vector)
-- Q: an ((n+d)m x (n+d)m) matrix of covariances, the covariance matrix of [A B].  Given?  Computed?

-- we want an (n x d) matrix X which solves, or "best solves", AX = B.

-- OLS: minimize ||AX - B||
ordinaryLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
ordinaryLS withConstant mA vB = do
  let mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (n, p) = LA.size mAwc
      dof    = n - p
  HU.checkVectorMatrix "b" "A" vB mAwc  -- vB <> mA is legal, b has same length as A has rows
  let vX  = mAwc <\> vB
      vU  = vB - (mAwc #> vX) -- residuals
      mse = (vU <.> vU) / realToFrac dof
      cov = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB Nothing vU
  return $ RE.RegressionResult (RE.estimates cov vX)
                               (realToFrac dof)
                               mse
                               rSq
                               aRSq
                               fStat
                               cov

weightedLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
weightedLS withConstant mA vB vW = do
  HU.checkEqualVectors "b" "w" vB vW
  let mW     = LA.diag vW
      vWB    = mW #> vB
      mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (_, p) = LA.size mAwc
  HU.checkVectorMatrix "b" "A" vB mAwc
  let mWA  = mW LA.<> mAwc
      vX   = mWA <\> vWB
      vU   = vB - (mAwc #> vX)
      vWU  = mW #> vU
      sumW = LA.sumElements vW
      effN = (sumW * sumW) / (vW <.> vW)
      mse  = effN / (effN - realToFrac p) * (vWU <.> vWU) / sumW
      cov  = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB (Just vW) vU
  return $ RE.RegressionResult (RE.estimates cov vX)
                               (effN - realToFrac p)
                               mse
                               rSq
                               aRSq
                               fStat
                               cov

totalLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
totalLS withConstant mA vB = do
  let mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (n, p) = LA.size mAwc
      dof    = realToFrac (n - p)
  HU.checkVectorMatrix "b" "Awc" vB mAwc
  let mAB      = mAwc LA.||| LA.asColumn vB
      (_, mV') = LA.rightSV mAB
--      gSV = sv V.! 0
--      tol = gSV * LA.eps
      sV22     = mV' LA.! p LA.! p
      vV12     = List.head $ LA.toColumns $ LA.subMatrix (0, p) (p, 1) mV' --LA.?? (LA.DropLast 1, LA.TakeLast 1)
      vX       = LA.scale (-1 / sV22) vV12 -- this is the TLS solution.  But in a shifted basis.??
--      mV2      = mV' LA.?? (LA.All, LA.TakeLast 1)
--      mABt     = mAB LA.<> mV2 LA.<> (LA.tr mV2)
--      mAt      = mABt LA.?? (LA.All, LA.DropLast 1)
      vBfit    = mAwc #> vX --(mAwc - mAt) #> vX -- this is confusing.
      vU       = vB - vBfit
--      (rSq, aRSq) = RE.goodnessOfFit p vB vU --(vB - mA #> vX)
      mse      = (vU <.> vU) / realToFrac (n - p)
      cov      = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB Nothing vU
  return $ RE.RegressionResult (RE.estimates cov vX) dof mse rSq aRSq fStat cov

weightedTLS
  :: Log.LogWithPrefixesLE effs
  => Bool
  -> Matrix R
  -> Vector R
  -> Vector R
  -> P.Sem effs (RE.RegressionResult R)
weightedTLS withConstant mA vB vW = do
  HU.checkEqualVectors "b" "w" vB vW
  let mW     = LA.diag vW
      vWB    = mW #> vB
      mAwc   = if withConstant then addBiasCol (LA.size vB) mA else mA -- add a constant, e.g., the b in y = mx + b
      (_, p) = LA.size mAwc
  HU.checkVectorMatrix "b" "Awc" vB mAwc
  let mWA      = mW LA.<> mAwc
      mWAB     = mWA LA.||| LA.asColumn vWB
      (_, mV') = LA.rightSV mWAB
--      gSV = sv V.! 0
--      tol = gSV * LA.eps
      sV22     = mV' LA.! p LA.! p
      vV12     = List.head $ LA.toColumns $ LA.subMatrix (0, p) (p, 1) mV' --LA.?? (LA.DropLast 1, LA.TakeLast 1)
      vX       = LA.scale (-1 / sV22) vV12 -- this is the WTLS solution.  But in a shifted basis.??
--      mV2      = mV' LA.?? (LA.All, LA.TakeLast 1)
--      mWABt    = mWAB LA.<> mV2 LA.<> (LA.tr mV2)
--      mWAt     = mWABt LA.?? (LA.All, LA.DropLast 1)
--      vBfit    = mAwc #> vX -- (mWA - mWAt) #> vX -- this is confusing
      vU       = vB - (mAwc #> vX)
      vWU      = mW #> vU
      sumW     = LA.sumElements vW
      effN     = (sumW * sumW) / (vW <.> vW)
      dof      = effN - realToFrac p
      mse      = effN / dof * (vWU <.> vWU) / sumW
--      (rSq, aRSq) = RE.goodnessOfFit p vWB vWU --(vB - mA #> vX)
      cov      = LA.scale mse (LA.inv $ LA.tr mAwc LA.<> mAwc)
  RE.FitStatistics rSq aRSq fStat <- RE.goodnessOfFit p vB (Just vW) vU
  return $ RE.RegressionResult (RE.estimates cov vX) dof mse rSq aRSq fStat cov

addBiasCol :: Int -> Matrix R -> Matrix R
addBiasCol rows mA =
  let colList = List.replicate rows 1
  in  if LA.size mA == (0, 0)
        then LA.matrix 1 colList
        else LA.col colList LA.||| mA