import matplotlib.pyplot as plt
import matplotlib.animation as animation
from time import sleep
import copy
import argparse


def animate_hanoi(n, frame_delay=750, src=0, dest=1, aux=2):
    colors = []
    for i in range(n):
        # Calculate the hue value
        hue = i / n
        # Convert HSL to RGB
        rgb = plt.cm.hsv(hue)[:3]  # Use HSV colormap and take RGB values
        colors.append(rgb)

    fig, ax = plt.subplots()
    fig.suptitle(f"Towers of Hanoi simulation for n={n}", fontsize=21)
    ax.set_xlim(-1, 3)
    ax.set_ylim(0, n + 1)
    ax.set_xticks([0, 1, 2])
    ax.set_xticklabels(['A', 'B', 'C'])
    ax.set_yticks([])

    frames = []

    def move_and_frame(src, dest):
        disk = towers[src].pop()
        towers[dest].append(disk)

        # Draw the disks
        frame = []
        for i, tower in enumerate(towers):
            for j, disk in enumerate(tower):
                frame += ax.bar(i, 1, bottom=j,
                                width=0.5, color=colors[disk-1]).patches
                frame.append(ax.text(i, j + 0.5, str(disk), ha='center',
                                     va='center', color='white', fontsize=16))

        frames.append(frame)

    def solve_game(n, src, dest, aux):
        if n == 1:
            move_and_frame(src, dest)
        else:
            solve_game(n - 1, src, aux, dest)
            move_and_frame(src, dest)
            solve_game(n - 1, aux, dest, src)

    # Create the towers
    towers = [[], [], []]
    for i in range(n):
        towers[src].append(n - i)

    # Frame the initial state
    move_and_frame(0, 0)

    # Solve it
    solve_game(n, src, dest, aux)

    ani = animation.ArtistAnimation(fig, frames, interval=frame_delay, blit=True,
                                    repeat=False)

    return ani


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="A simple program to animate the optimal solution for a Towers of Hanoi game with n disks")
    parser.add_argument("--disk-number", "-n", type=int,
                        help="Number of disks to simulate. Asks by default", default=0)
    parser.add_argument("--interval", "-i", type=int,
                        help="Interval (in ms) between frames", default=750)
    parser.add_argument("--output", "-o", type=str,
                        help="File to write the animation to (if specified)", default=None)
    args = parser.parse_args()
    if args.disk_number == 0:
        args.disk_number = int(input("❔ Number of disks in the game: "))
    ani = animate_hanoi(args.disk_number, args.interval)
    if args.output:
        ani.save(args.output)
        print("Animation saved succesfully")
    else:
        plt.show()