{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

module QLinear.Identity (e, Identity, HasIdentity (..)) where

import Data.Proxy
import qualified GHC.Natural as Natural
import GHC.TypeNats
import Internal.Matrix

type Identity n a = Matrix n n a

class HasIdentity a where
  zero :: a
  one :: a

instance (Num a) => HasIdentity a where
  zero :: a
zero = 0
  one :: a
one = 1

-- | Polymirphic identity matrix
--
-- Identity matrix can udjust to other matrix with known size. If size is unknown, just set it yourself
--
-- >>> e :: Identity 4 Int
-- [1,0,0,0]
-- [0,1,0,0]
-- [0,0,1,0]
-- [0,0,0,1]
-- >>> e ~+~ [matrix| 1 2; 3 4 |]
-- [2,2]
-- [3,5]
e :: forall n a. (KnownNat n, HasIdentity a) => Identity n a
e :: Identity n a
e = (Int, Int) -> [[a]] -> Identity n a
forall (m :: Nat) (n :: Nat) a. (Int, Int) -> [[a]] -> Matrix m n a
Matrix (Int
n, Int
n) ([[a]] -> Identity n a) -> [[a]] -> Identity n a
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> a -> a -> [[a]]
forall a. (Int, Int) -> a -> a -> [[a]]
finiteIdentityList (Int
n, Int
n) a
forall a. HasIdentity a => a
one a
forall a. HasIdentity a => a
zero
  where
    n :: Int
n = Natural -> Int
Natural.naturalToInt (Natural -> Int) -> Natural -> Int
forall a b. (a -> b) -> a -> b
$ Proxy n -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal (Proxy n
forall k (t :: k). Proxy t
Proxy @n)

infiniteIdentityList :: a -> a -> [[a]]
infiniteIdentityList :: a -> a -> [[a]]
infiniteIdentityList o :: a
o z :: a
z = [a] -> [[a]]
stream (a
o a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a]
forall a. a -> [a]
repeat a
z)
  where
    stream :: [a] -> [[a]]
stream seed :: [a]
seed = [a]
seed [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
stream (a
z a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
seed)

finiteIdentityList :: (Int, Int) -> a -> a -> [[a]]
finiteIdentityList :: (Int, Int) -> a -> a -> [[a]]
finiteIdentityList (m :: Int
m, n :: Int
n) o :: a
o z :: a
z = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ Int -> [[a]] -> [[a]]
forall a. Int -> [a] -> [a]
take Int
m ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ a -> a -> [[a]]
forall a. a -> a -> [[a]]
infiniteIdentityList a
o a
z