{-# 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 (
  (!),
  generate,
  mk0,
  mk1,
  mk2,
  mk3,
  mk4,
  mk5 )
import qualified Data.Vector.Fixed as V (
  and,
  foldl,
  fromList,
  head,
  ifoldl,
  ifoldr,
  imap,
  map,
  replicate,
  reverse,
  toList,
  zipWith )
import Data.Vector.Fixed.Cont ( Arity, arity )
import Linear.Vector ( Vec, delete )
import Naturals
import Normed ( Normed(..) )

-- We want the "max" that works on Ord, not the one that only works on
-- Bool/Integer from the Lattice class!
import NumericPrelude hiding ( (*), abs, max)
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.Field as Field ( C )
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, max)

-- | 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 Mat0 a = Mat Z Z a
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 Mat6 a = Mat N6 N6 a
type Mat7 a = Mat N7 N7 a

-- * Type synonyms for 1-by-n row "vectors".

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

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 synonyms for n-by-1 column "vectors".

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

type Col0 a = Col Z a
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
type Col6 a = Col N6 a
type Col7 a = Col N7 a
type Col8 a = Col N8 a
type Col9 a = Col N9 a
type Col10 a = Col N10 a
type Col11 a = Col N11 a
type Col12 a = Col N12 a
type Col13 a = Col N13 a
type Col14 a = Col N14 a
type Col15 a = Col N15 a
type Col16 a = Col N16 a
type Col17 a = Col N17 a
type Col18 a = Col N18 a
type Col19 a = Col N19 a
type Col20 a = Col N20 a
type Col21 a = Col N21 a
type Col22 a = Col N22 a
type Col23 a = Col N23 a
type Col24 a = Col N24 a
type Col25 a = Col N25 a
type Col26 a = Col N26 a
type Col27 a = Col N27 a
type Col28 a = Col N28 a
type Col29 a = Col N29 a
type Col30 a = Col N30 a
type Col31 a = Col N31 a
type Col32 a = Col N32 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 rows_one) == (Mat rows_two) =
    V.and $ V.zipWith comp rows_one rows_two
    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
--   >>> let r = cholesky m1
--   >>> frobenius_norm ((transpose r)*r - m1) < 1e-10
--   True
--   >>> is_upper_triangular r
--   True
--
--   >>> let k1 = [6, -3, 0, 0, 0, 0, 0] :: [Double]
--   >>> let k2 = [-3, 10.5, -7.5, 0, 0, 0, 0] :: [Double]
--   >>> let k3 = [0, -7.5, 12.5, 0, 0, 0, 0] :: [Double]
--   >>> let k4 = [0, 0, 0, 6, 0, 0, 0] :: [Double]
--   >>> let k5 = [0, 0, 0, 0, 6, 0, 0] :: [Double]
--   >>> let k6 = [0, 0, 0, 0, 0, 6, 0] :: [Double]
--   >>> let k7 = [0, 0, 0, 0, 0, 0, 15] :: [Double]
--   >>> let big_K = fromList [k1,k2,k3,k4,k5,k6,k7] :: Mat N7 N7 Double
--
--   >>> let e1 = [2.449489742783178,0,0,0,0,0,0] :: [Double]
--   >>> let e2 = [-1.224744871391589,3,0,0,0,0,0] :: [Double]
--   >>> let e3 = [0,-5/2,5/2,0,0,0,0] :: [Double]
--   >>> let e4 = [0,0,0,2.449489742783178,0,0,0] :: [Double]
--   >>> let e5 = [0,0,0,0,2.449489742783178,0,0] :: [Double]
--   >>> let e6 = [0,0,0,0,0,2.449489742783178,0] :: [Double]
--   >>> let e7 = [0,0,0,0,0,0,3.872983346207417] :: [Double]
--   >>> let expected = fromList [e1,e2,e3,e4,e5,e6,e7] :: Mat N7 N7 Double
--
--   >>> let r = cholesky big_K
--   >>> frobenius_norm (r - (transpose expected)) < 1e-12
--   True
--
cholesky :: forall m a. (Algebraic.C a, Arity m)
         => (Mat m m a) -> (Mat m m a)
