From 8925722e541d196aacc278b73f0ab13f5b5a03a5 Mon Sep 17 00:00:00 2001
From: jdrotleff <151064457+jpd-de@users.noreply.github.com>
Date: Thu, 6 Jun 2024 08:57:46 +0200
Subject: [PATCH] Added Step Back and Restart feature

This allows to jump back to the last frame or restart the visualization with a button.

Name changes:
Step: >>
Step back: <<
---
 znvis/visualizer/visualizer.py | 47 +++++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/znvis/visualizer/visualizer.py b/znvis/visualizer/visualizer.py
index ddfbe62..68d3a29 100644
--- a/znvis/visualizer/visualizer.py
+++ b/znvis/visualizer/visualizer.py
@@ -128,8 +128,9 @@ def _initialize_app(self):
         self.vis.reset_camera_to_default()
 
         # Add actions to the visualizer.
-        self.vis.add_action("Step", self._update_particles)
+        self.vis.add_action("<<", self._update_particles_back)
         self.vis.add_action("Play", self._continuous_trajectory)
+        self.vis.add_action(">>", self._update_particles)
         self.vis.add_action("Export Scene", self._export_scene)
         self.vis.add_action("Screenshot", self._take_screenshot)
         self.vis.add_action("Export Video", self._export_video)
@@ -553,6 +554,50 @@ def _update_particles(self, visualizer=None, step: int = None):
 
         visualizer.post_redraw()  # re-draw the window.
 
+    def _update_particles_back(self, visualizer=None, step: int = None):
+        """
+        Update the positions of the particles one step back (rewind-feature)
+
+        Parameters
+        ----------
+        step : int
+                Step to update to.
+
+        Returns
+        -------
+        Updates the positions of the particles in the box.
+        """
+        if visualizer is None:
+            visualizer = self.vis
+        if step is None:
+            if self.counter == self.number_of_steps - 1:
+                self.counter = self.number_of_steps - 1
+            else:
+                self.counter -= 1
+            step = self.counter
+
+        self._draw_particles(visualizer=visualizer)  # draw the particles.
+
+        # draw the vector field if it exists.
+        if self.vector_field is not None:
+            self._draw_vector_field(visualizer=visualizer) 
+
+        visualizer.post_redraw()  # re-draw the window.
+
+    def _restart_trajectory(self, visualizer=None):
+        if visualizer is None:
+            visualizer = self.vis
+        self.counter = 0
+
+        self._draw_particles(visualizer=visualizer)  # draw the particles.
+
+        # draw the vector field if it exists.
+        if self.vector_field is not None:
+            self._draw_vector_field(visualizer=visualizer) 
+
+        visualizer.post_redraw()  # re-draw the window.
+
+    
     def run_visualization(self):
         """
         Run the visualization.