{-# LANGUAGE BangPatterns #-}
module ListStats (mean, var, stDiv, nCr,  varM, qf,unStdNorm,unStdNormSplits, -- cov, corr,
                  stDivM, stNorm, erf, cn, erfInv, binomial, cdf, cdF, meanBy, stDivBy,
                  normal, gammaln,  adjustedMatrix, adjustedMatrixN, adjustedMatrixPN, normalP,
                  module ListAlgebra,
                 ) where

import Numeric
import Numeric.LinearAlgebra
import qualified Data.Vector as V (map, fromList)
import Data.List
import Control.DeepSeq (NFData)
import LINreg2 (foldl'Rnf )
import Data.Ord (comparing)
import ListAlgebra

----------------------------------------------------------------------------------------
--- mean, variance and correlations
----------------------------------------------------------------------------------------

insertZeros :: Num a =>  [a] -> [Int] -> [a]
insertZeros xs   []   = xs
insertZeros xs (a:as) =
        let   (front,back) = splitAt (a-1) xs  in
            (1 :front) ++ (0 : insertZeros back  as)

-- map a function then calculate the standard deviation
stDivBy :: (a -> Double) -> [a] -> Double
stDivBy f xs = (sumBy (sqr . (\a -> a - ns) . f) xs) / k
    where
       (sm,k) = foldl' (\(n,m) a -> (f a + n, m + 1)) (0,0) xs
       ns     = sm / k
       sqr x  = x * x

-- map a function then calculate the mean
meanBy :: (a -> Double) -> [a] -> Double
meanBy g xs = sm / k
   where (sm,k) = foldl' (\(n,m) a -> (g a + n, m + 1)) (0,0) xs

----
mean  :: (NFData a, Fractional a) => [a] -> a
mean     xs                    = fst suml / (snd suml - 1)
                    where suml = foldl'Rnf  kF (0, 1) xs
                          kF (!n , !r) b =  (n + b, 1 + r)
---
var  :: (NFData a, Floating a) => [a] -> a
var     xs                    =  (fst suml / snd suml)
                   where suml = foldl'Rnf  kF (0, 0) xs
                         kF (!n , !r) b = (n + (b - k)^2 , 1 + r)
                         k    = mean xs
-- variance of a list of values with the mean given
varM  :: (NFData a , Floating a) => a -> [a] -> a
varM     k  xs              =  fst suml / snd suml
                 where suml = foldl'Rnf  kF (0, 0) xs
                       kF (!n , !r) b = (n + (b - k)^2 , 1 + r)
--
--- the standard deviation without the Mean given
stDiv   :: (NFData a , Floating a) =>  [a] -> a
stDiv                       = sqrt . var

---- standard deviation with the mean given
stDivM        k             = sqrt . (varM k)

--- standardise by subtracting the mean and dividing by the standard deviation
--stNorm :: [Double] -> [Double]
stNorm       xs             = map (\ a -> (a - m) / std) xs
                 where m    = mean xs
                       std  = stDivM m xs
---
unStdNorm  xs ys  = map (\a -> std * a + m)  ys
                where  m    = mean xs
                       std  = stDivM m xs

unStdNormSplits  xs zs = unStdNorm (insertZeros xs zs)

---   selecting the 'best' sublist of lenght n of a list of values
bestN :: (Ord a , Ord b)  => (a -> b) -> Int -> [a] -> [a]
bestN  comp   lim  xs
    | length xs <= lim   =  xs
    | length back <= lim =  myBest lim comp xs
    | otherwise          =  min (sortBy (comparing comp) front) (bestN comp lim back)
    where
        (front, back) = splitAt lim xs
        myBest n  f   = myTake n .  sortBy (comparing f)
        mn            = lim `div` 2

--- computed the weighted average of the residuals of values, which represent
--- points on the y-axis plotted against x-values [1 .. n], where n is the
-- length of the list. The weight biased towards is determind as the
wgtAvgRd ::  (NFData a, Enum a, Fractional a) => Int -> [[a]]  -> a
wgtAvgRd    n   xs     = weightsSum / len
    where
        (weightsSum , len) = foldl'Rnf (getWeightedAvg (2 * n)) (0,0) xs
        getWeightedAvg k (wa,ln) ys  =   (wa + wVal , ln  + 1)
            where
                f     = fromIntegral
                len   = f $ length ys
                ks    = [1 .. len]
                sumSqrs  as bs = foldl'Rnf (\a (b,c) -> a + b * c) 0 (zip as bs)
                wVal  =  (sumSqrs ks ys) / (sumSqrs ks ks)   -- (f k / len)  *

--------------------------------------------------------------------------------------------
--  Defining the X matrix for the MML function
--------------------------------------------------------------------------------------------
-- We combine adjustMatrix, fillZeros and allPieces to
-- give the final adjusted matrix. This matrix is standardised, then transposed,
-- but the transposeition can be removed if we require that to occur later on
adjustedMatrix  xs   =  [( tS . adjustMatrix . fillZeroes_aux n 0 1) yss | yss <- allPieces xs]
                         where n  = length xs
                               tS =  fromLists . map (1:) . transpose . map stNorm

adjustedMatrixN m xs =  [tS $ adjustMatrix $ fillZeroes_aux n 0 1 yss | yss <- nPieces m xs]
                         where n  = length xs
                               tS = map (1:). transpose . map stNorm
adjustedMatrixPN k m xs = [(tS $ adjustMatrix $ fillZeroes_aux n 0 1 yss, cuts) |
                                (yss,cuts) <- nPiecesP k m Nothing xs, not $ null cuts]
                         where n  = length xs
                               tS = map (1:). transpose . map stNorm
--
-- test function
joins :: Eq a => [a] -> [[a]]
joins  xs | null xs = [] -- && (length (concat a) > 5)
          | otherwise = joins_aux xs ++ joins (tail xs)
    where
        joins_aux (x : y : ys) = ([x,y] :  (map (x :) $ joins_aux (y : ys)))
        joins_aux  _           = []
--------------------------------------------------------------------------------------------
-- calculating the binomial coefficient, and normal prob density
-------------------------------------------------------------------------------------------
nCr ::  (Integral a, Ord a) => a -> a -> a
nCr   n    r   | n < 0 || r < 0  =  error "can find combinations for negative numbers"
nCr   n   0                      =  1
nCr   0   _                      =  0
nCr   !n   !k                    =  nCr (n-1) (k-1) * n `div` k

--- binomial distribution
binomial :: (Integral a, Ord b, Fractional b, Ord a) => a -> a -> b -> b
binomial   n   r   p  | p > 1 || p < 0   = error "Invalid probability value" ---n nCr r *
binomial   n   k   p  | k >= 0 && k <= n = let  nck = fromIntegral (nCr n k) in nck * (p^k) * (1 - p)^(n - k)
                      | otherwise        = 0

-- value of x from a normal distribution with mean 'mn' and standard deviation "std"
normal x mn std = (1 / (std * sqrt (2 * pi))) * exp(- 0.5 * ((x - mn)^2) / (std ^2))
--normal x mn std = cdf (x + eps) mn std  - cdf (x - eps) mn std
    where
        eps = 0.0101095
        --z   = (x - mn) / std

-- calculating the normal with a given precision
normalP eps x mn std = cdf (x + eps) mn std  - cdf (x - eps) mn std
------------------------------------------------------------------------------------------------
--- approximating the error functions and its inverse
----------------------------------------------------------------------------------------------
erf	x	| x  <  0    =  -1 *  erf (- 1 * x)
		| otherwise  = 1 - t * exp k
    where
        t  = 1 / (1 + 0.5 * abs x)
	t3 = t * (- 0.82215223 +  t *  0.17087277)
        t2 = t * ( 0.27886807 +  t * (-1.13520398 +  t *  1.48851587 + t3))
        t4 = t * ( 0.37409196 +  t * ( 0.09678418 +  t *  (- 0.18628806) + t2))
	k  = - x * x -  1.26551223 + t *  1.00002368 +  t4
--

--cn  :: Int -> Double
cn	 0	=	1
cn	 n	=	foldr (\k  b -> b + (cn k  * cn (n - 1 - k)) / ((k + 1) * (2 * k + 1))) 0 [0.. n - 1]

-- inverse of the error function
--erfInv :: (Fractional a, Floating a, Enum a) => a -> a
erfInv	y | y <  -1 || y > 1  = error "erfInv argument out of range"
          | abs y == 1        = pol (-y* log (0.0))
          | y < - 0.7         = pol (- 1 * folD z (head c) (tail c) / folD z 1 d)
          | y <   0.7         = pol ( y *  folD (y^2) (head a) (tail a) / folD (y^2) 1 b)
          | otherwise         = pol (folD z1 (head c) (tail c)  / folD z1 1 d)
                  where a     = [0.886226899, -1.645349621,  0.914624893, -0.140543331]
                        b     = [-2.118377725,  1.442710462, -0.329097515,  0.012229801]
                        c     = [-1.970840454, -1.624906493,  3.429567803,  1.641345311]
                        d     = [3.543889200,  1.637067800]
                        z     = sqrt(- log ((1.0+y)/2.0))
                        z1    = sqrt(- log ((1.0+y)/2.0))
                        pol x = x - (erf(x) - y) / (2.0 /(sqrt pi)  * exp(-x*x))
                        folD  zz a ls = foldl'Rnf  (\a b -> a + zz * (fst b + snd b * zz)) a (zipp ls)

-- commulative density function
cdf   x mean std | std < 0      = error"Negative variance to CDF"
cdf   x mean std | otherwise    =  0.5 * (1 + erf ((x - mean) / (std * sqrt 2)))

-- the CDF for the standard normal distribution
cdF  n  =  cdf  n 0 1

-- the q function, which is 1 - CDF (x) (here given for the standard normal dist)
qf  x  =  1 - cdF x
--------------------------------------------------------------------------------------
--- The gamm function takel  from wikipedia
--------------------------------------------------------------------------------------
gammaln :: Double -> Double
gammaln xx = -tmp' + log(2.5066282746310005 * ser' / xx)
			where
			tmp' = (xx+5.5) - (xx+0.5)*log(xx+5.5)
			ser' = ser + sum (zipWith (/) cof [xx+1..])
			cof  = [76.18009172947146,-86.50532032941677,24.01409824083091,-1.231739572450155,0.001208650973866179,-0.000005395239384953]
			ser  = 1.000000000190015


-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    (
      Matrix(..)
    , Vector
    , fromList
    , fromVector
    , toVector
    , toList
    , dimension
    , center
    , multiply
    , multiplyV
    , transpose
    , power
    , norm
    , column
    , row
    , map
    , for
    , unsafeIndex
    , hasNaN
    , bounds
    , unsafeBounds
    ) where

import Prelude hiding (exponent, map, sum)
import Statistics.Function (for, square)
import Statistics.Matrix.Types
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Unboxed as U

-- | Convert from a row-major list.
fromList :: Int                 -- ^ Number of rows.
         -> Int                 -- ^ Number of columns.
         -> [Double]            -- ^ Flat list of values, in row-major order.
         -> Matrix
fromList r c = fromVector r c . U.fromList

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c 0 v
  where len    = U.length v

-- | Convert to a row-major flat vector.
toVector :: Matrix -> U.Vector Double
toVector (Matrix _ _ _ v) = v

-- | Convert to a row-major flat list.
toList :: Matrix -> [Double]
toList = U.toList . toVector

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _ _) = (r, c)

-- | Avoid overflow in the matrix.
avoidOverflow :: Matrix -> Matrix
avoidOverflow m@(Matrix r c e v)
  | center m > 1e140 = Matrix r c (e + 140) (U.map (* 1e-140) v)
  | otherwise        = m

