```{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE UndecidableInstances  #-}
module Numeric.Matrix.Internal
( MatrixTranspose (..)
, SquareMatrix (..)
, MatrixDeterminant (..)
, MatrixInverse (..)
, Matrix
, HomTransform4 (..)
, Mat22f, Mat23f, Mat24f
, Mat32f, Mat33f, Mat34f
, Mat42f, Mat43f, Mat44f
, Mat22d, Mat23d, Mat24d
, Mat32d, Mat33d, Mat34d
, Mat42d, Mat43d, Mat44d
, mat22, mat33, mat44
, (%*)
) where

import GHC.Base
import Numeric.DataFrame.Contraction        ((%*))
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Scalar.Internal
import Numeric.Vector.Internal

-- | Alias for DataFrames of rank 2
type Matrix (t :: l) (n :: k) (m :: k) = DataFrame t '[n,m]

class MatrixTranspose t (n :: k) (m :: k) where
-- | Transpose Mat
transpose :: Matrix t n m -> Matrix t m n

class SquareMatrix t (n :: Nat) where
-- | Mat with 1 on diagonal and 0 elsewhere
eye :: Matrix t n n
-- | Put the same value on the Mat diagonal, 0 otherwise
diag :: Scalar t -> Matrix t n n
-- | Sum of diagonal elements
trace :: Matrix t n n -> Scalar t

class MatrixDeterminant t (n :: Nat) where
-- | Determinant of  Mat
det :: Matrix t n n -> Scalar t

class MatrixInverse t (n :: Nat) where
-- | Matrix inverse
inverse :: Matrix t n n -> Matrix t n n

-- | Operations on 4x4 transformation matrices and vectors in homogeneous coordinates.
--   All angles are specified in radians.
--
--   Note: since version 2 of @easytensor@, DataFrames and matrices are row-major.
--         A good SIMD implementation may drastically improve performance
--         of 4D vector-matrix products of the form @v %* m@, but not so much
--         for products of the form @m %* v@ (due to memory layout).
--         Thus, all operations here assume the former form to benefit more from
--         SIMD in future.
class HomTransform4 t where
-- | Create a translation matrix from a vector.  The 4th coordinate is ignored.
--
--   If @p ! 3 == 1@ and @v ! 3 == 0@, then
--
--   > p %* translate4 v == p + v
--
translate4  :: Vector t 4 -> Matrix t 4 4
-- | Create a translation matrix from a vector.
--
--   If @p ! 3 == 1@, then
--
--   > p %* translate3 v == p + toHomVector v
--
translate3  :: Vector t 3 -> Matrix t 4 4
-- | Rotation matrix for a rotation around the X axis, angle is given in radians.
--   e.g. @p %* rotateX (pi/2)@ rotates point @p@ around @Ox@ by 90 degrees.
rotateX     :: t -> Matrix t 4 4
-- | Rotation matrix for a rotation around the Y axis, angle is given in radians.
--   e.g. @p %* rotateY (pi/2)@ rotates point @p@ around @Oy@ by 90 degrees.
rotateY     :: t -> Matrix t 4 4
-- | Rotation matrix for a rotation around the Z axis, angle is given in radians.
--   e.g. @p %* rotateZ (pi/2)@ rotates point @p@ around @Oz@ by 90 degrees.
rotateZ     :: t -> Matrix t 4 4
-- | Rotation matrix for a rotation around an arbitrary normalized vector
--   e.g. @p %* rotate (pi/2) v@ rotates point @p@ around @v@ by 90 degrees.
rotate      :: Vector t 3 -> t -> Matrix t 4 4
-- | Rotation matrix from the Euler angles roll (axis @Z@), yaw (axis @Y'@), and pitch (axis @X''@).
--   This order is known as Tait-Bryan angles (@Z-Y'-X''@ intrinsic rotations), or nautical angles, or Cardan angles.
--
--   > rotateEuler pitch yaw roll == rotateZ roll %* rotateY yaw %* rotateX pitch
--
--   https://en.wikipedia.org/wiki/Euler_angles#Conventions_2
rotateEuler :: t -- ^ pitch (axis @X''@)
-> t -- ^ yaw (axis @Y'@)
-> t -- ^ roll (axis @Z@)
-> Matrix t 4 4
-- | Create a transform matrix using up direction, camera position and a point to look at.
--   Just the same as GluLookAt.
lookAt      :: Vector t 3 -- ^ The up direction, not necessary unit length or perpendicular to the view vector
-> Vector t 3 -- ^ The viewers position
-> Vector t 3 -- ^ The point to look at
-> Matrix t 4 4
-- | A perspective symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
--   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
perspective :: t -- ^ Near plane clipping distance (always positive)
-> t -- ^ Far plane clipping distance (always positive)
-> t -- ^ Field of view of the y axis, in radians
-> t -- ^ Aspect ratio, i.e. screen's width\/height
-> Matrix t 4 4
-- | An orthogonal symmetric projection matrix. Right-handed coordinate system. (@x@ - right, @y@ - top)
--   http://en.wikibooks.org/wiki/GLSL_Programming/Vertex_Transformations
orthogonal  :: t -- ^ Near plane clipping distance
-> t -- ^ Far plane clipping distance
-> t -- ^ width
-> t -- ^ height
-> Matrix t 4 4
-- | Add one more dimension and set it to 1.
toHomPoint  :: Vector t 3 -> Vector t 4
-- | Add one more dimension and set it to 0.
toHomVector :: Vector t 3 -> Vector t 4
-- | Transform a homogenous vector or point into a normal 3D vector.
--   If the last coordinate is not zero, divide the rest by it.
fromHom     :: Vector t 4 -> Vector t 3

-- Type abbreviations

type Mat22f = Matrix Float 2 2
type Mat32f = Matrix Float 3 2
type Mat42f = Matrix Float 4 2
type Mat23f = Matrix Float 2 3
type Mat33f = Matrix Float 3 3
type Mat43f = Matrix Float 4 3
type Mat24f = Matrix Float 2 4
type Mat34f = Matrix Float 3 4
type Mat44f = Matrix Float 4 4

type Mat22d = Matrix Double 2 2
type Mat32d = Matrix Double 3 2
type Mat42d = Matrix Double 4 2
type Mat23d = Matrix Double 2 3
type Mat33d = Matrix Double 3 3
type Mat43d = Matrix Double 4 3
type Mat24d = Matrix Double 2 4
type Mat34d = Matrix Double 3 4
type Mat44d = Matrix Double 4 4

-- | Compose a 2x2D matrix
mat22 :: PrimBytes (t :: Type)
=> Vector t 2 -> Vector t 2 -> Matrix t 2 2
mat22 :: Vector t 2 -> Vector t 2 -> Matrix t 2 2
mat22 = (PrimBytes t, Dimensions '[2, 2]) => PackDF t '[2] 2 (Matrix t 2 2)
forall t (d :: Nat) (ds :: [Nat]).
(PrimBytes t, Dimensions (d : ds)) =>
PackDF t ds d (DataFrame t (d : ds))
packDF @_ @2 @'[2]
{-# INLINE mat22 #-}

-- | Compose a 3x3D matrix
mat33 :: PrimBytes (t :: Type)
=> Vector t 3 -> Vector t 3 -> Vector t 3 -> Matrix t 3 3
mat33 :: Vector t 3 -> Vector t 3 -> Vector t 3 -> Matrix t 3 3
mat33 = (PrimBytes t, Dimensions '[3, 3]) => PackDF t '[3] 3 (Matrix t 3 3)
forall t (d :: Nat) (ds :: [Nat]).
(PrimBytes t, Dimensions (d : ds)) =>
PackDF t ds d (DataFrame t (d : ds))
packDF @_ @3 @'[3]
{-# INLINE mat33 #-}

-- | Compose a 4x4D matrix
mat44 :: PrimBytes (t :: Type)
=> Vector t 4
-> Vector t 4
-> Vector t 4
-> Vector t 4
-> Matrix t 4 4
mat44 :: Vector t 4
-> Vector t 4 -> Vector t 4 -> Vector t 4 -> Matrix t 4 4
mat44 = (PrimBytes t, Dimensions '[4, 4]) => PackDF t '[4] 4 (Matrix t 4 4)
forall t (d :: Nat) (ds :: [Nat]).
(PrimBytes t, Dimensions (d : ds)) =>
PackDF t ds d (DataFrame t (d : ds))
packDF @_ @4 @'[4]
{-# INLINE mat44 #-}

instance ( KnownDim n, KnownDim m
, PrimArray t (Matrix t n m)
, PrimArray t (Matrix t m n)
) => MatrixTranspose t (n :: Nat) (m :: Nat) where
transpose :: Matrix t n m -> Matrix t m n
transpose Matrix t n m
df = case Matrix t n m -> Either t CumulDims
forall t a. PrimArray t a => a -> Either t CumulDims
uniqueOrCumulDims Matrix t n m
df of
Left t
a -> t -> Matrix t m n
forall t a. PrimArray t a => t -> a
a
Right CumulDims
_
| Word
wm <- KnownDim m => Word
forall k (n :: k). KnownDim n => Word
dimVal' @m
, Word
wn <- KnownDim n => Word
forall k (n :: k). KnownDim n => Word
dimVal' @n
, Int#
m <- case Word
wm of W# Word#
w -> Word# -> Int#
word2Int# Word#
w
, Int#
n <- case Word
wn of W# Word#
w -> Word# -> Int#
word2Int# Word#
w
-> let f :: (Int, Int) -> (# (Int, Int), t #)
f ( I# Int#
i,  I# Int#
j )
| Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
==# Int#
n) = (Int, Int) -> (# (Int, Int), t #)
f ( Int
0, Int# -> Int
I# (Int#
j Int# -> Int# -> Int#
+# Int#
1#) )
| Bool
otherwise         = (# ( Int# -> Int
I# (Int#
i Int# -> Int# -> Int#
+# Int#
1#), Int# -> Int
I# Int#
j )
, Int# -> Matrix t n m -> t
forall t a. PrimArray t a => Int# -> a -> t
ix# (Int#
i Int# -> Int# -> Int#
*# Int#
m Int# -> Int# -> Int#
+# Int#
j) Matrix t n m
df
#)
in case CumulDims
-> ((Int, Int) -> (# (Int, Int), t #))
-> (Int, Int)
-> (# (Int, Int), Matrix t m n #)
forall t a s.
PrimArray t a =>
CumulDims -> (s -> (# s, t #)) -> s -> (# s, a #)
gen# ([Word] -> CumulDims
CumulDims [Word
wmWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
wn, Word
wn, Word
1]) (Int, Int) -> (# (Int, Int), t #)
f (Int
0,Int
0) of (# (Int, Int)
_, Matrix t m n
r #) -> Matrix t m n
r

instance MatrixTranspose (t :: Type) (xn :: XNat) (xm :: XNat) where
transpose :: Matrix t xn xm -> Matrix t xm xn
transpose (XFrame (df :: DataFrame t ns))
| ((Dim y
D :: Dim n) :* (Dim y
D :: Dim m) :* TypedList Dim ys
U) <- Dimensions ns => TypedList Dim ns
forall k (ds :: [k]). Dimensions ds => Dims ds
dims @ns
, Dict (PrimBytes t)
Dict <- DataFrame t '[Head ns, Head (Tail ns)] -> Dict (PrimBytes t)
forall t (d :: Nat) (ds :: [Nat]).
KnownBackend t (d : ds) =>
DataFrame t (d : ds) -> Dict (PrimBytes t)
inferPrimElem DataFrame t ns
df
= DataFrame t '[y, y] -> Matrix t xm xn
forall l (ts :: l) (xns :: [XNat]) (ns :: [Nat]).
(All KnownDimType xns, FixedDims xns ns, Dimensions ns,
KnownBackends ts ns) =>
DataFrame ts ns -> DataFrame ts xns
XFrame (Matrix t y y -> DataFrame t '[y, y]
forall k k (t :: k) (n :: k) (m :: k).
MatrixTranspose t n m =>
Matrix t n m -> Matrix t m n
transpose DataFrame t ns
Matrix t y y
df :: Matrix t m n)
transpose _ = error "MatrixTranspose/transpose: impossible argument"
#endif

instance (KnownDim n, PrimArray t (Matrix t n n), Num t)
=> SquareMatrix t n where
eye :: Matrix t n n
eye
| Word
n <- KnownDim n => Word
forall k (n :: k). KnownDim n => Word
dimVal' @n
= let f :: Word -> (# Word, t #)
f Word
0 = (# Word
n    , t
1 #)
f Word
k = (# Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1, t
0 #)
in case CumulDims
-> (Word -> (# Word, t #)) -> Word -> (# Word, Matrix t n n #)
forall t a s.
PrimArray t a =>
CumulDims -> (s -> (# s, t #)) -> s -> (# s, a #)
gen# ([Word] -> CumulDims
CumulDims [Word
nWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
n, Word
n, Word
1]) Word -> (# Word, t #)
f Word
0 of
(# Word
_, Matrix t n n
r #) -> Matrix t n n
r
diag :: Scalar t -> Matrix t n n
diag Scalar t
se
| Word
n <- KnownDim n => Word
forall k (n :: k). KnownDim n => Word
dimVal' @n
, t
e <- Scalar t -> t
forall t. DataFrame t '[] -> t
unScalar Scalar t
se
= let f :: Word -> (# Word, t #)
f Word
0 = (# Word
n    , t
e #)
f Word
k = (# Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1, t
0 #)
in case CumulDims
-> (Word -> (# Word, t #)) -> Word -> (# Word, Matrix t n n #)
forall t a s.
PrimArray t a =>
CumulDims -> (s -> (# s, t #)) -> s -> (# s, a #)
gen# ([Word] -> CumulDims
CumulDims [Word
nWord -> Word -> Word
forall a. Num a => a -> a -> a
*Word
n, Word
n, Word
1]) Word -> (# Word, t #)
f Word
0 of
(# Word
_, Matrix t n n
r #) -> Matrix t n n
r
trace :: Matrix t n n -> Scalar t
trace Matrix t n n
df
| I# Int#
n <- Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
\$ KnownDim n => Word
forall k (n :: k). KnownDim n => Word
dimVal' @n
, Int#
n1 <- Int#
n Int# -> Int# -> Int#
+# Int#
1#
= let f :: Int# -> t
f Int#
0# = Int# -> Matrix t n n -> t
forall t a. PrimArray t a => Int# -> a -> t
ix# Int#
0# Matrix t n n
df
f Int#
k  = Int# -> Matrix t n n -> t
forall t a. PrimArray t a => Int# -> a -> t
ix# Int#
k  Matrix t n n
df t -> t -> t
forall a. Num a => a -> a -> a
+ Int# -> t
f (Int#
k Int# -> Int# -> Int#
-# Int#
n1)
in t -> Scalar t
forall t. t -> DataFrame t '[]
scalar (t -> Scalar t) -> t -> Scalar t
forall a b. (a -> b) -> a -> b
\$ Int# -> t
f (Int#
n Int# -> Int# -> Int#
*# Int#
n Int# -> Int# -> Int#
-# Int#
1#)
```