{-# OPTIONS -fglasgow-exts -fallow-undecidable-instances #-}
module Data.Matrix where

import Data.Array.Fast hiding ((!))
import qualified Data.Array.Fast as Arr
import Data.List hiding (transpose)

data Matrix a where
  Empty :: Matrix a
  Scalar :: a -> Matrix a
  Vector :: FArray a e => a Int e -> Matrix e
  VectorT :: FArray a e => a Int e -> Matrix e
  Matrix :: FArray a e => a (Int, Int) e -> Matrix e

instance Eq a => Eq (Matrix a) where
  Empty     == Empty     = True
  Scalar a  == Scalar b  = a==b
  Vector v  == Vector w  = elems v==elems w
  VectorT v == VectorT w = elems v==elems w
  Matrix m  == Matrix n  = elems m==elems n
  _         == _         = False

class Ix i => Indexable a i e | a -> i where
  (!) :: a e -> i -> e

instance (IArray a e, Ix i) => Indexable (a i) i e where
  (!) = (Arr.!)

instance Indexable Matrix (Int,Int) e where
  Empty ! _ = error "Can't index empty matrix"
  Scalar s ! (1,1) = s
  Vector v ! (r,1) = v!r
  VectorT v ! (1,c) = v!c
  Matrix a ! (m,n) = a!(m,n)
  _ ! _ = error "Matrix.!: Index out of range."

dim :: Matrix a -> (Int, Int)
dim Empty = (0,0)
dim (Scalar _) = (1,1)
dim (Vector v) = (snd (bounds v),1)
dim (VectorT v) = (1,snd (bounds v))
dim (Matrix m) = snd (bounds m)

instance (Show a) => Show (Matrix a) where
  show Empty = "matrix []"
  show (Scalar s) = "scalar "++show s
  show (Vector v) = "vector "++show (elems v)
  show (VectorT v) = "transpose (vector "++shows (elems v) ")"
  show (Matrix a) = "matrix "++show [ [ a!(i,j) | j <- [1..n] ]
                                      | i <- [1..m]]
    where (_,(m,n)) = bounds a

matmap :: (Num n, Num m, FArray b m) =>
          (n -> m) -> Matrix n -> Matrix m
matmap _ Empty = Empty
matmap f (Scalar s) = Scalar (f s)
matmap f (Vector v) = Vector (famap f v)
matmap f (VectorT v) = VectorT (famap f v)
matmap f (Matrix m) = Matrix (famap f m)

scalar :: (Num e) => e -> Matrix e
scalar = Scalar

vector :: (FArray a e, Num e) => [e] -> Matrix e
vector [] = Empty
vector [a] = Scalar a
vector xs = Vector $ listArray (1,length xs) xs

matrix :: (Num e, FArray a e) => [[e]] -> Matrix e
matrix [] = Empty
matrix [es] = transpose $ vector es
matrix ([]:_) = Empty
matrix es@([_]:_) = vector (map headAlone es)
matrix ms = Matrix (array ((1,1),(length ms,length (head ms)))
                          [((i,j),e) | (i,es) <- zip [1..] ms,
                                       (j,e)  <- zip [1..] es])

chkV :: Int -> e -> e
chkV 1 t = t
chkV _ _ = error "Vector index not 1"

assocMatrix :: (Num e, FArray a e) => (Int,Int) -> [((Int,Int),e)] -> Matrix e
assocMatrix (1,1) [((1,1),s)] = Scalar s
assocMatrix (1,1) [] = Scalar undefined
assocMatrix (1,1) _  = error "Bad initializers for scalar."
assocMatrix (0,_) [] = Empty
assocMatrix (0,_) _  = error "Non-empty initializer for empty matrix."
assocMatrix (_,0) [] = Empty
assocMatrix (_,0) _  = error "Non-empty initializer for empty matrix."
assocMatrix (1,n) vs =
  VectorT $ array (1,n) [ chkV i (j,v) | ((i,j),v) <- vs ]
assocMatrix (n,1) vs =
  Vector $ array (1,n) [ chkV j (i,v) | ((i,j),v) <- vs ]
assocMatrix (m,n) vs =
  Matrix $ array ((1,1),(m,n)) vs

accumMatrix :: (Num b, FArray a b) => (b -> e -> b) -> b ->
               (Int,Int) -> [((Int,Int),e)] -> Matrix b
accumMatrix f z (1,1) vs =
  Scalar $ foldl' f z [ chkV i $ chkV j $ v | ((i,j),v) <- vs]
accumMatrix _ _ (0,_) [] = Empty
accumMatrix _ _ (0,_) _  = error "Non-empty initializer for empty matrix."
accumMatrix _ _ (_,0) [] = Empty
accumMatrix _ _ (_,0) _  = error "Non-empty initializer for empty matrix."
accumMatrix f z (1,n) vs =
  VectorT $ accumArray f z (1,n) [ chkV i (j,v) | ((i,j),v) <- vs ]
accumMatrix f z (n,1) vs =
  Vector $ accumArray f z (1,n) [ chkV j (i,v) | ((i,j),v) <- vs ]
accumMatrix f z (m,n) vs =
  Matrix $ accumArray f z ((1,1),(m,n)) vs

transpose :: (Num e) => Matrix e -> Matrix e
transpose Empty = Empty
transpose (Scalar a) = Scalar a
transpose (Vector v) = VectorT v
transpose (VectorT v) = Vector v
transpose (Matrix a) = Matrix (array ((1,1),(n,m))
                                     [((j,i), e) | ((i,j),e) <- assocs a])
  where (_,(m,n)) = bounds a

-- headAlone takes the head of a singleton list, otherwise is an error.
headAlone :: [a] -> a
headAlone [a] = a
headAlone _   = error "Singleton list expected"

elementwise :: (Num m) =>
               (m -> n -> p) -> Matrix m -> Matrix n -> Matrix p
elementwise f (Scalar s)  m           = matmap (f s) m
elementwise f m           (Scalar s)  = matmap (flip f s) m
elementwise _ Empty       Empty       = Empty
elementwise f (Vector v)  (Vector w)  = Vector $ pwv f v w
elementwise f (VectorT v) (VectorT w) = VectorT $ pwv f v w
elementwise f (Matrix a) (Matrix b)
  | snd (bounds b) == (m,n) = Matrix $ array ((1,1),(m,n))
                                       [ ((i,j),f e (b!(i,j)))
                                             | ((i,j),e) <- assocs a ]
  where (_,(m,n)) = bounds a
elementwise _ _ _ = error "Matrix shapes do not match"

pwv :: (IArray a m, IArray b n, FArray c p) =>
       (m -> n -> p) -> a Int m -> b Int n -> c Int p
pwv f v w
  | n==m = array (1,n) [ (i,f e (w!i)) | (i,e) <- assocs v ]
  | otherwise = error "Vector dimensions do not match!"
  where (_,n) = bounds v
        (_,m) = bounds w

instance (Num e, FArray a e) => Num (Matrix e) where
  (+) = elementwise (+)
  (-) = elementwise (-)
  Empty     * _         = Empty
  _         * Empty     = Empty
  Scalar s  * m         = matmap (s*) m
  m         * Scalar s  = matmap (*s) m
  a         * b
    | n==p =
        accumMatrix (+) 0 (m,q)
          [ ((i,k), a!(i,j) * b!(j,k))
            | i <- [1..m], k <- [1..q], j <- [1..n]]
    | otherwise = error "LHS Matrix width not equal to RHS Matrix height"
    where (m,n) = dim a
          (p,q) = dim b
  negate = matmap negate
  abs = matmap abs
  signum = matmap signum
  fromInteger = scalar . fromInteger