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

-- | Implements decorated and ordered trees, and most operations acting only on these.
-- The decorated tree is the most fundamental part of the implementation, and the internal
-- structure of the tree is to the largest extent what determines the actual identity
-- of a given element of the free operad.

import Prelude hiding (mapM)
import Data.Foldable (Foldable, foldMap)
import Data.Traversable
import Data.List (sort, sortBy, intersperse, (\\))
import Control.Applicative
import Data.Ord
import Data.Monoid

-- * Decorated and ordered trees

-- | The fundamental tree data type used. Leaves carry labels - most often integral -
-- and these are expected to control, e.g., composition points in shuffle operad compositions.
-- The vertices carry labels, used for the ordering on trees and to distinguish different
-- basis corollas of the same arity.
data (Ord a, Show a) => PreDecoratedTree a b = DTLeaf !b |
DTVertex {
vertexType :: !a,
subTrees :: ![PreDecoratedTree a b]}

-- | The arity of a corolla
vertexArity :: (Ord a, Show a) => PreDecoratedTree a b -> Int
vertexArity t = length (subTrees t)

instance (Ord a, Show a) => Functor (PreDecoratedTree a) where
fmap f (DTLeaf b) = DTLeaf (f b)
fmap f (DTVertex t ts) = DTVertex t (map (fmap f) ts)

instance (Ord a, Show a) => Foldable (PreDecoratedTree a) where
foldMap f (DTLeaf b) = f b
foldMap f (DTVertex _ ts) = mconcat (map (foldMap f) ts)

instance (Ord a, Show a) => Traversable (PreDecoratedTree a) where
traverse f (DTLeaf b) = DTLeaf <\$> f b
traverse f (DTVertex t ts) = DTVertex t <\$> traverse (traverse f) ts

instance (Ord a, Show a, Show b) => PPrint (PreDecoratedTree a b) where
pp (DTLeaf x) = show x
pp (DTVertex t ts) = "m" ++ show t ++ "(" ++ concat (intersperse "," (map pp ts)) ++ ")"

-- | If a tree has trees as labels for its leaves, we can replace the leaves with the roots of
-- those label trees. Thus we may glue together trees, as required by the compositions.
glueTrees :: (Ord a, Show a) => PreDecoratedTree a (PreDecoratedTree a b) -> PreDecoratedTree a b
glueTrees (DTLeaf newRoot) = newRoot
glueTrees (DTVertex t ts) = DTVertex t (map glueTrees ts)

instance (Ord a, Show a) => Monad (PreDecoratedTree a) where
return = DTLeaf;
a >>= f = glueTrees (fmap f a)

-- | This is the fundamental datatype of the whole project. Monomials in a free operad
-- are decorated trees, and we build a type for decorated trees here. We require our
-- trees to have Int labels, limiting us to at most 2 147 483 647 leaf labels.
type DecoratedTree a = PreDecoratedTree a Int

-- | Monomial orderings on the free operad requires us to choose an ordering for the
-- trees. These are parametrized by types implementing the type class 'TreeOrdering',
-- and this is a data type for a tree carrying its comparison type. We call these
-- /ordered trees/.
data (Ord a, Show a, TreeOrdering t) => OrderedTree a t = OT (DecoratedTree a) t deriving (Eq, Show, Read)

instance (Ord a, Show a, TreeOrdering t) => PPrint (OrderedTree a t) where
pp (OT dect _) = pp dect

-- | Monomial ordering for trees. We require this to be a total well-ordering, compatible
instance (Ord a, Show a, TreeOrdering t) => Ord (OrderedTree a t) where
compare (OT s o1) (OT t _) = treeCompare o1 s t

-- | Building an ordered tree with 'PathLex' ordering from a decorated tree.
ot :: (Ord a, Show a, TreeOrdering t) => DecoratedTree a -> OrderedTree a t
ot t = OT t ordering

-- | Extracting the underlying tree from an ordered tree.
dt :: (Ord a, Show a, TreeOrdering t) => OrderedTree a t -> DecoratedTree a
dt (OT t _) = t

-- ** Monomial orderings on the free operad

