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

{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_HADDOCK prune #-}

-- |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 Prelude hiding (  (<*), (*>) )

import Control.Applicative hiding ( (<*), (*>) )
import Control.Monad (ap)
import qualified Data.List as L
import qualified Data.Set as S -- only needed for toSet

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


-- |Given a field type k and a basis type b, Vect k b is the type of the free k-vector space over b.
-- Elements (values) of Vect k b consist of k-linear combinations of elements (values) of b.
--
-- In order for Vect k b to be a vector space, it is necessary that k is a field (that is, an instance of Fractional).
-- In practice, we often relax this condition, and require that k is a ring (that is, an instance of Num). In that case,
-- Vect k b should more correctly be called (the type of) the free k-module over b.
--
-- Most of the code requires that b is an instance of Ord. This is primarily to enable us to simplify to a normal form.
newtype Vect k b = V [(b,k)] deriving (Eq,Ord)

instance (Show k, Eq k, 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

-- |Return the coefficient of the specified basis element in a vector
coeff :: (Num k, Eq b) => b -> Vect k b -> k
coeff b v = sum [k | (b',k) <- terms v, b' == b]

-- |Remove the term for a specified basis element from a vector
removeTerm :: (Eq k, Num k, Ord b) => b -> Vect k b -> Vect k b
removeTerm b (V ts) = V $ filter ((/=b) . fst) ts
-- v <-> coeff b v *> return b

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

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

-- |Addition of vectors (same as add)
(<+>) :: (Eq k, Num k, Ord b) => 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 :: (Eq k, Num k, Ord b) => [Vect k b] -> Vect k b
sumv = foldl (<+>) zerov

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

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

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

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

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

-- |Same as smultR. Mnemonic is \"multiply through (from the right)\"
(<*) :: (Eq k, 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 :: (Eq k, Num k, Ord b) => 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) )


-- |Given a field k, (Vect k) is a functor, the \"free k-vector space\" functor.
--
-- In the mathematical sense, this can be regarded as a functor from the category Set (of sets) to the category k-Vect
-- (of k-vector spaces). In Haskell, instead of Set we have Hask, the category of Haskell types. However, for our purposes
-- it is helpful to identify Hask with Set, by identifying a Haskell type with its set of inhabitants.
--
-- The type constructor (Vect k) gives the action of the functor on objects in the category,
-- taking a set (type) to a free k-vector space. fmap gives the action of the functor on arrows in the category,
-- taking a function between sets (types) to a linear map between vector spaces.
--
-- Note that if f is not order-preserving, then (fmap f) is not guaranteed to return results in normal form,
-- so it may be preferable to use (nf . fmap f).
instance Functor (Vect k) where
    -- lift a function on the basis to a function on the vector space
    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

-- From GHC 7.10, Monad has Applicative as a superclass, so we must define an instance.
-- It doesn't particularly make sense for Vect k.
-- (Although given Vect k b, we could represent the dual space as Vect k (b -> ()),
-- and then have a use for <*>.)
instance Num k => Applicative (Vect k) where
    pure = return
    -- pure b = V [(b,1)]
    (<*>) = ap
    -- V fs <*> V xs = V [(f x, a*b) | (f,a) <- fs, (x,b) <- xs]

-- |Given a field k, the type constructor (Vect k) is a monad, the \"free k-vector space monad\".
--
-- In order to understand this, it is probably easiest to think of a free k-vector space as a kind of container,
-- a bit like a list, except that order doesn't matter, and you're allowed arbitrary (even negative or fractional)
-- quantities of the basis elements in the container.
--
-- According to this way of thinking, return is the function that puts a basis element into the vector space (container).
--
-- Given a function f from the basis of one vector space to another vector space (a -> Vect k b),
-- bind (>>=) lifts it to a function (>>= f) from the first vector space to the second (Vect k a -> Vect k b).
--
-- Note that in general (>>= f) applied to a vector will not return a result in normal form,
-- so it is usually preferable to use (linear f) instead.
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]
    V ts >>= f = V $ [ (b,y*x) | (a,x) <- ts, let V us = f a, (b,y) <- us]
    -- Note that as we can't assume Ord a in the Monad instance, we need to call "nf" afterwards

-- |A linear map between vector spaces A and B can be defined by giving its action on the basis elements of A.
-- The action on all elements of A then follows by linearity.
--
-- If we have A = Vect k a, B = Vect k b, and f :: a -> Vect k b is a function from the basis elements of A into B,
-- then @linear f@ is the linear map that this defines by linearity.
linear :: (Eq k, Num k, Ord b) => (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 an element of the field k to an element of the trivial k-vector space
wrap :: (Eq k, Num k) => k -> Vect k ()
wrap 0 = zerov
wrap x = V [( (),x)]

-- |Unwrap an element of the trivial k-vector space to an element of the field k
unwrap :: Num k => Vect k () -> k
unwrap (V []) = 0
unwrap (V [( (),x)]) = x

-- |Given a finite vector space basis b, Dual b can be used to represent a basis for the dual vector space.
-- The intention is that for a given individual basis element b_i, (Dual b_i) represents the indicator function for b_i,
-- which takes b_i to 1 and all other basis elements to 0.
--
-- (Note that if the basis b is infinite, then Dual b may only represent a sub-basis of the dual vector space.)
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

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

zerof v = zerov

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


-- Lens
coeffLens :: (Ord b, Eq k, Num k, Functor f) => b -> (k -> f k) -> (Vect k b -> f (Vect k b))
coeffLens b = lens (coeff b) (setter b)
    where setter b = \(V ts) k -> (k *> return b) <+> (V $ filter ((/=b) . fst) ts)
          lens getter setter f a = fmap (setter a) (f (getter a))
-- Can be used with lens-family, for example
-- e1 ^. coeffLens (E 2) --> 0
-- e1 & coeffLens (E 2) .~ 2 --> e1+2e2
-- e1 & coeffLens (E 1) %~ (+2) --> 3e1