cholesky m = ifoldl2 f zero m
  where
    f :: Int -> Int -> (Mat m m a) -> a -> (Mat m m a)
    f i j cur_R _ = set_idx cur_R (i,j) (r cur_R i j)

    r :: (Mat m m a) -> Int -> Int -> a
    r cur_R i j
      | i == j = sqrt(m !!! (i,j) - sum [(cur_R !!! (k,i))^2 | k <- [0..i-1]])
      | i < j = (((m !!! (i,j))
                  - sum [(cur_R !!! (k,i)) NP.* (cur_R !!! (k,j))
                             | k <- [0..i-1]]))/(cur_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
--
is_upper_triangular' :: forall m n a.
                        (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 matrix =
  ifoldl2 f True matrix
  where
    f :: Int -> Int -> Bool -> a -> Bool
    f _ _ False _ = False
    f i j True x
      | i <= j = True
      -- use "less than or equal to" so zero is a valid epsilon
      | otherwise = abs x <= epsilon


-- | Returns True if the given matrix is upper-triangular, and False
--   otherwise. We don't delegate to the general
--   'is_upper_triangular'' here because it imposes additional
--   typeclass constraints throughout the library.
--
--   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
--
is_upper_triangular :: forall m n a.
                       (Eq a, Ring.C a, Arity m, Arity n)
                    => Mat m n a -> Bool
is_upper_triangular matrix =
  ifoldl2 f True matrix
  where
    f :: Int -> Int -> Bool -> a -> Bool
    f _ _ False _ = False
    f i j True x
      | i <= j = True
      | otherwise = x == 0



-- | Returns True if the given matrix is lower-triangular, and False
--   otherwise.
--
--   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 :: (Eq a,
                        Ring.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


-- | Delete the @i@th row and @j@th column from the matrix. The name
--   \"preminor\" is made up, but is meant to signify that this is
--   usually used in the computationof a minor. A minor is simply the
--   determinant of a preminor in that case.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> preminor m 0 0 :: Mat2 Int
--   ((5,6),(8,9))
--   >>> preminor m 1 1 :: Mat2 Int
--   ((1,3),(7,9))
--
preminor :: (Arity m, Arity n)
      => Mat (S m) (S n) a
      -> Int
      -> Int
      -> Mat m n a
preminor (Mat rows) i j = m
  where
    rows' = delete rows i
    m = Mat $ V.map ((flip delete) j) rows'


-- | Compute the i,jth minor of a @matrix@.
--
--   Examples:
--
--   >>> let m1 = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Double
--   >>> minor m1 1 1
--   -12.0
--
minor :: (Arity m, Determined (Mat m m) a)
      => Mat (S m) (S m) a
      -> Int
      -> Int
      -> a
minor matrix i j = determinant (preminor matrix i j)

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 = unscalar

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)

          determinant_recursive =
            sum [ (-1)^(toInteger j) NP.* (m' 0 j) NP.* (minor m 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 (S m) (S n) a
        -> Mat (S n) (S p) a
        -> Mat (S m) (S p) a
(*) m1 m2 = construct lambda
  where
    lambda i j = (transpose $ row m1 i) `dot` (column m2 j)



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

  (Mat rows_one) + (Mat rows_two) =
    Mat $ V.zipWith (V.zipWith (+)) rows_one rows_two

  (Mat rows_one) - (Mat rows_two) =
    Mat $ V.zipWith (V.zipWith (-)) rows_one rows_two

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


instance (Ring.C a, Arity m, Arity n, m ~ n) => Ring.C (Mat (S m) (S n) a) where
  -- The first * is ring multiplication, the second is matrix
  -- multiplication.
  one = identity_matrix
  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 (Absolute.C a,
          Algebraic.C a,
          ToRational.C a,
          Arity m)
         => Normed (Col m a) where
  -- | Generic p-norms for vectors in R^m that are represented as m-by-1
  --   matrices.
  --
  --   Examples:
  --
  --   >>> let v1 = vec2d (3,4)
  --   >>> norm_p 1 v1
  --   7.0
  --   >>> norm_p 2 v1
  --   5.0
  --
  --   >>> let v1 = vec2d (-1,1) :: Col2 Double
  --   >>> norm_p 1 v1 :: Double
  --   2.0
  --
  --   >>> let v1 = vec0d :: Col0 Double
  --   >>> norm v1
  --   0.0
  --
  norm_p p (Mat rows) =
    (root p') $ sum [fromRational' (toRational $ abs 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.foldl P.max 0) $ V.map (V.foldl P.max 0) 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 :: (Arity m, Arity n, Algebraic.C a, Ring.C a)
               => Mat m n a
               -> a
frobenius_norm matrix =
  sqrt $ element_sum2 $ squares
  where
    squares = map2 (^2) matrix


-- 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))
--
vec0d :: Col0 a
vec0d = Mat mk0

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))

-- Get the scalar value out of a 1x1 matrix.
unscalar :: Mat1 a -> a
unscalar (Mat rows) = V.head $ V.head rows


dot :: (Ring.C a, Arity m)
       => Col (S m) a
       -> Col (S m) a
       -> a
v1 `dot` v2 = element_sum2 $ zipwith2 (NP.*) v1 v2


-- | 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,
          m ~ S t,
          Arity t,
          ToRational.C a)
          => Col m a
          -> Col m 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 = element_sum2 . diagonal