-- | The type class that parametrizes types implementing tree orderings.
class (Eq t, Show t) => TreeOrdering t where
treeCompare :: (Ord a, Show a) => t -> DecoratedTree a -> DecoratedTree a -> Ordering
treeCompare o t1 t2 = comparePathSequence o (orderedPathSequence t1) (orderedPathSequence t2)
comparePathSequence :: (Ord a, Show a) => t -> ([[a]],Shuffle) -> ([[a]],Shuffle) -> Ordering
ordering :: t

-- | Finding the path sequences. cf. Dotsenko-Khoroshkin.
pathSequence :: (Ord a, Show a) => DecoratedTree a -> ([[a]],Shuffle)
pathSequence (DTLeaf j) = ([[]],[j])
pathSequence (DTVertex l ts) = let
pathSequences = map pathSequence \$! ts
paths = concatMap fst \$! pathSequences
leaves = concatMap snd \$! pathSequences
in (map (l:) \$! paths, leaves)

-- | Reordering the path sequences to mirror the actual leaf ordering.
orderedPathSequence :: (Ord a, Show a) => DecoratedTree a -> ([[a]],Shuffle)
orderedPathSequence t = (map fst . sortBy (comparing snd) \$ zip ps1 ps2, ps2) where (ps1, ps2) = pathSequence t

-- | Degree reverse lexicographic path sequence ordering.
data RPathLex = RPathLex deriving (Eq, Ord, Show, Read)

-- | Changes direction of an ordering.
reverseOrder :: Ordering -> Ordering
reverseOrder LT = GT
reverseOrder GT = LT
reverseOrder EQ = EQ

instance TreeOrdering RPathLex where
treeCompare o s t = if (nLeaves s) /= (nLeaves t) then comparing nLeaves s t
else if s == t then EQ
else comparePathSequence o (orderedPathSequence s) (orderedPathSequence t)
comparePathSequence _ (paths,perms) (patht,permt) = let
clS = zipWith (comparing length) paths patht
coS = zipWith compare paths patht
cS = zipWith (\comp1 comp2 -> if comp1 == EQ then comp2 else reverseOrder comp1) clS coS
in
if any (/= EQ) cS then head (filter (/=EQ) cS)
else compare perms permt
ordering = RPathLex

-- | Path lexicographic ordering. Orders trees first by lexicographic comparison on
-- the ordered path sequence, and then by lexicographic comparison on the leaf orderings.
data PathLex = PathLex deriving (Eq, Ord, Show, Read)

instance TreeOrdering PathLex where
treeCompare o s t = if (nLeaves s) /= (nLeaves t) then comparing nLeaves s t
else if s == t then EQ
else comparePathSequence o (orderedPathSequence s) (orderedPathSequence t)
comparePathSequence _ (paths,perms) (patht,permt) = let
clS = zipWith (comparing length) paths patht
coS = zipWith compare paths patht
cs = zipWith (\comp1 comp2 -> if comp1 == EQ then comp2 else comp1) clS coS
in
if any (/= EQ) cs then head (filter (/=EQ) cs)
else compare perms permt
ordering = PathLex

-- | Forest lexicographic ordering. Currently not implemented.
data ForestLex = ForestLex deriving (Eq, Ord, Show)

instance TreeOrdering ForestLex where
comparePathSequence = error "Forest lexicographic ordering is not yet implemented."
ordering = ForestLex

-- ** Utility functions on trees
--
-- Trees are represented rooted, and all operations act on a specific root, and may recurse from there.

-- | Build a single corolla in a decorated tree. Takes a list for labels for the leaves, and derives
-- the arity of the corolla from those. This, and the composition functions, form the preferred method
-- to construct trees.
corolla :: (Ord a, Show a) => a -> [Int] -> DecoratedTree a
corolla label leaflabels =
if null leaflabels
then error "The operadic Buchberger, and many other algorithms, require the absence of 0-ary operations."
else DTVertex label (map DTLeaf leaflabels)

-- | Build a single leaf.
leaf :: (Ord a, Show a) => Int -> DecoratedTree a
leaf n = DTLeaf n

