{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RebindableSyntax #-}

-- | Boxed matrices; that is, boxed m-vectors of boxed n-vectors. We
--   assume that the underlying representation is
--   Data.Vector.Fixed.Boxed.Vec for simplicity. It was tried in
--   generality and failed.
--
module Linear.Matrix
where

import Data.List (intercalate)

import Data.Vector.Fixed (
  (!),
  N1,
  N2,
  N3,
  N4,
  N5,
  S,
  Z,
  generate,
  mk1,
  mk2,
  mk3,
  mk4,
  mk5
  )
import qualified Data.Vector.Fixed as V (
  and,
  fromList,
  head,
  ifoldl,
  imap,
  map,
  maximum,
  replicate,
  toList,
  zipWith )
import Data.Vector.Fixed.Cont ( Arity, arity )
import Linear.Vector ( Vec, delete, element_sum )
import Normed ( Normed(..) )

import NumericPrelude hiding ( (*), abs )
import qualified NumericPrelude as NP ( (*) )
import qualified Algebra.Absolute as Absolute ( C )
import Algebra.Absolute ( abs )
import qualified Algebra.Additive as Additive ( C )
import qualified Algebra.Algebraic as Algebraic ( C )
import Algebra.Algebraic ( root )
import qualified Algebra.Ring as Ring ( C )
import qualified Algebra.Module as Module ( C )
import qualified Algebra.RealRing as RealRing ( C )
import qualified Algebra.ToRational as ToRational ( C )
import qualified Algebra.Transcendental as Transcendental ( C )
import qualified Prelude as P ( map )

-- | Our main matrix type.
data Mat m n a = (Arity m, Arity n) => Mat (Vec m (Vec n a))

-- Type synonyms for n-by-n matrices.
type Mat1 a = Mat N1 N1 a
type Mat2 a = Mat N2 N2 a
type Mat3 a = Mat N3 N3 a
type Mat4 a = Mat N4 N4 a
type Mat5 a = Mat N5 N5 a

-- | Type synonym for row vectors expressed as 1-by-n matrices.
type Row n a = Mat N1 n a

-- Type synonyms for 1-by-n row "vectors".
type Row1 a = Row N1 a
type Row2 a = Row N2 a
type Row3 a = Row N3 a
type Row4 a = Row N4 a
type Row5 a = Row N5 a

-- | Type synonym for column vectors expressed as n-by-1 matrices.
type Col n a = Mat n N1 a

-- Type synonyms for n-by-1 column "vectors".
type Col1 a = Col N1 a
type Col2 a = Col N2 a
type Col3 a = Col N3 a
type Col4 a = Col N4 a
type Col5 a = Col N5 a

-- We need a big column for Gaussian quadrature.
type N10 = S (S (S (S (S N5))))
type Col10 a = Col N10 a


instance (Eq a) => Eq (Mat m n a) where
  -- | Compare a row at a time.
  --
  --   Examples:
  --
  --   >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
  --   >>> let m2 = fromList [[1,2],[3,4]] :: Mat2 Int
  --   >>> let m3 = fromList [[5,6],[7,8]] :: Mat2 Int
  --   >>> m1 == m2
  --   True
  --   >>> m1 == m3
  --   False
  --
  (Mat rows1) == (Mat rows2) =
    V.and $ V.zipWith comp rows1 rows2
    where
      -- Compare a row, one column at a time.
      comp row1 row2 = V.and (V.zipWith (==) row1 row2)


instance (Show a) => Show (Mat m n a) where
  -- | Display matrices and vectors as ordinary tuples. This is poor
  --   practice, but these results are primarily displayed
  --   interactively and convenience trumps correctness (said the guy
  --   who insists his vector lengths be statically checked at
  --   compile-time).
  --
  --   Examples:
  --
  --   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
  --   >>> show m
  --   ((1,2),(3,4))
  --
  show (Mat rows) =
    "(" ++ (intercalate "," (V.toList row_strings)) ++ ")"
    where
      row_strings = V.map show_vector rows
      show_vector v1 =
        "(" ++ (intercalate "," element_strings) ++ ")"
        where
          v1l = V.toList v1
          element_strings = P.map show v1l


-- | Convert a matrix to a nested list.
toList :: Mat m n a -> [[a]]
toList (Mat rows) = map V.toList (V.toList rows)


-- | Create a matrix from a nested list.
fromList :: (Arity m, Arity n) => [[a]] -> Mat m n a
fromList vs = Mat (V.fromList $ map V.fromList vs)


-- | Unsafe indexing. Much faster than the safe indexing.
(!!!) :: (Arity m, Arity n) => Mat m n a -> (Int, Int) -> a
(!!!) (Mat rows) (i, j) = (rows ! i) ! j


-- | Safe indexing.
--
--   Examples:
--
--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> m !!? (-1,-1)
--   Nothing
--   >>>  m !!? (-1,0)
--   Nothing
--   >>>  m !!? (-1,1)
--   Nothing
--   >>>  m !!? (0,-1)
--   Nothing
--   >>>  m !!? (0,0)
--   Just 1
--   >>>  m !!? (0,1)
--   Just 2
--   >>>  m !!? (1,-1)
--   Nothing
--   >>>  m !!? (1,0)
--   Just 3
--   >>>  m !!? (1,1)
--   Just 4
--   >>>  m !!? (2,-1)
--   Nothing
--   >>>  m !!? (2,0)
--   Nothing
--   >>>  m !!? (2,1)
--   Nothing
--   >>>  m !!? (2,2)
--   Nothing
--
(!!?) :: (Arity m, Arity n) => Mat m n a -> (Int, Int) -> Maybe a
(!!?) matrix idx =
  ifoldl2 f Nothing matrix
  where
    f k l found cur = if (k,l) == idx then (Just cur) else found


-- | The number of rows in the matrix.
nrows :: forall m n a. (Arity m) => Mat m n a -> Int
nrows _ = arity (undefined :: m)


-- | The number of columns in the first row of the
--   matrix. Implementation stolen from Data.Vector.Fixed.length.
ncols :: forall m n a. (Arity n) => Mat m n a -> Int
ncols _ = arity (undefined :: n)


-- | Return the @i@th row of @m@ as a matrix. Unsafe.
row :: (Arity m, Arity n) => Mat m n a -> Int -> Row n a
row m i =
  construct lambda
  where
    lambda _ j = m !!! (i, j)


-- | Return the @j@th column of @m@ as a matrix. Unsafe.
column :: (Arity m, Arity n) => Mat m n a -> Int -> Col m a
column m j =
  construct lambda
  where
    lambda i _ = m !!! (i, j)


-- | Transpose @m@; switch it's columns and its rows. This is a dirty
--   implementation, but I don't see a better way.
--
--   TODO: Don't cheat with fromList.
--
--   Examples:
--
--   >>> let m = fromList [[1,2], [3,4]] :: Mat2 Int
--   >>> transpose m
--   ((1,3),(2,4))
--
transpose :: (Arity m, Arity n) => Mat m n a -> Mat n m a
transpose matrix =
  construct lambda
  where
    lambda i j = matrix !!! (j,i)


-- | Is @m@ symmetric?
--
--   Examples:
--
--   >>> let m1 = fromList [[1,2], [2,1]] :: Mat2 Int
--   >>> symmetric m1
--   True
--
--   >>> let m2 = fromList [[1,2], [3,1]] :: Mat2 Int
--   >>> symmetric m2
--   False
--
symmetric :: (Eq a, Arity m) => Mat m m 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.
--
--   Examples:
--
--   >>> let lambda i j = i + j
--   >>> construct lambda :: Mat3 Int
--   ((0,1,2),(1,2,3),(2,3,4))
--
construct :: forall m n a. (Arity m, Arity n)
          => (Int -> Int -> a) -> Mat m n a
construct lambda = Mat $ generate make_row
  where
    make_row :: Int -> Vec n a
    make_row i = generate (lambda i)


-- | Create an identity matrix with the right dimensions.
--
--   Examples:
--
--   >>> identity_matrix :: Mat3 Int
--   ((1,0,0),(0,1,0),(0,0,1))
--   >>> identity_matrix :: Mat3 Double
--   ((1.0,0.0,0.0),(0.0,1.0,0.0),(0.0,0.0,1.0))
--
identity_matrix :: (Arity m, Ring.C a) => Mat m m a
identity_matrix =
  construct (\i j -> if i == j then (fromInteger 1) else (fromInteger 0))


-- | Given a positive-definite matrix @m@, computes the
--   upper-triangular matrix @r@ with (transpose r)*r == m and all
--   values on the diagonal of @r@ positive.
--
--   Examples:
--
--   >>> let m1 = fromList [[20,-1], [-1,20]] :: Mat2 Double
--   >>> cholesky m1
--   ((4.47213595499958,-0.22360679774997896),(0.0,4.466542286825459))
--   >>> (transpose (cholesky m1)) * (cholesky m1)
--   ((20.000000000000004,-1.0),(-1.0,20.0))
--
cholesky :: forall m n a. (Algebraic.C a, Arity m, Arity n)
         => (Mat m n a) -> (Mat m n a)
cholesky m = construct r
  where
    r :: Int -> Int -> a
    r i j | i == j = sqrt(m !!! (i,j) - sum [(r k i)^2 | k <- [0..i-1]])
          | i < j =
              (((m !!! (i,j)) - sum [(r k i) NP.* (r k j) | k <- [0..i-1]]))/(r i i)
          | otherwise = 0


-- | Returns True if the given matrix is upper-triangular, and False
--   otherwise. The parameter @epsilon@ lets the caller choose a
--   tolerance.
--
--   Examples:
--
--   >>> let m = fromList [[1,1],[1e-12,1]] :: Mat2 Double
--   >>> is_upper_triangular m
--   False
--   >>> is_upper_triangular' 1e-10 m
--   True
--
--   TODO:
--
--     1. Don't cheat with lists.
--
is_upper_triangular' :: (Ord a, Ring.C a, Absolute.C a, Arity m, Arity n)
                    => a -- ^ The tolerance @epsilon@.
                    -> Mat m n a
                    -> Bool
is_upper_triangular' epsilon m =
  and $ concat results
  where
    results = [[ test i j | i <- [0..(nrows m)-1]] | j <- [0..(ncols m)-1] ]

    test :: Int -> Int -> Bool
    test i j
      | i <= j = True
      -- use "less than or equal to" so zero is a valid epsilon
      | otherwise = abs (m !!! (i,j)) <= epsilon


-- | Returns True if the given matrix is upper-triangular, and False
--   otherwise. A specialized version of 'is_upper_triangular\'' with
--   @epsilon = 0@.
--
--   Examples:
--
--   >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
--   >>> is_upper_triangular m
--   False
--
--   >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
--   >>> is_upper_triangular m
--   True
--
--   TODO:
--
--     1. The Ord constraint is too strong here, Eq would suffice.
--
is_upper_triangular :: (Ord a, Ring.C a, Absolute.C a, Arity m, Arity n)
                    => Mat m n a -> Bool
is_upper_triangular = is_upper_triangular' 0


-- | Returns True if the given matrix is lower-triangular, and False
--   otherwise. This is a specialized version of 'is_lower_triangular\''
--   with @epsilon = 0@.
--
--   Examples:
--
--   >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
--   >>> is_lower_triangular m
--   True
--
--   >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
--   >>> is_lower_triangular m
--   False
--
is_lower_triangular :: (Ord a,
                        Ring.C a,
                        Absolute.C a,
                        Arity m,
                        Arity n)
                    => Mat m n a
                    -> Bool
is_lower_triangular = is_upper_triangular . transpose


-- | Returns True if the given matrix is lower-triangular, and False
--   otherwise. The parameter @epsilon@ lets the caller choose a
--   tolerance.
--
--   Examples:
--
--   >>> let m = fromList [[1,1e-12],[1,1]] :: Mat2 Double
--   >>> is_lower_triangular m
--   False
--   >>> is_lower_triangular' 1e-12 m
--   True
--
is_lower_triangular' :: (Ord a,
                         Ring.C a,
                         Absolute.C a,
                         Arity m,
                         Arity n)
                    => a -- ^ The tolerance @epsilon@.
                    -> Mat m n a
                    -> Bool
