# MIT License
#
# Copyright (c) 2021 Tskit Developers
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
A Python version of the C AVL tree code for development purposes.

Based on Knuth's AVL tree code in TAOCP volume 3, adapted from
https://commandlinefanatic.com/cgi-bin/showarticle.cgi?article=art070

Note there is a bug in that Python translation which is missing
P.B = 0 at the end of A9.
"""
from __future__ import annotations

import dataclasses
from typing import Any

import numpy as np
import pytest


# The nodes of the tree are assumed to contain KEY, LLINK, and RLINK fields.
# We also have a new field
#
# B(P) = balance factor of NODE(P)
#
# the height of the right subtree minus the height of the left subtree; this field
# always contains either +1, 0, or -1.  A special header node also appears at the top
# of the tree, in location HEAD; the value of RLINK(HEAD) is a pointer to the root
# of the tree, and LLINK(HEAD) is used to keep track of the overall height of the tree.
# We assume that the tree is nonempty, namely that RLINK(HEAD) != ^.


@dataclasses.dataclass(eq=False)
class Node:
    key: Any = None
    llink: Node = None
    rlink: Node = None
    balance: int = 0

    def __str__(self):
        llink = None if self.llink is None else self.llink.key
        rlink = None if self.rlink is None else self.rlink.key
        return (
            f"Node(key={self.key}, balance={self.balance}, "
            f"llink={llink}, rlink={rlink})"
        )


# For convenience in description, the algorithm uses the notation LINK(a,P)
# as a synonym for LLINK(P) if a = -1, and for RLINK(P) if a = +1.


def get_link(a, P):
    if a == -1:
        return P.llink
    else:
        return P.rlink


def set_link(a, P, val):
    if a == -1:
        P.llink = val
    else:
        P.rlink = val


class AvlTree:
    def __init__(self):
        self.head = Node()
        self.size = 0
        self.height = 0

    @property
    def root(self):
        return self.head.rlink

    def __str__(self):
        stack = [(self.head, 0)]
        s = f"size = {self.size} height = {self.height}\n"
        while len(stack) > 0:
            node, depth = stack.pop()
            s += ("  " * depth) + f"KEY={node.key} B={node.balance}\n"
            for child in [node.llink, node.rlink]:
                if child is not None:
                    stack.append((child, depth + 1))
        return s

    def ordered_keys(self):
        """
        Return the keys in sorted order. This is done by an in-order
        traversal of the nodes.
        """

        def inorder(node):
            if node is not None:
                yield from inorder(node.llink)
                yield node.key
                yield from inorder(node.rlink)

        yield from inorder(self.root)

    def search(self, key):
        P = self.root
        while P is not None:
            if key == P.key:
                break
            elif key < P.key:
                P = P.llink
            else:
                P = P.rlink
        return P

    def __insert_empty(self, key):
        self.head.rlink = Node()
        self.head.rlink.key = key
        self.size = 1
        self.height = 1
        return self.head.rlink

    def __insert(self, K):
        # A1. [Initialize.] Set T <- HEAD, S <- P <- RLINK(HEAD).
        # (The pointer variable P will move down the tree; S will point
        # to the place where rebalancing may be necessary, and
        # T always points to the parent of S.)
        T = self.head
        S = P = self.head.rlink

        # A2. [Compare.] If K < KEY(P), go to A3; if K > KEY(P), go to A4; and if
        # K = KEY(P), the search terminates successfully.
        while True:
            if K == P.key:
                return P
            elif K < P.key:
                # A3. [Move left.] Set Q <- LLINK(P). If Q = ^, set Q <= AVAIL and
                # LLINK(P) <- Q and go to step A5. Otherwise if B(Q) != 0, set T <-
                # P and S <- Q. Finally set P <- Q and return to step A2.
                Q = P.llink
                if Q is None:
                    Q = Node()
                    P.llink = Q
                    break
            # A4. [Move right.] Set Q <- RLINK(P). If Q = ^, set Q <= AVAIL and
            # RLINK(P) <- Q and go to step A5. Otherwise if B(Q) != 0, set T <- P
            # and S <- Q. Finally set P <- Q and return to step A2.
            elif K > P.key:
                Q = P.rlink
                if Q is None:
                    Q = Node()
                    P.rlink = Q
                    break
            if Q.balance != 0:
                T = P
                S = Q
            P = Q
        # A5. [Insert.] Set KEY(Q) <- K, LLINK(Q) <- RLINK(Q) <- ^, and B(Q) <- 0.
        Q.key = K
        Q.llink = Q.rlink = None
        Q.balance = 0

        # A6. [Adjust balance factors.] If K < KEY(S) set a <- -1, otherwise set a
        # <- +1. Then set R <- P <- LINK(a,S), and repeatedly do the following
        # operations zero or more times until P = Q: If K < KEY(P) set B(P) <- -1
        # and P <- LLINK(P); if K > KEY(P), set B(P) <- +1 and P <- RLINK(P).
        if K < S.key:
            a = -1
        else:
            a = 1
        R = P = get_link(a, S)
        while P != Q:
            if K < P.key:
                P.balance = -1
                P = P.llink
            elif K > P.key:
                P.balance = 1
                P = P.rlink

        # A7. [Balancing act.] Several cases now arise:
        #
        #  i) If B(S) = 0, set B(S) <- a, LLINK(HEAD) <- LLINK(HEAD) + 1, and
        #  terminate the algorithm.
        #
        if S.balance == 0:
            S.balance = a
            self.height += 1

        # ii) If B(S) = -a, set B(S) <- 0 and terminate the algorithm.

        elif S.balance == -a:
            S.balance = 0

        # iii) If B(S) = a, go to step A8 if B(R) = a, to A9 if B(R) = -a.
        else:
            if R.balance == a:
                # A8. [Single rotation.] Set P <- R, LINK(a,S) <- LINK(-a,R),
                # LINK(-a,R) <- S,B(S) <- B(R) <- 0. Go to A10.
                P = R
                set_link(a, S, get_link(-a, R))
                set_link(-a, R, S)
                S.balance = R.balance = 0
            elif R.balance == -a:
                # A9. [Double rotation.] Set P <- LINK(-a,R),
                #  LINK(-a,R) <- LINK(a,P),LINK(a,P) <- R, LINK(a,S)
                #  <- LINK(-a,P), LINK(-a,P) <- S. Now set
                #
                #               { (-a,0), if B(P) =  a;
                #  (B(S),B(R))<-{ ( 0,0), if B(P) =  0;
                #               { ( 0,a), if B(P) = -a;
                #
                #  and then set B(P) <- 0
                P = get_link(-a, R)
                set_link(-a, R, get_link(a, P))
                set_link(a, P, R)
                set_link(a, S, get_link(-a, P))
                set_link(-a, P, S)
                if P.balance == a:
                    S.balance = -a
                    R.balance = 0
                elif P.balance == 0:
                    S.balance = 0
                    R.balance = 0
                else:
                    S.balance = 0
                    R.balance = a
                P.balance = 0

            # A10. [Finishing touch.] If S = RLINK(T) then set RLINK(T) <- P,
            # otherwise set LLINK(T) <- P.
            if S == T.rlink:
                T.rlink = P
            else:
                T.llink = P

        return Q

    def insert(self, key):
        if self.size == 0:
            return self.__insert_empty(key)
        return self.__insert(key)


class TestAvlTree:
    def verify_tree(self, tree):
        """
        Check that the tree fits the AVL tree properties.
        """
        # The height of a node is its maximum distance to a leaf
        node_height = {}

        def compute_height(node):
            if node is None:
                return 0
            val = 1 + max([compute_height(node.llink), compute_height(node.rlink)])
            node_height[node] = val
            return val

        compute_height(tree.head.rlink)
        assert tree.height == max(node_height.values())
        assert tree.height == node_height[tree.head.rlink]
        # print(tree)

        # The balance factor B is the height of the right subtree
        # minus the height of the left subtree
        stack = [tree.head.rlink]
        while len(stack) > 0:
            node = stack.pop()
            # print(node, node_height[node])
            assert node.balance in [-1, 0, 1]
            lheight = None
            if node.llink is not None:
                lheight = node_height[node.llink]
                stack.append(node.llink)
            rheight = None
            if node.rlink is not None:
                rheight = node_height[node.rlink]
                stack.append(node.rlink)
            if lheight is not None and rheight is not None:
                balance_factor = rheight - lheight
                assert node.balance == balance_factor
            elif lheight is None and rheight is None:
                assert node_height[node] == 1
                assert node.balance == 0
            else:
                # if one child is None, the height of this node must be 2
                assert node_height[node] == 2
                if lheight is None:
                    assert node.balance == 1
                else:
                    assert node.balance == -1

    def verify(self, keys):
        tree = AvlTree()
        key_set = set()
        for k in keys:
            node = tree.search(k)
            if k in key_set:
                assert node is not None
                assert node.key == k
            else:
                assert node is None
            node = tree.insert(k)
            key_set.add(k)
            self.verify_tree(tree)
            assert tree.search(k) is node
        for k in range(100):
            node = tree.search(k)
            if k in key_set:
                assert node is not None
                assert node.key == k
            else:
                assert node is None
        ordered_keys = list(tree.ordered_keys())
        assert ordered_keys == list(sorted(set(keys)))

        # Implement the inorder on an existing list to mimic C algorithm
        l2 = [None for _ in ordered_keys]

        def visit(node, index, out):
            if node is None:
                return index
            index = visit(node.llink, index, out)
            out[index] = node.key
            return visit(node.rlink, index + 1, out)

        visit(tree.root, 0, l2)
        assert l2 == ordered_keys

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_sequential(self, n):
        self.verify(range(n))

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_sequential_reversed(self, n):
        self.verify(range(n)[::-1])

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_random_integers(self, n):
        rng = np.random.RandomState(42)
        values = rng.randint(-100, 100, size=n)
        self.verify(values)

    @pytest.mark.parametrize("n", [0, 1, 10, 33, 64, 127, 133])
    def test_random_floats(self, n):
        rng = np.random.RandomState(42)
        values = rng.random(size=n)
        self.verify(values)
