"""
A good representation for partial orders, i.e. graphs between objects
in which some are 'smaller' than others but not all pairs of objects are
comparable.
"""


# A PartialOrder structure is organized as a tree with multiple
# parents and children per object.  A relation parent-to-children
# means that the parent is in front of the child in the specified order. There
# is such a relation every time an object is directly in front of another one.
# The tree is not transitive : if A <= B, and B <= C, then there is a
# parent-to-child link from A to B and from B to C but not from A to C.
# The tree cannot contain loops.


class PartialOrder:
    
    def __init__(self):
        self.root = Node(None)
        self.root.isroot = True
        self.objects = {}

    def add(self, obj):
        "Add 'obj' to the data structure."
        newnode = Node(obj)
        self.objects[obj] = newnode
        ins = Inserter(newnode, self.isbefore)
        ins.analysegraph(self.root)
        ins.updategraph(self.root)

    def enumerate(self):
        """List all objects in a compatible order
        (if a <= b then a is enumerated before b)."""
        lst = []
        self.root.predecessors(lst, {})
        return lst

    def __iter__(self):
        return iter(self.enumerate())

    def predecessors(self, obj):
        """List all objects before 'obj' in a compatible order
        (if a <= b then a is enumerated before b)."""
        node = self.objects[obj]
        lst = []
        node.predecessors(lst, {})
        return lst

    def isbefore(self, obj1, obj2):
        "Can be overridden."
        if obj1 <= obj2:
            return -1
        elif obj2 <= obj1:
            return 1
        else:
            return 0  # does not mean '==' but just uncomparable


class Node:

    def __init__(self, obj):
        self.obj = obj
        self.children = []
        self.parentcount = 0
        self.isroot = False

    def addchild(self, child):
        self.children.append(child)
        child.parentcount += 1

    def removechild(self, i):
        child = self.children[i]
        del self.children[i]
        child.parentcount -= 1

    def predecessors(self, result, seen):
        for f in self.children:
            if f not in seen:
                seen[f] = True
                f.predecessors(result, seen)
                result.append(f.obj)


class Inserter:

    def __init__(self, newnode, comparefn):
        self.newnode = newnode
        self.comparefn = comparefn
        self.parentsleft = {}
        self.behind = {}
        self.before = {}

    def analysegraph(self, node):
        # Analyses the current tree and makes relations from 'newnode'
        # to all objects behind it.  Relations from other objects to
        # 'newnode', as well as removal of extra relations that
        # non-transitivity forbid are handled by UpdateGraph.
        for f in node.children:
            if f not in self.parentsleft:
                # if AnalyseGraph never saw this child before,
                # mark it seen and set default flags
                self.parentsleft[f] = f.parentcount
            self.parentsleft[f] -= 1
            if self.parentsleft[f] == 0:
                # if we already visited all the parents of f
                if f.isroot:
                    compare = 1
                else:
                    compare = self.comparefn(self.newnode.obj, f.obj)
                if compare > 0:   # newnode is in front of f
                    self.newnode.addchild(f)   # add the relation newnode <= f
                    self.before[f] = True
                else:
                    self.analysegraph(f)
                    if compare < 0:  # newnode is behind f
                        self.behind[f] = True

    def updategraph(self, node, backchain=True):
        # Updates the tree. It creates "backchaining" (relations from an item
        # already in the tree to the new item) and it removes extra relations
        # that non-transitivity forbid.
        backchaindone = False
        for i in range(len(node.children)-1, -1, -1):
            f = node.children[i]
            if self.parentsleft.get(f) == 0:
                # if parentsleft>0, AnalyseGraph did not visit all parents
                # of f, which means that some parent was behind the new item,
                # and this means that f is indirectly behind the new item.
                # The link newnode <= f has already been built by
                # AnalyseGraph (directly or indirectly).
                del self.parentsleft[f]   # mark it as seen by updategraph
                if f in self.before:
                    # the new item is in front of f
                    if backchain:   # if we already backchained the new item,
                        # we can remove this old link, as there is a new
                        # indirect path to it through the new item, throught
                        # its backchain.
                        node.removechild(i)
                else:
                    # f and newnode are besides each other
                    # recursivity (we have to do backchaining in the sub-tree
                    # if the new item is behind f).
                    if self.updategraph(f, f in self.behind):
                        backchaindone = True
        if backchain and not backchaindone:
            node.addchild(self.newnode)  # backchaining
            backchaindone = True
        return backchaindone


if __name__ == '__main__':
    for n in range(50):
        po = PartialOrder()
        lst = range(20)
        import random; random.shuffle(lst)
        for item in lst:
            po.add(item)
        assert po.enumerate() == range(20)
        i = random.choice(lst)
        assert po.predecessors(i) == range(i)
    print "Ok."