# binary_trees.py
"""

Binary Trees

by Eric Riddoch
Oct 4, 2018

"""

# These imports are used in BST.draw().
import networkx as nx
from networkx.drawing.nx_agraph import graphviz_layout
from matplotlib import pyplot as plt
from timeit import default_timer as time
import numpy as np


class SinglyLinkedListNode:
    """A node with a value and a reference to the next node."""
    def __init__(self, data):
        self.value, self.next = data, None

class SinglyLinkedList:
    """A singly linked list with a head and a tail."""
    def __init__(self):
        self.head, self.tail = None, None

    def append(self, data):
        """Add a node containing the data to the end of the list."""
        n = SinglyLinkedListNode(data)
        if self.head is None:
            self.head, self.tail = n, n
        else:
            self.tail.next = n
            self.tail = n

    def iterative_find(self, data):
        """Search iteratively for a node containing the data.
        If there is no such node in the list, including if the list is empty,
        raise a ValueError.

        Returns:
            (SinglyLinkedListNode): the node containing the data.
        """
        current = self.head
        while current is not None:
            if current.value == data:
                return current
            current = current.next
        raise ValueError(str(data) + " is not in the list")

    def recursive_find(self, data):
        """Search recursively for the node containing the data.
        If there is no such node in the list, including if the list is empty,
        raise a ValueError.

        Returns:
            (SinglyLinkedListNode): the node containing the data.
        """
        def find(node):
            """recursive function that traverses the list in search of correct node"""
            if node == None: # are we at the end of the list?
                raise ValueError(f"{data} is not in the list!")
            if node.value == data: # did we find the node?
                return node
            
            return find(node.next) # look at the next node
        
        return find(self.head)
            
class BSTNode:
    """A node class for binary search trees. Contains a value, a
    reference to the parent node, and references to two child nodes.
    """
    def __init__(self, data):
        """Construct a new node and set the value attribute. The other
        attributes will be set when the node is added to a tree.
        """
        self.value = data
        self.prev = None        # A reference to this node's parent node.
        self.left = None        # self.left.value < self.value
        self.right = None       # self.value < self.right.value

