{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}

module Internal.Matrix where

import qualified Data.List as List
import Data.List.Split
import Data.Proxy
import Data.Tuple
import qualified GHC.Natural as Natural
import GHC.TypeNats
import Prelude hiding (length)

data Matrix (m :: Nat) (n :: Nat) a where
  Matrix :: forall m n a. (Int, Int) -> ![[a]] -> Matrix m n a

instance Show a => Show (Matrix m n a) where
  show :: Matrix m n a -> String
show (Matrix _ matrix :: [[a]]
matrix) = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
List.intercalate "\n" ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ ([a] -> String) -> [[a]] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> String
forall a. Show a => a -> String
show [[a]]
matrix

instance Functor (Matrix m n) where
  fmap :: (a -> b) -> Matrix m n a -> Matrix m n b
fmap f :: a -> b
f (Matrix size :: (Int, Int)
size a :: [[a]]
a) = (Int, Int) -> [[b]] -> Matrix m n b
forall (m :: Nat) (n :: Nat) a. (Int, Int) -> [[a]] -> Matrix m n a
Matrix (Int, Int)
size ([[b]] -> Matrix m n b) -> [[b]] -> Matrix m n b
forall a b. (a -> b) -> a -> b
$ ([a] -> [b]) -> [[a]] -> [[b]]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f) [[a]]
a

instance Applicative (Matrix m n) where
  pure :: a -> Matrix m n a
pure = (Int, Int) -> [[a]] -> Matrix m n a
forall (m :: Nat) (n :: Nat) a. (Int, Int) -> [[a]] -> Matrix m n a
Matrix (1, 1) ([[a]] -> Matrix m n a) -> (a -> [[a]]) -> a -> Matrix m n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> [[a]]) -> (a -> [a]) -> a -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Matrix size :: (Int, Int)
size fs :: [[a -> b]]
fs <*> :: Matrix m n (a -> b) -> Matrix m n a -> Matrix m n b
<*> (Matrix _ as :: [[a]]
as) =
    (Int, Int) -> [[b]] -> Matrix m n b
forall (m :: Nat) (n :: Nat) a. (Int, Int) -> [[a]] -> Matrix m n a
Matrix (Int, Int)
size ([[b]] -> Matrix m n b) -> [[b]] -> Matrix m n b
forall a b. (a -> b) -> a -> b
$ (([a -> b], [a]) -> [b]) -> [([a -> b], [a])] -> [[b]]
forall a b. (a -> b) -> [a] -> [b]
map (([a -> b] -> [a] -> [b]) -> ([a -> b], [a]) -> [b]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [a -> b] -> [a] -> [b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>)) ([([a -> b], [a])] -> [[b]]) -> [([a -> b], [a])] -> [[b]]
forall a b. (a -> b) -> a -> b
$ [[a -> b]] -> [[a]] -> [([a -> b], [a])]
forall a b. [a] -> [b] -> [(a, b)]
zip [[a -> b]]
fs [[a]]
as

instance (Eq a) => Eq (Matrix m n a) where
  Matrix _ a :: [[a]]
a == :: Matrix m n a -> Matrix m n a -> Bool
== Matrix _ b :: [[a]]
b = [[a]]
a [[a]] -> [[a]] -> Bool
forall a. Eq a => a -> a -> Bool
== [[a]]
b

type Vector n a = Matrix n 1 a

size :: Integral b => Matrix m n a -> (b, b)
size :: Matrix m n a -> (b, b)
size (Matrix (m :: Int
m, n :: Int
n) _) = (Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m, Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)

value :: Matrix m n a -> [[a]]
value :: Matrix m n a -> [[a]]
value (Matrix _ value :: [[a]]
value) = [[a]]
value