-- | Check whether a given root is a leaf.
isLeaf :: (Ord a, Show a) => DecoratedTree a -> Bool
isLeaf (DTLeaf _) = True
isLeaf _ = False

-- | Check whether a given root is a corolla.
isCorolla :: (Ord a, Show a) => DecoratedTree a -> Bool
isCorolla = not . isLeaf

-- | Change the leaves of a tree to take their values from a given list.
relabelLeaves :: (Ord a, Show a) => DecoratedTree a -> [b] -> PreDecoratedTree a b
relabelLeaves tree newLabels = fst \$ runState (mapM (\_ -> do; (x:xs) <- get; put xs; return x) tree) newLabels

-- | Find the permutation the leaf labeling ordains for inputs.
leafOrder :: (Ord a, Show a) => DecoratedTree a -> [Int]
leafOrder = foldMap (:[])

-- | Find the minimal leaf covering any given vertex.
minimalLeaf :: (Ord a, Show a, Ord b) => PreDecoratedTree a b -> b
minimalLeaf (DTLeaf lbl) = lbl
minimalLeaf vertex = minimum \$ map minimalLeaf (subTrees vertex)

-- | Compute the number of leaves of the entire tree covering a given vertex.
nLeaves :: (Ord a, Show a) => DecoratedTree a -> Int
nLeaves (DTLeaf _) = 1
nLeaves vertex = sum \$ map nLeaves (subTrees vertex)

-- | 'arityDegree' is one less than 'nLeaves'.
arityDegree :: (Ord a, Show a) => DecoratedTree a -> Int
arityDegree t = nLeaves t - 1

-- * Shuffles
-- Basic handling functions for building, recognizing and applying shuffle permutations.

-- | A shuffle is a special kind of sequence of integers.
type Shuffle = [Int]

-- | We need to recognize sorted sequences of integers.
isSorted :: (Ord a, Show a) => [a] -> Bool
isSorted xs = and \$ zipWith (<=) xs (tail xs)

-- | This tests whether a given sequence of integers really is a shuffle.
isShuffle :: Shuffle -> Bool
isShuffle as = (sort as) == [1..length as] &&
any (\k -> isShuffleIJ as k (length as-k)) [1..length as]

-- | This tests whether a given sequence of integers is an (i,j)-shuffle
isShuffleIJ :: Shuffle -> Int -> Int -> Bool
isShuffleIJ as i j = isSorted [a | a <- as, a <= i] &&
isSorted [a | a <- as, a > i] &&
length as == i+j

-- | This tests whether a given sequence of integers is admissible for a specific composition operation.
isShuffleIPQ :: Shuffle -> Int -> Int -> Bool
isShuffleIPQ as i p = let
initSegment = take i as
finalSegment = drop i as
upperTree = take p finalSegment
latterTree = drop p finalSegment
in initSegment == [1 .. length initSegment] &&
isSorted upperTree &&
isSorted latterTree

-- | This applies the resulting permutation from a shuffle to a set of elements
applyPerm :: Show a => Shuffle -> [a] -> [a]
applyPerm s is = --trace ("applyPerm " ++ show s ++ " " ++ show (length is) ++ "\n") \$
map (is!!) (map (subtract 1) s)

-- | Apply the permutation inversely to 'applyPerm'.
invApplyPerm :: Shuffle -> [a] -> [a]
invApplyPerm sh lst = map snd . sortBy (comparing fst) \$ zip sh lst

-- | Generate all subsets of length k from a given list.
kSubsets :: Int -> [Int] -> [[Int]]
kSubsets 0 _ = [[]]
kSubsets k ss = if length ss == k then [ss]
else if length ss < k then []
else do
map sort \$ (map ((head ss):) \$ kSubsets (k-1) (tail ss)) ++ kSubsets k (tail ss)

-- | Generates all shuffles from Sh_i(p,q).
allShuffles :: Int -> Int -> Int -> [Shuffle]
allShuffles i p q = if p<0 || q<0 || i<0 then error "Positive numbers, please!" else
do
let later = [i+1..i+p+q]
pS <- kSubsets p later
let qS = later \\ pS
return \$ [1..i] ++ pS ++ qS

```