From 05434d03207def5c8e30415e0555dc590f33cb16 Mon Sep 17 00:00:00 2001
From: Alon Druck <druckalon@gmail.com>
Date: Mon, 27 May 2024 12:22:35 +0000
Subject: [PATCH] Implement the suggested modifications

Signed-off-by: Alon Druck <alon.druck@ekumenlabs.com>
---
 beluga_tutorial/scripts/visualize.py | 47 ++++++++++++++--------------
 1 file changed, 23 insertions(+), 24 deletions(-)

diff --git a/beluga_tutorial/scripts/visualize.py b/beluga_tutorial/scripts/visualize.py
index 577d8c326..a5230a8ee 100755
--- a/beluga_tutorial/scripts/visualize.py
+++ b/beluga_tutorial/scripts/visualize.py
@@ -88,27 +88,14 @@ def plot_stages(yaml_data, axs, index):
     plt.draw()
 
 
-def on_key(event, yaml_data, axs, num_frames):
-    if not hasattr(on_key, "current_frame"):
-        on_key.current_frame = 0
-
-    if event.key == 'right':
-        on_key.current_frame = (on_key.current_frame + 1) % num_frames
-    elif event.key == 'left':
-        on_key.current_frame = (on_key.current_frame - 1) % num_frames
-
-    plot_stages(yaml_data, axs, on_key.current_frame)
-
-
 def main(argv=None) -> int:
+    """Run the entry point of the program."""
     parser = argparse.ArgumentParser(description=globals()['__doc__'])
 
     parser.add_argument(
-        '-p',
-        '--record-file-path',
+        'record_file_path',
         type=str,
         help='Absolute path to the record file generated by the beluga_tutorial example code',
-        required=True,
     )
 
     parser.add_argument(
@@ -135,20 +122,32 @@ def main(argv=None) -> int:
     fig, axs = plt.subplots(6, 1)
     num_frames = len(yaml_data['simulation_records'])
 
+    def plot_stages_update(current_frame: int) -> None:
+        """Update plots (Helper function)."""
+        plot_stages(yaml_data, axs, current_frame)
+
     if args.manual_control:
-        plot_stages(yaml_data, axs, 0)
-        fig.canvas.mpl_connect(
-            'key_press_event',
-            lambda event: on_key(event, yaml_data, axs, num_frames),
-        )
-    else:
+        current_frame = 0
+        plot_stages_update(current_frame)
+
+        def on_key(event):
+            """Handle key events (Callback function)."""
+            nonlocal current_frame
+
+            if event.key == 'right':
+                current_frame += 1
+            elif event.key == 'left':
+                current_frame -= 1
 
-        def update(current_frame):
-            plot_stages(yaml_data, axs, current_frame)
+            current_frame %= num_frames
+            plot_stages_update(current_frame)
 
+        fig.canvas.mpl_connect('key_press_event', on_key)
+
+    else:
         anim = animation.FuncAnimation(  # noqa: F841
             fig,
-            update,
+            plot_stages_update,
             frames=num_frames,
             blit=False,
             interval=args.interval_ms,