class BST:
    """Binary search tree data structure class.
    The root attribute references the first node in the tree.
    """
    def __init__(self):
        """Initialize the root attribute."""
        self.root = None

    def find(self, data):
        """Return the node containing the data. If there is no such node
        in the tree, including if the tree is empty, raise a ValueError.
        """

        # Define a recursive function to traverse the tree.
        def _step(current):
            """Recursively step through the tree until the node containing
            the data is found. If there is no such node, raise a Value Error.
            """
            if current is None:                     # Base case 1: dead end.
                raise ValueError(str(data) + " is not in the tree.")
            if data == current.value:               # Base case 2: data found!
                return current
            if data < current.value:                # Recursively search left.
                return _step(current.left)
            else:                                   # Recursively search right.
                return _step(current.right)

        # Start the recursion on the root of the tree.
        return _step(self.root)

    def insert(self, data):
        """Insert a new node containing the specified data.

        Raises:
            ValueError: if the data is already in the tree.

        Example:
            >>> tree = BST()                    |
            >>> for i in [4, 3, 6, 5, 7, 8, 1]: |            (4)
            ...     tree.insert(i)              |            / \
            ...                                 |          (3) (6)
            >>> print(tree)                     |          /   / \
            [4]                                 |        (1) (5) (7)
            [3, 6]                              |                  \
            [1, 5, 7]                           |                  (8)
            [8]                                 |
        """
        
        # is the tree empty?
        if self.root == None:
            self.root = BSTNode(data)
            return
        
        # Define a recursive function to traverse the tree.
        def _step(current):
            """Recursively step through the tree until the node containing
            the data is found. If there is no such node, raise a Value Error.
            """
            if data == current.value:               # Base case 1: data found!
                raise ValueError(f"{data} is already in the tree!")
                
            if data < current.value:                # Act on left.
                if current.left == None: # insert
                    current.left = BSTNode(data)
                    current.left.prev = current
                else:
                    return _step(current.left)
            
            if data > current.value:                # Act on right.
                if current.right == None: # insert
                    current.right = BSTNode(data)
                    current.right.prev = current
                else:
                    return _step(current.right)

        # perform insertion:
        _step(self.root)

            
    def remove(self, data):
        """Remove the node containing the specified data.

        Raises:
            ValueError: if there is no node containing the data, including if
                the tree is empty.

        Examples:
            >>> print(12)                       | >>> print(t3)
            [6]                                 | [5]
            [4, 8]                              | [3, 6]
            [1, 5, 7, 10]                       | [1, 4, 7]
            [3, 9]                              | [8]
            >>> for x in [7, 10, 1, 4, 3]:      | >>> for x in [8, 6, 3, 5]:
            ...     t1.remove(x)                | ...     t3.remove(x)
            ...                                 | ...
            >>> print(t1)                       | >>> print(t3)
            [6]                                 | [4]
            [5, 8]                              | [1, 7]
            [9]                                 |
                                                | >>> print(t4)
            >>> print(t2)                       | [5]
            [2]                                 | >>> t4.remove(1)
            [1, 3]                              | ValueError: <message>
            >>> for x in [2, 1, 3]:             | >>> t4.remove(5)
            ...     t2.remove(x)                | >>> print(t4)
            ...                                 | []
            >>> print(t2)                       | >>> t4.remove(5)
            []                                  | ValueError: <message>
        """
        
        def get_pred(node):
            """returns the immediate predecessor as a node"""

            # start at the left child
            left = node.left

            # return the rightmost child of left
            while left.right != None:
                left = left.right

            return left
        
        def swap(n1, n2):
            temp = n1.value
            n1.value = n2.value
            n2.value = temp
        
        def _step(target):
            # case 1: target is a leaf
            if target.left == None and target.right == None:
                # A: is it the root?
                if target == self.root:
                    self.root = None
                # B: is it a left child?
                elif target == target.prev.left:
                    target.prev.left = None
                    target.prev = None
                # C: is it a right child?
                elif target == target.prev.right:
                    target.prev.right = None
                    target = None

            # case 2: target has two children
            elif target.left != None and target.right != None:
                # get the immediate predecessor
                pred = get_pred(target)
                # swap the two values
                swap(target, pred)
                # remove the predecessor
                _step(pred)

            # case 3: target has one child
            elif target.left != None or target.right != None:
                # A: is it the root?
                if target == self.root:
                    if target.left != None:
                        self.root = target.left
                        self.root.prev = None
                    else:
                        self.root = target.right
                        self.root.prev = None

                # B: is target a left child?
                elif target.prev.left == target:
                    # get the child!
                    child = None
                    if target.left != None:
                        child = target.left
                    else:
                        child = target.right
                    
                    if target.left != None:
                        target.prev.left = child
                        target.left.prev = target.prev
                    else:
                        target.prev.left = child
                        target.right.prev = target.prev
                        
                # C: is target a right child?
                elif target.prev.right == target:
                    # get the child!
                    child = None
                    if target.left != None:
                        child = target.left
                    else:
                        child = target.right

                    if target.left != None:
                        target.prev.right = child
                        target.left.prev = target.prev
                    else:
                        target.prev.right = child
                        target.right.prev = target.prev

        target = self.find(data)
        _step(target)

        
    def __str__(self):
        """String representation: a hierarchical view of the BST.

        Example:  (3)
                  / \     '[3]          The nodes of the BST are printed
                (2) (5)    [2, 5]       by depth levels. Edges and empty
                /   / \    [1, 4, 6]'   nodes are not printed.
              (1) (4) (6)
        """
        if self.root is None:                       # Empty tree
            return "[]"
        out, current_level = [], [self.root]        # Nonempty tree
        while current_level:
            next_level, values = [], []
            for node in current_level:
                values.append(node.value)
                for child in [node.left, node.right]:
                    if child is not None:
                        next_level.append(child)
            out.append(values)
            current_level = next_level
        return "\n".join([str(x) for x in out])

    def draw(self):
        """Use NetworkX and Matplotlib to visualize the tree."""
        if self.root is None:
            return

        # Build the directed graph.
        G = nx.DiGraph()
        G.add_node(self.root.value)
        nodes = [self.root]
        while nodes:
            current = nodes.pop(0)
            for child in [current.left, current.right]:
                if child is not None:
                    G.add_edge(current.value, child.value)
                    nodes.append(child)

        # Plot the graph. This requires graphviz_layout (pygraphviz).
        nx.draw(G, pos=graphviz_layout(G, prog="dot"), arrows=True,
                with_labels=True, node_color="C1", font_size=8)
        plt.show()


