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

-- | Implements the operad element storage using a class that tries to delay all comparisons as long
-- as possible, by maintaining the initial term of any operad element in a separate storage. 

module Math.Operad.PolyBag where

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

-- | The type carrying operadic elements. An element in an operad is the leading monomial tree, its coefficient,
-- and a list of all other elements stored as (tree, coefficient) pairs. 
data (Show a, Ord a, Num n, TreeOrdering t) => OperadElement a n t = PB (OrderedTree a t) n [(OrderedTree a t,n)] deriving (Ord, Eq, Show, Read)

instance (Show a, Ord a, Num n, TreeOrdering t) => Num (OperadElement a n t) where
  a@(PB ma ca baga) + b@(PB mb cb bagb)
      | ma > mb = PB ma ca (baga ++ ((mb,cb):bagb))
      | ma < mb = b + a
      | ca+cb /= 0 = PB ma (ca+cb) (baga++bagb)
      | otherwise = let
        combinedMap = Map.fromListWith (+) (baga ++ bagb)
        maybeSum = Map.maxViewWithKey combinedMap
   in
     if isNothing maybeSum then PB ma 0 [] 
     else let
         ((mP,cP),mapP) = fromJust maybeSum
    in
      PB mP cP (Map.toList mapP)
  (*) = undefined
  negate pb = (-1) .*. pb
  abs = undefined
  signum = undefined
  fromInteger = undefined 

-- | Collapse the storage, removing duplicates from the list carrying the tail of the element.
collate :: (Show a, Ord a, Num n, TreeOrdering t) => OperadElement a n t -> OperadElement a n t
collate  = fromList . toList

-- | Given a list of (tree,coefficient)-pairs, reconstruct the corresponding operad element.
fromList :: (TreeOrdering t, Num n, Ord a, Show a) => [(OrderedTree a t,n)] -> OperadElement a n t
fromList lst = fromMaybe (PB (ot $ leaf 1) 0 []) $ do
    ((mP,cP),mapP) <- Map.maxViewWithKey (Map.fromList lst)
    return $ PB mP cP (Map.toList mapP)

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

-- | 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 (PB m c bag) = collate (PB (f m) c (map (first f) bag))

-- | 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 (PB m c bag) = foldr f [] ((m,c):bag)

instance (Ord a, Show a, Num n, TreeOrdering t) => PPrint (OperadElement a n t) where
      pp m = if str == "" then "0" else str 
          where str = foldMonomials (\(k,a) pstr -> pstr ++ "\n+" ++ show a ++ "*" ++ pp k) m

-- | Extract all occurring monomial trees from an operad element.
getTrees :: (Ord a, Show a, TreeOrdering t, Num n) =>
            OperadElement a n t -> [OrderedTree a t]
getTrees (PB m _ bag) = nub $ m : (map fst bag)

-- | Scalar multiplication.
(.*.) :: (Show a, Ord a, Num n, TreeOrdering t) => n -> OperadElement a n t -> OperadElement a n t
0 .*. (PB ma _ _) = PB ma 0 []
x .*. (PB ma ca baga) = PB ma (x*ca) (map (\(m,c) -> (m,x*c)) baga)



-- ** 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 = PB (ot dect) 1 [] -- oe $ Map.singleton (OT dt o) 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 = PB (ot dect) n [] -- oe $ Map.singleton (OT dt o) 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 = PB (ot $ leaf 1) 0 [] -- oe (Map.empty)

-- | Check whether an element is equal to 0. 
isZero :: (Ord a, Show a, TreeOrdering t, Num n) => OperadElement a n t -> Bool
isZero m = 0 == leadingCoefficient m -- Map.null 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 (PB m c _) = (m,c) -- Map.findMax $ m --  (t, m Map.! t) where t = maximum $ Map.keys m --

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