-- Copyright (c) 2010, David Amos. All rights reserved.

{-# LANGUAGE NoMonomorphismRestriction #-}

-- |A module defining the type and operations of free k-vector spaces over a basis b (for a field k)
module Math.Algebras.VectorSpace where

import qualified Data.List as L
import qualified Data.Set as S -- only needed for toSet

-- toSet = S.toList . S.fromList

infixr 7 *>
infixl 7 <*
infixl 6 <+>, <->


-- |Given a field type k (ie a Fractional instance), Vect k b is the type of the free k-vector space over the basis type b.
-- Elements of Vect k b consist of k-linear combinations of elements of b.
newtype Vect k b = V [(b,k)] deriving (Eq,Ord)

instance (Num k, Show b) => Show (Vect k b) where
    show (V []) = "0"
    show (V ts) = concatWithPlus $ map showTerm ts
        where showTerm (b,x) | show b == "1" = show x
                             | show x == "1" = show b
                             | show x == "-1" = "-" ++ show b
                             | otherwise = (if isAtomic (show x) then show x else "(" ++ show x ++ ")") ++
                                           show b
                                           -- (if ' ' `notElem` show b then show b else "(" ++ show b ++ ")")
                                           -- if we put this here we miss the two cases above
              concatWithPlus (t1:t2:ts) = if head t2 == '-'
                                          then t1 ++ concatWithPlus (t2:ts)
                                          else t1 ++ '+' : concatWithPlus (t2:ts)
              concatWithPlus [t] = t
              isAtomic (c:cs) = isAtomic' cs
              isAtomic' ('^':'-':cs) = isAtomic' cs
              isAtomic' ('+':cs) = False
              isAtomic' ('-':cs) = False
              isAtomic' (c:cs) = isAtomic' cs
              isAtomic' [] = True

terms (V ts) = ts

coeff b v = sum [k | (b',k) <- terms v, b' == b]

-- Deprecated
zero = V []

-- |The zero vector
zerov :: Vect k b
zerov = V []

-- |Addition of vectors
add :: (Ord b, Num k) => Vect k b -> Vect k b -> Vect k b
add (V ts) (V us) = V $ addmerge ts us

-- |Addition of vectors (same as add)
(<+>) :: (Ord b, Num k) => Vect k b -> Vect k b -> Vect k b
(<+>) = add

addmerge ((a,x):ts) ((b,y):us) =
    case compare a b of
    LT -> (a,x) : addmerge ts ((b,y):us)
    EQ -> if x+y == 0 then addmerge ts us else (a,x+y) : addmerge ts us
    GT -> (b,y) : addmerge ((a,x):ts) us
addmerge ts [] = ts
addmerge [] us = us

-- |Sum of a list of vectors
sumv :: (Ord b, Num k) => [Vect k b] -> Vect k b
sumv = foldl (<+>) zerov

-- |Negation of vector
neg :: (Num k) => Vect k b -> Vect k b
neg (V ts) = V $ map (\(b,x) -> (b,-x)) ts

-- |Subtraction of vectors
(<->) :: (Ord b, Num k) => Vect k b -> Vect k b -> Vect k b
(<->) u v = u <+> neg v

-- |Scalar multiplication (on the left)
smultL :: (Num k) => k -> Vect k b -> Vect k b
smultL 0 _ = zero -- V []
smultL k (V ts) = V [(ei,k*xi) | (ei,xi) <- ts]

-- |Same as smultL. Mnemonic is "multiply through (from the left)"
(*>) :: (Num k) => k -> Vect k b -> Vect k b
(*>) = smultL

-- |Scalar multiplication on the right
smultR :: (Num k) => Vect k b -> k -> Vect k b
smultR _ 0 = zero -- V []
smultR (V ts) k = V [(ei,xi*k) | (ei,xi) <- ts]

-- |Same as smultR. Mnemonic is "multiply through (from the right)"
(<*) :: (Num k) => Vect k b -> k -> Vect k b
(<*) = smultR

-- same as return
-- injection of basis elt into vector space
-- inject b = V [(b,1)]

-- same as fmap
-- liftFromBasis f (V ts) = V [(f b, x) | (b, x) <- ts]
-- if f is not order-preserving, then you need to call nf afterwards

-- |Convert an element of Vect k b into normal form. Normal form consists in having the basis elements in ascending order,
-- with no duplicates, and all coefficients non-zero
nf :: (Ord b, Num k) => Vect k b -> Vect k b
nf (V ts) = V $ nf' $ L.sortBy compareFst ts where
    nf' ((b1,x1):(b2,x2):ts) =
        case compare b1 b2 of
        LT -> if x1 == 0 then nf' ((b2,x2):ts) else (b1,x1) : nf' ((b2,x2):ts)
        EQ -> if x1+x2 == 0 then nf' ts else nf' ((b1,x1+x2):ts)
        GT -> error "nf': not pre-sorted"
    nf' [(b,x)] = if x == 0 then [] else [(b,x)]
    nf' [] = []
    compareFst (b1,x1) (b2,x2) = compare b1 b2
    -- compareFst = curry ( uncurry compare . (fst *** fst) )


-- lift a function on the basis to a function on the vector space
instance Functor (Vect k) where
    fmap f (V ts) = V [(f b, x) | (b,x) <- ts]
-- Note that if f is not order-preserving, then we need to call "nf" afterwards

instance Num k => Monad (Vect k) where
    return a = V [(a,1)]
    V ts >>= f = V $ concat [ [(b,y*x) | let V us = f a, (b,y) <- us] | (a,x) <- ts]
    -- Note that as we can't assume Ord a in the Monad instance, we need to call "nf" afterwards

linear :: (Ord b, Num k) => (a -> Vect k b) -> Vect k a -> Vect k b
linear f v = nf $ v >>= f

newtype EBasis = E Int deriving (Eq,Ord)

instance Show EBasis where show (E i) = "e" ++ show i

e i = return $ E i
e1 = e 1
e2 = e 2
e3 = e 3

-- dual (E i) = E (-i)


-- |Trivial k is the field k considered as a k-vector space. In maths, we would not normally make a distinction here,
-- but in the code, we need this if we want to be able to put k as one side of a tensor product.
type Trivial k = Vect k ()

wrap :: Num k => k -> Vect k ()
wrap 0 = zero
wrap x = V [( (),x)]

unwrap :: Num k => Vect k () -> k
unwrap (V []) = 0
unwrap (V [( (),x)]) = x

-- |Given a finite vector space basis b, Dual b represents a basis for the dual vector space. (If b is infinite, then Dual b is only a sub-basis.)
newtype Dual b = Dual b deriving (Eq,Ord)

instance Show basis => Show (Dual basis) where
    show (Dual b) = show b ++ "'"


e' i = return $ Dual $ E i
e1' = e' 1
e2' = e' 2
e3' = e' 3

dual :: Vect k b -> Vect k (Dual b)
dual = fmap Dual


(f <<+>> g) v = f v <+> g v

zerof v = zerov

sumf fs = foldl (<<+>>) zerof fs