```-- 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.

#ifndef USE_OLDMAP
#else
import qualified Data.Map as Map
import Data.Map (Map)
#endif
import Data.Maybe

#ifndef USE_OLDMAP
type MonomialMap a t n = Map a t n
#else
type MonomialMap a t n = Map (OrderedTree a t) n
#endif

-- | 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 (MonomialMap 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 -> MonomialMap 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.mapKeysWith (+) f 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

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