Refactorization

animation.py

@@ -2,26 +2,10 @@ import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from time import sleep
 import copy
-
-TOWER_NAMES = ['A', 'B', 'C']
+import argparse
 
 
-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):
+def animate_hanoi(n, frame_delay=750, src=0, dest=1, aux=2):
     colors = []
     for i in range(n):
         # Calculate the hue value
@@ -30,52 +14,71 @@ def animate_hanoi(n, frame_delay=0.5, src=0, dest=1, aux=2):
         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_xticklabels(['A', 'B', 'C'])
     ax.set_yticks([])
 
-    plot_elements = []
+    frames = []
 
-    def update(frame):
-        move, towers = states[frame]
-
-        # Clear previous plot
-        for el in plot_elements:
-            el.remove()
-        plot_elements.clear()
+    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):
-                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))
+                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))
 
-        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
+        frames.append(frame)
 
-    ani = animation.FuncAnimation(
-        fig, update, frames=len(states), repeat=False)
-    plt.show()
+    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__":
-    n = int(input("❔ Number of disks in the game: "))
-    d = int(input("❔ Delay between frames (in seconds): "))
-    animate_hanoi(n, d)
+    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()