#!/usr/bin/env python
# coding:utf-8
# Author:  Mozman
# Purpose: treemixin provides top level functions for binary trees
# Created: 03.05.2010

from .iterator import TreeIterator
from .walker import Walker
from .treeslice import TreeSlice


class TreeMixin(object):
    """
    Abstract-Base-Class for the pure Python Trees: BinaryTree, AVLTree and RBTree
    Mixin-Class for the Cython based Trees: FastBinaryTree, FastAVLTree, FastRBTree

    The TreeMixin Class
    ===================

    T has to implement following properties
    ---------------------------------------

    count -- get node count

    T has to implement following methods
    ------------------------------------

    get_walker(...)
        get a tree walker object, provides methods to traverse the tree see walker.py

    insert(...)
        insert(key, value) <==> T[key] = value, insert key into T

    remove(...)
        remove(key) <==> del T[key], remove key from T

    clear(...)
        T.clear() -> None.  Remove all items from T.

    Methods defined here
    --------------------

    * __contains__(k) -> True if T has a key k, else False, O(log(n))
    * __delitem__(y) <==> del T[y], del T[s:e], O(log(n))
    * __getitem__(y) <==> T[y], T[s:e], O(log(n))
    * __iter__() <==> iter(T)
    * __len__() <==> len(T), O(1)
    * __max__() <==> max(T), get max item (k,v) of T, O(log(n))
    * __min__() <==> min(T), get min item (k,v) of T, O(log(n))
    * __and__(other) <==> T & other, intersection
    * __or__(other) <==> T | other, union
    * __sub__(other) <==> T - other, difference
    * __xor__(other) <==> T ^ other, symmetric_difference
    * __repr__() <==> repr(T)
    * __setitem__(k, v) <==> T[k] = v, O(log(n))
    * clear() -> None, Remove all items from T, , O(n)
    * copy() -> a shallow copy of T, O(n*log(n))
    * discard(k) -> None, remove k from T, if k is present, O(log(n))
    * get(k[,d]) -> T[k] if k in T, else d, O(log(n))
    * is_empty() -> True if len(T) == 0, O(1)
    * items([reverse]) -> generator for (k, v) items of T, O(n)
    * keys([reverse]) -> generator for keys of T, O(n)
    * values([reverse]) -> generator for values of  T, O(n)
    * pop(k[,d]) -> v, remove specified key and return the corresponding value, O(log(n))
    * popitem() -> (k, v), remove and return some (key, value) pair as a 2-tuple, O(log(n))
    * setdefault(k[,d]) -> T.get(k, d), also set T[k]=d if k not in T, O(log(n))
    * update(E) -> None.  Update T from dict/iterable E, O(E*log(n))

    slicing by keys

    * itemslice(s, e) -> generator for (k, v) items of T for s <= key < e, O(n)
    * keyslice(s, e) -> generator for keys of T for s <= key < e, O(n)
    * valueslice(s, e) -> generator for values of T for s <= key < e, O(n)
    * T[s:e] -> TreeSlice object, with keys in range s <= key < e, O(n)
    * del T[s:e] -> remove items by key slicing, for s <= key < e, O(n)

    if 's' is None or T[:e] TreeSlice/iterator starts with value of min_key()
    if 'e' is None or T[s:] TreeSlice/iterator ends with value of max_key()
    T[:] is a TreeSlice which represents the whole tree.

    TreeSlice is a tree wrapper with range check, and contains no references
    to objects, deleting objects in the associated tree also deletes the object
    in the TreeSlice.

    * TreeSlice[k] -> get value for key k, raises KeyError if k not exists in range s:e
    * TreeSlice[s1:e1] -> TreeSlice object, with keys in range s1 <= key < e1

      * new lower bound is max(s, s1)
      * new upper bound is min(e, e1)

    TreeSlice methods:

    * items() -> generator for (k, v) items of T, O(n)
    * keys() -> generator for keys of T, O(n)
    * values() -> generator for values of  T, O(n)
    * __iter__ <==> keys()
    * __repr__ <==> repr(T)
    * __contains__(key)-> True if TreeSlice has a key k, else False, O(log(n))

    prev/succ operations

    * prev_item(key) -> get (k, v) pair, where k is predecessor to key, O(log(n))
    * prev_key(key) -> k, get the predecessor of key, O(log(n))
    * succ_item(key) -> get (k,v) pair as a 2-tuple, where k is successor to key, O(log(n))
    * succ_key(key) -> k, get the successor of key, O(log(n))

    traverse tree

    * treeiter([rtype, reverse]) -> extended TreeIterator (has prev, succ, goto, ... methods)
    * foreach(f, [order]) -> visit all nodes of tree and call f(k, v) for each node, O(n)

    Heap methods

    * max_item() -> get largest (key, value) pair of T, O(log(n))
    * max_key() -> get largest key of T, O(log(n))
    * min_item() -> get smallest (key, value) pair of T, O(log(n))
    * min_key() -> get smallest key of T, O(log(n))
    * pop_min() -> (k, v), remove item with minimum key, O(log(n))
    * pop_max() -> (k, v), remove item with maximum key, O(log(n))
    * nlargest(i[,pop]) -> get list of i largest items (k, v), O(i*log(n))
    * nsmallest(i[,pop]) -> get list of i smallest items (k, v), O(i*log(n))

    Set methods (using frozenset)

    * intersection(t1, t2, ...) -> Tree with keys *common* to all trees
    * union(t1, t2, ...) -> Tree with keys from *either* trees
    * difference(t1, t2, ...) -> Tree with keys in T but not any of t1, t2, ...
    * symmetric_difference(t1) -> Tree with keys in either T and t1  but not both
    * issubset(S) -> True if every element in T is in S
    * issuperset(S) -> True if every element in S is in T
    * isdisjoint(S) ->  True if T has a null intersection with S

    Classmethods

    * fromkeys(S[,v]) -> New tree with keys from S and values equal to v.

    """

    def get_walker(self):
        return Walker(self)

    def __repr__(self):
        """ x.__repr__(...) <==> repr(x) """
        tpl = "%s({%s})" % (self.__class__.__name__, '%s')
        return tpl % ", ".join(
            ("%r: %r" % item for item in list(self.items())))

    def copy(self):
        """ T.copy() -> get a shallow copy of T. """
        tree = self.__class__()
        self.foreach(tree.insert, order=-1)
        return tree
    __copy__ = copy

    def __contains__(self, key):
        """ k in T -> True if T has a key k, else False """
        try:
            self.get_value(key)
            return True
        except KeyError:
            return False

    def __len__(self):
        """ x.__len__() <==> len(x) """
        return self.count

    def __min__(self):
        """ x.__min__() <==> min(x) """
        return self.min_item()

    def __max__(self):
        """ x.__max__() <==> max(x) """
        return self.max_item()

    def __and__(self, other):
        """ x.__and__(other) <==> self & other """
        return self.intersection(other)

    def __or__(self, other):
        """ x.__or__(other) <==> self | other """
        return self.union(other)

    def __sub__(self, other):
        """ x.__sub__(other) <==> self - other """
        return self.difference(other)

    def __xor__(self, other):
        """ x.__xor__(other) <==> self ^ other """
        return self.symmetric_difference(other)

    def discard(self, key):
        """ x.discard(k) -> None, remove k from T, if k is present """
        try:
            self.remove(key)
        except KeyError:
            pass

    def __del__(self):
        self.clear()

    def is_empty(self):
        """ x.is_empty() -> False if T contains any items else True"""
        return self.count == 0

    def keys(self, reverse=False):
        """ T.iterkeys([reverse]) -> an iterator over the keys of T, in ascending
        order if reverse is True, iterate in descending order, reverse defaults
        to False
        """
        return (item[0] for item in self.items(reverse))
    __iter__ = keys

    def __reversed__(self):
        return self.keys(reverse=True)

    def values(self, reverse=False):
        """ T.values([reverse]) -> an iterator over the values of T, in ascending order
        if reverse is True, iterate in descending order, reverse defaults to False
        """
        return (item[1] for item in self.items(reverse))

    def items(self, reverse=False):
        """ T.items([reverse]) -> an iterator over the (key, value) items of T,
        in ascending order if reverse is True, iterate in descending order,
        reverse defaults to False
        """
        if self.is_empty():
            return
        node = self.get_walker()
        direction = 1 if reverse else 0
        other = 1 - direction
        go_down = True
        while True:
            if node.has_child(direction) and go_down:
                node.push()
                node.down(direction)
            else:
                yield node.item
                if node.has_child(other):
                    node.down(other)
                    go_down = True
                else:
                    if node.stack_is_empty():
                        return  # all done
                    node.pop()
                    go_down = False

    def treeiter(self, rtype='key', reverse=False):
        """ T.treeiter([rtype, reverse]) -> TreeIterator,
        rtype in ('key', 'value', 'item').
        """
        return TreeIterator(self, rtype, reverse)

    def __getitem__(self, key):
        """ x.__getitem__(y) <==> x[y] """
        if isinstance(key, slice):
            return TreeSlice(self, key.start, key.stop)
        else:
            return self.get_value(key)

    def __setitem__(self, key, value):
        """ x.__setitem__(i, y) <==> x[i]=y """
        if isinstance(key, slice):
            raise ValueError('setslice is not supported')
        self.insert(key, value)

    def __delitem__(self, key):
        """ x.__delitem__(y) <==> del x[y] """
        if isinstance(key, slice):
            self.delitems(self.keyslice(key.start, key.stop))
        else:
            self.remove(key)

    def delitems(self, keys):
        """ T.delitems(keys) -> remove all items by keys
        """
        # convert generator to a set, because the content of the
        # tree will be modified!
        for key in frozenset(keys):
            self.remove(key)

    def keyslice(self, startkey, endkey):
        """ T.keyslice(startkey, endkey) -> key iterator:
        startkey <= key < endkey.
        """
        return (item[0] for item in self.itemslice(startkey, endkey))

    def itemslice(self, startkey, endkey):
        """ T.itemslice(s, e) -> item iterator: s <= key < e.

        if s is None: start with min element -> T[:e]
        if e is None: end with max element -> T[s:]
        T[:] -> all items

        """
        if self.is_empty():
            return

        if startkey is None:
            # no lower bound
            def can_go_left(): return node.has_left() and visit_left
        else:
            # don't visit subtrees without keys in search range
            def can_go_left(): return node.key > startkey and node.has_left() and visit_left

        if endkey is None:
            # no upper bound
            def can_go_right(): return node.has_right()
        else:
            # don't visit subtrees without keys in search range
            def can_go_right(): return node.key < endkey and node.has_right()

        if (startkey, endkey) == (None, None):
            def key_in_range(): return True
        elif startkey is None:
            def key_in_range(): return node.key < endkey
        elif endkey is None:
            def key_in_range(): return node.key >= startkey
        else:
            def key_in_range(): return startkey <= node.key < endkey

        node = self.get_walker()
        visit_left = True
        while True:
            if can_go_left():
                node.push()
                node.go_left()
            else:
                if key_in_range():
                    yield node.item
                if can_go_right():
                    node.go_right()
                    visit_left = True
                else:
                    if node.stack_is_empty():
                        return
                    node.pop()
                    # left side is already done
                    visit_left = False

    def valueslice(self, startkey, endkey):
        """ T.valueslice(startkey, endkey) -> value iterator:
        startkey <= key < endkey.
        """
        return (item[1] for item in self.itemslice(startkey, endkey))

    def get_value(self, key):
        node = self.root
        while node is not None:
            if key == node.key:
                return node.value
            elif key < node.key:
                node = node.left
            else:
                node = node.right
        raise KeyError(str(key))

    def __getstate__(self):
        return dict(list(self.items()))

    def __setstate__(self, state):
        # note for myself: this is called like __init__, so don't use clear()
        # to remove existing data!
        self._root = None
        self._count = 0
        self.update(state)

    def setdefault(self, key, default=None):
        """ T.setdefault(k[,d]) -> T.get(k,d), also set T[k]=d if k not in T """
        try:
            return self.get_value(key)
        except KeyError:
            self.insert(key, default)
            return default

    def update(self, *args):
        """ T.update(E) -> None. Update T from E : for (k, v) in E: T[k] = v """
        for items in args:
            try:
                generator = list(items.items())
            except AttributeError:
                generator = iter(items)

            for key, value in generator:
                self.insert(key, value)

    @classmethod
    def fromkeys(cls, iterable, value=None):
        """ T.fromkeys(S[,v]) -> New tree with keys from S and values equal to v.

        v defaults to None.
        """
        tree = cls()
        for key in iterable:
            tree.insert(key, value)
        return tree

    def get(self, key, default=None):
        """ T.get(k[,d]) -> T[k] if k in T, else d.  d defaults to None. """
        try:
            return self.get_value(key)
        except KeyError:
            return default

    def pop(self, key, *args):
        """ T.pop(k[,d]) -> v, remove specified key and return the corresponding value
        If key is not found, d is returned if given, otherwise KeyError is raised
        """
        if len(args) > 1:
            raise TypeError(
                "pop expected at most 2 arguments, got %d" %
                (1 + len(args)))
        try:
            value = self.get_value(key)
            self.remove(key)
            return value
        except KeyError:
            if len(args) == 0:
                raise
            else:
                return args[0]

    def popitem(self):
        """ T.popitem() -> (k, v), remove and return some (key, value) pair as a
        2-tuple; but raise KeyError if T is empty
        """
        if self.is_empty():
            raise KeyError("popitem(): tree is empty")
        walker = self.get_walker()
        walker.goto_leaf()
        result = walker.item
        self.remove(walker.key)
        return result

    def foreach(self, func, order=0):
        """ Visit all tree nodes and process key, value.

        parm func: function(key, value)
        param int order: inorder = 0, preorder = -1, postorder = +1

        """
        def _traverse():
            if order == -1:
                func(node.key, node.value)
            if node.has_left():
                node.push()
                node.go_left()
                _traverse()
                node.pop()
            if order == 0:
                func(node.key, node.value)
            if node.has_right():
                node.push()
                node.go_right()
                _traverse()
                node.pop()
            if order == +1:
                func(node.key, node.value)

        node = self.get_walker()
        _traverse()

    def min_item(self):
        """ Get item with min key of tree, raises ValueError if tree is empty. """
        if self.count == 0:
            raise ValueError("Tree is empty")
        node = self._root
        while node.left is not None:
            node = node.left
        return (node.key, node.value)

    def pop_min(self):
        """ T.pop_min() -> (k, v), remove item with minimum key, raise ValueError
        if T is empty.
        """
        item = self.min_item()
        self.remove(item[0])
        return item

    def min_key(self):
        """ Get min key of tree, raises ValueError if tree is empty. """
        key, value = self.min_item()
        return key

    def prev_item(self, key):
        """ Get predecessor (k,v) pair of key, raises KeyError if key is min key
        or key does not exist.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.prev_item(key)

    def succ_item(self, key):
        """ Get successor (k,v) pair of key, raises KeyError if key is max key
        or key does not exist.
        """
        if self.count == 0:
            raise KeyError("Tree is empty")
        walker = self.get_walker()
        return walker.succ_item(key)

    def prev_key(self, key):
        """ Get predecessor to key, raises KeyError if key is min key
        or key does not exist.
        """
        key, value = self.prev_item(key)
        return key

    def succ_key(self, key):
        """ Get successor to key, raises KeyError if key is max key
        or key does not exist.
        """
        key, value = self.succ_item(key)
        return key

    def max_item(self):
        """ Get item with max key of tree, raises ValueError if tree is empty. """
        if self.count == 0:
            raise ValueError("Tree is empty")
        node = self._root
        while node.right is not None:
            node = node.right
        return (node.key, node.value)

    def pop_max(self):
        """ T.pop_max() -> (k, v), remove item with maximum key, raise ValueError
        if T is empty.
        """
        item = self.max_item()
        self.remove(item[0])
        return item

    def max_key(self):
        """ Get max key of tree, raises ValueError if tree is empty. """
        key, value = self.max_item()
        return key

    def nsmallest(self, n, pop=False):
        """ T.nsmallest(n) -> get list of n smallest items (k, v).
        If pop is True, remove items from T.
        """
        if pop:
            return [self.pop_min() for _ in range(min(len(self), n))]
        else:
            items = list(self.items())
            return [next(items) for _ in range(min(len(self), n))]

    def nlargest(self, n, pop=False):
        """ T.nlargest(n) -> get list of n largest items (k, v).
        If pop is True remove items from T.
        """
        if pop:
            return [self.pop_max() for _ in range(min(len(self), n))]
        else:
            items = self.items(reverse=True)
            return [next(items) for _ in range(min(len(self), n))]

    def intersection(self, *trees):
        """ x.intersection(t1, t2, ...) -> Tree, with keys *common* to all trees
        """
        thiskeys = frozenset(list(self.keys()))
        sets = _build_sets(trees)
        rkeys = thiskeys.intersection(*sets)
        return self.__class__(((key, self.get(key)) for key in rkeys))

    def union(self, *trees):
        """ x.union(t1, t2, ...) -> Tree with keys from *either* trees
        """
        thiskeys = frozenset(list(self.keys()))
        rkeys = thiskeys.union(*_build_sets(trees))
        return self.__class__(((key, self.get(key)) for key in rkeys))

    def difference(self, *trees):
        """ x.difference(t1, t2, ...) -> Tree with keys in T but not any of t1,
        t2, ...
        """
        thiskeys = frozenset(list(self.keys()))
        rkeys = thiskeys.difference(*_build_sets(trees))
        return self.__class__(((key, self.get(key)) for key in rkeys))

    def symmetric_difference(self, tree):
        """ x.symmetric_difference(t1) -> Tree with keys in either T and t1  but
        not both
        """
        thiskeys = frozenset(list(self.keys()))
        rkeys = thiskeys.symmetric_difference(frozenset(list(tree.keys())))
        return self.__class__(((key, self.get(key)) for key in rkeys))

    def issubset(self, tree):
        """ x.issubset(tree) -> True if every element in x is in tree """
        thiskeys = frozenset(list(self.keys()))
        return thiskeys.issubset(frozenset(list(tree.keys())))

    def issuperset(self, tree):
        """ x.issubset(tree) -> True if every element in tree is in x """
        thiskeys = frozenset(list(self.keys()))
        return thiskeys.issuperset(frozenset(list(tree.keys())))

    def isdisjoint(self, tree):
        """ x.isdisjoint(S) ->  True if x has a null intersection with tree """
        thiskeys = frozenset(list(self.keys()))
        return thiskeys.isdisjoint(frozenset(list(tree.keys())))


def _build_sets(trees):
    return [frozenset(list(tree.keys())) for tree in trees]
