"""AVL Binary Search Tree module."""

import math
import shutil


class BinaryTree(object):
    """
    AVL Binary search tree.
    """

    def __init__(self, iterable=None, autobalance=True):
        """Initialize bst with root and size."""
        self.root = None
        self._size = 0
        self.autobalance = autobalance
        self.rotations = 0
        if iterable:
            if isinstance(iterable, range):
                iterable = treegen(iterable)
            try:
                for val in iterable:
                    self.insert(val)
            except:
                raise Exception

    def insert(self, val):
        """Insert a new node into the bst."""
        cur = self.root
        if cur is None:
            self.root = BinaryTreeNode(val)
            self._size += 1
        else:
            while True:
                if val == cur.val:
                    return
                child, left_or_right = cur.left_or_right(val)
                if child is None:
                    self._size += 1
                    new_node = BinaryTreeNode(val, parent=cur)
                    setattr(cur, left_or_right, new_node)
                    self._bubble_up_depth_from(cur)
                    if self.autobalance:
                        r_root = self._find_unbalanced(new_node)
                        if r_root:
                            self._rebalance(r_root)
                    break
                else:
                    cur = child

    def search(self, val):
        """Return node with value val if it exists, otherwise None."""
        cur_node = self.root
        while cur_node:
            if cur_node.val == val:
                return cur_node
            cur_node, trash = cur_node.left_or_right(val)

    def delete(self, val, error=False):
        """Delete a node and reorganize tree as needed."""
        to_d = self.search(val)
        if error and to_d is None:
            raise ValueError('Not in tree.')
        if to_d:
            replacement = None
            check_from = to_d.parent
            self._size -= 1
            if to_d.is_leaf():
                to_d.set_parents_child(None)
            else:
                children = to_d.children()
                if len(children) == 1:
                    child = children[0]
                    child.parent = to_d.parent
                    to_d.set_parents_child(child)
                    replacement = child
                else:
                    lmost = self.node_furthest('left', from_=to_d.right)
                    replacement = lmost
                    if lmost.parent is to_d:
                        check_from = lmost
                    else:
                        check_from = lmost.parent
                    if lmost.right:
                        lmost.right.parent = lmost.parent
                    lmost.set_parents_child(lmost.right)
                    if to_d.right:
                        to_d.right.parent = lmost
                    if to_d.left:
                        to_d.left.parent = lmost
                    lmost.right = to_d.right
                    lmost.left = to_d.left

                    to_d.set_parents_child(lmost)
                    lmost.parent = to_d.parent
            if to_d.parent is None:
                self.root = replacement
            if check_from:
                self._bubble_up_depth_from(check_from)
            if self.autobalance:
                r_root = self._find_unbalanced(check_from)
                if r_root:
                    self._rebalance(r_root)

    def contains(self, val):
        """Return whether val in bst."""
        return bool(self.search(val))

    def balance(self, from_='root'):
        """Return left vs right balance from a node on the bst."""
        if from_ == 'root':
            from_ = self.root
        if from_ is None:
            return 0
        left_depth = from_.left.depth if from_.left else 0
        right_depth = from_.right.depth if from_.right else 0
        return left_depth - right_depth

    def pre_order(self, start='root', attr='val'):
        """Return a generator of pre-order traversal."""
        if start == 'root':
            start = self.root
        if start:
            yield getattr(start, attr) if attr else start
            for node in start.children():
                for child in self.pre_order(start=node, attr=attr):
                    yield child

    def post_order(self, start='root', attr='val'):
        """Return a generator of post-order traversal."""
        if start == 'root':
            start = self.root
        if start:
            for node in start.children():
                for child in self.post_order(start=node, attr=attr):
                    yield child
            yield getattr(start, attr) if attr else start

    def in_order(self, start='root', attr='val'):
        """Return a generator of tree's nodes in order."""
        if start == 'root':
            start = self.root
        if start:
            for child in self.in_order(start=start.left, attr=attr):
                yield child
            yield getattr(start, attr) if attr else start
            for child in self.in_order(start=start.right, attr=attr):
                yield child

    def breadth_first(self, start='root', attr='val'):
        """Return generator of breadth first traversal of tree rooted at root."""
        if start == 'root':
            start = self.root
        parent_queue = []
        if start:
            parent_queue.append(start)
        while parent_queue:
            parent = parent_queue.pop(0)
            yield getattr(parent, attr) if attr else parent
            parent_queue.extend(parent.children())

    def _rebalance(self, node):
        """Rebalance a tree starting at node."""
        self.rotations += 1
        if self.balance(from_=node) < 0:
            if self.balance(from_=node.right) <= 0:
                r_root = self._lr(node)
            else:
                r_root = self._rlr(node)
        else:
            if self.balance(from_=node.left) >= 0:
                r_root = self._rr(node)
            else:
                r_root = self._lrr(node)
        self._reset_depths_from(r_root)
        self._bubble_up_depth_from(r_root)
        next_node = self._find_unbalanced(node)
        if next_node:
            self._rebalance(next_node)

    def _find_unbalanced(self, node):
        """Bubble up from a node and check for unbalanced trees."""
        while node:
            if abs(self.balance(from_=node)) > 1:
                return node
            node = node.parent

    def _lr(self, old_root):
        r"""
        Perform a left rotation (lr) around the old_root (old_root).

             3                      5
              \                    / \
               5        -->       3   9
              / \                  \
             4   9                  4
        """
        new_root = old_root.right
        old_root.right = new_root.left
        if new_root.left:
            new_root.left.parent = old_root
        if old_root.parent is None:
            self.root = new_root
        old_root.set_parents_child(new_root)
        new_root.parent = old_root.parent
        new_root.left = old_root
        old_root.parent = new_root
        return new_root

    def _rr(self, old_root):
        r"""
        Perform a right rotation (rl) around the rotation_root (r_root).

                  7              5
                 /              / \
                5      -->     1   7
               / \                /
              1   6              6
        """
        new_root = old_root.left
        old_root.left = new_root.right