class AVL(BST):
    """Adelson-Velsky Landis binary search tree data structure class.
    Rebalances after insertion when needed.
    """
    def insert(self, data):
        """Insert a node containing the data into the tree, then rebalance."""
        BST.insert(self, data)      # Insert the data like usual.
        n = self.find(data)
        while n:                    # Rebalance from the bottom up.
            n = self._rebalance(n).prev

    def remove(*args, **kwargs):
        """Disable remove() to keep the tree in balance."""
        raise NotImplementedError("remove() is disabled for this class")

    def _rebalance(self,n):
        """Rebalance the subtree starting at the specified node."""
        balance = AVL._balance_factor(n)
        if balance == -2:                                   # Left heavy
            if AVL._height(n.left.left) > AVL._height(n.left.right):
                n = self._rotate_left_left(n)                   # Left Left
            else:
                n = self._rotate_left_right(n)                  # Left Right
        elif balance == 2:                                  # Right heavy
            if AVL._height(n.right.right) > AVL._height(n.right.left):
                n = self._rotate_right_right(n)                 # Right Right
            else:
                n = self._rotate_right_left(n)                  # Right Left
        return n

    @staticmethod
    def _height(current):
        """Calculate the height of a given node by descending recursively until
        there are no further child nodes. Return the number of children in the
        longest chain down.
                                    node | height
        Example:  (c)                  a | 0
                  / \                  b | 1
                (b) (f)                c | 3
                /   / \                d | 1
              (a) (d) (g)              e | 0
                    \                  f | 2
                    (e)                g | 0
        """
        if current is None:     # Base case: the end of a branch.
            return -1           # Otherwise, descend down both branches.
        return 1 + max(AVL._height(current.right), AVL._height(current.left))

    @staticmethod
    def _balance_factor(n):
        return AVL._height(n.right) - AVL._height(n.left)

    def _rotate_left_left(self, n):
        temp = n.left
        n.left = temp.right
        if temp.right:
            temp.right.prev = n
        temp.right = n
        temp.prev = n.prev
        n.prev = temp
        if temp.prev:
            if temp.prev.value > temp.value:
                temp.prev.left = temp
            else:
                temp.prev.right = temp
        if n is self.root:
            self.root = temp
        return temp

    def _rotate_right_right(self, n):
        temp = n.right
        n.right = temp.left
        if temp.left:
            temp.left.prev = n
        temp.left = n
        temp.prev = n.prev
        n.prev = temp
        if temp.prev:
            if temp.prev.value > temp.value:
                temp.prev.left = temp
            else:
                temp.prev.right = temp
        if n is self.root:
            self.root = temp
        return temp

    def _rotate_left_right(self, n):
        temp1 = n.left
        temp2 = temp1.right
        temp1.right = temp2.left
        if temp2.left:
            temp2.left.prev = temp1
        temp2.prev = n
        temp2.left = temp1
        temp1.prev = temp2
        n.left = temp2
        return self._rotate_left_left(n)

    def _rotate_right_left(self, n):
        temp1 = n.right
        temp2 = temp1.left
        temp1.left = temp2.right
        if temp2.right:
            temp2.right.prev = temp1
        temp2.prev = n
        temp2.right = temp1
        temp1.prev = temp2
        n.right = temp2
        return self._rotate_right_right(n)


def benchmark():
    """Compare the build and search times of the SinglyLinkedList, BST, and
    AVL classes. For search times, use SinglyLinkedList.iterative_find(),
    BST.find(), and AVL.find() to search for 5 random elements in each
    structure. Plot the number of elements in the structure versus the build
    and search times. Use log scales where appropriate.
    """
    
    # get the lines in english.txt
    with open('english.txt', 'r') as infile:
        lines = infile.readlines()
    
    # create a domain
    domain = 2**np.arange(3, 11)
    
    times_build_linked  = []
    times_build_avl   = []
    times_build_btree = []
    
    times_linked = []
    times_avl    = []
    times_btree  = []
    
    for i in domain:        
        # get a random set of lines of size 2**i
        choices = np.random.choice(lines, i, replace=False)
        
        ####################################################
        #                  BUILD THE TREES                 #
        ####################################################
        
        # create a BST, SinglyLinkedList, and AVL Tree
        btree = BST()
        avl = AVL()
        linked = SinglyLinkedList()
        
        # Linked List
        start = time()
        for choice in choices:
            linked.append(choice)
        times_build_linked.append(time() - start)
        
        # BST
        start = time()
        for choice in choices:
            btree.insert(choice)
        times_build_btree.append(time() - start)
        
        # AVL
        start = time()
        for choice in choices:
            avl.insert(choice)
        times_build_avl.append(time() - start)

        
        ####################################################
        #                  FIND THE NODES                  #
        ####################################################
        
        # from the nodes, randomely choose 5 of them
        five_choices = np.random.choice(choices, 5, replace=False)
        
        # Linked List
        start = time()
        for choice in five_choices:
            linked.iterative_find(choice)
        times_linked.append(time() - start)
        
        # BST            
        start = time()
        for choice in five_choices:
            btree.find(choice)
        times_btree.append(time() - start)
        
        # AVL
        start = time()
        for choice in five_choices:
            avl.find(choice)
        times_avl.append(time() - start)
        
    fig, axes = plt.subplots(1,2)
    
    # Plot build times
    axes[0].loglog(domain, times_build_linked, ".-", basex=2, basey=10, label="Linked List")
    axes[0].loglog(domain, times_build_avl, ".-", basex=2, basey=10, label="AVL")
    axes[0].loglog(domain, times_build_btree, ".-", basex=2, basey=10, label="BST")
    axes[0].set_title("Build the Structure")
    axes[0].legend()
    
    # Plot find times
    axes[1].loglog(domain, times_linked, ".-", basex=2, basey=10, label="Linked List")
    axes[1].loglog(domain, times_avl, ".-", basex=2, basey=10, label="AVL")
    axes[1].loglog(domain, times_btree, ".-", basex=2, basey=10, label="BST")
    axes[1].set_title("find() functions")
    axes[1].legend()
    
    plt.tight_layout()
    plt.show()     
        
            
        
        
    
    
