from .types import Parenthood, NumericalNodeProperty, Node, Edge


class UFDS:
    """
    Union-find data structure for felzenszwalb algorithm.
    Each unionFind instance X maintains a family of disjoint sets of
    hashable objects, supporting the following two methods:
    - X[item] returns a name for the set containing the given item.
      Each set is named by an arbitrarily-chosen one of its members; as
      long as the set remains unchanged it will keep the same name. If
      the item is not yet part of a set in X, a new singleton set is
      created for it.
    - X.union(item1, item2, ...) merges the sets containing each item
      into a single larger set.  If any item is not yet part of a set
      in X, it is added to X as one of the members of the merged set.
      Union-find data structure. Based on Josiah Carlson's code,
      http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/215912
      with significant additional changes by D. Eppstein.
      http://www.ics.uci.edu/~eppstein/PADS/UnionFind.py
    """

    def __init__(self):
        """
        Create a new empty union-find structure.
        ****************************************
        """
        # UFDS permit to compute the hierarchy
        # it is responsible to keep track of:
        # the parent is crucial, but it only guarantee the tree structure
        # should not be used outside of the object
        self._parent: Parenthood = {}
        # size is subjective, it is only used to fuse objects, no real quantitative interest
        # should not be used outside of the object
        self._size: NumericalNodeProperty = {}

    def get_root(self, node: Node) -> Node:
        """
        Find and return the name of the set containing the node.
        ********************************************************
        """
        # check for previously unknown object
        # if unknown, create an entry in the parent dico
        # then, size is set to 1
        # intra is set to 0
        if node not in self._parent:
            self._parent[node] = node
            self._size[node] = 1
            return node
        # find path of objects leading to the root
        path = [node]
        root = self._parent[node]
        while root != path[-1]:
            path.append(root)
            root = self._parent[root]
        # compress the path and return
        for ancestor in path:
            self._parent[ancestor] = root
        return root

    def __iter__(self):
        """
        Iterate through all items ever found or unioned by this structure.
        ******************************************************************
        """
        return iter(self.parents)

    def union(self, edge: Edge):
        """
        Find the sets containing the objects and merge them all.
        ********************************************************
        """
        n1, n2 = edge
        roots = [self.get_root(n1), self.get_root(n2)]
        # Find the heaviest root according to its weight.
        heaviest = max(roots, key=lambda r: self._size[r])
        for r in roots:
            if r != heaviest:
                self._size[heaviest] += self._size[r]
                self._parent[r] = heaviest