module Matrix
    ( Z
    , S
    , N
    , Tuple (..)
    , tupleToList
    , tupleFromList
    , AdditiveMonoid (..)
    , scale
    , Vector (..)
    , Matrix (..)
    , matrixFromColumnVectors
    , matrixFromRowVectors
    , columnVectors
    , rowVectors
    , act
    , mult
    , diag
    , scalar
    , id
    , transpose
    , adj
    , det
    , inv
    , rank
    ) where

import Prelude hiding (id,zip,zipWith)
import Control.Monad

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

newtype Flip f x y = Flip{ unFlip :: f y x }

class Functor t => Zippable t where
    zip     :: t a -> t b -> t (a,b)
    zipWith :: (a -> b -> c) -> t a -> t b -> t c
    zip           = zipWith (,)
    zipWith f a b = fmap (uncurry f) (zip a b)

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

data Z
data S x

class N n where
    pr :: t Z -> (forall m. N m => t m -> t (S m)) -> t n

instance N Z where
    pr z s = z

instance N n => N (S n) where
    pr z s = s (pr z s)

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

data Tuple n x where
    Nil  :: Tuple Z x
    Cons :: N n => x -> Tuple n x -> Tuple (S n) x

instance Functor (Tuple n) where
    fmap f Nil = Nil
    fmap f (Cons x xs) = Cons (f x) (fmap f xs)

instance Zippable (Tuple n) where
    zipWith f Nil Nil = Nil
    zipWith f (Cons x xs) (Cons y ys) = Cons (f x y) (zipWith f xs ys)

tupleToList :: Tuple n x -> [x]
tupleToList Nil = []
tupleToList (Cons x xs) = x : tupleToList xs

newtype F2 x n = F2{ unF2 :: [x] -> Maybe (Tuple n x) }

tupleFromList :: N n => [x] -> Maybe (Tuple n x)
tupleFromList = unF2 (pr z s)
    where z = F2 z'
              where z' [] = Just Nil
                    z'  _ = Nothing
          s (F2 f) = F2 s'
              where s' (x:xs) = do ys <- f xs
                                   return (Cons x ys)
                    s' [] = Nothing

replTuple :: N n => x -> Tuple n x
replTuple x = unFlip (pr (Flip Nil) (Flip . Cons x . unFlip))

sumTuple :: (Num x) => Tuple n x -> x
sumTuple = sum . tupleToList

transposeTuple :: N n => Tuple m (Tuple n x) -> Tuple n (Tuple m x)
transposeTuple Nil         = replTuple Nil
transposeTuple (Cons x xs) = zipWith Cons x (transposeTuple xs)

newtype F a b n = F{ unF :: b -> Tuple n a }