-- | Matrix-matrix multiplication. Matrices must be of compatible
-- sizes (/note: not checked/).
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ e1 _) m2@(Matrix _ c2 e2 _) =
  Matrix r1 c2 (e1 + e2) $ U.generate (r1*c2) go
  where
    go t = sum $ U.zipWith (*) (row m1 i) (column m2 j)
      where (i,j) = t `quotRem` c2

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Raise matrix to /n/th power. Power must be positive
-- (/note: not checked).
power :: Matrix -> Int -> Matrix
power mat 1 = mat
power mat n = avoidOverflow res
  where
    mat2 = power mat (n `quot` 2)
    pow  = multiply mat2 mat2
    res | odd n     = multiply pow mat
        | otherwise = pow

-- | Element in the center of matrix (not corrected for exponent).
center :: Matrix -> Double
center mat@(Matrix r c _ _) =
    unsafeBounds U.unsafeIndex mat (r `quot` 2) (c `quot` 2)

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sum . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c _ v) i = U.backpermute v $ U.enumFromStepN i c r
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c _ v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

map :: (Double -> Double) -> Matrix -> Matrix
map f (Matrix r c e v) = Matrix r c e (U.map f v)

-- | Indicate whether any element of the matrix is @NaN@.
hasNaN :: Matrix -> Bool
hasNaN = U.any isNaN . toVector

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

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

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 e _) = Matrix c0 r0 e . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r


import Control.Monad.ST (runST)
import Criterion.Main
import Data.Complex
import Statistics.Sample
import Statistics.Transform
import Statistics.Correlation.Pearson
import System.Random.MWC
import qualified Data.Vector.Unboxed as U

-- Test sample
sample :: U.Vector Double
sample = runST $ flip uniformVector 10000 =<< create

-- Weighted test sample
sampleW :: U.Vector (Double,Double)
sampleW = U.zip sample (U.reverse sample)

-- Comlex vector for FFT tests
sampleC :: U.Vector (Complex Double)
sampleC = U.zipWith (:+) sample (U.reverse sample)