is_lower_triangular' epsilon = (is_upper_triangular' epsilon) . transpose


-- | Returns True if the given matrix is triangular, and False
--   otherwise.
--
--   Examples:
--
--   >>> let m = fromList [[1,0],[1,1]] :: Mat2 Int
--   >>> is_triangular m
--   True
--
--   >>> let m = fromList [[1,2],[0,3]] :: Mat2 Int
--   >>> is_triangular m
--   True
--
--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> is_triangular m
--   False
--
is_triangular :: (Ord a,
                  Ring.C a,
                  Absolute.C a,
                  Arity m,
                  Arity n)
              => Mat m n a
              -> Bool
is_triangular m = is_upper_triangular m || is_lower_triangular m


-- | Return the (i,j)th minor of m.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> minor m 0 0 :: Mat2 Int
--   ((5,6),(8,9))
--   >>> minor m 1 1 :: Mat2 Int
--   ((1,3),(7,9))
--
minor :: (m ~ S r,
          n ~ S t,
          Arity r,
          Arity t)
      => Mat m n a
      -> Int
      -> Int
      -> Mat r t a
minor (Mat rows) i j = m
  where
    rows' = delete rows i
    m = Mat $ V.map ((flip delete) j) rows'


class (Eq a, Ring.C a) => Determined p a where
  determinant :: (p a) -> a

