import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Arc
from matplotlib.collections import PatchCollection
import seaborn as sns
sns.set()

from check import *
from methods import compute_fibonacci_sequence


# -----------------------------------------------------------------------------
# Golden spiral approximation
# -----------------------------------------------------------------------------
def plot_golden_spiral_approximation(n):
    # check value of the argument
    if n < 1:
        raise ValueError("Argument n should be larger than 1")
    # compute F_i with i from 0 to n
    sequence = compute_fibonacci_sequence(n)

    # initialize variables
    x_ref, y_ref = 0, 0
    x_min, x_max, y_min, y_max = 0, 0, 0, 0
    angle = 180
    center = (1, 1)

    # initialize figure
    fig, ax = plt.subplots(1, figsize=(16, 16))
    rectangles = []
    for i, side in enumerate(sequence):
        if i == 0:
            previous_side = side
            continue
        # add a square of side F_i
        rectangles.append(Rectangle([x_ref, y_ref], side, side))
        # annotate the square if not too small
        if i > n - 8:
            ax.annotate(side, xy=(x_ref + 0.45 * side, y_ref + 0.45 * side), fontsize=10)
        # add the corresponding arc of the spiral
        ax.add_patch(
            Arc(center, 2 * side, 2 * side,
                angle=angle, theta1=0, theta2=90,
                edgecolor="black", antialiased=True)
            )

        # update variables
        angle += 90
        angle %= 360
        if i % 4 == 1:
            if i == 1:
                x_max, y_max = 1, 1
            else:
                y_min -= side
            x_ref += side
            center = (x_ref, y_ref + side + previous_side)
        elif i % 4 == 2:
            x_max += side
            x_ref -= previous_side
            y_ref += side
            center = (x_ref, y_ref)
        elif i % 4 == 3:
            y_max += side
            x_ref -= (side + previous_side)
            y_ref -= previous_side
            center = (x_ref + side + previous_side, y_ref)
        else:
            x_min -= side
            y_ref -= (side + previous_side)
            center = (x_ref + side + previous_side, y_ref + side + previous_side)
        previous_side = side

    # color the squares depending on their size and add them to the figure
    col = PatchCollection(rectangles, alpha=.95, edgecolor="black")
    col.set(array=np.asarray(range(n + 1)), cmap="Blues")
    ax.add_collection(col)

    # customize the axes
    ax.set_aspect("equal", "box")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)
    plt.tight_layout()
    plt.savefig("plots/golden_spiral_approximation.png")
    plt.show()


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------
if __name__ == "__main__":

    n = 16
    check_fibonacci_argument(n)

    plot_golden_spiral_approximation(n)