-- | Zip together two matrices.
--
--   TODO: don't cheat with construct (map V.zips instead).
--
--   Examples:
--
--   >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
--   >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
--   >>> zip2 m1 m2
--   (((1,1)),((1,2)),((1,3)))
--
--   >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> let m2 = fromList [[1,1],[1,1]] :: Mat2 Int
--   >>> zip2 m1 m2
--   (((1,1),(2,1)),((3,1),(4,1)))
--
zip2 :: (Arity m, Arity n) => Mat m n a -> Mat m n b -> Mat m n (a,b)
zip2 m1 m2 =
  construct lambda
  where
    lambda i j = (m1 !!! (i,j), m2 !!! (i,j))


-- | Zip together three matrices.
--
--   TODO: don't cheat with construct (map V.zips instead).
--
--   Examples:
--
--   >>> let m1 = fromList [[1],[1],[1]] :: Col3 Int
--   >>> let m2 = fromList [[1],[2],[3]] :: Col3 Int
--   >>> let m3 = fromList [[4],[5],[6]] :: Col3 Int
--   >>> zip2three m1 m2 m3
--   (((1,1,4)),((1,2,5)),((1,3,6)))
--
--   >>> let m1 = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> let m2 = fromList [[1,1],[1,1]] :: Mat2 Int
--   >>> let m3 = fromList [[8,2],[6,3]] :: Mat2 Int
--   >>> zip2three m1 m2 m3
--   (((1,1,8),(2,1,2)),((3,1,6),(4,1,3)))
--
zip2three :: (Arity m, Arity n)
          => Mat m n a
          -> Mat m n a
          -> Mat m n a
          -> Mat m n (a,a,a)
zip2three m1 m2 m3 =
  construct lambda
  where
    lambda i j = (m1 !!! (i,j), m2 !!! (i,j), m3 !!! (i,j))


-- | Zip together two matrices using the supplied function.
--
--   Examples:
--
--   >>> let c1 = fromList [[1],[2],[3]] :: Col3 Integer
--   >>> let c2 = fromList [[4],[5],[6]] :: Col3 Integer
--   >>> zipwith2 (^) c1 c2
--   ((1),(32),(729))
--
zipwith2 :: (Arity m, Arity n)
           => (a -> b -> c)
           -> Mat m n a
           -> Mat m n b
           -> Mat m n c
zipwith2 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. The fold occurs
--   from top-left to bottom-right.
--
--   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


-- | Left fold over the entries of a matrix (top-left to bottom-right).
--
foldl2 :: forall a b m n.
          (b -> a -> b)
        -> b
        -> Mat m n a
        -> b
foldl2 f initial matrix =
  -- Use the index fold but ignore the index arguments.
  let g _ _ = f in ifoldl2 g initial matrix


