{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

module FixedMatrix
where

import FixedVector as FV
import qualified Data.Vector.Fixed as V
import Data.Vector.Fixed.Internal (arity)

type Mat v w a = Vn v (Vn w a)
type Mat2 a = Mat Vec2D Vec2D a
type Mat3 a = Mat Vec3D Vec3D a
type Mat4 a = Mat Vec4D Vec4D a

-- | Convert a matrix to a nested list.
toList :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> [[a]]
toList m = Prelude.map FV.toList (FV.toList m)

-- | Create a matrix from a nested list.
fromList :: (V.Vector v (Vn w a), V.Vector w a) => [[a]] -> Mat v w a
fromList vs = FV.fromList $ Prelude.map FV.fromList vs


-- | Unsafe indexing.
(!) :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> (Int, Int) -> a
(!) m (i, j) = (row m i) FV.! j

-- | Safe indexing.
(!?) :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a
                                               -> (Int, Int)
                                               -> Maybe a
(!?) m (i, j)
  | i < 0 || j < 0 = Nothing
  | i > FV.length m = Nothing
  | otherwise = if j > FV.length (row m j)
                then Nothing
                else Just $ (row m j) FV.! j


-- | The number of rows in the matrix. Implementation stolen from
--   Data.Vector.Fixed.Length.
nrows :: forall v w a. (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> Int
nrows _ = arity (undefined :: V.Dim v)

-- | The number of columns in the first row of the
--   matrix. Implementation stolen from
ncols :: forall v w a. (V.Vector v (Vn w a), V.Vector w a) => Mat v w a -> Int
ncols _ = arity (undefined :: V.Dim w)

-- | Return the @i@th row of @m@. Unsafe.
row :: (V.Vector v (Vn w a), V.Vector w a) => Mat v w a
                                           -> Int
                                           -> Vn w a
row m i = m FV.! i


-- | Return the @j@th column of @m@. Unsafe.
column :: (V.Vector v a, V.Vector v (Vn w a), V.Vector w a) => Mat v w a
                                                            -> Int
                                                            -> Vn v a
column m j =
  FV.map (element j) m
  where
    element = flip (FV.!)


-- | Transose @m@; switch it's columns and its rows. This is a dirty
--   implementation.. it would be a little cleaner to use imap, but it
--   doesn't seem to work.
transpose :: (V.Vector v (Vn w a),
              V.Vector w (Vn v a),
              V.Vector v a,
              V.Vector w a)
             => Mat v w a
             -> Mat w v a
transpose m = FV.fromList column_list
  where
    column_list = [ column m i | i <- [0..(ncols m)-1] ]

-- | Is @m@ symmetric?
symmetric :: (V.Vector v (Vn w a),
              V.Vector w a,
              v ~ w,
              V.Vector w Bool,
              Eq a)
             => Mat v w a
             -> Bool
symmetric m =
  m == (transpose m)


-- | Construct a new matrix from a function @lambda@. The function
--   @lambda@ should take two parameters i,j corresponding to the
--   entries in the matrix. The i,j entry of the resulting matrix will
--   have the value returned by lambda i j.
construct :: forall v w a.
             (V.Vector v (Vn w a),
              V.Vector w a)
             => (Int -> Int -> a)
             -> Mat v w a
construct lambda = rows
  where
    -- The arity trick is used in Data.Vector.Fixed.length.
    imax = (arity (undefined :: V.Dim w)) - 1
    jmax = (arity (undefined :: V.Dim w)) - 1
    row' i = FV.fromList [ lambda i j | j <- [0..jmax] ]
    rows = FV.fromList [ row' i | i <- [0..imax] ]