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

TOWER_NAMES = ['A', 'B', 'C']


def move_and_save(src, dest, towers, states):
    disk = towers[src].pop()
    towers[dest].append(disk)
    states.append(((disk, src, dest), copy.deepcopy(towers)))


def solve_game(n, src, dest, aux, towers, states):
    if n == 1:
        move_and_save(src, dest, towers, states)
    else:
        solve_game(n - 1, src, aux, dest, towers, states)
        move_and_save(src, dest, towers, states)
        solve_game(n - 1, aux, dest, src, towers, states)


def animate_hanoi(n, frame_delay=0.5, 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)

    # Create the towers
    towers = [[], [], []]
    for i in range(n):
        towers[src].append(n - i)
    # Solve it
    states = [(None, copy.deepcopy(towers))]
    solve_game(n, src, dest, aux, towers, states)

    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(TOWER_NAMES)
    ax.set_yticks([])

    plot_elements = []

    def update(frame):
        move, towers = states[frame]

        # Clear previous plot
        for el in plot_elements:
            el.remove()
        plot_elements.clear()

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

        sleep(frame_delay)  # Delay the frame
        if move:
            print(
                f"🟦 Move {frame}: Disk ({move[0]}) on tower {TOWER_NAMES[move[1]]} to tower {TOWER_NAMES[move[2]]}")
        return plot_elements

    ani = animation.FuncAnimation(
        fig, update, frames=len(states), repeat=False)
    plt.show()


if __name__ == "__main__":
    n = int(input("❔ Number of disks in the game: "))
    d = int(input("❔ Delay between frames (in seconds): "))
    animate_hanoi(n, d)