-- | Fold over the entire matrix passing the coordinates @i@ and @j@
--   (of the row/column) to the accumulation function. The fold occurs
--   from bottom-right to top-left.
--
--   The order of the arguments in the supplied function are different
--   from those in V.ifoldr; we keep them similar to ifoldl2.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> ifoldr2 (\i j cur _ -> cur + i + j) 0 m
--   18
--
ifoldr2 :: forall a b m n.
           (Int -> Int -> b -> a -> b)
        -> b
        -> Mat m n a
        -> b
ifoldr2 f initial (Mat rows) =
  V.ifoldr row_function initial rows
  where
    -- | Swap the order of arguments in @f@ so that it agrees with the
    --   @f@ passed to ifoldl2.
    g :: Int -> Int -> a -> b -> b
    g w x y z = f w x z y

    row_function :: Int -> Vec n a -> b -> b
    row_function idx r rowinit = V.ifoldr (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)


-- | Reverse the order of elements in a matrix.
--
--   Examples:
--
--   >>> let m1 = fromList [[1,2,3]] :: Row3 Int
--   >>> reverse2 m1
--   ((3,2,1))
--
--   >>> let m1 = vec3d (1,2,3 :: Int)
--   >>> reverse2 m1
--   ((3),(2),(1))
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> reverse2 m
--   ((9,8,7),(6,5,4),(3,2,1))
--
reverse2 :: (Arity m, Arity n) => Mat m n a -> Mat m n a
reverse2 (Mat rows) = Mat $ V.reverse $ V.map V.reverse rows


-- | Unsafely set the (i,j) element of the given matrix.
--
--   Examples:
--
--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
--   >>> set_idx m (1,1) 17
--   ((1,2,3),(4,17,6),(7,8,9))
--
set_idx :: forall m n a.
           (Arity m, Arity n)
        => Mat m n a
        -> (Int, Int)
        -> a
        -> Mat m n a
set_idx matrix (i,j) newval =
  imap2 updater matrix
  where
    updater :: Int -> Int -> a -> a
    updater k l existing =
      if k == i && l == j
      then newval
      else existing


-- | Compute the i,jth cofactor of the given @matrix@. This simply
--   premultiplues the i,jth minor by (-1)^(i+j).
cofactor :: (Arity m, Determined (Mat m m) a)
         => Mat (S m) (S m) a
         -> Int
         -> Int
         -> a
cofactor matrix i j =
  (-1)^(toInteger i + toInteger j) NP.* (minor matrix i j)


-- | Compute the inverse of a matrix using cofactor expansion
--   (generalized Cramer's rule).
--
--   Examples:
--
--   >>> let m1 = fromList [[37,22],[17,54]] :: Mat2 Double
--   >>> let e1 = [54/1624, -22/1624] :: [Double]
--   >>> let e2 = [-17/1624, 37/1624] :: [Double]
--   >>> let expected = fromList [e1, e2] :: Mat2 Double
--   >>> let actual = inverse m1
--   >>> frobenius_norm (actual - expected) < 1e-12
--   True
--
inverse :: (Arity m,
            Determined (Mat (S m) (S m)) a,
            Determined (Mat m m) a,
            Field.C a)
        => Mat (S m) (S m) a
        -> Mat (S m) (S m) a
inverse matrix =
  (1 / (determinant matrix)) *> (transpose $ construct lambda)
  where
    lambda i j = cofactor matrix i j



-- | Retrieve the rows of a matrix as a column matrix. If the given
--   matrix is m-by-n, the result would be an m-by-1 column whose
--   entries are 1-by-n row matrices.
--
--   Examples:
--
--   >>> let m = fromList [[1,2],[3,4]] :: Mat2 Int
--   >>> (rows2 m) !!! (0,0)
--   ((1,2))
--   >>> (rows2 m) !!! (1,0)
--   ((3,4))
--
rows2 :: (Arity m, Arity n)
      => Mat m n a
      -> Col m (Row n a)
rows2 (Mat rows) =
  Mat $ V.map (mk1. Mat . mk1) rows



-- | Sum the elements of a matrix.
--
--   Examples:
--
--   >>> let m = fromList [[1,-1],[3,4]] :: Mat2 Int
--   >>> element_sum2 m
--   7
--
element_sum2 :: (Arity m, Arity n, Additive.C a) => Mat m n a -> a
element_sum2 = foldl2 (+) zero