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

module Linear.Matrix
where

import Data.List (intercalate)

import Data.Vector.Fixed (
  Dim,
  Vector
  )
import qualified Data.Vector.Fixed as V (
  Fun(..),
  N1,
  and,
  eq,
  foldl,
  fromList,
  length,
  map,
  maximum,
  replicate,
  toList,
  zipWith
  )
import Data.Vector.Fixed.Internal (Arity, arity, S, Dim)
import Linear.Vector
import Normed

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

data Mat v w a = (Vector v (w a), Vector w a) => Mat (v (w a))
type Mat1 a = Mat D1 D1 a
type Mat2 a = Mat D2 D2 a
type Mat3 a = Mat D3 D3 a
type Mat4 a = Mat D4 D4 a

-- We can't just declare that all instances of Vector are instances of
-- Eq unfortunately. We wind up with an overlapping instance for
-- w (w a).
instance (Eq a, Vector v Bool, Vector w Bool) => Eq (Mat v w 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, Vector v String, Vector w String) => Show (Mat v w 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 v w a -> [[a]]
toList (Mat rows) = map V.toList (V.toList rows)

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


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

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


-- | The number of rows in the matrix.
nrows :: Mat v w a -> Int
nrows (Mat rows) = V.length rows

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

-- | Return the @i@th row of @m@. Unsafe.
row :: Mat v w a -> Int -> w a
row (Mat rows) i = rows ! i


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


-- | Transpose @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.
--
--   TODO: Don't cheat with fromList.
--
--   Examples:
--
--   >>> let m = fromList [[1,2], [3,4]] :: Mat2 Int
--   >>> transpose m
--   ((1,3),(2,4))
--
transpose :: (Vector w (v a),
              Vector v a,
              Vector w a)
             => Mat v w a
             -> Mat w v a
transpose m = Mat $ V.fromList column_list
  where
    column_list = [ column m i | i <- [0..(ncols m)-1] ]


-- | 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 :: (Vector v (w a),
              Vector w a,
              v ~ w,
              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.
--
--   TODO: Don't cheat with fromList.
--
--   Examples:
--
--   >>> let lambda i j = i + j
--   >>> construct lambda :: Mat3 Int
--   ((0,1,2),(1,2,3),(2,3,4))
--
construct :: forall v w a.
             (Vector v (w a),
              Vector w a)
             => (Int -> Int -> a)
             -> Mat v w a
construct lambda = Mat rows
  where
    -- The arity trick is used in Data.Vector.Fixed.length.
    imax = (arity (undefined :: Dim v)) - 1
    jmax = (arity (undefined :: Dim w)) - 1
    row' i = V.fromList [ lambda i j | j <- [0..jmax] ]
    rows = V.fromList [ row' i | i <- [0..imax] ]

-- | 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 a v w.
            (Algebraic.C a,
             Vector v (w a),
             Vector w a,
             Vector v a)
            => (Mat v w a)
            -> (Mat v w 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


-- | Matrix multiplication. Our 'Num' instance doesn't define one, and
--   we need additional restrictions on the result type anyway.
--
--   Examples:
--
--   >>> let m1 = fromList [[1,2,3], [4,5,6]]  :: Mat D2 D3 Int
--   >>> let m2 = fromList [[1,2],[3,4],[5,6]] :: Mat D3 D2 Int
--   >>> m1 * m2
--   ((22,28),(49,64))
--
infixl 7 *
(*) :: (Ring.C a,
         Vector v a,
         Vector w a,
         Vector z a,
         Vector v (z a))
        => Mat v w a
        -> Mat w z a
        -> Mat v z 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,
          Vector v (w a),
          Vector w a)
         => Additive.C (Mat v w 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,
          Vector v (w a),
          Vector w a,
          v ~ w)
         => Ring.C (Mat v w a) where
  -- The first * is ring multiplication, the second is matrix
  -- multiplication.
  m1 * m2 = m1 * m1


instance (Ring.C a,
          Vector v (w a),
          Vector w a)
         => Module.C a (Mat v w 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,
          Vector v (w a),
          Vector w a,
          Vector v a,
          Vector v [a])
         => Normed (Mat v w a) where
  -- | Generic p-norms. The usual norm in R^n is (norm_p 2). We treat
  --   all matrices as big vectors.
  --
  --   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. We don't use V.maximum here because it
  --   relies on a type constraint that the vector be non-empty and I
  --   don't know how to pattern match it away.
  --
  --   Examples:
  --
  --   >>> let v1 = vec3d (1,5,2)
  --   >>> norm_infty v1
  --   5
  --
  norm_infty m@(Mat rows)
    | nrows m == 0 || ncols m == 0 = 0
    | otherwise =
        fromRational' $ toRational $
        P.maximum $ V.toList $ V.map (P.maximum . V.toList) rows





-- 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 h = 0.5 :: Double
--   >>> let g1 (Mat (D2 (D1 x) (D1 y))) = 1.0 + h*exp(-(x^2))/(1.0 + y^2)
--   >>> let g2 (Mat (D2 (D1 x) (D1 y))) = 0.5 + h*atan(x^2 + y^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)
--
vec2d :: (a,a) -> Mat D2 D1 a
vec2d (x,y) = Mat (D2 (D1 x) (D1 y))

vec3d :: (a,a,a) -> Mat D3 D1 a
vec3d (x,y,z) = Mat (D3 (D1 x) (D1 y) (D1 z))

vec4d :: (a,a,a,a) -> Mat D4 D1 a
vec4d (w,x,y,z) = Mat (D4 (D1 w) (D1 x) (D1 y) (D1 z))

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

dot :: (RealRing.C a,
        Dim w ~ V.N1,
        Dim v ~ S n,
        Vector v a,
        Vector w a,
        Vector w (v a),
        Vector w (w a))
       => Mat v w a
       -> Mat v w 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,
          Dim w ~ V.N1,
          Dim v ~ S n,
          Vector w (w a),
          Vector v [a],
          Vector v a,
          Vector w a,
          Vector v (w a),
          Vector w (v a),
          ToRational.C a)
          => Mat v w a
          -> Mat v w a
          -> a
angle v1 v2 =
  acos theta
  where
   theta = (recip norms) NP.* (v1 `dot` v2)
   norms = (norm v1) NP.* (norm v2)