-- Copyright 2009 Mikael Vejdemo Johansson <mik@stanford.edu>
-- Released under a BSD license

-- | Implements the operad element storage using the Haskell native Data.Map storage type.

module Math.Operad.MapOperad where

import qualified Data.Map as Map
import Data.Map (Map)
import Data.Maybe
import Control.Arrow
import Math.Operad.OrderedTree
import Math.Operad.PPrint

-- | The type carrying operadic elements. An element in an operad is an associative array 
-- with keys being labeled trees and values being their coefficients. 
newtype (Ord a, Show a, TreeOrdering t) => OperadElement a n t = OE (Map (OrderedTree a t) n) deriving (Eq, Ord, Show, Read)

instance (Ord a, Show a, Show n, TreeOrdering t) => PPrint (OperadElement a n t) where
      pp (OE m) = if str == "" then "0" else str 
          where str = Map.foldWithKey (\k a result -> result ++ "\n+" ++ show a ++ "*" ++ pp k) "" m

-- | Extracting the internal structure of the an element of the free operad.
extractMap :: (Ord a, Show a, TreeOrdering t) => OperadElement a n t -> Map (OrderedTree a t) n
extractMap (OE m) = m

-- | Arithmetic in the operad.
instance (Ord a, Show a, Num n, TreeOrdering t) => Num (OperadElement a n t) where
    (OE m) + (OE n) = OE $ Map.filter (/=0) $ Map.unionWith (+) m n
    negate on = (-1).*.on
    (OE m) * (OE n) = OE $ Map.filter (/=0) $ Map.intersectionWith (*) m n
    abs _ = undefined
    signum _ = undefined
    fromInteger _ = undefined

-- | Scalar multiplication in the operad.
(.*.) :: (Ord a, Show a, Eq n, Show n, Num n, TreeOrdering t) => n -> OperadElement a n t -> OperadElement a n t
nn .*. (OE m) = OE $ Map.map (nn*) m

-- | Apply a function to each monomial tree in the operad element.
mapMonomials :: (Show a, Ord a, Show b, Ord b, Num n, TreeOrdering s, TreeOrdering t) =>
                (OrderedTree a s -> OrderedTree b t) -> OperadElement a n s -> OperadElement b n t
mapMonomials f (OE m) = OE $ Map.fromListWith (+) $ map (first f) (Map.toList m)

-- | Fold a function over all monomial trees in an operad element, collating the results in a list.
foldMonomials :: (Show a, Ord a, Num n, TreeOrdering t) => 
                 ((OrderedTree a t,n) -> [b] -> [b]) -> OperadElement a n t -> [b]
foldMonomials f (OE m) = Map.foldWithKey (curry f) [] m

-- | Given a list of (tree,coefficient)-pairs, reconstruct the corresponding operad element.
fromList :: (TreeOrdering t, Show a, Ord a, Num n) => [(OrderedTree a t,n)] -> OperadElement a n t
fromList = OE . Map.filter (/=0) . Map.fromListWith (+)

-- | Given an operad element, extract a list of (tree, coefficient) pairs. 
toList :: (TreeOrdering t, Show a, Ord a) => OperadElement a n t -> [(OrderedTree a t, n)]
toList (OE m) = Map.toList m

-- ** Handling polynomials in the free operad

-- | Construct an element in the free operad from its internal structure. Use this instead of the constructor.
oe :: (Ord a, Show a, TreeOrdering t, Num n) => [(OrderedTree a t, n)] -> OperadElement a n t
oe = fromList

-- | Construct a monomial in the free operad from a tree and a tree ordering. It's coefficient will be 1.
oet :: (Ord a, Show a, TreeOrdering t, Num n) => DecoratedTree a -> OperadElement a n t
oet dect = oe $ [(OT dect ordering, 1)]

-- | Construct a monomial in the free operad from a tree, a tree ordering and a coefficient.
oek :: (Ord a, Show a, TreeOrdering t, Num n) => DecoratedTree a -> n -> OperadElement a n t
oek dect n = oe $ [(OT dect ordering, n)]

-- | Return the zero of the corresponding operad, with type appropriate to the given element.
-- Can be given an appropriately casted undefined to construct a zero.
zero :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t 
zero = oe [(ot $ leaf 1, 0)]

-- | Check whether an element is equal to 0. 
isZero :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> Bool
isZero (OE m) = Map.null $ Map.filter (/=0) m

-- | Extract the leading term of an operad element. 
leadingTerm :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> (OrderedTree a t, n)
leadingTerm (OE om) = fromMaybe (ot (leaf 0), 0) $ do 
                        ((m,c),_) <- Map.maxViewWithKey om
                        return (m,c) 

-- | Extract the ordered tree for the leading term of an operad element.
leadingOMonomial :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> OrderedTree a t
leadingOMonomial = fst . leadingTerm

-- | Extract the tree for the leading term of an operad element.
leadingMonomial :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> DecoratedTree a
leadingMonomial = dt .leadingOMonomial

-- | Extract the leading coefficient of an operad element.
leadingCoefficient :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> n
leadingCoefficient = snd . leadingTerm

-- | Extract all occurring monomial trees from an operad element.
getTrees :: (TreeOrdering t, Show a, Ord a) => OperadElement a n t -> [OrderedTree a t]
getTrees (OE m) = Map.keys m