instance (Eq a, Ring.C a) => Determined (Mat (S Z) (S Z)) a where
  determinant (Mat rows) = (V.head . V.head) rows

instance (Ord a,
          Ring.C a,
          Absolute.C a,
          Arity n,
          Determined (Mat (S n) (S n)) a)
         => Determined (Mat (S (S n)) (S (S n))) a where
  -- | The recursive definition with a special-case for triangular matrices.
  --
  --   Examples:
  --
  --   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
  --   >>> determinant m
  --   -1
  --
  determinant m
    | is_triangular m = product [ m !!! (i,i) | i <- [0..(nrows m)-1] ]
    | otherwise = determinant_recursive
        where
          m' i j = m !!! (i,j)

          det_minor i j = determinant (minor m i j)

          determinant_recursive =
            sum [ (-1)^(toInteger j) NP.* (m' 0 j) NP.* (det_minor 0 j)
              | j <- [0..(ncols m)-1] ]



-- | Matrix multiplication.
--
--   Examples:
--
--   >>> let m1 = fromList [[1,2,3], [4,5,6]]  :: Mat N2 N3 Int
--   >>> let m2 = fromList [[1,2],[3,4],[5,6]] :: Mat N3 N2 Int
--   >>> m1 * m2
--   ((22,28),(49,64))
--
infixl 7 *
(*) :: (Ring.C a, Arity m, Arity n, Arity p)
        => Mat m n a
        -> Mat n p a
        -> Mat m p a
(*) m1 m2 = construct lambda
  where
    lambda i j =
      sum [(m1 !!! (i,k)) NP.* (m2 !!! (k,j)) | k <- [0..(ncols m1)-1] ]



instance (Ring.C a, Arity m, Arity n) => Additive.C (Mat m n a) where

  (Mat rows1) + (Mat rows2) =
    Mat $ V.zipWith (V.zipWith (+)) rows1 rows2

  (Mat rows1) - (Mat rows2) =
    Mat $ V.zipWith (V.zipWith (-)) rows1 rows2

  zero = Mat (V.replicate $ V.replicate (fromInteger 0))


instance (Ring.C a, Arity m, Arity n, m ~ n) => Ring.C (Mat m n a) where
  -- The first * is ring multiplication, the second is matrix
  -- multiplication.
  m1 * m2 = m1 * m2


instance (Ring.C a, Arity m, Arity n) => Module.C a (Mat m n a) where
  -- We can multiply a matrix by a scalar of the same type as its
  -- elements.
  x *> (Mat rows) = Mat $ V.map (V.map (NP.* x)) rows


instance (Algebraic.C a,
          ToRational.C a,
          Arity m)
         => Normed (Mat (S m) N1 a) where
  -- | Generic p-norms for vectors in R^n that are represented as nx1
  --   matrices.
  --
  --   Examples:
  --
  --   >>> let v1 = vec2d (3,4)
  --   >>> norm_p 1 v1
  --   7.0
  --   >>> norm_p 2 v1
  --   5.0
  --
  norm_p p (Mat rows) =
    (root p') $ sum [fromRational' (toRational x)^p' | x <- xs]
    where
      p' = toInteger p
      xs = concat $ V.toList $ V.map V.toList rows

  -- | The infinity norm.
  --
  --   Examples:
  --
  --   >>> let v1 = vec3d (1,5,2)
  --   >>> norm_infty v1
  --   5
  --
  norm_infty (Mat rows) =
    fromRational' $ toRational $ V.maximum $ V.map V.maximum rows


-- | Compute the Frobenius norm of a matrix. This essentially treats
--   the matrix as one long vector containing all of its entries (in
--   any order, it doesn't matter).
--
--   Examples:
--
--   >>> let m = fromList [[1, 2, 3],[4,5,6],[7,8,9]] :: Mat3 Double
--   >>> frobenius_norm m == sqrt 285
--   True
--
--   >>> let m = fromList [[1, -1, 1],[-1,1,-1],[1,-1,1]] :: Mat3 Double
--   >>> frobenius_norm m == 3
--   True
--
frobenius_norm :: (Algebraic.C a, Ring.C a) => Mat m n a -> a
frobenius_norm (Mat rows) =
  sqrt $ element_sum $ V.map row_sum rows
  where
    -- | Square and add up the entries of a row.
    row_sum = element_sum . V.map (^2)


-- Vector helpers. We want it to be easy to create low-dimension
-- column vectors, which are nx1 matrices.

-- | Convenient constructor for 2D vectors.
--
--   Examples:
--
--   >>> import Roots.Simple
--   >>> let fst m = m !!! (0,0)
--   >>> let snd m = m !!! (1,0)
--   >>> let h = 0.5 :: Double
--   >>> let g1 m = 1.0 + h NP.* exp(-((fst m)^2))/(1.0 + (snd m)^2)
--   >>> let g2 m = 0.5 + h NP.* atan((fst m)^2 + (snd m)^2)
--   >>> let g u = vec2d ((g1 u), (g2 u))
--   >>> let u0 = vec2d (1.0, 1.0)
--   >>> let eps = 1/(10^9)
--   >>> fixed_point g eps u0
--   ((1.0728549599342185),(1.0820591495686167))
--
vec1d :: (a) -> Col1 a
vec1d (x) = Mat (mk1 (mk1 x))

vec2d :: (a,a) -> Col2 a
vec2d (x,y) = Mat (mk2 (mk1 x) (mk1 y))

vec3d :: (a,a,a) -> Col3 a
vec3d (x,y,z) = Mat (mk3 (mk1 x) (mk1 y) (mk1 z))

vec4d :: (a,a,a,a) -> Col4 a
vec4d (w,x,y,z) = Mat (mk4 (mk1 w) (mk1 x) (mk1 y) (mk1 z))

vec5d :: (a,a,a,a,a) -> Col5 a
vec5d (v,w,x,y,z) = Mat (mk5 (mk1 v) (mk1 w) (mk1 x) (mk1 y) (mk1 z))

-- Since we commandeered multiplication, we need to create 1x1
-- matrices in order to multiply things.
scalar :: a -> Mat1 a
scalar x = Mat (mk1 (mk1 x))

dot :: (RealRing.C a, n ~ N1, m ~ S t, Arity t)
       => Mat m n a
       -> Mat m n a
       -> a
v1 `dot` v2 = ((transpose v1) * v2) !!! (0, 0)


-- | The angle between @v1@ and @v2@ in Euclidean space.
--
--   Examples:
--
--   >>> let v1 = vec2d (1.0, 0.0)
--   >>> let v2 = vec2d (0.0, 1.0)
--   >>> angle v1 v2 == pi/2.0
--   True
--
angle :: (Transcendental.C a,
          RealRing.C a,
          n ~ N1,
          m ~ S t,
          Arity t,
          ToRational.C a)
          => Mat m n a
          -> Mat m n a
          -> a
angle v1 v2 =
  acos theta
  where
   theta = (recip norms) NP.* (v1 `dot` v2)
   norms = (norm v1) NP.* (norm v2)


-- | Retrieve the diagonal elements of the given matrix as a \"column
--   vector,\" i.e. a m-by-1 matrix. We require the matrix to be
--   square to avoid ambiguity in the return type which would ideally
--   have dimension min(m,n) supposing an m-by-n matrix.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> diagonal m
--   ((1),(5),(9))
--
diagonal :: (Arity m) => Mat m m a -> Col m a
diagonal matrix =
  construct lambda
  where
    lambda i _ = matrix !!! (i,i)


-- | Given a square @matrix@, return a new matrix of the same size
--   containing only the on-diagonal entries of @matrix@. The
--   off-diagonal entries are set to zero.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> diagonal_part m
--   ((1,0,0),(0,5,0),(0,0,9))
--
diagonal_part :: (Arity m, Ring.C a)
         => Mat m m a
         -> Mat m m a
diagonal_part matrix =
  construct lambda
  where
    lambda i j = if i == j then matrix !!! (i,j) else 0


-- | Given a square @matrix@, return a new matrix of the same size
--   containing only the on-diagonal and below-diagonal entries of
--   @matrix@. The above-diagonal entries are set to zero.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> lt_part m
--   ((1,0,0),(4,5,0),(7,8,9))
--
lt_part :: (Arity m, Ring.C a)
        => Mat m m a
        -> Mat m m a
lt_part matrix =
  construct lambda
  where
    lambda i j = if i >= j then matrix !!! (i,j) else 0


-- | Given a square @matrix@, return a new matrix of the same size
--   containing only the below-diagonal entries of @matrix@. The on-
--   and above-diagonal entries are set to zero.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> lt_part_strict m
--   ((0,0,0),(4,0,0),(7,8,0))
--
lt_part_strict :: (Arity m, Ring.C a)
        => Mat m m a
        -> Mat m m a
lt_part_strict matrix =
  construct lambda
  where
    lambda i j = if i > j then matrix !!! (i,j) else 0


-- | Given a square @matrix@, return a new matrix of the same size
--   containing only the on-diagonal and above-diagonal entries of
--   @matrix@. The below-diagonal entries are set to zero.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> ut_part m
--   ((1,2,3),(0,5,6),(0,0,9))
--
ut_part :: (Arity m, Ring.C a)
        => Mat m m a
        -> Mat m m a
ut_part = transpose . lt_part . transpose


-- | Given a square @matrix@, return a new matrix of the same size
--   containing only the above-diagonal entries of @matrix@. The on-
--   and below-diagonal entries are set to zero.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> ut_part_strict m
--   ((0,2,3),(0,0,6),(0,0,0))
--
ut_part_strict :: (Arity m, Ring.C a)
        => Mat m m a
        -> Mat m m a
ut_part_strict = transpose . lt_part_strict . transpose


-- | Compute the trace of a square matrix, the sum of the elements
--   which lie on its diagonal. We require the matrix to be
--   square to avoid ambiguity in the return type which would ideally
--   have dimension min(m,n) supposing an m-by-n matrix.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> trace m
--   15
--
trace :: (Arity m, Ring.C a) => Mat m m a -> a
trace matrix =
  let (Mat rows) = diagonal matrix
  in
    element_sum $ V.map V.head rows


-- | Zip together two column matrices.
--
--   Examples:
--
--   >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
--   >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
--   >>> colzip m1 m2
--   (((1,1)),((1,2)),((1,3)))
--
colzip :: Arity m => Col m a -> Col m a -> Col m (a,a)
colzip c1 c2 =
  construct lambda
  where
    lambda i j = (c1 !!! (i,j), c2 !!! (i,j))


-- | Zip together two column matrices using the supplied function.
--
--   Examples:
--
--   >>> let c1 = fromList [[1],[2],[3]] :: Col3 Integer
--   >>> let c2 = fromList [[4],[5],[6]] :: Col3 Integer
--   >>> colzipwith (^) c1 c2
--   ((1),(32),(729))
--
colzipwith :: Arity m
           => (a -> a -> b)
           -> Col m a
           -> Col m a
           -> Col m b
colzipwith f c1 c2 =
  construct lambda
  where
    lambda i j = f (c1 !!! (i,j)) (c2 !!! (i,j))


-- | Map a function over a matrix of any dimensions.
--
--   Examples:
--
--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> map2 (^2) m
--   ((1,4),(9,16))
--
map2 :: (a -> b) -> Mat m n a -> Mat m n b
map2 f (Mat rows) =
  Mat $ V.map g rows
  where
    g = V.map f


-- | Fold over the entire matrix passing the coordinates @i@ and @j@
--   (of the row/column) to the accumulation function.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>>  ifoldl2 (\i j cur _ -> cur + i + j) 0 m
--   18
--
ifoldl2 :: forall a b m n.
           (Int -> Int -> b -> a -> b)
        -> b
        -> Mat m n a
        -> b
ifoldl2 f initial (Mat rows) =
  V.ifoldl row_function initial rows
  where
    -- | The order that we need this in (so that @g idx@ makes sense)
    --   is a little funny. So that we don't need to pass weird
    --   functions into ifoldl2, we swap the second and third
    --   arguments of @f@ calling the result @g@.
    g :: Int -> b -> Int -> a -> b
    g w x y = f w y x

    row_function :: b -> Int -> Vec n a -> b
    row_function rowinit idx r = V.ifoldl (g idx) rowinit r


-- | Map a function over a matrix of any dimensions, passing the
--   coordinates @i@ and @j@ to the function @f@.
--
--   Examples:
--
--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> imap2 (\i j _ -> i+j) m
--   ((0,1),(1,2))
--
imap2 :: (Int -> Int -> a -> b) -> Mat m n a -> Mat m n b
imap2 f (Mat rows) =
  Mat $ V.imap g rows
  where
    g i = V.imap (f i)