{-# LANGUAGE DeriveGeneric #-}

{-|
Module      : HABQTlib.Data

This module contains data types and helper functions for working with quantum
state vectors and density matrices.
-}
module HABQTlib.Data
  ( Dim
  , Rank
  , NumberOfParticles
  , QBitNum
  , Weight
  , MHMCiter
  , OptIter
  , OutputVerb(..)
  , DensityMatrix(..)
  , truncateRank
  , PureStateVector(..)
  , pureStateLikelihood
  , svToDM
  , WeighedDensityMatrix(..)
  , mkWDM
  , mkWDM1
  , WeighedPureStateVector(..)
  , (<+>)
  , fidelity
  , fidelityDM
  , validSV
  , validDM
  , validPartNum
  , validRank
  , validMHMCiter
  , validOptIter
  , validQBitN
  ) where

import Control.Newtype.Generics (Newtype, unpack)
import Data.Bool.HT (select)
import Data.Complex
import Data.Validation
import GHC.Generics (Generic)
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra (Matrix)

-- | Dimension of Hilbert space.
type Dim = Int

-- | Rank of mixed state.
type Rank = Int

-- | Number of particles per rank.
type NumberOfParticles = Int

-- | Number of quantum bits.
type QBitNum = Int

-- | Number of MHMC iterations to perform when resampling.
type MHMCiter = Int

-- | Number of optimisation steps to perform when searching for optimal
-- measurment.
type OptIter = Int

-- | Weight associated with a particle.
type Weight = Double

-- | Output verbosity settings.
data OutputVerb
  = NoOutput -- ^ No stdout output
  | FidOutput -- ^ Only output fidelities and weights of hierarchical mean estimates
  | FullOutput -- ^ Full output, including resampling diagnostic information
  deriving (Eq, Show, Ord)

-- | Density matrix are stored as hmatrix matrices of complex doubles.
newtype DensityMatrix = DensityMatrix
  { getDensityMatrix :: Matrix (Complex Double)
  } deriving (Eq, Show, Generic)

-- | Pure state vectors are stored as hmatrix matrices of complex doubles.
-- Such matrices only have one column.
newtype PureStateVector = PureStateVector
  { getStateVector :: Matrix (Complex Double)
  } deriving (Eq, Show, Generic)

instance Newtype DensityMatrix

instance Newtype PureStateVector

-- | Check whether a pure state vector is properly normed.
validSV :: PureStateVector -> Validation [String] PureStateVector
validSV =
  validate
    ["State vector must have unit norm."]
    (\x -> abs (1 - LA.norm_2 (unpack x)) < 1e-12)

-- | Verify that density matrix is Hermitian and has trace 1.
validDM :: DensityMatrix -> Validation [String] DensityMatrix
validDM =
  let traceU :: LA.Matrix (Complex Double) -> Bool
      traceU dm =
        (abs (1 - (magnitude . LA.sumElements . LA.takeDiag) dm) < 1e-12)
      hermU dm = (LA.norm_2 (dm - LA.tr dm) < 1e-6)
      both = ((&&) <$> traceU <*> hermU) . unpack
      dmM = ["Density matrix must be Hermitian and have trace of 1."]
   in validate dmM both

qnM :: [String]
qnM = ["Number of quantum bits must be a positive integer."]

-- | Verify that number of quantum bits is positive.
validQBitN :: QBitNum -> Validation [String] QBitNum
validQBitN = validate qnM (> 0)

miM :: [String]
miM = ["Number of MHMC iterations must be a positive integer."]

-- | Verify that number of MHMC iterations is a positive integer.
validMHMCiter :: MHMCiter -> Validation [String] MHMCiter
validMHMCiter = validate miM (> 0)

pnM :: [String]
pnM = ["Number of particles per rank must be a positive integer."]

-- | Verify that particle number is a positive integer.
validPartNum :: NumberOfParticles -> Validation [String] NumberOfParticles
validPartNum = validate pnM (> 0)

rM :: [String]
rM = ["Rank must be a positive integer."]

-- | Verify that rank is a positive integer. Setting rank to be higher than
-- the dimension of space creates poinless performance overhead, but isn't
-- prevented by validation.
validRank :: Rank -> Validation [String] Rank
validRank = validate rM (> 0)

oiM :: [String]
oiM = ["Number of POVM optimisation iterations must be a positive integer."]

-- | Verify that number of POVM optimisation iterations is positive.
validOptIter :: OptIter -> Validation [String] OptIter
validOptIter = validate oiM (> 0)

-- | Weighed density matrix where weight is stored separately as first
-- coordinate of a tuple.
newtype WeighedDensityMatrix = WeighedDensityMatrix
  { getWDM :: (Weight, DensityMatrix)
  } deriving (Eq, Show, Generic)

-- | A shorter alias for curried 'WeighedDensityMatrix' constructor.
mkWDM :: Weight -> DensityMatrix -> WeighedDensityMatrix
mkWDM w dm = WeighedDensityMatrix (w, dm)

-- | Alias for @'mkWDM' 1@.
mkWDM1 :: DensityMatrix -> WeighedDensityMatrix
mkWDM1 = mkWDM 1

-- | Weighed state vector where weight is stored separately as first coordinate
-- of a tuple.
newtype WeighedPureStateVector = WeighedPureStateVector
  { getWSV :: (Weight, PureStateVector)
  } deriving (Eq, Show, Generic)

instance Newtype WeighedDensityMatrix

instance Newtype WeighedPureStateVector

-- | Fidelity (probability of measurement) between pure states.
fidelity :: PureStateVector -> PureStateVector -> Double
fidelity (PureStateVector sv1) (PureStateVector sv2) =
  let ips = magnitude (LA.atIndex (LA.tr sv1 LA.<> sv2) (0, 0)) ^ (2 :: Int)
   in select ips [(ips < 0, 0), (ips > 1, 1)]

-- | Fidelity (probability of measurement) between mixed states.
fidelityDM :: DensityMatrix -> DensityMatrix -> Double
fidelityDM (DensityMatrix dm1) (DensityMatrix dm2) =
  let (u1, s1) = LA.leftSV dm1
      (u2, s2) = LA.leftSV dm2
      ss1 = LA.cmap (\x -> sqrt x :+ 0) s1
      ss2 = LA.cmap (\x -> sqrt x :+ 0) s2
   in LA.norm_nuclear
        (u1 LA.<> LA.diag ss1 LA.<> LA.tr u1 LA.<> u2 LA.<> LA.diag ss2 LA.<>
         LA.tr u2) ^
      (2 :: Int)

-- | Calculate the density matrix of a given pure state.
svToDM :: PureStateVector -> DensityMatrix
svToDM (PureStateVector sv) = DensityMatrix $ sv LA.<> LA.tr sv

infix 8 <+>

-- | Given two weighed density matrixes, compute their mixture. Associative
-- operation.
(<+>) :: WeighedDensityMatrix -> WeighedDensityMatrix -> WeighedDensityMatrix
wdm0 <+> wdm1 =
  WeighedDensityMatrix
    (w0 + w1, DensityMatrix $ LA.scale c0 dm0 + LA.scale c1 dm1)
  where
    up = fmap unpack . unpack
    (w0, dm0) = up wdm0
    (w1, dm1) = up wdm1
    cs = LA.fromList [w0 :+ 0, w1 :+ 0]
    csn = LA.scale (1 / LA.sumElements cs) cs
    c0 = csn LA.! 0
    c1 = csn LA.! 1

-- | Probability of obtaining a measurement result when projecting a system in
-- mixed state determined by a density matrix onto a pure state.
pureStateLikelihood :: PureStateVector -> DensityMatrix -> Double
pureStateLikelihood (PureStateVector sv) (DensityMatrix dm) =
  let p =
        LA.magnitude . LA.sumElements . LA.takeDiag $ LA.tr sv LA.<> dm LA.<> sv
   in select p [(p < 0, 0), (p > 1, 1)]

-- | Set smallest eigenvalues of a weighed density matrix to zero until
-- specified rank is reached.
truncateRank :: Rank -> WeighedDensityMatrix -> WeighedDensityMatrix
truncateRank targetRank (WeighedDensityMatrix (w, DensityMatrix dm)) =
  let (u, s, _) = LA.svd dm
      st = LA.real $ LA.subVector 0 targetRank s
      ut = u LA.¿ [0 .. (targetRank - 1)]
      stn = LA.scale (1 / LA.sumElements st) st
   in WeighedDensityMatrix
        (w, DensityMatrix $ ut LA.<> LA.diag stn LA.<> LA.tr ut)


module Seed
  ( Seed
  , createSeed
  , arrowsAt
  , drawSeed
  )
where

import           Arrow                          ( Arrow
                                                , arrange
                                                , endpoint
                                                , updateArrow
                                                )
import qualified Fourier
import           Data.Complex
import           Graphics.Gloss

-- A seed represents an initial state of arrow frequencies, directions, and magnitudes
-- Any complex function can be traced out by changing this initial state.
-- It also contains some drawing and scaling information.
data Seed = Seed { _arrows :: [(Arrow, Integer)], _spaceScale :: Float, _timeScale :: Float, _trailLength :: Float }

createSeed :: (Float -> Complex Float) -> Int -> Float -> Float -> Float -> Seed
createSeed f n spaceScale timeScale trailLength = Seed
  { _arrows      = arrows
  , _spaceScale  = spaceScale
  , _timeScale   = timeScale
  -- measure trail length by percentage of a time-scaled second
  , _trailLength = trailLength
  }
 where
  -- order arrow frequencies 0, 1, -1, 2, -2, 3, -3, etc.
  freqs             = take n (0 : concat [ [x, -x] | x <- [1 ..] ])
  startingPositions = map (Fourier.coefficient f) freqs
  arrows            = zip startingPositions freqs

arrowsAt :: Seed -> Float -> [Arrow]
arrowsAt seed t = map (\(arr, freq) -> updateArrow arr freq t) (_arrows seed)

drawSeed :: Seed -> Float -> Picture
drawSeed seed seconds = scale s s $ Pictures [vectors, tip, trail]
 where
  s         = _spaceScale seed
  t         = _timeScale seed * seconds
  positions = arrange $ arrowsAt seed t
  (x, y)    = last positions
  -- draw stuff
  vectors   = color (light aquamarine) . line $ positions
  tip       = color red . translate x y $ circleSolid (5 / _spaceScale seed)
  trail     = color white . line $ map (endpoint . arrowsAt seed) ts
  -- the trail draws through points with a step size of 1/1000 of a (time-scaled) second
  tf        = truncateTo 3 t
  ti        = max 0 (tf - _trailLength seed)
  ts        = [ti, ti + 0.001 .. t]

-- truncate a real number to the given number of decimal places
truncateTo :: (RealFrac a) => Integer -> a -> a
truncateTo n = (/ 10 ^ n) . fromIntegral . truncate . (* 10 ^ n)


-- | Internal utilities.

module Math.Probably.Utils where

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Vector.Storable as V
import Math.Probably.Types
import Numeric.LinearAlgebra
import Statistics.Distribution
import Statistics.Distribution.Normal

lookupDefault  :: Ord k => a -> k -> Map k a -> a
lookupDefault d k m = fromMaybe d (Map.lookup k m)

-- | A spherical Gaussian distribution.
sphereGauss :: ContinuousParams -> ContinuousParams -> Double -> Double
sphereGauss xs m sd = product $ zipWith density normalDists xsAsList where
  xsAsList    = toList xs
  meanAsList  = toList m
  normalDists = map (`normalDistr` sd) meanAsList

-- | Scalar-vector multiplication.
(.*) :: Double -> ContinuousParams -> ContinuousParams
z .* xs = mapVector (* z) xs

-- | Scalar-vector subtraction.
(.-) :: ContinuousParams -> ContinuousParams -> ContinuousParams
xs .- ys = zipVectorWith (-) xs ys

-- | Scalar-vector addition.
(.+) :: ContinuousParams -> ContinuousParams -> ContinuousParams
xs .+ ys = zipVectorWith (+) xs ys

-- | The leapfrog integrator.
leapfrog :: Gradient -> Particle -> Double -> Particle
leapfrog glTarget (q, r) e = (qf, rf) where
  rm = adjustMomentum glTarget e (q, r)
  qf = adjustPosition e (rm, q)
  rf = adjustMomentum glTarget e (qf, rm)

-- | Adjust momentum according to a half-leapfrog step.
adjustMomentum :: Gradient -> Double -> Particle -> ContinuousParams
adjustMomentum glTarget e (t, r) = r .+ ((e / 2) .* glTarget t)

-- | Adjust position according to a half-leapfrog step.
adjustPosition :: Double -> Particle -> ContinuousParams
adjustPosition e (r, t) = t .+ (e .* r)

-- | A target augmented by momentum auxilliary variables.
auxilliaryTarget :: (ContinuousParams -> Double) -> Particle -> Double
auxilliaryTarget lTarget (t, r) =
  lTarget t - 0.5 * innerProduct r r

innerProduct :: ContinuousParams -> ContinuousParams -> Double
innerProduct xs ys = V.sum $ V.zipWith (*) xs ys

indicate :: Integral a => Bool -> a
indicate True  = 1
indicate False = 0

fi :: (Integral a, Num b) => a -> b
fi = fromIntegral


{-# LANGUAGE FlexibleContexts #-}
module Legacy.QuantumCircuit where

import Data.Complex
import Data.Either
import Data.List
import Data.Function
import System.Random
import Control.Monad.IO.Class

import Helper

import Tensor

-- applyGate :: (QuantumRegister q Num RealFloat, RealFloat a) => q (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Index
--   -> q (Complex a) -- ^ State vector
--   -> Either String (q (Complex a))
-- applyGate x i y
--     | not (isSquare x) = Left "Invalid Gate Tensor"
--     | not (isSquare y) = Left "Invalid Statevector"
--     | i < 0 || i >= rank y = Left "Index out of bounds"
--     | otherwise = Right $ norm $ build (rank y) [(i, x)] pauliId * y

-- applyGateAll :: (QuantumRegister q Num RealFloat, RealFloat a) => q (Complex a) -> q (Complex a) -> Either String (q (Complex a))
-- applyGateAll x y
--     | not (isSquare x) = Left "Invalid Gate Tensor"
--     | not (isSquare y) = Left "Invalid Statevector"
--     | otherwise = Right $ norm $ build (rank y) [] x * y

-- applyControl :: (QuantumRegister q, RealFloat a) => q (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Control qubit index
--   -> Int -- ^ Apply qubit index
--   -> q (Complex a) -- ^ State vector
--   -> Either String (q (Complex a))
-- applyControl x ctl i y
--     | not (isSquare x) = Left "Invalid gate tensor"
--     | not (isSquare y) = Left "Invalid state vector"
--     | i < 0 || i >= rank y = Left "Application qubit index out of bounds"
--     | ctl < 0 || ctl >= rank y = Left "Control qubit index out of bounds"
--     | ctl == i = Left "Control and application qubit cannot be the same"
--     | otherwise = Right $ norm $ (build (rank y) [(ctl, mask0)] pauliId
--         + build (rank y) (sortBy (compare `on` fst) [(ctl, mask1), (i, x)]) pauliId) * y

-- collapse :: (RealFloat a, QuantumRegister q) => Int -- ^ Index of the qubit
--         -> Bool -- ^ Collapsed State
--         -> q (Complex a) -- ^ Qubits to collapse
--         -> Either String (q (Complex a)) -- ^ Resulting Tensor
-- collapse i state q
--     | i < 0 || i >= rank q = Left "Index out of bounds"
--     | not (isSquare q) = Left "Invalid state vector"
--     | state = norm <$> applyGate mask1 i q
--     | otherwise = norm <$> applyGate mask0 i q

-- measure i v = do
--     a <- randomRIO (0.0, 1.0)
--     if a < fst p then
--         return $ fmap (\x -> (False, x)) (collapse i False =<< v)
--     else
--         return $ fmap (\x -> (True, x)) (collapse i True =<< v)
--     where
--         prob = getProb i =<< v
--         p = fromRight (0.0, 0.0) prob


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Gamma
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The gamma distribution.  This is a continuous probability
-- distribution with two parameters, /k/ and ϑ. If /k/ is
-- integral, the distribution represents the sum of /k/ independent
-- exponentially distributed random variables, each of which has a
-- mean of ϑ.

module Statistics.Distribution.Gamma
    (
      GammaDistribution
    -- * Constructors
    , gammaDistr
    , improperGammaDistr
    -- * Accessors
    , gdShape
    , gdScale
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<$>), (<*>))
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_NaN, m_neg_inf)
import Numeric.SpecFunctions (incompleteGamma, invIncompleteGamma, logGamma, digamma)
import Statistics.Distribution.Poisson.Internal as Poisson
import qualified Statistics.Distribution as D
import qualified System.Random.MWC.Distributions as MWC

-- | The gamma distribution.
data GammaDistribution = GD {
      gdShape :: {-# UNPACK #-} !Double -- ^ Shape parameter, /k/.
    , gdScale :: {-# UNPACK #-} !Double -- ^ Scale parameter, ϑ.
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON GammaDistribution
instance ToJSON GammaDistribution

instance Binary GammaDistribution where
    put (GD x y) = put x >> put y
    get = GD <$> get <*> get

-- | Create gamma distribution. Both shape and scale parameters must
-- be positive.
gammaDistr :: Double            -- ^ Shape parameter. /k/
           -> Double            -- ^ Scale parameter, ϑ.
           -> GammaDistribution
gammaDistr k theta
  | k     <= 0 = error $ msg ++ "shape must be positive. Got " ++ show k
  | theta <= 0 = error $ msg ++ "scale must be positive. Got " ++ show theta
  | otherwise  = improperGammaDistr k theta
    where msg = "Statistics.Distribution.Gamma.gammaDistr: "

-- | Create gamma distribution. This constructor do not check whether
--   parameters are valid
improperGammaDistr :: Double            -- ^ Shape parameter. /k/
                   -> Double            -- ^ Scale parameter, ϑ.
                   -> GammaDistribution
improperGammaDistr = GD

instance D.Distribution GammaDistribution where
    cumulative = cumulative

instance D.ContDistr GammaDistribution where
    density    = density
    logDensity (GD k theta) x
      | x <= 0    = m_neg_inf
      | otherwise = log x * (k - 1) - (x / theta) - logGamma k - log theta * k
    quantile   = quantile

instance D.Variance GammaDistribution where
    variance (GD a l) = a * l * l

instance D.Mean GammaDistribution where
    mean (GD a l) = a * l

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

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

instance D.MaybeEntropy GammaDistribution where
  maybeEntropy (GD a l)
    | a > 0 && l > 0 =
      Just $
      a
      + log l
      + logGamma a
      + (1-a) * digamma a
    | otherwise = Nothing

instance D.ContGen GammaDistribution where
    genContVar (GD a l) = MWC.gamma a l

density :: GammaDistribution -> Double -> Double
density (GD a l) x
  | a < 0 || l <= 0   = m_NaN
  | x <= 0            = 0
  | a == 0            = if x == 0 then m_pos_inf else 0
  | x == 0            = if a < 1 then m_pos_inf else if a > 1 then 0 else 1/l
  | a < 1             = Poisson.probability (x/l) a * a / x
  | otherwise         = Poisson.probability (x/l) (a-1) / l

cumulative :: GammaDistribution -> Double -> Double
cumulative (GD k l) x
  | x <= 0    = 0
  | otherwise = incompleteGamma k (x/l)

quantile :: GammaDistribution -> Double -> Double
quantile (GD k l) p
  | p == 0         = 0
  | p == 1         = 1/0
  | p > 0 && p < 1 = l * invIncompleteGamma k p
  | otherwise      =
    error $ "Statistics.Distribution.Gamma.quantile: p must be in [0,1] range. Got: "++show p


{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE EmptyDataDecls             #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FunctionalDependencies     #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE Rank2Types                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE ViewPatterns               #-}

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

{- |
Module      :  Numeric.LinearAlgebra.Static
Copyright   :  (c) Alberto Ruiz 2014
License     :  BSD3
Stability   :  experimental

Experimental interface with statically checked dimensions.

See code examples at http://dis.um.es/~alberto/hmatrix/static.html.

-}

module Numeric.LinearAlgebra.Static(
    -- * Vector
       ℝ, R,
    vec2, vec3, vec4, (&), (#), split, headTail,
    vector,
    linspace, range, dim,
    -- * Matrix
    L, Sq, build,
    row, col, (|||),(===), splitRows, splitCols,
    unrow, uncol,
    tr,
    eye,
    diag,
    blockAt,
    matrix,
    -- * Complex
    ℂ, C, M, Her, her, 𝑖,
    toComplex,
    fromComplex,
    complex,
    real,
    imag,
    sqMagnitude,
    magnitude,
    -- * Products
    (<>),(#>),(<.>),
    -- * Linear Systems
    linSolve, (<\>),
    -- * Factorizations
    svd, withCompactSVD, svdTall, svdFlat, Eigen(..),
    withNullspace, withOrth, qr, chol,
    -- * Norms
    Normed(..),
    -- * Random arrays
    Seed, RandDist(..),
    randomVector, rand, randn, gaussianSample, uniformSample,
    -- * Misc
    mean, meanCov,
    Disp(..), Domain(..),
    withVector, withMatrix, exactLength, exactDims,
    toRows, toColumns, withRows, withColumns,
    Sized(..), Diag(..), Sym, sym, mTm, unSym, (<·>)
) where

import           Control.Arrow               ((***))
import qualified Data.Bifunctor              as BF (first)
import           Data.Proxy                  (Proxy (..))
import           Data.Type.Equality          ((:~:) (Refl))
import           GHC.TypeLits
import           Internal.Static
import           Numeric.LinearAlgebra       hiding (C, Konst (..), R, build, chol, col,
                                              complex, diag, disp, dot, eig, eigSH,
                                              eigenvalues, eigenvaluesSH, fromComplex,
                                              fromList, gaussianSample, linspace, mTm,
                                              magnitude, matrix, meanCov, qr, rand, randn,
                                              randomVector, range, real, row, size, svd,
                                              sym, takeDiag, toColumns, toComplex, toRows,
                                              unSym, uniformSample, vector, ( #> ), (<.>),
                                              (<>), (<\>), (===), (|||))
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as LA
import           Text.Printf
#if MIN_VERSION_base(4,11,0)
import           Prelude                     hiding ((<>))
#endif

ud1 :: R n -> Vector ℝ
ud1 (R (Dim v)) = v

infixl 4 &
(&) :: forall n . KnownNat n
    => R n -> ℝ -> R (n+1)
u & x = u # (konst x :: R 1)

infixl 4 #
(#) :: forall n m . (KnownNat n, KnownNat m)
    => R n -> R m -> R (n+m)
(R u) # (R v) = R (vconcat u v)

vec2 :: ℝ -> ℝ -> R 2
vec2 a b = R (gvec2 a b)

vec3 :: ℝ -> ℝ -> ℝ -> R 3
vec3 a b c = R (gvec3 a b c)

vec4 :: ℝ -> ℝ -> ℝ -> ℝ -> R 4
vec4 a b c d = R (gvec4 a b c d)

vector :: KnownNat n => [ℝ] -> R n
vector = fromList

matrix :: (KnownNat m, KnownNat n) => [ℝ] -> L m n
matrix = fromList

linspace :: forall n . KnownNat n => (ℝ,ℝ) -> R n
linspace (a,b) = v
  where
    v = mkR (LA.linspace (size v) (a,b))

range :: forall n . KnownNat n => R n
range = v
  where
    v = mkR (LA.linspace d (1,fromIntegral d))
    d = size v

dim :: forall n . KnownNat n => R n
dim = v
  where
    v = mkR (scalar (fromIntegral $ size v))

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

ud2 :: L m n -> Matrix ℝ
ud2 (L (Dim (Dim x))) = x

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

diag :: KnownNat n => R n -> Sq n
diag = diagR 0

eye :: KnownNat n => Sq n
eye = diag 1

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

blockAt :: forall m n . (KnownNat m, KnownNat n) =>  ℝ -> Int -> Int -> Matrix Float -> L m n
blockAt x r c a = res
  where
    z = scalar x
    z1 = LA.konst x (r,c)
    z2 = LA.konst x (max 0 (m'-(ra+r)), max 0 (n'-(ca+c)))
    ra = min (rows a) . max 0 $ m'-r
    ca = min (cols a) . max 0 $ n'-c
    sa = subMatrix (0,0) (ra, ca) a
    (m',n') = size res
    res = mkL $ fromBlocks [[z1,z,z],[z,sa,z],[z,z,z2]]

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

row :: R n -> L 1 n
row = mkL . asRow . ud1

--col :: R n -> L n 1
col v = tr . row $ v

unrow :: L 1 n -> R n
unrow = mkR . head . LA.toRows . ud2

--uncol :: L n 1 -> R n
uncol v = unrow . tr $ v

infixl 2 ===
(===) :: (KnownNat r1, KnownNat r2, KnownNat c) => L r1 c -> L r2 c -> L (r1+r2) c
a === b = mkL (extract a LA.=== extract b)

infixl 3 |||
-- (|||) :: (KnownNat r, KnownNat c1, KnownNat c2) => L r c1 -> L r c2 -> L r (c1+c2)
a ||| b = tr (tr a === tr b)

type Sq n  = L n n
--type CSq n = CL n n

type GL = forall n m . (KnownNat n, KnownNat m) => L m n
type GSq = forall n . KnownNat n => Sq n

isKonst :: forall m n . (KnownNat m, KnownNat n) => L m n -> Maybe (ℝ,(Int,Int))
isKonst s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

isKonstC :: forall m n . (KnownNat m, KnownNat n) => M m n -> Maybe (ℂ,(Int,Int))
isKonstC s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

infixr 8 <>
(<>) :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n
(<>) = mulR

infixr 8 #>
(#>) :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
(#>) = appR

infixr 8 <·>
(<·>) :: KnownNat n => R n -> R n -> ℝ
(<·>) = dotR

infixr 8 <.>
(<.>) :: KnownNat n => R n -> R n -> ℝ
(<.>) = dotR

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

class Diag m d | m -> d
  where
    takeDiag :: m -> d

instance KnownNat n => Diag (L n n) (R n)
  where
    takeDiag x = mkR (LA.takeDiag (extract x))

instance KnownNat n => Diag (M n n) (C n)
  where
    takeDiag x = mkC (LA.takeDiag (extract x))

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

toComplex :: KnownNat n => (R n, R n) -> C n
toComplex (r,i) = mkC $ LA.toComplex (ud1 r, ud1 i)

fromComplex :: KnownNat n => C n -> (R n, R n)
fromComplex (C (Dim v)) = let (r,i) = LA.fromComplex v in (mkR r, mkR i)

complex :: KnownNat n => R n -> C n
complex r = mkC $ LA.toComplex (ud1 r, LA.konst 0 (size r))

real :: KnownNat n => C n -> R n
real = fst . fromComplex

imag :: KnownNat n => C n -> R n
imag = snd . fromComplex

sqMagnitude :: KnownNat n => C n -> R n
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2

magnitude :: KnownNat n => C n -> R n
magnitude = sqrt . sqMagnitude

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

linSolve :: (KnownNat m, KnownNat n) => L m m -> L m n -> Maybe (L m n)
linSolve (extract -> a) (extract -> b) = fmap mkL (LA.linearSolve a b)

(<\>) :: (KnownNat m, KnownNat n, KnownNat r) => L m n -> L m r -> L n r
(extract -> a) <\> (extract -> b) = mkL (a LA.<\> b)

svd :: (KnownNat m, KnownNat n) => L m n -> (L m m, R n, L n n)
svd (extract -> m) = (mkL u, mkR s', mkL v)
  where
    (u,s,v) = LA.svd m
    s' = vjoin [s, z]
    z = LA.konst 0 (max 0 (cols m - LA.size s))

svdTall :: (KnownNat m, KnownNat n, n <= m) => L m n -> (L m n, R n, L n n)
svdTall (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

svdFlat :: (KnownNat m, KnownNat n, m <= n) => L m n -> (L m m, R m, L n m)
svdFlat (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

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

class Eigen m l v | m -> l, m -> v
  where
    eigensystem :: m -> (l,v)
    eigenvalues :: m -> l

newtype Sym n = Sym (Sq n) deriving Show

sym :: KnownNat n => Sq n -> Sym n
sym m = Sym $ (m + tr m)/2

mTm :: (KnownNat m, KnownNat n) => L m n -> Sym n
mTm x = Sym (tr x <> x)

unSym :: Sym n -> Sq n
unSym (Sym x) = x

𝑖 :: Sized ℂ s c => s
𝑖 = konst iC

newtype Her n = Her (M n n)

her :: KnownNat n => M n n -> Her n
her m = Her $ (m + LA.tr m)/2

instance (KnownNat n) => Disp (Sym n)
  where
    disp n (Sym x) = do
        let a = extract x
        let su = LA.dispf n a
        printf "Sym %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance (KnownNat n) => Disp (Her n)
  where
    disp n (Her x) = do
        let a = extract x
        let su = LA.dispcf n a
        printf "Her %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance KnownNat n => Eigen (Sym n) (R n) (L n n)
  where
    eigenvalues (Sym (extract -> m)) =  mkR . LA.eigenvaluesSH . LA.trustSym $ m
    eigensystem (Sym (extract -> m)) = (mkR l, mkL v)
      where
        (l,v) = LA.eigSH . LA.trustSym $ m

instance KnownNat n => Eigen (Sq n) (C n) (M n n)
  where
    eigenvalues (extract -> m) = mkC . LA.eigenvalues $ m
    eigensystem (extract -> m) = (mkC l, mkM v)
      where
        (l,v) = LA.eig m

chol :: KnownNat n => Sym n -> Sq n
chol (extract . unSym -> m) = mkL $ LA.chol $ LA.trustSym m

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

withNullspace
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => L n k -> z)
    -> z
withNullspace (LA.nullspace . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withOrth
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k. (KnownNat k) => L n k -> z)
    -> z
withOrth (LA.orth . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withCompactSVD
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => (L m k, R k, L n k) -> z)
    -> z
withCompactSVD (LA.compactSVD . extract -> (u,s,v)) f =
    case someNatVal $ fromIntegral $ LA.size s of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL u :: L m k, mkR s :: R k, mkL v :: L n k)

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

qr :: (KnownNat m, KnownNat n) => L m n -> (L m m, L m n)
qr (extract -> x) = (mkL q, mkL r)
  where
    (q,r) = LA.qr x

-- use qrRaw?

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

split :: forall p n . (KnownNat p, KnownNat n, p<=n) => R n -> (R p, R (n-p))
split (extract -> v) = ( mkR (subVector 0 p' v) ,
                         mkR (subVector p' (LA.size v - p') v) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

headTail :: (KnownNat n, 1<=n) => R n -> (ℝ, R (n-1))
headTail = ((!0) . extract *** id) . split

splitRows :: forall p m n . (KnownNat p, KnownNat m, KnownNat n, p<=m) => L m n -> (L p n, L (m-p) n)
splitRows (extract -> x) = ( mkL (takeRows p' x) ,
                             mkL (dropRows p' x) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

splitCols :: forall p m n. (KnownNat p, KnownNat m, KnownNat n, KnownNat (n-p), p<=n) => L m n -> (L m p, L m (n-p))
splitCols = (tr *** tr) . splitRows . tr

toRows :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R n]
toRows (LA.toRows . extract -> vs) = map mkR vs

withRows
    :: forall n z . KnownNat n
    => [R n]
    -> (forall m . KnownNat m => L m n -> z)
    -> z
withRows (LA.fromRows . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.rows m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkL m :: L m n)

toColumns :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R m]
toColumns (LA.toColumns . extract -> vs) = map mkR vs

withColumns
    :: forall m z . KnownNat m
    => [R m]
    -> (forall n . KnownNat n => L m n -> z)
    -> z
withColumns (LA.fromColumns . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.cols m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy n)) -> f (mkL m :: L m n)

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

build
  :: forall m n . (KnownNat n, KnownNat m)
    => (ℝ -> ℝ -> ℝ)
    -> L m n
build f = r
  where
    r = mkL $ LA.build (size r) f

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

withVector
    :: forall z
     . Vector ℝ
    -> (forall n . (KnownNat n) => R n -> z)
    -> z
withVector v f =
    case someNatVal $ fromIntegral $ LA.size v of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkR v :: R m)

-- | Useful for constraining two dependently typed vectors to match each
-- other in length when they are unknown at compile-time.
exactLength
    :: forall n m . (KnownNat n, KnownNat m)
    => R m
    -> Maybe (R n)
exactLength v = do
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy m)
    return $ mkR (unwrap v)

withMatrix
    :: forall z
     . Matrix ℝ
    -> (forall m n . (KnownNat m, KnownNat n) => L m n -> z)
    -> z
withMatrix a f =
    case someNatVal $ fromIntegral $ rows a of
       Nothing -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) ->
           case someNatVal $ fromIntegral $ cols a of
               Nothing -> error "static/dynamic mismatch"
               Just (SomeNat (_ :: Proxy n)) ->
                  f (mkL a :: L m n)

-- | Useful for constraining two dependently typed matrices to match each
-- other in dimensions when they are unknown at compile-time.
exactDims
    :: forall n m j k . (KnownNat n, KnownNat m, KnownNat j, KnownNat k)
    => L m n
    -> Maybe (L j k)
exactDims m = do
    Refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy j)
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy k)
    return $ mkL (unwrap m)

randomVector
    :: forall n . KnownNat n
    => Seed
    -> RandDist
    -> R n
randomVector s d = mkR (LA.randomVector s d
                          (fromInteger (natVal (Proxy :: Proxy n)))
                       )

rand
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
rand = mkL <$> LA.rand (fromInteger (natVal (Proxy :: Proxy m)))
                       (fromInteger (natVal (Proxy :: Proxy n)))

randn
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
randn = mkL <$> LA.randn (fromInteger (natVal (Proxy :: Proxy m)))
                         (fromInteger (natVal (Proxy :: Proxy n)))

gaussianSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n
    -> Sym n
    -> L m n
gaussianSample s (extract -> mu) (Sym (extract -> sigma)) =
    mkL $ LA.gaussianSample s (fromInteger (natVal (Proxy :: Proxy m)))
                            mu (LA.trustSym sigma)

uniformSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n    -- ^ minimums of each row
    -> R n    -- ^ maximums of each row
    -> L m n
uniformSample s (extract -> mins) (extract -> maxs) =
    mkL $ LA.uniformSample s (fromInteger (natVal (Proxy :: Proxy m)))
                           (zip (LA.toList mins) (LA.toList maxs))

meanCov
    :: forall m n . (KnownNat m, KnownNat n, 1 <= m)
    => L m n
    -> (R n, Sym n)
meanCov (extract -> vs) = mkR *** (Sym . mkL . LA.unSym) $ LA.meanCov vs

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

class Domain field vec mat | mat -> vec field, vec -> mat field, field -> mat vec
  where
    mul :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => mat m k -> mat k n -> mat m n
    app :: forall m n . (KnownNat m, KnownNat n) => mat m n -> vec n -> vec m
    dot :: forall n . (KnownNat n) => vec n -> vec n -> field
    cross :: vec 3 -> vec 3 -> vec 3
    diagR ::  forall m n k . (KnownNat m, KnownNat n, KnownNat k) => field -> vec k -> mat m n
    dvmap :: forall n. KnownNat n => (field -> field) -> vec n -> vec n
    dmmap :: forall n m. (KnownNat m, KnownNat n) => (field -> field) -> mat n m -> mat n m
    outer :: forall n m. (KnownNat m, KnownNat n) => vec n -> vec m -> mat n m
    zipWithVector :: forall n. KnownNat n => (field -> field -> field) -> vec n -> vec n -> vec n
    det :: forall n. KnownNat n => mat n n -> field
    invlndet :: forall n. KnownNat n => mat n n -> (mat n n, (field, field))
    expm :: forall n. KnownNat n => mat n n -> mat n n
    sqrtm :: forall n. KnownNat n => mat n n -> mat n n
    inv :: forall n. KnownNat n => mat n n -> mat n n

instance Domain ℝ R L
  where
    mul = mulR
    app = appR
    dot = dotR
    cross = crossR
    diagR = diagRectR
    dvmap = mapR
    dmmap = mapL
    outer = outerR
    zipWithVector = zipWithR
    det = detL
    invlndet = invlndetL
    expm = expmL
    sqrtm = sqrtmL
    inv = invL

instance Domain ℂ C M
  where
    mul = mulC
    app = appC
    dot = dotC
    cross = crossC
    diagR = diagRectC
    dvmap = mapC
    dmmap = mapM'
    outer = outerC
    zipWithVector = zipWithC
    det = detM
    invlndet = invlndetM
    expm = expmM
    sqrtm = sqrtmM
    inv = invM

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

mulR :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n

mulR (isKonst -> Just (a,(_,k))) (isKonst -> Just (b,_)) = konst (a * b * fromIntegral k)

mulR (isDiag -> Just (0,a,_)) (isDiag -> Just (0,b,_)) = diagR 0 (mkR v :: R k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulR (isDiag -> Just (0,a,_)) (extract -> b) = mkL (asColumn a * takeRows (LA.size a) b)

mulR (extract -> a) (isDiag -> Just (0,b,_)) = mkL (takeColumns (LA.size b) a * asRow b)

mulR a b = mkL (extract a LA.<> extract b)

appR :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
appR (isDiag -> Just (0, w, _)) v = mkR (w * subVector 0 (LA.size w) (extract v))
appR m v                          = mkR (extract m LA.#> extract v)

dotR :: KnownNat n => R n -> R n -> ℝ
dotR (extract -> u) (extract -> v) = LA.dot u v

crossR :: R 3 -> R 3 -> R 3
crossR (extract -> x) (extract -> y) = vec3 z1 z2 z3
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerR :: (KnownNat m, KnownNat n) => R n -> R m -> L n m
outerR (extract -> x) (extract -> y) = mkL (LA.outer x y)

mapR :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapR f (unwrap -> v) = mkR (LA.cmap f v)

zipWithR :: KnownNat n => (ℝ -> ℝ -> ℝ) -> R n -> R n -> R n
zipWithR f (extract -> x) (extract -> y) = mkR (LA.zipVectorWith f x y)

mapL :: (KnownNat n, KnownNat m) => (ℝ -> ℝ) -> L n m -> L n m
mapL f = overMatL' (LA.cmap f)

detL :: KnownNat n => Sq n -> ℝ
detL = LA.det . unwrap

invlndetL :: KnownNat n => Sq n -> (L n n, (ℝ, ℝ))
invlndetL = BF.first mkL . LA.invlndet . unwrap

expmL :: KnownNat n => Sq n -> Sq n
expmL = overMatL' LA.expm

sqrtmL :: KnownNat n => Sq n -> Sq n
sqrtmL = overMatL' LA.sqrtm

invL :: KnownNat n => Sq n -> Sq n
invL = overMatL' LA.inv

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

mulC :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => M m k -> M k n -> M m n

mulC (isKonstC -> Just (a,(_,k))) (isKonstC -> Just (b,_)) = konst (a * b * fromIntegral k)

mulC (isDiagC -> Just (0,a,_)) (isDiagC -> Just (0,b,_)) = diagR 0 (mkC v :: C k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulC (isDiagC -> Just (0,a,_)) (extract -> b) = mkM (asColumn a * takeRows (LA.size a) b)

mulC (extract -> a) (isDiagC -> Just (0,b,_)) = mkM (takeColumns (LA.size b) a * asRow b)

mulC a b = mkM (extract a LA.<> extract b)

appC :: (KnownNat m, KnownNat n) => M m n -> C n -> C m
appC (isDiagC -> Just (0, w, _)) v = mkC (w * subVector 0 (LA.size w) (extract v))
appC m v                           = mkC (extract m LA.#> extract v)

dotC :: KnownNat n => C n -> C n -> ℂ
dotC (extract -> u) (extract -> v) = LA.dot u v

crossC :: C 3 -> C 3 -> C 3
crossC (extract -> x) (extract -> y) = mkC (LA.fromList [z1, z2, z3])
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerC :: (KnownNat m, KnownNat n) => C n -> C m -> M n m
outerC (extract -> x) (extract -> y) = mkM (LA.outer x y)

mapC :: KnownNat n => (ℂ -> ℂ) -> C n -> C n
mapC f (unwrap -> v) = mkC (LA.cmap f v)

zipWithC :: KnownNat n => (ℂ -> ℂ -> ℂ) -> C n -> C n -> C n
zipWithC f (extract -> x) (extract -> y) = mkC (LA.zipVectorWith f x y)

mapM' :: (KnownNat n, KnownNat m) => (ℂ -> ℂ) -> M n m -> M n m
mapM' f = overMatM' (LA.cmap f)

detM :: KnownNat n => M n n -> ℂ
detM = LA.det . unwrap

invlndetM :: KnownNat n => M n n -> (M n n, (ℂ, ℂ))
invlndetM = BF.first mkM . LA.invlndet . unwrap

expmM :: KnownNat n => M n n -> M n n
expmM = overMatM' LA.expm

sqrtmM :: KnownNat n => M n n -> M n n
sqrtmM = overMatM' LA.sqrtm

invM :: KnownNat n => M n n -> M n n
invM = overMatM' LA.inv

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

diagRectR :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℝ -> R k -> L m n
diagRectR x v
    | m' == 1 = mkL (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = matrix []
  where
    r = mkL (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

diagRectC :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℂ -> C k -> M m n
diagRectC x v
    | m' == 1 = mkM (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = fromList []
  where
    r = mkM (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

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

mean :: (KnownNat n, 1<=n) => R n -> ℝ
mean v = v <·> (1/dim)

test :: (Bool, IO ())
test = (ok,info)
  where
    ok =   extract (eye :: Sq 5) == ident 5
           && (unwrap .unSym) (mTm sm :: Sym 3) == tr ((3><3)[1..]) LA.<> (3><3)[1..]
           && unwrap (tm :: L 3 5) == LA.matrix 5 [1..15]
           && thingS == thingD
           && precS == precD
           && withVector (LA.vector [1..15]) sumV == sumElements (LA.fromList [1..15])

    info = do
        print $ u
        print $ v
        print (eye :: Sq 3)
        print $ ((u & 5) + 1) <·> v
        print (tm :: L 2 5)
        print (tm <> sm :: L 2 3)
        print thingS
        print thingD
        print precS
        print precD
        print $ withVector (LA.vector [1..15]) sumV
        splittest

    sumV w = w <·> konst 1

    u = vec2 3 5

    𝕧 x = vector [x] :: R 1

    v = 𝕧 2 & 4 & 7

    tm :: GL
    tm = lmat 0 [1..]

    lmat :: forall m n . (KnownNat m, KnownNat n) => ℝ -> [ℝ] -> L m n
    lmat z xs = r
      where
        r = mkL . reshape n' . LA.fromList . take (m'*n') $ xs ++ repeat z
        (m',n') = size r

    sm :: GSq
    sm = lmat 0 [1..]

    thingS = (u & 1) <·> tr q #> q #> v
      where
        q = tm :: L 10 3

    thingD = vjoin [ud1 u, 1] LA.<.> tr m LA.#> m LA.#> ud1 v
      where
        m = LA.matrix 3 [1..30]

    precS = (1::Float) + (2::Float) * ((1 :: R 3) * (u & 6)) <·> konst 2 #> v
    precD = 1 + 2 * vjoin[ud1 u, 6] LA.<.> LA.konst 2 (LA.size (ud1 u) +1, LA.size (ud1 v)) LA.#> ud1 v

splittest
    = do
    let v = range :: R 7
        a = snd (split v) :: R 4
    print $ a
    print $ snd . headTail . snd . headTail $ v
    print $ first (vec3 1 2 3)
    print $ second (vec3 1 2 3)
    print $ third (vec3 1 2 3)
    print $ (snd $ splitRows eye :: L 4 6)
 where
    first v = fst . headTail $ v
    second v = first . snd . headTail $ v
    third v = first . snd . headTail . snd . headTail $ v

instance (KnownNat n', KnownNat m') => Testable (L n' m')
  where
    checkT _ = test

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

instance KnownNat n => Normed (R n)
  where
    norm_0 v = norm_0 (extract v)
    norm_1 v = norm_1 (extract v)
    norm_2 v = norm_2 (extract v)
    norm_Inf v = norm_Inf (extract v)

instance (KnownNat m, KnownNat n) => Normed (L m n)
  where
    norm_0 m = norm_0 (extract m)
    norm_1 m = norm_1 (extract m)
    norm_2 m = norm_2 (extract m)
    norm_Inf m = norm_Inf (extract m)

mkSym f = Sym . f . unSym
mkSym2 f x y = Sym (f (unSym x) (unSym y))

instance KnownNat n =>  Num (Sym n)
  where
    (+) = mkSym2 (+)
    (*) = mkSym2 (*)
    (-) = mkSym2 (-)
    abs = mkSym abs
    signum = mkSym signum
    negate = mkSym negate
    fromInteger = Sym . fromInteger

instance KnownNat n => Fractional (Sym n)
  where
    fromRational = Sym . fromRational
    (/) = mkSym2 (/)

instance KnownNat n => Floating (Sym n)
  where
    sin   = mkSym sin
    cos   = mkSym cos
    tan   = mkSym tan
    asin  = mkSym asin
    acos  = mkSym acos
    atan  = mkSym atan
    sinh  = mkSym sinh
    cosh  = mkSym cosh
    tanh  = mkSym tanh
    asinh = mkSym asinh
    acosh = mkSym acosh
    atanh = mkSym atanh
    exp   = mkSym exp
    log   = mkSym log
    sqrt  = mkSym sqrt
    (**)  = mkSym2 (**)
    pi    = Sym pi

instance KnownNat n => Additive (Sym n) where
    add = (+)

instance KnownNat n => Transposable (Sym n) (Sym n) where
    tr  = id
    tr' = id

instance KnownNat n => Transposable (Her n) (Her n) where
    tr          = id
    tr' (Her m) = Her (tr' m)


{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE PatternSynonyms      #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wall #-}
-- | This module is wrongly named.
module VectorSpace (
    LinMap (.., LI),
    lmul,
    HasDim(Dim, dimDict),
    toRawMatrix,
    evalL,
    L (..),
    linear,
    VectorSpace (..),
    toVector,
    fromVector,
) where

import Data.Constraint       (Dict (..), withDict, (:-))
import Data.Proxy            (Proxy (..))
import GHC.TypeLits
import Overloaded.Categories

import qualified Control.Category      as C
import qualified Data.Constraint.Nat   as C
import qualified Numeric.LinearAlgebra as L

-- import qualified Numeric.LinearAlgebra.Static as LS

data LinMap a b where
    LZ :: LinMap a b
    LD :: Double -> LinMap a a
    LH :: LinMap a b -> LinMap a c -> LinMap a (b, c)
    LV :: LinMap a c -> LinMap b c -> LinMap (a, b) c
    LA :: LinMap a b -> LinMap a b -> LinMap a b

deriving instance Show (LinMap a b)

pattern LI :: forall a b. () => (b ~ a) => LinMap a b
pattern LI = LD 1

lmul :: Double -> LinMap a b -> LinMap a b
lmul _ LZ       = LZ
lmul k (LD x)   = LD (k * x)
lmul k (LH f g) = LH (lmul k f) (lmul k g)
lmul k (LV f g) = LV (lmul k f) (lmul k g)
lmul k (LA f g) = LA (lmul k f) (lmul k g)

lcomp :: LinMap b c -> LinMap a b -> LinMap a c
lcomp LZ       _        = LZ
lcomp _        LZ       = LZ
lcomp (LD k)   h        = lmul k h
lcomp h        (LD k)   = lmul k h
lcomp (LA f g) h        = LA (lcomp f h) (lcomp g h)
lcomp f        (LA g h) = LA (lcomp f g) (lcomp f h)
lcomp (LH f g) h        = LH (lcomp f h) (lcomp g h)
lcomp h        (LV f g) = LV (lcomp h f) (lcomp h g)
lcomp (LV f g) (LH u v) = LA (lcomp f u) (lcomp g v)

instance Category LinMap where
    id  = LI
    (.) = lcomp

instance CategoryWith1 LinMap where
    type Terminal LinMap = ()
    terminal = LZ

instance CartesianCategory LinMap where
    type Product LinMap = (,)
    proj1  = LV C.id LZ
    proj2  = LV LZ C.id
    fanout = LH

instance CategoryWith0 LinMap where
    type Initial LinMap = ()
    initial = LZ

instance CocartesianCategory LinMap where
    type Coproduct LinMap = (,)
    inl   = LH C.id LZ
    inr   = LH LZ C.id
    fanin = LV

instance BicartesianCategory LinMap where
    distr = LH
        (LH (LV (LV LI LZ) LZ) (LV LZ LI))
        (LH (LV (LV LZ LI) LZ) (LV LZ LI))

newtype L a b = L (forall r. LinMap r a -> LinMap r b)

lfst :: LinMap a (b, c) -> LinMap a b
lfst (LA f g) = LA (lfst f) (lfst g)
lfst (LH f _) = f
lfst (LV f g) = LV (lfst f) (lfst g)
lfst LZ       = LZ
lfst (LD k)   = LV (LD k) LZ

lsnd :: LinMap a (b, c) -> LinMap a c
lsnd (LH _ g) = g
lsnd (LA f g) = LA (lsnd f) (lsnd g)
lsnd (LV f g) = LV (lsnd f) (lsnd g)
lsnd LZ       = LZ
lsnd (LD k)   = LV LZ (LD k)

linitial :: LinMap r () -> LinMap r a
linitial _ = LZ

linear :: Double -> L a a
linear k = L $ lmul k

-- lmult :: Double -> Double -> LinMap r (a, a) -> LinMap r a
-- lmult x y (LH f g) = LA (LK y f) (LK x g)
-- lmult x y (LV f g) = LV (lmult x y f) (lmult x y g)
-- lmult x y (LA f g) = LA (lmult x y f) (lmult x y g)
-- lmult x y (LK k f) = LK k (lmult x y f)
-- lmult _ _ LZ       = LZ
-- lmult x y LI       = LV (LK y LI) (LK x LI)

instance Category L where
    id = L id
    L f . L g = L (f . g)

instance CategoryWith1 L where
    type Terminal L = ()

    terminal = L (\_ -> LZ)

instance CartesianCategory L where
    type Product L = (,)

    proj1 = L lfst
    proj2 = L lsnd

    fanout (L f) (L g) = L $ \x -> LH (f x) (g x)

instance CategoryWith0 L where
    type Initial L = ()

    initial = L linitial

-- Is this correct?
instance CocartesianCategory L where
    type Coproduct L = (,)

    inl = L $ \f -> LH f LZ
    inr = L $ \g -> LH LZ g

    fanin (L f) (L g) = L $ \x -> LA (f (lfst x)) (g (lsnd x))

class HasDim a where
    type Dim a :: Nat

    dimDict :: Proxy a -> Dict (KnownNat (Dim a))

    splitPair :: (a ~ (b, c)) => (Dict (HasDim b), Dict (HasDim c))
    splitPair = error "impossible: splitPair"

instance HasDim () where
    type Dim () = 0
    dimDict _ = Dict

instance HasDim Double where
    type Dim Double = 1
    dimDict _ = Dict

instance (HasDim a, HasDim b) => HasDim (a, b) where
    type Dim (a, b) = Dim a + Dim b

    dimDict _ =
        withDimDict (Proxy :: Proxy a) $
        withDimDict (Proxy :: Proxy b) $
        withDict (C.plusNat :: (KnownNat (Dim a), KnownNat (Dim b)) :- KnownNat (Dim a + Dim b))
        Dict

    splitPair = (Dict, Dict)

withDimDict :: HasDim a => Proxy a -> (KnownNat (Dim a) => r) -> r
withDimDict p = withDict (dimDict p)

dim :: forall a. HasDim a => Proxy a -> Int
dim p = withDimDict p $ fromInteger $ natVal (Proxy :: Proxy (Dim a))

toRawMatrix :: forall a b. (HasDim a, HasDim b) => LinMap a b -> L.Matrix Double
toRawMatrix LZ       = (dim (Proxy :: Proxy a) L.>< dim (Proxy :: Proxy b)) (repeat 0)
toRawMatrix (LD k)   = L.scale k (L.ident (dim (Proxy :: Proxy a)))
toRawMatrix (LA f g) = L.add (toRawMatrix f) (toRawMatrix g)
toRawMatrix (LH f g) = go splitPair f g where
    go :: (Dict (HasDim x), Dict (HasDim y)) -> LinMap a x -> LinMap a y -> L.Matrix Double
    go (Dict, Dict) f' g' = toRawMatrix f' L.||| toRawMatrix g'
toRawMatrix (LV f g) = go splitPair f g where
    go :: (Dict (HasDim x), Dict (HasDim y)) -> LinMap x b -> LinMap y b -> L.Matrix Double
    go (Dict, Dict) f' g' = toRawMatrix f' L.=== toRawMatrix g'

evalL :: (HasDim a, HasDim b) => L a b -> L.Matrix Double
evalL (L f) = toRawMatrix (f (LD 1))

-- toStaticMatrix :: forall a b. (HasDim a, HasDim b) => LinMap a b -> LS.L (Dim a) (Dim b)
-- toStaticMatrix LZ =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) 0
-- toStaticMatrix LI =
--     withDimDict (Proxy :: Proxy a) LS.eye
-- toStaticMatrix (LA f g) =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) $
--     L.add (toStaticMatrix f) (toStaticMatrix g)
-- toStaticMatrix (LK k f) =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) $
--     toStaticMatrix f LS.<> LS.diag (LS.konst k)
-- toStaticMatrix (LH f g) = go splitPair f g where
--     go :: forall x y. (x,y) ~ b => (Dict (HasDim x), Dict (HasDim y)) -> LinMap a x -> LinMap a y -> LS.L (Dim a) (Dim x + Dim y)
--     go (Dict, Dict) f' g' =
--         withDimDict (Proxy :: Proxy a) $
--         withDimDict (Proxy :: Proxy b) $
--         withDimDict (Proxy :: Proxy x) $
--         withDimDict (Proxy :: Proxy y) $
--         toStaticMatrix f' LS.||| toStaticMatrix g'
-- toStaticMatrix (LV f g) = go splitPair f g where
--     go :: forall x y. (x,y) ~ a => (Dict (HasDim x), Dict (HasDim y)) -> LinMap x b -> LinMap y b -> LS.L (Dim x + Dim y) (Dim b)
--     go (Dict, Dict) f' g' =
--         withDimDict (Proxy :: Proxy a) $
--         withDimDict (Proxy :: Proxy b) $
--         withDimDict (Proxy :: Proxy x) $
--         withDimDict (Proxy :: Proxy y) $
--         toStaticMatrix f' LS.=== toStaticMatrix g'

-------------------------------------------------------------------------------
-- Vector space
-------------------------------------------------------------------------------

class HasDim a => VectorSpace a where
    toVector' :: a -> [Double] -> [Double]

    fromVector' :: [Double] -> (a -> [Double] -> r) -> r

toVector :: VectorSpace a => a -> [Double]
toVector x = toVector' x []

fromVector :: VectorSpace a => [Double] -> a
fromVector ds = fromVector' ds const

instance VectorSpace Double where
    toVector' d = (d :)

    fromVector' []     k = k 0 []
    fromVector' (d:ds) k = k d ds

instance (VectorSpace a, VectorSpace b) => VectorSpace (a, b) where
    toVector' (a, b) = toVector' a . toVector' b

    fromVector' xs k =
        fromVector' xs $ \a ys ->
        fromVector' ys $ \b zs ->
        k (a, b) zs


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.ContainerBoot
-- Copyright   :  (c) Alberto Ruiz 2010
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable
--
-- Module to avoid cyclyc dependencies.
--
-----------------------------------------------------------------------------

module Numeric.ContainerBoot (
    -- * Basic functions
    ident, diag, ctrans,
    -- * Generic operations
    Container(..),
    -- * Matrix product and related functions
    Product(..),
    mXm,mXv,vXm,
    outer, kronecker,
    -- * Element conversion
    Convert(..),
    Complexable(),
    RealElement(),

    RealOf, ComplexOf, SingleOf, DoubleOf,

    IndexOf,
    module Data.Complex,
    -- * Experimental
    build', konst'
) where

import Data.Packed
import Data.Packed.ST as ST
import Numeric.Conversion
import Data.Packed.Internal
import Numeric.GSL.Vector
import Data.Complex
import Control.Monad(ap)

import Numeric.LinearAlgebra.LAPACK(multiplyR,multiplyC,multiplyF,multiplyQ)

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

type family IndexOf (c :: * -> *)

type instance IndexOf Vector = Int
type instance IndexOf Matrix = (Int,Int)

type family ArgOf (c :: * -> *) a

type instance ArgOf Vector a = a -> a
type instance ArgOf Matrix a = a -> a -> a

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

-- | Basic element-by-element functions for numeric containers
class (Complexable c, Fractional e, Element e) => Container c e where
    -- | create a structure with a single element
    scalar      :: e -> c e
    -- | complex conjugate
    conj        :: c e -> c e
    scale       :: e -> c e -> c e
    -- | scale the element by element reciprocal of the object:
    --
    -- @scaleRecip 2 (fromList [5,i]) == 2 |> [0.4 :+ 0.0,0.0 :+ (-2.0)]@
    scaleRecip  :: e -> c e -> c e
    addConstant :: e -> c e -> c e
    add         :: c e -> c e -> c e
    sub         :: c e -> c e -> c e
    -- | element by element multiplication
    mul         :: c e -> c e -> c e
    -- | element by element division
    divide      :: c e -> c e -> c e
    equal       :: c e -> c e -> Bool
    --
    -- element by element inverse tangent
    arctan2     :: c e -> c e -> c e
    --
    -- | cannot implement instance Functor because of Element class constraint
    cmap        :: (Element b) => (e -> b) -> c e -> c b
    -- | constant structure of given size
    konst       :: e -> IndexOf c -> c e
    -- | create a structure using a function
    --
    -- Hilbert matrix of order N:
    --
    -- @hilb n = build (n,n) (\\i j -> 1/(i+j+1))@
    build       :: IndexOf c -> (ArgOf c e) -> c e
    --build       :: BoundsOf f -> f -> (ContainerOf f) e
    --
    -- | indexing function
    atIndex     :: c e -> IndexOf c -> e
    -- | index of min element
    minIndex    :: c e -> IndexOf c
    -- | index of max element
    maxIndex    :: c e -> IndexOf c
    -- | value of min element
    minElement  :: c e -> e
    -- | value of max element
    maxElement  :: c e -> e
    -- the C functions sumX/prodX are twice as fast as using foldVector
    -- | the sum of elements (faster than using @fold@)
    sumElements :: c e -> e
    -- | the product of elements (faster than using @fold@)
    prodElements :: c e -> e

    -- | A more efficient implementation of @cmap (\\x -> if x>0 then 1 else 0)@
    --
    -- @> step $ linspace 5 (-1,1::Double)
    -- 5 |> [0.0,0.0,0.0,1.0,1.0]@

    step :: RealElement e => c e -> c e

    -- | Element by element version of @case compare a b of {LT -> l; EQ -> e; GT -> g}@.
    --
    -- Arguments with any dimension = 1 are automatically expanded:
    --
    -- @> cond ((1>\<4)[1..]) ((3>\<1)[1..]) 0 100 ((3>\<4)[1..]) :: Matrix Double
    -- (3><4)
    -- [ 100.0,   2.0,   3.0,  4.0
    -- ,   0.0, 100.0,   7.0,  8.0
    -- ,   0.0,   0.0, 100.0, 12.0 ]@

    cond :: RealElement e
         => c e -- ^ a
         -> c e -- ^ b
         -> c e -- ^ l
         -> c e -- ^ e
         -> c e -- ^ g
         -> c e -- ^ result

    -- | Find index of elements which satisfy a predicate
    --
    -- @> find (>0) (ident 3 :: Matrix Double)
    -- [(0,0),(1,1),(2,2)]@

    find :: (e -> Bool) -> c e -> [IndexOf c]

    -- | Create a structure from an association list
    --
    -- @> assoc 5 0 [(2,7),(1,3)] :: Vector Double
    -- 5 |> [0.0,3.0,7.0,0.0,0.0]@

    assoc :: IndexOf c        -- ^ size
          -> e                -- ^ default value
          -> [(IndexOf c, e)] -- ^ association list
          -> c e              -- ^ result

    -- | Modify a structure using an update function
    --
    -- @> accum (ident 5) (+) [((1,1),5),((0,3),3)] :: Matrix Double
    -- (5><5)
    --  [ 1.0, 0.0, 0.0, 3.0, 0.0
    --  , 0.0, 6.0, 0.0, 0.0, 0.0
    --  , 0.0, 0.0, 1.0, 0.0, 0.0
    --  , 0.0, 0.0, 0.0, 1.0, 0.0
    --  , 0.0, 0.0, 0.0, 0.0, 1.0 ]@

    accum :: c e              -- ^ initial structure
          -> (e -> e -> e)    -- ^ update function
          -> [(IndexOf c, e)] -- ^ association list
          -> c e              -- ^ result

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

instance Container Vector Float where
    scale = vectorMapValF Scale
    scaleRecip = vectorMapValF Recip
    addConstant = vectorMapValF AddConstant
    add = vectorZipF Add
    sub = vectorZipF Sub
    mul = vectorZipF Mul
    divide = vectorZipF Div
    equal u v = dim u == dim v && maxElement (vectorMapF Abs (sub u v)) == 0.0
    arctan2 = vectorZipF ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = id
    cmap = mapVector
    atIndex = (@>)
    minIndex     = round . toScalarF MinIdx
    maxIndex     = round . toScalarF MaxIdx
    minElement  = toScalarF Min
    maxElement  = toScalarF Max
    sumElements  = sumF
    prodElements = prodF
    step = stepF
    find = findV
    assoc = assocV
    accum = accumV
    cond = condV condF

instance Container Vector Double where
    scale = vectorMapValR Scale
    scaleRecip = vectorMapValR Recip
    addConstant = vectorMapValR AddConstant
    add = vectorZipR Add
    sub = vectorZipR Sub
    mul = vectorZipR Mul
    divide = vectorZipR Div
    equal u v = dim u == dim v && maxElement (vectorMapR Abs (sub u v)) == 0.0
    arctan2 = vectorZipR ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = id
    cmap = mapVector
    atIndex = (@>)
    minIndex     = round . toScalarR MinIdx
    maxIndex     = round . toScalarR MaxIdx
    minElement  = toScalarR Min
    maxElement  = toScalarR Max
    sumElements  = sumR
    prodElements = prodR
    step = stepD
    find = findV
    assoc = assocV
    accum = accumV
    cond = condV condD

instance Container Vector (Complex Double) where
    scale = vectorMapValC Scale
    scaleRecip = vectorMapValC Recip
    addConstant = vectorMapValC AddConstant
    add = vectorZipC Add
    sub = vectorZipC Sub
    mul = vectorZipC Mul
    divide = vectorZipC Div
    equal u v = dim u == dim v && maxElement (mapVector magnitude (sub u v)) == 0.0
    arctan2 = vectorZipC ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = conjugateC
    cmap = mapVector
    atIndex = (@>)
    minIndex     = minIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    maxIndex     = maxIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    minElement  = ap (@>) minIndex
    maxElement  = ap (@>) maxIndex
    sumElements  = sumC
    prodElements = prodC
    step = undefined -- cannot match
    find = findV
    assoc = assocV
    accum = accumV
    cond = undefined -- cannot match

instance Container Vector (Complex Float) where
    scale = vectorMapValQ Scale
    scaleRecip = vectorMapValQ Recip
    addConstant = vectorMapValQ AddConstant
    add = vectorZipQ Add
    sub = vectorZipQ Sub
    mul = vectorZipQ Mul
    divide = vectorZipQ Div
    equal u v = dim u == dim v && maxElement (mapVector magnitude (sub u v)) == 0.0
    arctan2 = vectorZipQ ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = conjugateQ
    cmap = mapVector
    atIndex = (@>)
    minIndex     = minIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    maxIndex     = maxIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    minElement  = ap (@>) minIndex
    maxElement  = ap (@>) maxIndex
    sumElements  = sumQ
    prodElements = prodQ
    step = undefined -- cannot match
    find = findV
    assoc = assocV
    accum = accumV
    cond = undefined -- cannot match

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

instance (Container Vector a) => Container Matrix a where
    scale x = liftMatrix (scale x)
    scaleRecip x = liftMatrix (scaleRecip x)
    addConstant x = liftMatrix (addConstant x)
    add = liftMatrix2 add
    sub = liftMatrix2 sub
    mul = liftMatrix2 mul
    divide = liftMatrix2 divide
    equal a b = cols a == cols b && flatten a `equal` flatten b
    arctan2 = liftMatrix2 arctan2
    scalar x = (1><1) [x]
    konst v (r,c) = reshape c (konst v (r*c))
    build = buildM
    conj = liftMatrix conj
    cmap f = liftMatrix (mapVector f)
    atIndex = (@@>)
    minIndex m = let (r,c) = (rows m,cols m)
                     i = (minIndex $ flatten m)
                 in (i `div` c,i `mod` c)
    maxIndex m = let (r,c) = (rows m,cols m)
                     i = (maxIndex $ flatten m)
                 in (i `div` c,i `mod` c)
    minElement = ap (@@>) minIndex
    maxElement = ap (@@>) maxIndex
    sumElements = sumElements . flatten
    prodElements = prodElements . flatten
    step = liftMatrix step
    find = findM
    assoc = assocM
    accum = accumM
    cond = condM

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

-- | Matrix product and related functions
class Element e => Product e where
    -- | matrix product
    multiply :: Matrix e -> Matrix e -> Matrix e
    -- | dot (inner) product
    dot        :: Vector e -> Vector e -> e
    -- | sum of absolute value of elements (differs in complex case from @norm1@)
    absSum     :: Vector e -> RealOf e
    -- | sum of absolute value of elements
    norm1      :: Vector e -> RealOf e
    -- | euclidean norm
    norm2      :: Vector e -> RealOf e
    -- | element of maximum magnitude
    normInf    :: Vector e -> RealOf e

instance Product Float where
    norm2      = toScalarF Norm2
    absSum     = toScalarF AbsSum
    dot        = dotF
    norm1      = toScalarF AbsSum
    normInf    = maxElement . vectorMapF Abs
    multiply = multiplyF

instance Product Double where
    norm2      = toScalarR Norm2
    absSum     = toScalarR AbsSum
    dot        = dotR
    norm1      = toScalarR AbsSum
    normInf    = maxElement . vectorMapR Abs
    multiply = multiplyR

instance Product (Complex Float) where
    norm2      = toScalarQ Norm2
    absSum     = toScalarQ AbsSum
    dot        = dotQ
    norm1      = sumElements . fst . fromComplex . vectorMapQ Abs
    normInf    = maxElement . fst . fromComplex . vectorMapQ Abs
    multiply = multiplyQ

instance Product (Complex Double) where
    norm2      = toScalarC Norm2
    absSum     = toScalarC AbsSum
    dot        = dotC
    norm1      = sumElements . fst . fromComplex . vectorMapC Abs
    normInf    = maxElement . fst . fromComplex . vectorMapC Abs
    multiply = multiplyC

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

-- synonym for matrix product
mXm :: Product t => Matrix t -> Matrix t -> Matrix t
mXm = multiply

-- matrix - vector product
mXv :: Product t => Matrix t -> Vector t -> Vector t
mXv m v = flatten $ m `mXm` (asColumn v)

-- vector - matrix product
vXm :: Product t => Vector t -> Matrix t -> Vector t
vXm v m = flatten $ (asRow v) `mXm` m

{- | Outer product of two vectors.

@\> 'fromList' [1,2,3] \`outer\` 'fromList' [5,2,3]
(3><3)
 [  5.0, 2.0, 3.0
 , 10.0, 4.0, 6.0
 , 15.0, 6.0, 9.0 ]@
-}
outer :: (Product t) => Vector t -> Vector t -> Matrix t
outer u v = asColumn u `multiply` asRow v

{- | Kronecker product of two matrices.

@m1=(2><3)
 [ 1.0,  2.0, 0.0
 , 0.0, -1.0, 3.0 ]
m2=(4><3)
 [  1.0,  2.0,  3.0
 ,  4.0,  5.0,  6.0
 ,  7.0,  8.0,  9.0
 , 10.0, 11.0, 12.0 ]@

@\> kronecker m1 m2
(8><9)
 [  1.0,  2.0,  3.0,   2.0,   4.0,   6.0,  0.0,  0.0,  0.0
 ,  4.0,  5.0,  6.0,   8.0,  10.0,  12.0,  0.0,  0.0,  0.0
 ,  7.0,  8.0,  9.0,  14.0,  16.0,  18.0,  0.0,  0.0,  0.0
 , 10.0, 11.0, 12.0,  20.0,  22.0,  24.0,  0.0,  0.0,  0.0
 ,  0.0,  0.0,  0.0,  -1.0,  -2.0,  -3.0,  3.0,  6.0,  9.0
 ,  0.0,  0.0,  0.0,  -4.0,  -5.0,  -6.0, 12.0, 15.0, 18.0
 ,  0.0,  0.0,  0.0,  -7.0,  -8.0,  -9.0, 21.0, 24.0, 27.0
 ,  0.0,  0.0,  0.0, -10.0, -11.0, -12.0, 30.0, 33.0, 36.0 ]@
-}
kronecker :: (Product t) => Matrix t -> Matrix t -> Matrix t
kronecker a b = fromBlocks
              . splitEvery (cols a)
              . map (reshape (cols b))
              . toRows
              $ flatten a `outer` flatten b

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

class Convert t where
    real    :: Container c t => c (RealOf t) -> c t
    complex :: Container c t => c t -> c (ComplexOf t)
    single  :: Container c t => c t -> c (SingleOf t)
    double  :: Container c t => c t -> c (DoubleOf t)
    toComplex   :: (Container c t, RealElement t) => (c t, c t) -> c (Complex t)
    fromComplex :: (Container c t, RealElement t) => c (Complex t) -> (c t, c t)

instance Convert Double where
    real = id
    complex = comp'
    single = single'
    double = id
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert Float where
    real = id
    complex = comp'
    single = id
    double = double'
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert (Complex Double) where
    real = comp'
    complex = id
    single = single'
    double = id
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert (Complex Float) where
    real = comp'
    complex = id
    single = id
    double = double'
    toComplex = toComplex'
    fromComplex = fromComplex'

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

type family RealOf x

type instance RealOf Double = Double
type instance RealOf (Complex Double) = Double

type instance RealOf Float = Float
type instance RealOf (Complex Float) = Float

type family ComplexOf x

type instance ComplexOf Double = Complex Double
type instance ComplexOf (Complex Double) = Complex Double

type instance ComplexOf Float = Complex Float
type instance ComplexOf (Complex Float) = Complex Float

type family SingleOf x

type instance SingleOf Double = Float
type instance SingleOf Float  = Float

type instance SingleOf (Complex a) = Complex (SingleOf a)

type family DoubleOf x

type instance DoubleOf Double = Double
type instance DoubleOf Float  = Double

type instance DoubleOf (Complex a) = Complex (DoubleOf a)

type family ElementOf c

type instance ElementOf (Vector a) = a
type instance ElementOf (Matrix a) = a

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

class Build f where
    build' :: BoundsOf f -> f -> ContainerOf f

type family BoundsOf x

type instance BoundsOf (a->a) = Int
type instance BoundsOf (a->a->a) = (Int,Int)

type family ContainerOf x

type instance ContainerOf (a->a) = Vector a
type instance ContainerOf (a->a->a) = Matrix a

instance (Element a, Num a) => Build (a->a) where
    build' = buildV

instance (Element a, Num a) => Build (a->a->a) where
    build' = buildM

buildM (rc,cc) f = fromLists [ [f r c | c <- cs] | r <- rs ]
    where rs = map fromIntegral [0 .. (rc-1)]
          cs = map fromIntegral [0 .. (cc-1)]

buildV n f = fromList [f k | k <- ks]
    where ks = map fromIntegral [0 .. (n-1)]

----------------------------------------------------
-- experimental

class Konst s where
    konst' :: Element e => e -> s -> ContainerOf' s e

type family ContainerOf' x y

type instance ContainerOf' Int a = Vector a
type instance ContainerOf' (Int,Int) a = Matrix a

instance Konst Int where
    konst' = constantD

instance Konst (Int,Int) where
    konst' k (r,c) = reshape c $ konst' k (r*c)

--------------------------------------------------------
-- | conjugate transpose
ctrans :: (Container Vector e, Element e) => Matrix e -> Matrix e
ctrans = liftMatrix conj . trans

-- | Creates a square matrix with a given diagonal.
diag :: (Num a, Element a) => Vector a -> Matrix a
diag v = diagRect 0 v n n where n = dim v

-- | creates the identity matrix of given dimension
ident :: (Num a, Element a) => Int -> Matrix a
ident n = diag (constantD 1 n)

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

findV p x = foldVectorWithIndex g [] x where
    g k z l = if p z then k:l else l

findM p x = map ((`divMod` cols x)) $ findV p (flatten x)

assocV n z xs = ST.runSTVector $ do
        v <- ST.newVector z n
        mapM_ (\(k,x) -> ST.writeVector v k x) xs
        return v

assocM (r,c) z xs = ST.runSTMatrix $ do
        m <- ST.newMatrix z r c
        mapM_ (\((i,j),x) -> ST.writeMatrix m i j x) xs
        return m

accumV v0 f xs = ST.runSTVector $ do
        v <- ST.thawVector v0
        mapM_ (\(k,x) -> ST.modifyVector v k (f x)) xs
        return v

accumM m0 f xs = ST.runSTMatrix $ do
        m <- ST.thawMatrix m0
        mapM_ (\((i,j),x) -> ST.modifyMatrix m i j (f x)) xs
        return m

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

condM a b l e t = reshape (cols a'') $ cond a' b' l' e' t'
  where
    args@(a'':_) = conformMs [a,b,l,e,t]
    [a', b', l', e', t'] = map flatten args

condV f a b l e t = f a' b' l' e' t'
  where
    [a', b', l', e', t'] = conformVs [a,b,l,e,t]


{-# LANGUAGE AllowAmbiguousTypes, DataKinds, RankNTypes, TypeFamilies,
             TypeOperators #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
{-# OPTIONS_GHC -fconstraint-solver-iterations=16 #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Shaped tensor-based implementation of Convolutional Neural Network
-- for classification of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistCnnShaped where

import Prelude

import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shape
import qualified Data.Array.ShapedS as OS
import           Data.Proxy (Proxy)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, type (+), type (<=), type Div)
import qualified Numeric.LinearAlgebra as HM

-- until stylish-haskell accepts NoStarIsType
import qualified GHC.TypeLits

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, varS)
import HordeAd.Tool.MnistData

convMnistLayerS
  :: forall kheight_minus_1 kwidth_minus_1 out_channels
            in_height in_width in_channels batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width
     , KnownNat in_channels, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1  -- wrongly reported as redundant
     , DualMonad d r m )
  => DualNumber d (OS.Array '[ out_channels, in_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[batch_size, in_channels, in_height, in_width] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> m (DualNumber d (OS.Array '[ batch_size, out_channels
                                , (in_height + kheight_minus_1) `Div` 2
                                , (in_width + kwidth_minus_1) `Div` 2 ] r))
convMnistLayerS ker x bias = do
  let yConv = conv24 ker x
      replicateBias
        :: DualNumber d (OS.Array '[] r)
        -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                   , in_width + kwidth_minus_1 ] r)
      replicateBias = konstS . fromS0
      biasStretched = ravelFromListS
                      $ replicate (valueOf @batch_size)
                      $ mapS replicateBias bias
        -- TODO: this is weakly typed; add and use replicateS instead
        -- or broadcastS or stretchS, possibly with transposeS?
  yRelu <- reluAct $ yConv + biasStretched
  maxPool24 @1 @2 yRelu

convMnistTwoS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width in_channels batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , in_channels ~ 1
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => OS.Array '[batch_size, in_channels, in_height, in_width] r
  -- All below is the type of all paramters of this nn. The same is reflected
  -- in the length function below and read from variables further down.
  -> DualNumber d (OS.Array '[ out_channels, in_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> DualNumber d (OS.Array '[ out_channels, out_channels
                             , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
  -> DualNumber d (OS.Array '[out_channels] r)
  -> DualNumber d (OS.Array '[ num_hidden
                             , out_channels
                                 GHC.TypeLits.*
                                   (((in_height + kheight_minus_1) `Div` 2
                                     + kheight_minus_1) `Div` 2)
                                 GHC.TypeLits.*
                                   (((in_width + kwidth_minus_1) `Div` 2
                                     + kwidth_minus_1) `Div` 2)
                             ] r)
  -> DualNumber d (OS.Array '[num_hidden] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel, num_hidden] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel] r)
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
convMnistTwoS x ker1 bias1 ker2 bias2
              weigthsDense biasesDense weigthsReadout biasesReadout = do
  t1 <- convMnistLayerS ker1 (constant x) bias1
  t2 <- convMnistLayerS ker2 t1 bias2
  let m1 = mapS reshapeS t2
      m2 = transpose2S m1
      denseLayer = weigthsDense <>$ m2 + asColumnS biasesDense
  denseRelu <- reluAct denseLayer
  returnLet $ weigthsReadout <>$ denseRelu + asColumnS biasesReadout

convMnistLenS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width )
  => Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> Proxy in_height
  -> Proxy in_width
  -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
convMnistLenS _ _ _ _ _ _ =
  ( 0
  , []
  , []
  , [ Data.Array.Shape.shapeT @'[ out_channels, 1
                                , kheight_minus_1 + 1, kwidth_minus_1 + 1 ]
    , Data.Array.Shape.shapeT @'[out_channels]
    , Data.Array.Shape.shapeT @'[ out_channels, out_channels
                                , kheight_minus_1 + 1, kwidth_minus_1 + 1 ]
    , Data.Array.Shape.shapeT @'[out_channels]
    , Data.Array.Shape.shapeT @'[ num_hidden
                                , out_channels
                                    GHC.TypeLits.*
                                      ((in_height + kheight_minus_1) `Div` 2
                                       + kheight_minus_1) `Div` 2
                                    GHC.TypeLits.*
                                      ((in_width + kwidth_minus_1) `Div` 2
                                       + kheight_minus_1) `Div` 2
                                ]
    , Data.Array.Shape.shapeT @'[num_hidden]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel, num_hidden]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel]
    ]
  )

convMnistS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => OS.Array '[batch_size, 1, in_height, in_width] r
  -> DualNumberVariables d r
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
convMnistS x variables = do
  let ker1 = varS variables 0
      bias1 = varS variables 1
      ker2 = varS variables 2
      bias2 = varS variables 3
      weigthsDense = varS variables 4
      biasesDense = varS variables 5
      weigthsReadout = varS variables 6
      biasesReadout = varS variables 7
  convMnistTwoS @kheight_minus_1 @kwidth_minus_1 @num_hidden @out_channels
                x ker1 bias1 ker2 bias2
                weigthsDense biasesDense weigthsReadout biasesReadout

convMnistLossFusedS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width batch_size d r m.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width, KnownNat batch_size
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , DualMonad d r m )
  => Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> ( OS.Array '[batch_size, in_height, in_width] r
     , OS.Array '[batch_size, SizeMnistLabel] r )
  -> DualNumberVariables d r
  -> m (DualNumber d r)
convMnistLossFusedS _ _ _ _ (glyphS, labelS) variables = do
  let xs :: OS.Array '[batch_size, 1, in_height, in_width] r
      xs = OS.reshape glyphS
  result <- convMnistS @kheight_minus_1 @kwidth_minus_1
                       @num_hidden @out_channels
                       xs variables
  let targets2 = HM.tr $ HM.reshape (valueOf @SizeMnistLabel)
                       $ OS.toVector labelS
  vec <- lossSoftMaxCrossEntropyL targets2 (fromS2 result)
  returnLet $ scale (recip $ fromIntegral (valueOf @batch_size :: Int))
            $ sumElements0 vec

-- For simplicity, testing is performed in mini-batches of 1.
-- See RNN for testing done in batches.
convMnistTestS
  :: forall kheight_minus_1 kwidth_minus_1 num_hidden out_channels
            in_height in_width r.
     ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
     , KnownNat num_hidden, KnownNat out_channels
     , KnownNat in_height, KnownNat in_width
     , 1 <= kheight_minus_1
     , 1 <= kwidth_minus_1
     , IsScalar 'DModeGradient r )
  => Proxy r
  -> Proxy kheight_minus_1
  -> Proxy kwidth_minus_1
  -> Proxy num_hidden
  -> Proxy out_channels
  -> [( OS.Array '[in_height, in_width] r
      , OS.Array '[SizeMnistLabel] r )]
  -> Domains r
  -> r
convMnistTestS _ _ _ _ _ inputs parameters =
  let matchesLabels :: ( OS.Array '[in_height, in_width] r
                       , OS.Array '[SizeMnistLabel] r )
                    -> Bool
      matchesLabels (glyph, label) =
        let tx :: OS.Array '[1, 1, in_height, in_width] r
            tx = OS.reshape glyph
            nn :: DualNumberVariables 'DModeGradient r
               -> DualMonadValue r (DualNumber 'DModeGradient (OS.Array '[SizeMnistLabel, 1] r))
            nn = convMnistS @kheight_minus_1 @kwidth_minus_1
                            @num_hidden @out_channels
                            tx
            value = primalValue nn parameters
        in V.maxIndex (OS.toVector value) == V.maxIndex (OS.toVector label)
  in fromIntegral (length (filter matchesLabels inputs))
     / fromIntegral (length inputs)


import Foreign.Marshal.Array
import Data.Complex
import Foreign.Storable.Complex
import FFTW

main = do
    inA  <- fftwAllocComplex 1024
    outA <- fftwAllocComplex 1024

    plan <- planDFT1d 1024 inA outA Forward fftwEstimate

    pokeArray inA $ map (:+ 0) [0..1023]
    execute plan
    res <- peekArray 1024 outA

    fftwFree inA
    fftwFree outA

    print res