-- Simple benchmark for functions from Statistics.Sample
main :: IO ()
main =
  defaultMain
  [ bgroup "sample"
    [ bench "range"            $ nf (\x -> range x)            sample
      -- Mean
    , bench "mean"             $ nf (\x -> mean x)             sample
    , bench "meanWeighted"     $ nf (\x -> meanWeighted x)     sampleW
    , bench "harmonicMean"     $ nf (\x -> harmonicMean x)     sample
    , bench "geometricMean"    $ nf (\x -> geometricMean x)    sample
      -- Variance
    , bench "variance"         $ nf (\x -> variance x)         sample
    , bench "varianceUnbiased" $ nf (\x -> varianceUnbiased x) sample
    , bench "varianceWeighted" $ nf (\x -> varianceWeighted x) sampleW
      -- Correlation
    , bench "pearson"          $ nf (\x -> pearson (U.reverse sample) x) sample
    , bench "pearson'"          $ nf (\x -> pearson' (U.reverse sample) x) sample
    , bench "pearsonFast"      $ nf (\x -> pearsonFast (U.reverse sample) x) sample
      -- Other
    , bench "stdDev"           $ nf (\x -> stdDev x)           sample
    , bench "skewness"         $ nf (\x -> skewness x)         sample
    , bench "kurtosis"         $ nf (\x -> kurtosis x)         sample
      -- Central moments
    , bench "C.M. 2"           $ nf (\x -> centralMoment 2 x)  sample
    , bench "C.M. 3"           $ nf (\x -> centralMoment 3 x)  sample
    , bench "C.M. 4"           $ nf (\x -> centralMoment 4 x)  sample
    , bench "C.M. 5"           $ nf (\x -> centralMoment 5 x)  sample
    ]
  , bgroup "FFT"
    [ bgroup "fft"
      [ bench  (show n) $ whnf fft   (U.take n sampleC) | n <- fftSizes ]
    , bgroup "ifft"
      [ bench  (show n) $ whnf ifft  (U.take n sampleC) | n <- fftSizes ]
    , bgroup "dct"
      [ bench  (show n) $ whnf dct   (U.take n sample)  | n <- fftSizes ]
    , bgroup "dct_"
      [ bench  (show n) $ whnf dct_  (U.take n sampleC) | n <- fftSizes ]
    , bgroup "idct"
      [ bench  (show n) $ whnf idct  (U.take n sample)  | n <- fftSizes ]
    , bgroup "idct_"
      [ bench  (show n) $ whnf idct_ (U.take n sampleC) | n <- fftSizes ]
    ]
  ]

fftSizes :: [Int]
fftSizes = [32,128,512,2048]


{-# LANGUAGE CPP                                      #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE PolyKinds                                #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fno-warn-orphans                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType #-}
#endif

-- |
-- Module      : Numeric.LinearAlgebra.Static.Backprop
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- A wrapper over "Numeric.LinearAlgebra.Static" (type-safe vector and
-- matrix operations based on blas/lapack) that allows its operations to
-- work with <https://hackage.haskell.org/package/backprop backprop>.  Also
-- provides orphan instances of 'Backprop' for types in
-- "Numeric.LinearAlgebra.Static".
--
-- In short, these functions are "lifted" to work with 'BVar's.
--
-- Using 'evalBP' will run the original operation:
--
-- @
-- 'evalBP' :: (forall s. 'Reifies' s 'W'. 'BVar' s a -> 'BVar' s b) -> a -> b
-- @
--
-- But using 'gradBP' or 'backprop' will give you the gradient:
--
-- @
-- 'gradBP' :: (forall s. 'Reifies' s 'W'. 'BVar' s a -> 'BVar' s b) -> a -> a
-- @
--
-- These can act as a drop-in replacement to the API of
-- "Numeric.LinearAlgebra.Static".  Just change your imports, and your
-- functions are automatically backpropagatable.  Useful types are all
-- re-exported.
--
-- Also contains 'sumElements' 'BVar' operation.
--
-- Formulas for gradients come from the following papers:
--
--     * https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
--     * http://www.dtic.mil/dtic/tr/fulltext/u2/624426.pdf
--     * http://www.cs.cmu.edu/~zkolter/course/15-884/linalg-review.pdf
--     * https://arxiv.org/abs/1602.07527
--
-- Some functions are notably unlifted:
--
--     * 'H.svd': I can't find any resources that allow you to backpropagate
--     if the U and V matrices are used!  If you find one, let me know, or
--     feel free to submit a PR!  Because of this, Currently only a version
--     that exports only the singular values is exported.
--     * 'H.svdTall', 'H.svdFlat': Not sure where to start for these
--     * 'qr': Same story.
--     https://github.com/tensorflow/tensorflow/issues/6504 might yield
--     a clue?
--     * 'H.her': No 'Num' instance for 'H.Her' makes this impossible at
--     the moment with the current backprop API
--     * 'H.exmp': Definitely possible, but I haven't dug deep enough to
--     figure it out yet!  There is a description here
--     https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf but it
--     requires some things I am not familiar with yet.  Feel free to
--     submit a PR!
--     * 'H.sqrtm': Also likely possible.  Maybe try to translate
--     http://people.cs.umass.edu/~smaji/projects/matrix-sqrt/ ?  PRs
--     welcomed!
--     * 'H.linSolve': Haven't figured out where to start!
--     * 'H.</>': Same story
--     * Functions returning existential types, like 'H.withNullSpace',
--     'H.withOrth', 'H.withRows', etc.; not quite sure what the best way
--     to handle these are at the moment.
--     * 'H.withRows' and 'H.withColumns' made "type-safe", without
--     existential types, with 'fromRows' and 'fromColumns'.

module Numeric.LinearAlgebra.Static.Backprop (
  -- * Vector
    H.R
  , H.ℝ
  , vec2
  , vec3
  , vec4
  , (&)
  , (#)
  , split
  , headTail
  , vector
  , linspace
  , H.range
  , H.dim
  -- * Matrix
  , H.L
  , H.Sq
  , row
  , col
  , (|||)
  , (===)
  , splitRows
  , splitCols
  , unrow
  , uncol
  , tr
  , H.eye
  , diag
  , matrix
  -- * Complex
  , H.ℂ
  , H.C
  , H.M
  , H.𝑖
  , toComplex
  , fromComplex
  , complex
  , real
  , imag
  , sqMagnitude
  , magnitude
  -- * Products
  , (<>)
  , (#>)
  , (<.>)
  -- * Factorizations
  , svd
  , svd_
  , H.Eigen
  , eigensystem
  , eigenvalues
  , chol
  -- * Norms
  , H.Normed
  , norm_0
  , norm_1V
  , norm_1M
  , norm_2V
  , norm_2M
  , norm_InfV
  , norm_InfM
  -- * Misc
  , mean
  , meanCov
  , meanL
  , cov
  , H.Disp(..)
  -- ** Domain
  , H.Domain
  , mul
  , app
  , dot
  , cross
  , diagR
  , vmap
  , vmap'
  , dvmap
  , mmap
  , mmap'
  , dmmap
  , outer
  , zipWithVector
  , zipWithVector'
  , dzipWithVector
  , det
  , invlndet
  , lndet
  , inv
  -- ** Conversions
  , toRows
  , toColumns
  , fromRows
  , fromColumns
  -- ** Misc Operations
  , konst
  , sumElements
  , extractV
  , extractM
  , create
  , H.Diag
  , takeDiag
  , H.Sym
  , sym
  , mTm
  , unSym
  , (<·>)
  -- * Backprop types re-exported
  -- | Re-exported for convenience.
  --
  -- @since 0.1.1.0
  , BVar
  , Backprop
  , Reifies
  , W
  ) where

import           Data.Bifunctor
import           Data.Coerce
import           Data.Functor.Identity
import           Data.Maybe
import           Data.Proxy
import           Data.Vinyl                          (Rec(..))
import           Foreign.Storable
import           GHC.TypeLits
import           Lens.Micro hiding                   ((&))
import           Numeric.Backprop
import           Numeric.Backprop.Class
import           Unsafe.Coerce
import qualified Data.Vector                         as V
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Generic.Sized           as SVG
import qualified Data.Vector.Sized                   as SV
import qualified Data.Vector.Storable.Sized          as SVS
import qualified Numeric.Backprop.Explicit           as BE
import qualified Numeric.LinearAlgebra               as HU
import qualified Numeric.LinearAlgebra.Static        as H
import qualified Numeric.LinearAlgebra.Static.Vector as H
import qualified Prelude.Backprop                    as B

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

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

instance Backprop (H.C n) where
    zero = zeroNum
    add = addNum
    one = oneNum

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

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

instance KnownNat n => Backprop (H.Sym n) where
    zero = zeroNum
    add = addNum
    one = oneNum

vec2
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 2)
vec2 = isoVar2 H.vec2 (\(H.rVec->v) -> (SVS.index v 0, SVS.index v 1))
{-# INLINE vec2 #-}

vec3
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 3)
vec3 = isoVar3 H.vec3 (\(H.rVec->v) -> (SVS.index v 0, SVS.index v 1, SVS.index v 2))
{-# INLINE vec3 #-}

vec4
    :: Reifies s W
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R 4)
vec4 vX vY vZ vW = isoVarN
    (\(Identity x :& Identity y :& Identity z :& Identity w :& RNil) -> H.vec4 x y z w)
    (\(H.rVec->v) -> Identity (SVS.index v 0)
                  :& Identity (SVS.index v 1)
                  :& Identity (SVS.index v 2)
                  :& Identity (SVS.index v 3)
                  :& RNil
    )
    (vX :& vY :& vZ :& vW :& RNil)
{-# INLINE vec4 #-}

(&) :: (KnownNat n, 1 <= n, KnownNat (n + 1), Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
    -> BVar s (H.R (n + 1))
(&) = isoVar2 (H.&) (\(H.split->(dxs,dy)) -> (dxs, fst (H.headTail dy)))
infixl 4 &
{-# INLINE (&) #-}

(#) :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R m)
    -> BVar s (H.R (n + m))
(#) = isoVar2 (H.#) H.split
infixl 4 #
{-# INLINE (#) #-}

split
    :: forall p n s. (KnownNat p, KnownNat n, p <= n, Reifies s W)
    => BVar s (H.R n)
    -> (BVar s (H.R p), BVar s (H.R (n - p)))
split v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.split (uncurry (H.#)) v
    {-# NOINLINE t #-}
{-# INLINE split #-}

headTail
    :: (Reifies s W, KnownNat n, 1 <= n)
    => BVar s (H.R n)
    -> (BVar s H.ℝ, BVar s (H.R (n - 1)))
headTail v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.headTail
                   (\(d, dx) -> (H.konst d :: H.R 1) H.# dx)
                   v
    {-# NOINLINE t #-}
{-# INLINE headTail #-}

-- | Potentially extremely bad for anything but short lists!!!
vector
    :: forall n s. (KnownNat n, Reifies s W)
    => SV.Vector n (BVar s H.ℝ)
    -> BVar s (H.R n)
vector = BE.isoVar afSV
            (H.vecR . SVG.convert) (SVG.convert . H.rVec)
       . collectVar
{-# INLINE vector #-}

linspace
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s H.ℝ
    -> BVar s H.ℝ
    -> BVar s (H.R n)
linspace = liftOp2 . op2 $ \l u ->
    ( H.linspace (l, u)
    , \d -> let n1 = fromInteger $ natVal (Proxy @n) - 1
                dDot = ((H.range - 1) H.<.> d) / n1
                dSum = HU.sumElements . H.extract $ d
            in  (dSum - dDot, dDot)
    )
{-# INLINE linspace #-}

row :: Reifies s W
    => BVar s (H.R n)
    -> BVar s (H.L 1 n)
row = isoVar H.row H.unrow
{-# INLINE row #-}

col :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.L n 1)
col = isoVar H.col H.uncol
{-# INLINE col #-}

(|||) :: (KnownNat c, KnownNat r1, KnownNat (r1 + r2), Reifies s W)
      => BVar s (H.L c r1)
      -> BVar s (H.L c r2)
      -> BVar s (H.L c (r1 + r2))
(|||) = isoVar2 (H.|||) H.splitCols
infixl 3 |||
{-# INLINE (|||) #-}

(===) :: (KnownNat c, KnownNat r1, KnownNat (r1 + r2), Reifies s W)
      => BVar s (H.L r1        c)
      -> BVar s (H.L r2        c)
      -> BVar s (H.L (r1 + r2) c)
(===) = isoVar2 (H.===) H.splitRows
infixl 2 ===
{-# INLINE (===) #-}

splitRows
    :: forall p m n s. (KnownNat p, KnownNat m, KnownNat n, p <= m, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L p n), BVar s (H.L (m - p) n))
splitRows v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.splitRows (uncurry (H.===)) v
    {-# NOINLINE t #-}
{-# INLINE splitRows #-}

splitCols
    :: forall p m n s. (KnownNat p, KnownNat m, KnownNat n, KnownNat (n - p), p <= n, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L m p), BVar s (H.L m (n - p)))
splitCols v = (t ^^. _1, t ^^. _2)
  where
    t = isoVar H.splitCols (uncurry (H.|||)) v
    {-# NOINLINE t #-}
{-# INLINE splitCols #-}

unrow
    :: (KnownNat n, Reifies s W)
    => BVar s (H.L 1 n)
    -> BVar s (H.R n)
unrow = isoVar H.unrow H.row
{-# INLINE unrow #-}

uncol
    :: (KnownNat n, Reifies s W)
    => BVar s (H.L n 1)
    -> BVar s (H.R n)
uncol = isoVar H.uncol H.col
{-# INLINE uncol #-}

tr  :: (HU.Transposable m mt, HU.Transposable mt m, Backprop m, Reifies s W)
    => BVar s m
    -> BVar s mt
tr = isoVar H.tr H.tr
{-# INLINE tr #-}

diag
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.Sq n)
diag = liftOp1 . op1 $ \x -> (H.diag x, H.takeDiag)
{-# INLINE diag #-}

-- | Potentially extremely bad for anything but short lists!!!
matrix
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => [BVar s H.ℝ]
    -> BVar s (H.L m n)
matrix = maybe (error "matrix: invalid number of elements")
               ( isoVar (H.vecL . SVG.convert . runABP) (ABP . SVG.convert . H.lVec)
               . collectVar
               . ABP
               )
       . SV.fromList @(m * n)
{-# INLINE matrix #-}

-- | Matrix product
(<>)
    :: (KnownNat m, KnownNat k, KnownNat n, Reifies s W)
    => BVar s (H.L m k)
    -> BVar s (H.L k n)
    -> BVar s (H.L m n)
(<>) = mul
infixr 8 <>
{-# INLINE (<>) #-}

-- | Matrix-vector product
(#>)
    :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
    -> BVar s (H.R m)
(#>) = app
infixr 8 #>
{-# INLINE (#>) #-}

-- | Dot product
(<.>)
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s H.ℝ
(<.>) = dot
infixr 8 <.>
{-# INLINE (<.>) #-}

-- | Can only get the singular values, for now.  Let me know if you find an
-- algorithm that can compute the gradients based on differentials for the
-- other matricies!
--
svd :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
svd = liftOp1 . op1 $ \x ->
    let (u, σ, v) = H.svd x
    in  ( σ
        , \(dΣ :: H.R n) -> (u H.<> H.diagR 0 dΣ) H.<> H.tr v
                -- must manually associate because of bug in diagR in
                -- hmatrix-0.18.2.0
        )
{-# INLINE svd #-}

-- | Version of 'svd' that returns the full SVD, but if you attempt to find
-- the gradient, it will fail at runtime if you ever use U or V.
svd_
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.L m m), BVar s (H.R n), BVar s (H.L n n))
svd_ r = (t ^^. _1, t ^^. _2, t ^^. _3)
  where
    o :: Op '[H.L m n] (H.L m m, H.R n, H.L n n)
    o = op1 $ \x ->
        let msv@(u, _, v) = H.svd x
        in  ( msv
            , \(dU, dΣ, dV) ->
                    if H.norm_0 dU == 0 && H.norm_0 dV == 0
                      then (u H.<> H.diagR 0 dΣ) H.<> H.tr v
                      else error "svd_: Cannot backprop if U and V are used."
            )
    {-# INLINE o #-}
    t = liftOp1 o r
    {-# NOINLINE t #-}
{-# INLINE svd_ #-}

helpEigen :: KnownNat n => H.Sym n -> (H.R n, H.L n n, H.L n n, H.L n n)
helpEigen x = (l, v, H.inv v, H.tr v)
  where
    (l, v) = H.eigensystem x
{-# INLINE helpEigen #-}

-- | /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
eigensystem
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> (BVar s (H.R n), BVar s (H.L n n))
eigensystem u = (t ^^. _1, t ^^. _2)
  where
    o :: Op '[H.Sym n] (H.R n, H.L n n)
    o = op1 $ \x ->
        let (l, v, vInv, vTr) = helpEigen x
            lRep = H.rowsL . SV.replicate $ l
            fMat = (1 - H.eye) * (lRep - H.tr lRep)
        in  ( (l, v)
            , \(dL, dV) -> unsafeCoerce $
                       H.tr vInv
                  H.<> (H.diag dL + fMat * (vTr H.<> dV))
                  H.<> vTr
            )
    {-# INLINE o #-}
    t = liftOp1 o u
    {-# NOINLINE t #-}
{-# INLINE eigensystem #-}

-- | /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
eigenvalues
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.R n)
eigenvalues = liftOp1 . op1 $ \x ->
    let (l, _, vInv, vTr) = helpEigen x
    in  ( l
        , \dL -> unsafeCoerce $
                 H.tr vInv H.<> H.diag dL H.<> vTr
        )
{-# INLINE eigenvalues #-}

-- | Algorithm from https://arxiv.org/abs/1602.07527
--
-- The paper also suggests a potential imperative algorithm that might
-- help.  Need to benchmark to see what is best.
--
-- /NOTE/ The gradient is not necessarily symmetric!  The gradient is not
-- meant to be retireved directly; insteadl, 'eigenvalues' is meant to be
-- used as a part of a larger computation, and the gradient as an
-- intermediate step.
chol
    :: forall n s. (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.Sq n)
chol = liftOp1 . op1 $ \x ->
    let l = H.chol x
        lInv = H.inv l
        phi :: H.Sq n
        phi = H.build $ \i j -> case compare i j of
                                  LT -> 1
                                  EQ -> 0.5
                                  GT -> 0
    in  ( l
        , \dL -> let s = H.tr lInv H.<> (phi * (H.tr l H.<> dL)) H.<> lInv
                 in  unsafeCoerce $ s + H.tr s - H.eye * s
        )
{-# INLINE chol #-}

-- | Number of non-zero items
norm_0
    :: (H.Normed a, Backprop a, Reifies s W)
    => BVar s a
    -> BVar s H.ℝ
norm_0 = liftOp1 . op1 $ \x -> (H.norm_0 x, const (zero x))
{-# INLINE norm_0 #-}

-- | Sum of absolute values
norm_1V
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_1V = liftOp1 . op1 $ \x -> (H.norm_1 x, (* signum x) . H.konst)
{-# INLINE norm_1V #-}

-- | Maximum 'H.norm_1' of columns
norm_1M
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_1M = liftOp1 . op1 $ \x ->
    let n = H.norm_1 x
    in  (n, \d -> let d' = H.konst d
                  in  H.colsL
                    . SV.map (\c -> if H.norm_1 c == n
                                      then d' * signum c
                                      else 0
                             )
                    . H.lCols
                    $ x
        )
{-# INLINE norm_1M #-}

-- | Square root of sum of squares
--
-- Be aware that gradient diverges when the norm is zero
norm_2V
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_2V = liftOp1 . op1 $ \x ->
    let n = H.norm_2 x
    in (n, \d -> x * H.konst (d / n))
{-# INLINE norm_2V #-}

-- | Maximum singular value
norm_2M
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_2M = liftOp1 . op1 $ \x ->
    let n = H.norm_2 x
        (head.H.toColumns->u1,_,head.H.toColumns->v1) = H.svd x
    in (n, \d -> H.konst d * (u1 `H.outer` v1))
{-# INLINE norm_2M #-}

-- | Maximum absolute value
norm_InfV
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
norm_InfV = liftOp1 . op1 $ \x ->
    let n :: H.ℝ
        n = H.norm_Inf x
    in  (n, \d -> H.vecR
                . SVS.map (\e -> if abs e == n
                                   then signum e * d
                                   else 0
                          )
                . H.rVec
                $ x
        )
{-# ANN norm_InfV "HLint: ignore Use camelCase" #-}
{-# INLINE norm_InfV #-}

-- | Maximum 'H.norm_1' of rows
norm_InfM
    :: (KnownNat n, KnownNat m, Reifies s W)
    => BVar s (H.L n m)
    -> BVar s H.ℝ
norm_InfM = liftOp1 . op1 $ \x ->
    let n = H.norm_Inf x
    in  (n, \d -> let d' = H.konst d
                  in  H.rowsL
                    . SV.map (\c -> if H.norm_1 c == n
                                      then d' * signum c
                                      else 0
                             )
                    . H.lRows
                    $ x
        )
{-# ANN norm_InfM "HLint: ignore Use camelCase" #-}
{-# INLINE norm_InfM #-}

mean
    :: (KnownNat n, 1 <= n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s H.ℝ
mean = liftOp1 . op1 $ \x -> (H.mean x, H.konst . (/ H.norm_0 x))
{-# INLINE mean #-}

gradCov
    :: forall m n. (KnownNat m, KnownNat n)
    => H.L m n
    -> H.R n
    -> H.Sym n
    -> H.L m n
gradCov x μ dσ = H.rowsL
               . SV.map (subtract (dDiffsSum / m))
               . H.lRows
               $ dDiffs
  where
    diffs = H.rowsL . SV.map (subtract μ) . H.lRows $ x
    dDiffs = H.konst (2/n) * (diffs H.<> H.tr (H.unSym dσ))
    dDiffsSum = sum . H.toRows $ dDiffs
    m = fromIntegral $ natVal (Proxy @m)
    n = fromIntegral $ natVal (Proxy @n)
{-# INLINE gradCov #-}

-- | Mean and covariance.  If you know you only want to use one or the
-- other, use 'meanL' or 'cov'.
meanCov
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> (BVar s (H.R n), BVar s (H.Sym n))
meanCov v = (t ^^. _1, t ^^. _2)
  where
    m = fromInteger $ natVal (Proxy @m)
    t = ($ v) . liftOp1 . op1 $ \x ->
        let ms@(μ, _) = H.meanCov x
        in  ( ms
            , \(dμ, dσ) ->
                let gradMean = H.rowsL
                             . SV.replicate
                             $ (dμ / H.konst m)
                in  gradMean + gradCov x μ dσ
            )
    {-# NOINLINE t #-}
{-# INLINE meanCov #-}

-- | 'meanCov', but if you know you won't use the covariance.
meanL
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.R n)
meanL = liftOp1 . op1 $ \x ->
    ( fst (H.meanCov x)
    , H.rowsL . SV.replicate . (/ H.konst m)
    )
  where
    m = fromInteger $ natVal (Proxy @m)
{-# INLINE meanL #-}

-- | 'cov', but if you know you won't use the covariance.
cov
    :: forall m n s. (KnownNat n, KnownNat m, 1 <= m, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.Sym n)
cov = liftOp1 . op1 $ \x ->
    let (μ, σ) = H.meanCov x
    in  (σ, gradCov x μ)
{-# INLINE cov #-}

mul :: ( KnownNat m
       , KnownNat k
       , KnownNat n
       , H.Domain field vec mat
       , Backprop (mat m k)
       , Backprop (mat k n)
       , HU.Transposable (mat m k) (mat k m)
       , HU.Transposable (mat k n) (mat n k)
       , Reifies s W
       )
    => BVar s (mat m k)
    -> BVar s (mat k n)
    -> BVar s (mat m n)
mul = liftOp2 . op2 $ \x y ->
    ( x `H.mul` y
    , \d -> (d `H.mul` H.tr y, H.tr x `H.mul` d)
    )
{-# INLINE mul #-}

app :: ( KnownNat m
       , KnownNat n
       , H.Domain field vec mat
       , HU.Transposable (mat m n) (mat n m)
       , Backprop (mat m n)
       , Backprop (vec n)
       , Reifies s W
       )
    => BVar s (mat m n)
    -> BVar s (vec n)
    -> BVar s (vec m)
app = liftOp2 . op2 $ \xs y ->
    ( xs `H.app` y
    , \d -> (d `H.outer` y, H.tr xs `H.app` d)
    )
{-# INLINE app #-}

dot :: ( KnownNat n
       , H.Domain field vec mat
       , H.Sized field (vec n) d
       , Num (vec n)
       , Backprop (vec n)
       , Reifies s W
       )
    => BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s field
dot = liftOp2 . op2 $ \x y ->
    ( x `H.dot` y
    , \d -> let d' = H.konst d
            in  (d' * y, x * d')
    )
{-# INLINE dot #-}

cross
    :: ( H.Domain field vec mat
       , Reifies s W
       , Backprop (vec 3)
       )
    => BVar s (vec 3)
    -> BVar s (vec 3)
    -> BVar s (vec 3)
cross = liftOp2 . op2 $ \x y ->
    ( x `H.cross` y
    , \d -> (y `H.cross` d, d `H.cross` x)
    )
{-# INLINE cross #-}

-- | Create matrix with diagonal, and fill with default entries
diagR
    :: forall m n k field vec mat s.
       ( H.Domain field vec mat
       , Num (vec k)
       , Num (mat m n)
       , KnownNat m
       , KnownNat n
       , KnownNat k
       , HU.Container HU.Vector field
       , H.Sized field (mat m n) HU.Matrix
       , H.Sized field (vec k) HU.Vector
       , Backprop field
       , Backprop (vec k)
       , Reifies s W
       )
    => BVar s field             -- ^ default value
    -> BVar s (vec k)           -- ^ diagonal
    -> BVar s (mat m n)
diagR = liftOp2 . op2 $ \c x ->
    ( H.diagR c x
    , \d -> ( HU.sumElements . H.extract $ H.diagR 1 (0 :: vec k) * d
            , fromJust . H.create . HU.takeDiag . H.extract $ d
            )
    )
{-# INLINE diagR #-}

-- | Note: if possible, use the potentially much more performant 'vmap''.
vmap
    :: (KnownNat n, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
vmap f = isoVar (H.vecR . SVG.convert @V.Vector . runABP)
                (ABP . SVG.convert . H.rVec)
       . B.fmap f
       . isoVar (ABP . SVG.convert . H.rVec) (H.vecR . SVG.convert . runABP)
{-# INLINE vmap #-}

-- | 'vmap', but potentially more performant.  Only usable if the mapped
-- function does not depend on any external 'BVar's.
vmap'
    :: ( Num (vec n)
       , Storable field
       , H.Sized field (vec n) HU.Vector
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
vmap' f = liftOp1 . op1 $ bimap (fromJust . H.create . VG.convert)
                                ((*) . fromJust . H.create . VG.convert)
                        . V.unzip
                        . V.map (backprop f)
                        . VG.convert
                        . H.extract
{-# INLINE vmap' #-}

-- TODO: Can be made more efficient if backprop exports
-- a custom-total-derivative version

-- | Note: Potentially less performant than 'vmap''.
dvmap
    :: ( KnownNat n
       , H.Domain field vec mat
       , Num (vec n)
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
dvmap f = liftOp1 . op1 $ \x ->
    ( H.dvmap (evalBP f) x
    , (H.dvmap (gradBP f) x *)
    )
{-# INLINE dvmap #-}

-- | Note: if possible, use the potentially much more performant 'mmap''.
mmap
    :: (KnownNat n, KnownNat m, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.L n m)
    -> BVar s (H.L n m)
mmap f = isoVar (H.vecL . SVG.convert @V.Vector . runABP)
                (ABP . SVG.convert . H.lVec)
       . B.fmap f
       . isoVar (ABP . SVG.convert . H.lVec) (H.vecL . SVG.convert . runABP)
{-# INLINE mmap #-}

-- | 'mmap', but potentially more performant.  Only usable if the mapped
-- function does not depend on any external 'BVar's.
mmap'
    :: forall n m mat field s.
       ( KnownNat m
       , Num (mat n m)
       , Backprop (mat n m)
       , Backprop field
       , H.Sized field (mat n m) HU.Matrix
       , HU.Element field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (mat n m)
    -> BVar s (mat n m)
mmap' f = liftOp1 . op1 $ bimap (fromJust . H.create . HU.reshape m . VG.convert)
                                ((*) . fromJust . H.create . HU.reshape m . VG.convert)
                        . V.unzip
                        . V.map (backprop f)
                        . VG.convert
                        . HU.flatten
                        . H.extract
  where
    m :: Int
    m = fromInteger $ natVal (Proxy @m)
{-# INLINE mmap' #-}

-- | Note: Potentially less performant than 'mmap''.
dmmap
    :: ( KnownNat n
       , KnownNat m
       , H.Domain field vec mat
       , Num (mat n m)
       , Backprop (mat n m)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field)
    -> BVar s (mat n m)
    -> BVar s (mat n m)
dmmap f = liftOp1 . op1 $ \x ->
    ( H.dmmap (evalBP f) x
    , (H.dmmap (gradBP f) x *)
    )
{-# INLINE dmmap #-}

outer
    :: ( KnownNat m
       , KnownNat n
       , H.Domain field vec mat
       , HU.Transposable (mat n m) (mat m n)
       , Backprop (vec n)
       , Backprop (vec m)
       , Reifies s W
       )
    => BVar s (vec n)
    -> BVar s (vec m)
    -> BVar s (mat n m)
outer = liftOp2 . op2 $ \x y ->
    ( x `H.outer` y
    , \d -> ( d `H.app` y
            , H.tr d `H.app` x)
    )
{-# INLINE outer #-}

-- | Note: if possible, use the potentially much more performant
-- 'zipWithVector''.
zipWithVector
    :: (KnownNat n, Reifies s W)
    => (BVar s H.ℝ -> BVar s H.ℝ -> BVar s H.ℝ)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s (H.R n)
zipWithVector f x y = isoVar (H.vecR . SVG.convert . runABP)
                             (ABP . SVG.convert . H.rVec)
                    $ B.liftA2 @(ABP (SV.Vector _)) f (iv x) (iv y)
  where
    iv = isoVar (ABP . SVG.convert . H.rVec) (H.vecR . SVG.convert . runABP)
{-# INLINE zipWithVector #-}

zipWithVector'
    :: ( Num (vec n)
       , Backprop (vec n)
       , Storable field
       , Backprop field
       , H.Sized field (vec n) HU.Vector
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s (vec n)
zipWithVector' f = liftOp2 . op2 $ \(VG.convert.H.extract->x) (VG.convert.H.extract->y) ->
    let (z, dx, dy) = V.unzip3 $ V.zipWith (\x' -> retup . backprop2 f x') x y
    in  ( fromJust (H.create (VG.convert z))
        , \d -> ( d * fromJust (H.create (VG.convert dx))
                , d * fromJust (H.create (VG.convert dy))
                )
        )
  where
    retup (x, (y, z)) = (x, y, z)
{-# INLINE zipWithVector' #-}

-- | A version of 'zipWithVector'' that is potentially less performant but
-- is based on 'H.zipWithVector' from 'H.Domain'.
dzipWithVector
    :: ( KnownNat n
       , H.Domain field vec mat
       , Num (vec n)
       , Backprop (vec n)
       , Backprop field
       , Reifies s W
       )
    => (forall s'. Reifies s' W => BVar s' field -> BVar s' field -> BVar s' field)
    -> BVar s (vec n)
    -> BVar s (vec n)
    -> BVar s (vec n)
dzipWithVector f = liftOp2 . op2 $ \x y ->
    ( H.zipWithVector (evalBP2 f) x y
    , \d -> let dx = H.zipWithVector (\x' -> fst . gradBP2 f x') x y
                dy = H.zipWithVector (\x' -> snd . gradBP2 f x') x y
            in  (d * dx, d * dy)
    )
{-# INLINE dzipWithVector #-}

det :: ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s field
det = liftOp1 . op1 $ \x ->
    let xDet = H.det x
        xInv = H.inv x
    in  ( xDet, \d -> H.konst (d * xDet) * H.tr xInv )
{-# INLINE det #-}

-- | The inverse and the natural log of the determinant together.  If you
-- know you don't need the inverse, it is best to use 'lndet'.
invlndet
    :: forall n mat field vec d s.
       ( KnownNat n
       , Num (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Backprop field
       , Backprop (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> (BVar s (mat n n), (BVar s field, BVar s field))
invlndet v = (t ^^. _1, (t ^^. _2, t ^^. _3))
  where
    o :: Op '[mat n n] (mat n n, field, field)
    o = op1 $ \x ->
      let (i,(ldet, s)) = H.invlndet x
          iTr           = H.tr i
      in  ( (i, ldet, s)
          , \(dI, dLDet, _) ->
                let gradI    = - iTr `H.mul` dI `H.mul` iTr
                    gradLDet = H.konst dLDet * H.tr i
                in  gradI + gradLDet
          )
    {-# INLINE o #-}
    t = liftOp1 o v
    {-# NOINLINE t #-}
{-# INLINE invlndet #-}

-- | The natural log of the determinant.
lndet
    :: forall n mat field vec d s.
       ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , H.Sized field (mat n n) d
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s field
lndet = liftOp1 . op1 $ \x ->
          let (i,(ldet,_)) = H.invlndet x
          in  (ldet, (* H.tr i) . H.konst)
{-# INLINE lndet #-}

inv :: ( KnownNat n
       , Num (mat n n)
       , Backprop (mat n n)
       , H.Domain field vec mat
       , HU.Transposable (mat n n) (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s (mat n n)
inv = liftOp1 . op1 $ \x ->
    let xInv   = H.inv x
        xInvTr = H.tr xInv
    in  ( xInv, \d -> - xInvTr `H.mul` d `H.mul` xInvTr )
{-# INLINE inv #-}

toRows
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> SV.Vector m (BVar s (H.R n))
toRows = runABP . sequenceVar
       . isoVar (coerce (H.lRows @m)) (coerce H.rowsL)
{-# INLINE toRows #-}

toColumns
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> SV.Vector n (BVar s (H.R m))
toColumns = runABP . sequenceVar
          . isoVar (coerce (H.lCols @_ @n)) (coerce H.colsL)
{-# INLINE toColumns #-}

fromRows
    :: forall m n s. (KnownNat m, KnownNat n, Reifies s W)
    => SV.Vector m (BVar s (H.R n))
    -> BVar s (H.L m n)
fromRows = isoVar (coerce H.rowsL) (coerce (H.lRows @m))
         . collectVar . ABP
{-# INLINE fromRows #-}

fromColumns
    :: forall m n s. (KnownNat n, KnownNat m, Reifies s W)
    => SV.Vector n (BVar s (H.R m))
    -> BVar s (H.L m n)
fromColumns = isoVar (coerce H.colsL) (coerce (H.lCols @_ @n))
            . collectVar . ABP
{-# INLINE fromColumns #-}

konst
    :: forall t s d q.
     ( H.Sized t s d
     , HU.Container d t
     , Backprop t
     , Reifies q W
     )
    => BVar q t
    -> BVar q s
konst = liftOp1 . op1 $ \x ->
    ( H.konst x
    , HU.sumElements . H.extract
    )
{-# INLINE konst #-}

sumElements
    :: forall t s d q.
     ( H.Sized t s d
     , HU.Container d t
     , Backprop s
     , Reifies q W
     )
    => BVar q s
    -> BVar q t
sumElements = liftOp1 . op1 $ \x ->
    ( HU.sumElements . H.extract $ x
    , H.konst
    )
{-# INLINE sumElements #-}

-- | If there are extra items in the total derivative, they are dropped.
-- If there are missing items, they are treated as zero.
extractV
    :: forall t s q.
       ( H.Sized t s HU.Vector
       , HU.Konst t Int HU.Vector
       , HU.Container HU.Vector t
       , Backprop s
       , Reifies q W
       )
    => BVar q s
    -> BVar q (HU.Vector t)
extractV = liftOp1 . op1 $ \x ->
    let n = H.size x
    in  ( H.extract x
        , \d -> let m  = HU.size d
                    m' = case compare n m of
                            LT -> HU.subVector 0 n d
                            EQ -> d
                            GT -> HU.vjoin [d, HU.konst 0 (n - m)]
                in  fromJust . H.create $ m'
        )
{-# INLINE extractV #-}

-- | If there are extra items in the total derivative, they are dropped.
-- If there are missing items, they are treated as zero.
extractM
    :: forall t s q.
       ( H.Sized t s HU.Matrix
       , Backprop s
       , HU.Konst t (Int, Int) HU.Matrix
       , HU.Container HU.Matrix t
       , Num (HU.Matrix t)
       , Reifies q W
       )
    => BVar q s
    -> BVar q (HU.Matrix t)
extractM = liftOp1 . op1 $ \x ->
    let (xI,xJ) = H.size x
    in  ( H.extract x
        , \d -> let (dI,dJ) = HU.size d
                    m' = case (compare xI dI, compare xJ dJ) of
                           (LT, LT) -> d HU.?? (HU.Take xI, HU.Take xJ)
                           (LT, EQ) -> d HU.?? (HU.Take xI, HU.All)
                           (LT, GT) -> d HU.?? (HU.Take xI, HU.All)
                                HU.||| HU.konst 0 (xI, xJ - dJ)
                           (EQ, LT) -> d HU.?? (HU.All    , HU.Take xJ)
                           (EQ, EQ) -> d
                           (EQ, GT) -> d HU.?? (HU.All, HU.All)
                                HU.||| HU.konst 0 (xI, xJ - dJ)
                           (GT, LT) -> d HU.?? (HU.All, HU.Take xJ)
                                HU.=== HU.konst 0 (xI - dI, xJ)
                           (GT, EQ) -> d HU.?? (HU.All, HU.All)
                                HU.=== HU.konst 0 (xI - dI, xJ)
                           (GT, GT) -> HU.fromBlocks
                              [[d,0                            ]
                              ,[0,HU.konst 0 (xI - dI, xJ - dJ)]
                              ]
                in  fromJust . H.create $ m'
        )
{-# INLINE extractM #-}

create
    :: (H.Sized t s d, Backprop s, Num (d t), Backprop (d t), Reifies q W)
    => BVar q (d t)
    -> Maybe (BVar q s)
create = sequenceVar . isoVar H.create (maybe 0 H.extract)
{-# INLINE create #-}

takeDiag
    :: ( KnownNat n
       , H.Diag (mat n n) (vec n)
       , H.Domain field vec mat
       , Num field
       , Backprop (mat n n)
       , Reifies s W
       )
    => BVar s (mat n n)
    -> BVar s (vec n)
takeDiag = liftOp1 . op1 $ \x ->
    ( H.takeDiag x
    , H.diagR 0
    )
{-# INLINE takeDiag #-}

-- |
-- \[
-- \frac{1}{2} (M + M^T)
-- \]
sym :: (KnownNat n, Reifies s W)
    => BVar s (H.Sq n)
    -> BVar s (H.Sym n)
sym = liftOp1 . op1 $ \x ->
    ( H.sym x
    , H.unSym . H.sym . H.unSym
    )
{-# INLINE sym #-}

-- |
-- \[
-- M^T M
-- \]
mTm :: (KnownNat m, KnownNat n, Reifies s W)
    => BVar s (H.L m n)
    -> BVar s (H.Sym n)
mTm = liftOp1 . op1 $ \x ->
    ( H.mTm x
    , \d -> 2 * (x H.<> H.unSym d)
    )
{-# INLINE mTm #-}

-- | Warning: the gradient is going necessarily symmetric, and so is /not/
-- meant to be used directly.  Rather, it is meant to be used in the middle
-- (or at the end) of a longer computation.
unSym
    :: (KnownNat n, Reifies s W)
    => BVar s (H.Sym n)
    -> BVar s (H.Sq n)
unSym = isoVar H.unSym unsafeCoerce
{-# INLINE unSym #-}

-- | Unicode synonym for '<.>>'
(<·>)
    :: (KnownNat n, Reifies s W)
    => BVar s (H.R n)
    -> BVar s (H.R n)
    -> BVar s H.ℝ
(<·>) = dot
infixr 8 <·>
{-# INLINE (<·>) #-}

afSV :: Backprop a => BE.AddFunc (SV.Vector n a)
afSV = BE.AF (SV.zipWith add)
{-# INLINE afSV #-}

-- support for complex types
toComplex :: (KnownNat n, Reifies s W)
          => BVar s (H.R n) -> BVar s (H.R n) -> BVar s (H.C n)
toComplex = isoVar2 (curry H.toComplex) H.fromComplex
{-# INLINE toComplex #-}

fromComplex :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> (BVar s (H.R n), BVar s (H.R n))
fromComplex c = let ri = isoVar H.fromComplex H.toComplex c in (ri ^^. _1, ri ^^. _2)
{-# INLINE fromComplex #-}

complex :: (KnownNat n, Reifies s W)
       => BVar s (H.R n) -> BVar s (H.C n)
complex = isoVar H.complex H.real
{-# INLINE complex #-}

real :: (KnownNat n, Reifies s W)
     => BVar s (H.C n) -> BVar s (H.R n)
real = isoVar H.real (\r -> H.toComplex (r, H.konst 0))
{-# INLINE real #-}

imag :: (KnownNat n, Reifies s W)
     => BVar s (H.C n) -> BVar s (H.R n)
imag = isoVar H.imag (\r -> H.toComplex (H.konst 0, r))
{-# INLINE imag #-}

sqMagnitude :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> BVar s (H.R n)
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2
{-# INLINE sqMagnitude #-}

magnitude :: (KnownNat n, Reifies s W)
            => BVar s (H.C n) -> BVar s (H.R n)
magnitude c = sqrt $ sqMagnitude c
{-# INLINE magnitude #-}


-- some tests of the interface for pure
-- computations with inplace updates

import Numeric.LinearAlgebra
import Data.Packed.ST
import Data.Packed.Convert

import Data.Array.Unboxed
import Data.Array.ST
import Control.Monad.ST
import Control.Monad

main = sequence_[
    print test1,
    print test2,
    print test3,
    print test4,
    test5,
    test6,
    print test7,
    test8,
    test0]

-- helper functions
vector l = fromList l :: Vector Double
norm v = pnorm PNorm2 v

-- hmatrix vector and matrix
v = vector [1..10]
m = (5><10) [1..50::Double]

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

-- vector creation by in-place updates on a copy of the argument
test1 = fun v

fun :: Element t => Vector t -> Vector t
fun x = runSTVector $ do
    a <- thawVector x
    mapM_ (flip (modifyVector a) (+57)) [0 .. dim x `div` 2 - 1]
    return a

-- another example: creation of an antidiagonal matrix from a list
test2 = antiDiag 5 8 [1..] :: Matrix Double

antiDiag :: (Element b) => Int -> Int -> [b] -> Matrix b
antiDiag r c l = runSTMatrix $ do
    m <- newMatrix 0 r c
    let d = min r c - 1
    sequence_ $ zipWith (\i v -> writeMatrix m i (c-1-i) v) [0..d] l
    return m

-- using vector or matrix functions on mutable objects requires freezing:
test3 = g1 v

g1 x = runST $ do
    a <- thawVector x
    writeVector a (dim x -1) 0
    b <- freezeVector a
    return (norm b)

--  another possibility:
test4 = g2 v

g2 x = runST $ do
    a <- thawVector x
    writeVector a (dim x -1) 0
    t <- liftSTVector norm a
    return t

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

-- haskell arrays
hv = listArray (0,9) [1..10::Double]
hm = listArray ((0,0),(4,9)) [1..50::Double]

-- conversion from standard Haskell arrays
test5 = do
    print $ norm (vectorFromArray hv)
    print $ norm v
    print $ rcond (matrixFromArray hm)
    print $ rcond m

-- conversion to mutable ST arrays
test6 = do
    let y = clearColumn m 1
    print y
    print (matrixFromArray y)

clearColumn x c = runSTUArray $ do
    a <- mArrayFromMatrix x
    forM_ [0..rows x-1] $ \i->
        writeArray a (i,c) (0::Double)
    return a

-- hmatrix functions applied to mutable ST arrays
test7 = unitary (listArray (1,4) [3,5,7,2] :: UArray Int Double)

unitary v = runSTUArray $ do
    a <- thaw v
    n <- norm `fmap` vectorFromMArray a
    b <- mapArray (/n) a
    return b

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

-- (just to check that they are not affected)
test0 = do
    print v
    print m
    --print hv
    --print hm

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

histogram n ds = runSTVector $ do
    h <- newVector (0::Double) n -- number of bins
    let inc k = modifyVector h k (+1)
    mapM_ inc ds
    return h

-- check that newVector is really called with a fresh new array
histoCheck ds = runSTVector $ do
    h <- newVector (0::Double) 15 -- > constant for this test
    let inc k = modifyVector h k (+1)
    mapM_ inc ds
    return h

hc = fromList [1 .. 15::Double]

-- check that thawVector creates a new array
histoCheck2 ds = runSTVector $ do
    h <- thawVector hc
    let inc k = modifyVector h k (+1)
    mapM_ inc ds
    return h

test8 = do
    let ds = [0..14]
    print $ histogram 15 ds
    print $ histogram 15 ds
    print $ histogram 15 ds
    print $ histoCheck ds
    print $ histoCheck ds
    print $ histoCheck ds
    print $ histoCheck2 ds
    print $ histoCheck2 ds
    print $ histoCheck2 ds
    putStrLn "----------------------"


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
module DataRep where

import Prelude hiding ((.), id, curry, uncurry)
import Control.Exception (assert)
import Data.List
import Data.Proxy
import qualified Data.Vector.Generic as VG
import Foreign.Storable
import qualified Numeric.LinearAlgebra as HM

import qualified Base
import Base hiding (Scalable (..))

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

class (Fractional (Scalar a), Storable (Scalar a), HM.Numeric (Scalar a), Additive a) => VectorSpace a where
  type Scalar a
  dim :: Proxy a -> Int
  toVector :: a -> HM.Vector (Scalar a)
  fromVector :: HM.Vector (Scalar a) -> a

  scale :: Scalar a -> a -> a
  scale a = fromVector . HM.scale a . toVector

linComb :: VectorSpace a => [(a, Scalar a)] -> a
linComb xs = foldl' (.+.) zero [scale s a | (a,s) <- xs]

instance VectorSpace Double where
  type Scalar Double = Double
  scale = (*)
  dim _ = 1
  toVector = VG.singleton
  fromVector = (VG.! 0)

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => VectorSpace (a, b) where
  type Scalar (a, b) = Scalar a
  scale s (a,b) = (scale s a, scale s b)
  dim _ = dim (Proxy :: Proxy a) + dim (Proxy :: Proxy b)
  toVector (a,b) = toVector a <> toVector b
  fromVector v =
    case VG.splitAt (dim (Proxy :: Proxy a)) v of
      (v1, v2) -> (fromVector v1, fromVector v2)

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

newtype LinMap s a b = LinMap (HM.Matrix s)

{-
infixr 0 ⊸

data a ⊸ b where
  LinMap :: (Scalar a ~ Scalar b) => !(HM.Matrix (Scalar a)) -> a ⊸ b

と書けると格好良いが、スカラー型の異なるベクトル空間の直積とかで困ったことになる。
-}

instance Category (LinMap s) where
  type Obj (LinMap s) a = (VectorSpace a, Scalar a ~ s)

  id :: forall a. Obj (LinMap s) a => LinMap s a a
  id = LinMap $ HM.ident $ dim (Proxy :: Proxy a)

  LinMap m1 . LinMap m2 = LinMap (m1 HM.<> m2)

instance ToFun (LinMap s) where
  toFun (LinMap m) a = fromVector (m HM.#> toVector a)

instance (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => Additive (LinMap s a b) where
  zero = LinMap $ HM.konst 0 (m, n)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)
  LinMap m1 .+. LinMap m2 = LinMap $ HM.add m1 m2

instance (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => VectorSpace (LinMap s a b) where
  type Scalar (LinMap s a b) = s
  scale s (LinMap m) = LinMap (HM.scale s m)
  dim _ = dim (Proxy :: Proxy a) * dim (Proxy :: Proxy b)
  toVector (LinMap m) = HM.flatten m
  fromVector v = LinMap (HM.reshape (dim (Proxy :: Proxy b)) v)

instance Monoidal (LinMap s) where
  LinMap m1 >< LinMap m2 = LinMap $ HM.diagBlock [m1, m2]
  monObj = ObjR

instance Cartesian (LinMap s) where
  exl :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s (a,b) a
  exl = LinMap $ HM.ident m HM.||| HM.konst 0 (m, n)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  exr :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s (a,b) b
  exr = LinMap $ HM.konst 0 (n, m) HM.||| HM.ident n
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  dup :: forall a. Obj (LinMap s) a => LinMap s a (a,a)
  dup = LinMap $
      HM.ident n
      HM.===
      HM.ident n
    where
      n = dim (Proxy :: Proxy a)

instance Cocartesian (LinMap s) where
  inl :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s a (a,b)
  inl = LinMap $
      HM.ident m
      HM.===
      HM.konst 0 (n, m)
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  inr :: forall a b. (Obj (LinMap s) a, Obj (LinMap s) b) => LinMap s b (a,b)
  inr = LinMap $
      HM.konst 0 (n, m)
      HM.===
      HM.ident m
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)

  jam :: forall a. (Obj (LinMap s) a) => LinMap s (a,a) a
  jam = LinMap $ HM.ident n HM.||| HM.ident n
    where
      n = dim (Proxy :: Proxy a)

instance (VectorSpace s, Scalar s ~ s) => Base.Scalable (LinMap s) s where
  scale s = assert (dim (Proxy :: Proxy s) == 1) $ LinMap ((1 HM.>< 1) [s])

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

-- | Dual vector space
newtype Dual a = Dual (LinMap (Scalar a) a (Scalar a))

instance VectorSpace a => Additive (Dual a) where
  zero = Dual $ LinMap $ HM.konst 0 (n, 1)
    where
      n = dim (Proxy :: Proxy a)
  Dual (LinMap m1) .+. Dual (LinMap m2) = Dual $ LinMap $ HM.add m1 m2

instance VectorSpace a => VectorSpace (Dual a) where
  type Scalar (Dual a) = Scalar a
  scale s (Dual (LinMap m)) = Dual (LinMap (HM.scale s m))
  dim _ = dim (Proxy :: Proxy a)
  toVector (Dual (LinMap m)) = HM.flatten m
  fromVector v = Dual (LinMap (HM.reshape 1 v))

toDual :: VectorSpace a => a -> Dual a
toDual a = Dual $ LinMap $ HM.asRow (toVector a)

toDualMap :: forall a b s. (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => LinMap s a b -> LinMap s (Dual b) (Dual a)
toDualMap (LinMap m) = LinMap (HM.tr' m)

fromDual :: forall a. VectorSpace a => Dual a -> a
fromDual (Dual (LinMap m)) = fromVector $ HM.flatten m

fromDualMap :: forall a b s. (VectorSpace a, VectorSpace b, Scalar a ~ s, Scalar b ~ s) => LinMap s (Dual b) (Dual a) -> LinMap s a b
fromDualMap (LinMap m) = LinMap (HM.tr' m)

instance (VectorSpace u, VectorSpace s, Scalar u ~ s, Scalar s ~ s) => Base.HasDot (LinMap s) s u where
  dot x = case toDual x of Dual f -> f
  undot f = fromDual (Dual f)

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

testDual = (toFun f (2,1) == 7, toFun f2 (2,1) == 7)
  where
    f :: LinMap Double (Double, Double) Double
    f = LinMap $ HM.asRow (VG.fromList [2,3])
    f2 = fromDualMap (toDualMap f)

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

data a :⊗ b where
  TensorProd :: (Scalar a ~ s, Scalar b ~ s) => HM.Matrix s -> a :⊗ b

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => Additive (a :⊗ b) where
  zero = TensorProd (HM.konst 0 (m, n))
    where
      m = dim (Proxy :: Proxy a)
      n = dim (Proxy :: Proxy b)
  TensorProd m1 .+. TensorProd m2 = TensorProd $ m1 `HM.add` m2

instance (VectorSpace a, VectorSpace b, Scalar a ~ Scalar b) => VectorSpace (a :⊗ b) where
  type Scalar (a :⊗ b) = Scalar a
  scale s (TensorProd m) = TensorProd $ HM.scale s m
  dim _ = dim (Proxy :: Proxy a) * dim (Proxy :: Proxy b)
  toVector (TensorProd m) = HM.flatten m
  fromVector v = TensorProd (HM.reshape (dim (Proxy :: Proxy b)) v)

curry
  :: forall a b c s. (VectorSpace a, VectorSpace b, VectorSpace c, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s)
  => LinMap s (a :⊗ b) c -> LinMap s a (LinMap s b c)
curry (LinMap m) = LinMap $ HM.reshape (nb*nc) $ HM.flatten m
  where
    nb = dim (Proxy :: Proxy b)
    nc = dim (Proxy :: Proxy c)

uncurry
  :: forall a b c s. (VectorSpace a, VectorSpace b, VectorSpace c, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s)
  => LinMap s a (LinMap s b c) -> LinMap s (a :⊗ b) c
uncurry (LinMap m) = LinMap $ HM.reshape nc $ HM.flatten m
  where
    nc = dim (Proxy :: Proxy c)

mapTensor
  :: forall a b c d s.
     (VectorSpace a, VectorSpace b, VectorSpace c, VectorSpace d, Scalar a ~ s, Scalar b ~ s, Scalar c ~ s, Scalar d ~ s)
  => LinMap s a b -> LinMap s c d -> LinMap s (a :⊗ c) (b :⊗ d)
mapTensor (LinMap m1) (LinMap m2) = LinMap $ ((na*nc) HM.>< (nb*nd)) [(m1 `HM.atIndex` (i, k)) * (m2 `HM.atIndex` (j, l)) | i <- [0..na-1], j <- [0..nc-1], k <- [0..nb-1], l <- [0..nd-1]]
  where
    na = dim (Proxy :: Proxy a)
    nb = dim (Proxy :: Proxy b)
    nc = dim (Proxy :: Proxy c)
    nd = dim (Proxy :: Proxy d)

TensorProd test_Tensor = toFun h t
  where
    f, g :: LinMap Double Double Double
    f = Base.scale 2
    g = Base.scale 3
    h :: LinMap Double (Double :⊗ Double) (Double :⊗ Double)
    h = mapTensor f g

    t :: Double :⊗ Double
    t = TensorProd $ (1 HM.>< 1) [4]


module STCR2Z2T0S0KoffkaCross where

import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           DFT.Plan
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           STC.PowerMethod
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoDecayStr:taoReversalStr:taoCornerStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFilePath:numIterationStr:writeSourceFlagStr:dStr:wStr:sigmaStr:thetaStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      taoDecay = read taoDecayStr :: Double
      taoReversal = read taoReversalStr :: Double
      taoCorner = read taoCornerStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      d = read dStr :: Int
      w = read wStr :: Int
      sigma = read sigmaStr :: Double
      theta = read thetaStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z2T0S0KoffkaCross"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0ReversalCornerRadial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          taoDecay
          taoReversal
          taoCorner
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  let xs =
        [ R2S1RPPoint (d, w, theta, 1)
        , R2S1RPPoint (d, -w, theta, 1)
        , R2S1RPPoint (w, d, theta + 90, 1)
        , R2S1RPPoint (-w, d, theta + 90, 1)
        , R2S1RPPoint (-d, w, theta + 180, 1)
        , R2S1RPPoint (-d, -w, theta + 180, 1)
        , R2S1RPPoint (w, -d, theta + 270, 1)
        , R2S1RPPoint (-w, -d, theta + 270, 1)
        ]
      bias =
        -- computeBiasR2T0S0Gaussian
        --   numPoint
        --   numPoint
        --   theta0Freqs
        --   scale0Freqs
        --   90
        --   sigma
        --   xs
        computeS $
        R.zipWith
          (+)
          (computeBiasR2T0S0Gaussian
             numPoint
             numPoint
             theta0Freqs
             scale0Freqs
             90
             sigma
             xs)
          (computeBiasR2T0S0Gaussian
             numPoint
             numPoint
             theta0Freqs
             scale0Freqs
             (-90)
             sigma
             xs)
      eigenVec =
        computeInitialEigenVectorR2T0S0
          numPoint
          numPoint
          theta0Freqs
          scale0Freqs
          thetaFreqs
          scaleFreqs
          xs
  powerMethodR2Z2T0S0Bias
    plan
    folderPath
    numPoint
    numPoint
    numOrientation
    thetaFreqs
    theta0Freqs
    numScale
    scaleFreqs
    scale0Freqs
    arrR2Z2T0S0
    numIteration
    writeSourceFlag
    -- (printf "_%d" (round r :: Int))
    (printf
       "_%d_%d_%d_%d_%d_%.2f_%.2f_%d_%d"
       numPoint
       (round maxScale :: Int)
       (round taoDecay :: Int)
       (round taoReversal :: Int)
       (round taoCorner :: Int)
       thetaSigma
       scaleSigma
       d
       w)
    0.5
    bias
    eigenVec


{-# LANGUAGE GADTs #-}
module Math.IRT.MLE.Internal.Generic
  ( MLEResult (..)
  , generic_mleEst
  , logLike
  ) where

import Numeric.AD (Mode, Scalar)
import Numeric.AD.Halley

import Statistics.Distribution (Distribution, ContDistr)

import Math.IRT.Fisher
import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.MLE.Internal.MLEResult

generic_mleEst :: (Distribution d, ContDistr d, DensityDeriv d, LogLikelihood d) => [Bool] -> [d] -> Int -> Double -> MLEResult
generic_mleEst rs params steps vTheta =
    let est    = (!! steps) $ extremumNoEq (logLike rs params) vTheta
        fisher = fisherInfoObserved est rs params
    in case fisher of
         (FisherInfo _ t s) -> MLEResult est t s

logLike :: (Mode a, Floating a, Scalar a ~ Double, Distribution d, LogLikelihood d) => [Bool] -> [d] -> a -> a
logLike responses params vTheta =
    let logLik   = zipWith (\a b -> logLikelihood a b vTheta) responses params
        bmePrior = 1
    in sum logLik + log bmePrior


{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE ViewPatterns      #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs      #-}

module Main where

import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.Devel          ( MatrixOrder(..)
                                            , matrixFromVector)
import Numeric.LinearAlgebra.HMatrix
-- import Control.Lens
import Control.Applicative
import qualified Data.ByteString.Lazy as BL
import Data.Csv
import Data.List
import qualified Data.Vector          as V
import qualified Data.Vector.Storable as VS
import GHC.Generics (Generic)
-- import Unsafe.Coerce

import TinyMath             as Tiny

matrixFromVector' :: VS.Storable a => MatrixOrder -> Int -> Int -> V.Vector [a] -> Matrix a
matrixFromVector' o r c = matrixFromVector o r c . VS.fromList . concat . V.toList

identity :: Int -> Matrix Double
identity = ident

pause :: IO ()
pause = do
  putStrLn "\n Program paused. Press enter to continue."
  _ <- getChar
  return ()

class Decodable a where
  decode' :: HasHeader -> BL.ByteString -> V.Vector a

instance Decodable [Double] where
  decode' :: HasHeader -> BL.ByteString -> V.Vector [Double]
  decode' h f = d
    where
      d = case decode NoHeader f :: Either String (V.Vector [Double]) of
        Left err -> V.empty
        Right xs -> xs

instance Decodable (Double, Double) where
  decode' :: HasHeader -> BL.ByteString -> V.Vector (Double, Double)
  decode' h f = d
    where
      d = case decode NoHeader f :: Either String (V.Vector (Double, Double)) of
        Left err -> V.empty
        Right xs -> xs

-- instance VS.Vector Integer where
--   konst :: Integer -> (Int, Int) -> Matrix Integer
--   konst = undefined

--computeCost :: Matrix a -> Matrix a -> a
computeCost :: Matrix Double -> Matrix Double -> Double
computeCost xs theta = cost
  where
    m    = fromIntegral . fst . size $ xs   :: Double
    x    = xs ¿ [0,1]
    y    = xs ¿ [2]
    hyp  = mul x theta                      :: Matrix Double
    temp = (hyp - y)^2                      :: Matrix Double
    cost = ((2*m)**(-1) * sumElements temp ) :: Double

-- app xs (VS.fromList [1.0,1.0,1.0])
-- ((2*m)**(-1)) * (sumElements $ (xs :: Matrix Double) ¿ [1])

gradientDescent :: Matrix a -> Matrix a -> Double -> Integer -> Double
gradientDescent = undefined

-- ===================== Main                   ====================

main :: IO ()
main = do
--  ==================== Part 1: Basic Function ====================
  putStrLn "Running warmUpExercise ..."
  putStrLn "5x5 Identity Matrix:"
  print (identity 5)
  pause

-- ======================= Part 2: Plotting =======================
  putStrLn "Plotting Data..."
  f <- BL.readFile "./data/ex1data1.txt"
  let
    d = decode' NoHeader f
    in Tiny.plot $ V.toList d

  pause

-- =================== Part 3: Cost and Gradient descent ===================
  let
    -- | Some gradient descent settings
    iterations = 1500 :: Integer
    alpha      = 0.1  :: Double

    d     = decode' NoHeader f :: V.Vector [Double]
    mtx   = matrixFromVector' ColumnMajor 2 (V.length d) d
    m     = V.length d
    -- | a Matrix with a column of 1's preadded.
    -- | initialize fittig parameters
    xs    = konst 1 (m, 1) ||| tr mtx      -- | Add a column of ones to x
    theta = konst 0 (2,1)      :: Matrix Double -- | initialize fitting parameters

    -- | compute and display initial cost
    cost  = computeCost xs theta
  print cost

  let
    theta'= gradientDescent xs theta alpha iterations

    -- m =
  return ()


{-# LANGUAGE CPP              #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies     #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Tests.Properties
Copyright   :  (c) Alberto Ruiz 2008
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

Testing properties.

-}

module Test.Numeric.LinearAlgebra.Tests.Properties (
    dist, (|~|), (~~), (~:), Aprox((:~)), (~=),
    zeros, ones,
    square,
    unitary,
    hermitian,
    wellCond,
    positiveDefinite,
    upperTriang,
    upperHessenberg,
    luProp,
    invProp,
    pinvProp,
    detProp,
    nullspaceProp,
--    bugProp,
    svdProp1, svdProp1a, svdProp1b, svdProp2, svdProp3, svdProp4,
    svdProp5a, svdProp5b, svdProp6a, svdProp6b, svdProp7,
    eigProp, eigSHProp, eigProp2, eigSHProp2,
    geigProp,
    qrProp, rqProp, rqProp1, rqProp2, rqProp3,
    hessProp,
    schurProp1, schurProp2,
    cholProp, exactProp,
    expmDiagProp,
    multProp1, multProp2,
    subProp,
    linearSolveProp, linearSolvePropH, linearSolveProp2,

    -- Binary properties
    vectorBinaryRoundtripProp
  , staticVectorBinaryRoundtripProp
  , matrixBinaryRoundtripProp
  , staticMatrixBinaryRoundtripProp
  , staticVectorBinaryFailProp
) where

import           Numeric.LinearAlgebra.HMatrix hiding (Testable)
import qualified Numeric.LinearAlgebra.Static  as Static
import           Test.QuickCheck

import           Data.Binary
import           Data.Binary.Get               (runGet)
import           Data.Either                   (isLeft)
#if MIN_VERSION_base(4,11,0)
import           Prelude                       hiding ((<>))
#endif

import           Debug.Trace

(~=) :: Float -> Float -> Bool
a ~= b = abs (a - b) < 1e-10

trivial :: Testable a => Bool -> a -> Property
trivial = (`classify` "trivial")

-- relative error
dist :: (Num a, Normed a) => a -> a -> Float
dist = relativeError norm_Inf

infixl 4 |~|
a |~| b = a :~1~: b
--a |~| b = dist a b < 10^^(-10)

a ~~ b = fromList a |~| fromList b

data Aprox a = (:~) a Int
-- (~:) :: (Normed a, Num a) => Aprox a -> a -> Bool
a :~n~: b = dist a b < 10^^(-n)

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

square m = rows m == cols m

-- orthonormal columns
orthonormal m = tr m <> m |~| ident (cols m)

unitary m = square m && orthonormal m

hermitian m = square m && m |~| tr m

wellCond m = rcond m > 1/100

positiveDefinite m = minimum (toList e) > 0
    where (e,_v) = eigSH m

upperTriang m = rows m == 1 || down == z
    where down = fromList $ concat $ zipWith drop [1..] (toLists (tr m))
          z = konst 0 (size down)

upperHessenberg m = rows m < 3 || down == z
    where down = fromList $ concat $ zipWith drop [2..] (toLists (tr m))
          z = konst 0 (size down)

zeros (r,c) = reshape c (konst 0 (r*c))

ones (r,c) = zeros (r,c) + 1

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

luProp m =
  -- trace ("x: " ++ show m)
  -- trace ("m: " ++ show m ++ " =~ " ++ show (p <> l <> u) ++ " ... " ++ show (m |~| p <> l <> u))
  -- trace ("f (det p) |~| f s: " ++ show (f (det p)) ++ " = " ++ show  (f s) ++ " ... " ++ show (f (det p) |~| f s))
  m |~| p <> l <> u && f (det p) |~| f s
    where (l,u,p,s) = lu m
          f x = fromList [x]

invProp m = m <> inv m |~| ident (rows m)

pinvProp m =  m <> p <> m |~| m
           && p <> m <> p |~| p
           && hermitian (m<>p)
           && hermitian (p<>m)
    where p = pinv m

detProp m =
  trace ("d1: " ++ show (s d1, s d2))
  s d1 |~| s d2
    where d1 = det m
          d2 = det' * det q
          det' = product $ toList $ takeDiag r
          (q,r) = qr m
          s x = fromList [x]

nullspaceProp m = null nl `trivial` (null nl || m <> n |~| zeros (r,c)
                                     && orthonormal n)
    where n = nullspaceSVD (Left (1*peps)) m (rightSV m)
          nl = toColumns n
          r = rows m
          c = cols m - rank m

------------------------------------------------------------------
{-
-- testcase for nonempty fpu stack
-- uncommenting unitary' signature eliminates the problem
bugProp m = m |~| u <> real d <> tr v && unitary' u && unitary' v
    where (u,d,v) = svd m
          -- unitary' :: (Num (Vector t), Field t) => Matrix t -> Bool
          unitary' a = unitary a
-}
------------------------------------------------------------------

-- fullSVD
svdProp1 m = m |~| u <> real d <> tr v && unitary u && unitary v
  where
    (u,s,v) = svd m
    d = diagRect 0 s (rows m) (cols m)

svdProp1a svdfun m = m |~| u <> real d <> tr v && unitary u && unitary v
  where
    (u,s,v) = svdfun m
    d = diagRect 0 s (rows m) (cols m)

svdProp1b svdfun m = unitary u && unitary v
  where
    (u,_,v) = svdfun m

-- thinSVD
svdProp2 thinSVDfun m
    =  m |~| u <> diag (real s) <> tr v
    && orthonormal u && orthonormal v
    && size s == min (rows m) (cols m)
  where
    (u,s,v) = thinSVDfun m

-- compactSVD
svdProp3 m = (m |~| u <> real (diag s) <> tr v
             && orthonormal u && orthonormal v)
  where
    (u,s,v) = compactSVD m

svdProp4 m' = m |~| u <> real (diag s) <> tr v
           && orthonormal u && orthonormal v
           && (size s == r || r == 0 && size s == 1)
  where
    (u,s,v) = compactSVD m
    m = fromBlocks [[m'],[m']]
    r = rank m'

svdProp5a m = all (s1|~|) [s3,s5] where
    s1       = singularValues (m :: Matrix Float)
--  s2       = svRd m
    (_,s3,_) = svd m
--  (_,s4,_) = svdRd m
    (_,s5,_) = thinSVD m
--  (_,s6,_) = thinSVDRd m

svdProp5b m = all (s1|~|) [s3,s5] where
    s1       = singularValues (m :: Matrix (Complex Float))
--  s2       = svCd m
    (_,s3,_) = svd m
--  (_,s4,_) = svdCd m
    (_,s5,_) = thinSVD m
--  (_,s6,_) = thinSVDCd m

svdProp6a m = s |~| s' && v |~| v' && s |~| s'' && u |~| u'
  where
    (u,s,v) = svd (m :: Matrix Float)
    (s',v') = rightSV m
    (u',s'') = leftSV m

svdProp6b m = s |~| s' && v |~| v' && s |~| s'' && u |~| u'
  where
    (u,s,v) = svd (m :: Matrix (Complex Float))
    (s',v') = rightSV m
    (u',s'') = leftSV m

svdProp7 m = s |~| s' && u |~| u' && v |~| v' && s |~| s'''
  where
    (u,s,v) = svd m
    (s',v') = rightSV m
    (u',_s'') = leftSV m
    s''' = singularValues m

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

eigProp m = complex m <> v |~| v <> diag s
    where (s, v) = eig m

eigSHProp m = m <> v |~| v <> real (diag s)
              && unitary v
              && m |~| v <> real (diag s) <> tr v
    where (s, v) = eigSH' m

eigProp2 m = fst (eig m) |~| eigenvalues m

eigSHProp2 m = fst (eigSH' m) |~| eigenvaluesSH' m

geigProp a b = a' <> v <> diag betas' |~| b' <> v <> diag alphas
    where (alphas, betas, v) = geig a b
          betas' = complex betas
          a' = complex a
          b' = complex b

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

qrProp m = q <> r |~| m && unitary q && upperTriang r
    where (q,r) = qr m

rqProp m = r <> q |~| m && unitary q && upperTriang' r
    where (r,q) = rq m

rqProp1 m = r <> q |~| m
    where (r,q) = rq m

rqProp2 m = unitary q
    where (_r,q) = rq m

rqProp3 m = upperTriang' r
    where (r,_q) = rq m

upperTriang' r = upptr (rows r) (cols r) * r |~| r
    where upptr f c = build (f,c) $ \r' c' -> if r'-t > c' then 0 else 1
              where t = fromIntegral (f-c)

hessProp m = m |~| p <> h <> tr p && unitary p && upperHessenberg h
    where (p,h) = hess m

schurProp1 m = m |~| u <> s <> tr u && unitary u && upperTriang s
    where (u,s) = schur m

schurProp2 m = m |~| u <> s <> tr u && unitary u && upperHessenberg s -- fixme
    where (u,s) = schur m

cholProp m = m |~| tr c <> c && upperTriang c
    where c = chol (trustSym m)

exactProp m = chol (trustSym m) == chol (trustSym (m+0))

expmDiagProp m = expm (logm m) :~ 7 ~: complex m
    where logm = matFunc log

-- reference multiply
mulH a b = fromLists [[ doth ai bj | bj <- toColumns b] | ai <- toRows a ]
    where doth u v = sum $ zipWith (*) (toList u) (toList v)

multProp1 p (a,b) = (a <> b) :~p~: (mulH a b)

multProp2 p (a,b) = (tr (a <> b)) :~p~: (tr b <> tr a)

linearSolveProp f m = f m m |~| ident (rows m)

linearSolvePropH f m = f m (unSym m) |~| ident (rows (unSym m))

linearSolveProp2 f (a,x) = not wc `trivial` (not wc || a <> f a b |~| b)
    where q = min (rows a) (cols a)
          b = a <> x
          wc = rank a == q

subProp m = m == (conj . tr . fromColumns . toRows) m

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

vectorBinaryRoundtripProp :: Vector Float -> Bool
vectorBinaryRoundtripProp vec = decode (encode vec) == vec

staticVectorBinaryRoundtripProp :: Static.R 5 -> Bool
staticVectorBinaryRoundtripProp vec =
  let
    decoded = decode (encode vec) :: Static.R 500
  in
    Static.extract decoded == Static.extract vec

matrixBinaryRoundtripProp :: Matrix Float -> Bool
matrixBinaryRoundtripProp mat = decode (encode mat) == mat

staticMatrixBinaryRoundtripProp :: Static.L 100 200 -> Bool
staticMatrixBinaryRoundtripProp mat =
  let
    decoded = decode (encode mat) :: Static.L 100 200
  in
    (Static.extract decoded) == (Static.extract mat)

staticVectorBinaryFailProp :: Static.R 20 -> Bool
staticVectorBinaryFailProp vec =
  let
    decoded = runGet get (encode vec) :: Either String (Static.R 50)
  in
    isLeft decoded


{-# LANGUAGE ConstrainedClassMethods #-}
{-# LANGUAGE ConstraintKinds         #-}
{-# LANGUAGE DefaultSignatures       #-}
{-# LANGUAGE FlexibleContexts        #-}
{-# LANGUAGE FlexibleInstances       #-}
{-# LANGUAGE FunctionalDependencies  #-}
{-# LANGUAGE KindSignatures          #-}
{-# LANGUAGE MultiParamTypeClasses   #-}
{-# LANGUAGE TypeFamilies            #-}
{-# LANGUAGE ViewPatterns            #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Diagrams.Coordinates.Isomorphic
-- Copyright   :  (C) 2015 Christopher Chalmers
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Christopher Chalmers
-- Stability   :  experimental
-- Portability :  non-portable
--
-- This module defines a class for coordinates that are (loosely)
-- isomorphic to the standard spaces ('V2' and 'V3'). This allows plots
-- to accept more data types for plot data.
--
----------------------------------------------------------------------------
module Diagrams.Coordinates.Isomorphic
  ( -- * Type constraints
    HasIndexedBasis, Euclidean

    -- * Vector like
  , VectorLike (..)
  , V2Like, V3Like

    -- * Point like
  , PointLike (..)
  , P2Like, P3Like
  )
  where

import           Control.Lens
import           Data.Complex
import           Data.Typeable

import           Diagrams.Prelude

type HasIndexedBasis v = (HasBasis v, TraversableWithIndex (E v) v)

-- | Umbrella class giving everything needed for working in the space. This is
--   basically 'V2' or 'V3' from "linear".
type Euclidean (v :: * -> *) = (HasLinearMap v, HasIndexedBasis v, Metric v)

-- vector like ---------------------------------------------------------

-- | Provides an 'Iso'' between @a@ and @v n@. This is normally used to
--   convert between the data type you're already using, @a@, and diagram's
--   native form, @v n@.
class (Euclidean v, Typeable v) => VectorLike v n a | a -> v n where
  -- | Isomorphism from @Point v n@ to something 'PointLike' @a@.
  --
  -- >>> V2 3 5 ^. vectorLike :: (Int, Int)
  -- (3,5)
  vectorLike :: Iso' (v n) a

  -- | Isomorphism from something 'PointLike' @a@ to @Point v n@.
  --
  -- >>> ((3, 5) :: (Int, Int)) ^. unvectorLike
  -- V2 3 5
  unvectorLike :: Iso' a (v n)
  unvectorLike = from vectorLike
  {-# INLINE unvectorLike #-}

instance VectorLike V2 n (V2 n) where
  vectorLike = id
  {-# INLINE vectorLike #-}

type V2Like = VectorLike V2

instance n ~ m => VectorLike V2 n (n, m) where
  vectorLike = iso unr2 r2
  {-# INLINE vectorLike #-}

instance VectorLike V2 n (Complex n) where
  vectorLike = iso (\(V2 x y) -> x :+ y)
                   (\(i :+ j) -> V2 i j)
  {-# INLINE vectorLike #-}

type V3Like = VectorLike V3

instance VectorLike V3 n (V3 n) where
  vectorLike = id
  {-# INLINE vectorLike #-}

instance (n ~ m, m ~ o) => VectorLike V3 n (n, m, o) where
  vectorLike = iso unr3 r3
  {-# INLINE vectorLike #-}

-- point like ----------------------------------------------------------

-- | Provides an 'Iso'' between @a@ and @'Point' v n@. This is normally used to
--   convert between the data type you're already using, @a@, and diagram's
--   native form, @'Point' v n@.
class (Euclidean v, Typeable v) => PointLike v n a | a -> v n where
  -- | Isomorphism from @'Point' v n@ to something 'PointLike' @a@.
  --
  -- >>> mkP2 3 5 ^. pointLike :: (Int, Int)
  -- (3,5)
  pointLike :: Iso' (Point v n) a

  -- | Isomorphism from something 'PointLike' @a@ to @Point v n@.
  --
  -- >>> ((3, 5) :: (Int, Int)) ^. unpointLike
  -- P (V2 3 5)
  unpointLike :: Iso' a (Point v n)
  unpointLike = from pointLike
  {-# INLINE unpointLike #-}

-- | Things that are isomorphic to points in R2.
type P2Like = PointLike V2

instance (Euclidean v, Typeable v) => PointLike v n (Point v n) where
  pointLike = id

instance PointLike V2 n (V2 n) where
  pointLike = _Point
  {-# INLINE pointLike #-}

instance n ~ m => PointLike V2 n (n, m) where
  pointLike = iso unp2 p2
  {-# INLINE pointLike #-}

instance PointLike V2 n (Complex n) where
  pointLike = iso (\(unp2 -> (x,y)) -> x :+ y)
                  (\(i :+ j)        -> p2 (i,j))
  {-# INLINE pointLike #-}

type P3Like = PointLike V3

instance (n ~ m, m ~ o) => PointLike V3 n (n, m, o) where
  pointLike = iso unp3 p3
  {-# INLINE pointLike #-}


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}
module Sound.Random
(
    -- * Randoms
    Randoms(..)
    -- * Samples
    , randomSamples
    -- * Vectors
    , vrandom
    , vrandoms
    -- * Tables
    , trandoms
    -- * Spectrums
    , vwhitespectrum
    , twhitespectrum
    -- * Reexports
    , Ran.StdGen
    , Ran.mkStdGen
)
where

import qualified Prelude as P
import qualified Data.Bits as B
import qualified Data.Complex as Cp

import qualified System.Random as Ran

import qualified Data.Vector.Unboxed as Vu

import Sound.Class
import Sound.StreamVector
import Sound.Table

class Randoms g v f where
    randoms :: g -> f v
    randomRs :: v -> v -> g -> f v

instance (Ran.Random a, FromList f) => Randoms Ran.StdGen a f where
    randoms g = fromList (error "L.randoms: list must be infinite") $ Ran.randoms g
    randomRs r s g = fromList (error "L.randoms: list must be infinite") $ Ran.randomRs (r,s) g

-- | Random sample values (random numbers between -1 and 1).
randomSamples :: (Randoms g v f, Functor f, Fractional v) => g -> f v
randomSamples = fmap (\ x -> x + x - 1) . randoms

vrandom :: forall g a. (Randoms g a [], Ran.Random a, Vu.Unbox a) => Int -> Ran.StdGen -> Vu.Vector a
vrandom size gen = vfromlist $ take size (randoms gen :: [a])

vrandoms :: forall g a. (Randoms g a [], Ran.Random a, Fractional a, Vu.Unbox a) => Int -> g -> Vu.Vector a
vrandoms size gen = Vu.map (\ x -> x + x - 1) . vfromlist $ take size (randoms gen :: [a])

trandoms :: forall g a. (Randoms g a [], Ran.Random a, Fractional a, Vu.Unbox a) => Lgsize -> g -> Tab a
trandoms lgsize gen =
    mktab $ Vu.fromList $ take n (randomSamples gen :: [a])
    where
        n = B.shiftL 1 lgsize

-- | The size must be even.
vwhitespectrum :: (Ran.Random a, RealFloat a, Vu.Unbox a) => Int -> Ran.StdGen -> Vu.Vector (Cp.Complex a)
vwhitespectrum n gen =
    vfromlist $ 0 : map (\ x -> Cp.cis (2 * pi * x)) (P.tail phases ++ reverse phases)
    where
        -- x k = x ((n - k) mod n); k = 0, 1, 2, ... n - 1
        phases = take h $ Ran.randoms gen
        h = div n 2

twhitespectrum :: (Ran.Random a, RealFloat a, Vu.Unbox a) => Lgsize -> Ran.StdGen -> Tab (Cp.Complex a)
twhitespectrum lgsize gen =
    mktab $ vwhitespectrum size gen
    where
        size = B.shiftL 1 lgsize