unfoldTuple :: N n => (b -> (a,b)) -> (b -> Tuple n a)
unfoldTuple f = unF (pr z s)
    where z = F (const Nil)
          s (F p) = F s'
              where s' b = Cons x (p b')
                        where (x,b') = f b

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

class AdditiveMonoid x where
    zero :: x
    add  :: x -> x -> x

scale :: (Num x, Functor f) => x -> f x -> f x
scale s = fmap (*s)

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

newtype Vector m x = Vector{ unVector :: Tuple m x }

instance Functor (Vector m) where
    fmap f (Vector x) = Vector $ fmap f x

instance Zippable (Vector m) where
    zipWith f (Vector a) (Vector b) = Vector (zipWith f a b)

instance (N n, Num x) => AdditiveMonoid (Vector n x) where
    zero = Vector (replTuple 0)
    add (Vector x) (Vector y) = Vector $ zipWith (+) x y

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

newtype Matrix m n x = Matrix{ unMatrix :: Tuple m (Tuple n x) }

instance Functor (Matrix m n) where
    fmap f = Matrix . fmap (fmap f) . unMatrix

instance Zippable (Matrix m n) where
    zipWith f (Matrix a) (Matrix b) = Matrix (zipWith (zipWith f) a b)

instance (N m, N n, Num x) => AdditiveMonoid (Matrix m n x) where
    zero = Matrix (replTuple (replTuple 0))
    add (Matrix x) (Matrix y) = Matrix $ zipWith (zipWith (+)) x y

matrixFromColumnVectors :: (N m, N n) => Tuple n (Vector m x) -> Matrix m n x
matrixFromColumnVectors = transpose . matrixFromRowVectors

matrixFromRowVectors :: (N m, N n) => Tuple m (Vector n x) -> Matrix m n x
matrixFromRowVectors = Matrix . fmap unVector

columnVectors :: (N m, N n) => Matrix m n x -> Tuple n (Vector m x)
columnVectors = rowVectors . transpose

rowVectors :: (N m, N n) => Matrix m n x -> Tuple m (Vector n x)
rowVectors = fmap Vector . unMatrix

act :: Num x => Matrix m n x -> Vector n x -> Vector m x
act (Matrix m) (Vector v) =
    Vector $ fmap (\row -> sumTuple (zipWith (*) row v)) m

mult :: (Num x, N m, N n, N o) => Matrix m n x -> Matrix n o x -> Matrix m o x
mult (Matrix a) (Matrix b) = Matrix (fmap f a)
    where bT = transposeTuple b
          f aRow = fmap (\bCol -> sumTuple (zipWith (*) bCol aRow)) bT

diag :: (N n, Num x) => Tuple n x -> Matrix n n x
diag xs = Matrix (f xs)
    where f :: (N n, Num x) => Tuple n x -> Tuple n (Tuple n x)
          f Nil         = Nil
          f (Cons x xs) = Cons (Cons x (replTuple 0)) (fmap (Cons 0) (f xs))

scalar :: (N n, Num x) => x -> Matrix n n x
scalar = diag . replTuple

id :: (N n, Num x) => Matrix n n x
id = scalar 1

transpose :: N n => Matrix m n x -> Matrix n m x
transpose = Matrix . transposeTuple . unMatrix

adj :: (N n, Num x) => Matrix n n x -> Matrix n n x
adj a@(Matrix table) =
    case table of
    Nil      -> Matrix Nil
    Cons _ _ -> transpose (zipWith (*) signs (fmap det (minors a)))
    where signs :: (N m, N n, Num x) => Matrix m n x
          signs = Matrix $ unfoldTuple (\x -> (f x, not x)) True
              where f = unfoldTuple (\y -> (if y then 1 else -1, not y))

minors :: (N m, N n) =>
          Matrix (S m) (S n) x -> Matrix (S m) (S n) (Matrix m n x)
minors = Matrix . fmap (fmap Matrix . transposeTuple . fmap f) . f . unMatrix
    where f :: Tuple (S n) x -> Tuple (S n) (Tuple n x)
          f (Cons x xs) = case xs of
                          Nil      -> Cons xs Nil
                          Cons _ _ -> Cons xs (fmap (Cons x) (f xs))

detFromAdjunction :: (N n, Num x) => Matrix n n x -> Matrix n n x -> x
detFromAdjunction a@(Matrix x) adjA =
    case x of
    Nil -> 1
    Cons _ _ -> case a `mult` adjA of
                Matrix rows -> case rows of
                               Cons row _ -> case row of
                                             Cons x _ -> x

det :: (N n, Num x) => Matrix n n x -> x
det a = detFromAdjunction a (adj a)

inv :: (N n, Fractional x) => Matrix n n x -> Maybe (Matrix n n x)
inv a = if detA /= 0 then Just (fmap (/ detA) adjA) else Nothing
    where adjA = adj a
          detA = detFromAdjunction a adjA

rank :: (N m, N n, Num x) => Matrix m n x -> Int
rank = rank' . unMatrix

rank' :: (N m, N n, Num x) => Tuple m (Tuple n x) -> Int
rank' Nil = 0
rank' t@(Cons row _) =
    case row of
    Nil -> 0
    Cons _ _ -> case f t of
                Just t' -> rank' t' + 1
                Nothing -> 0
    where f :: (N m, N n, Num x) =>
               Tuple (S m) (Tuple (S n) x) -> Maybe (Tuple m (Tuple n x))
          f t@(Cons _ _) =
              case pivot t of
              Just (Cons row rows) -> Just (fmap (g row) rows)
              Nothing -> Nothing
          g :: (N n, Num x) => Tuple (S n) x -> Tuple (S n) x -> Tuple n x
          g (Cons x xs) (Cons y ys) = zipWith (\y' x' -> x * y' - y * x') xs ys

pivot :: (N m, N n, Num x) =>
         Tuple (S m) (Tuple (S n) x) -> Maybe (Tuple (S m) (Tuple (S n) x))
pivot t =
    case pickupRow t of
    Just (row, rows) -> Just (Cons row rows)
    Nothing -> case pickupRow (transposeTuple t) of
               Just (col, cols) -> Just (transposeTuple (Cons col cols))
               Nothing -> Nothing

pickupRow :: (N m, N n, Num x)
          => Tuple (S m) (Tuple (S n) x)
          -> Maybe (Tuple (S n) x, Tuple m (Tuple (S n) x))
pickupRow (Cons row@(Cons x _) rows) =
    if x /= 0
    then Just (row, rows)
    else case rows of
         Nil -> Nothing
         Cons _ _ ->  case pickupRow rows of
                      Nothing -> Nothing
                      Just (row', rows') -> Just (row', Cons row rows')

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