From 027207105504a7ee7344fc74e6792a5d8495faf5 Mon Sep 17 00:00:00 2001
From: jbrhm <jrabraham1904@gmail.com>
Date: Sat, 6 Jul 2024 16:42:14 -0400
Subject: [PATCH] New State Machine Visualizer

---
 launch/autonomy.launch       |  2 ++
 scripts/visualizer.py        | 31 ++++++++++++++++---------------
 src/navigation/navigation.py |  4 +++-
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/launch/autonomy.launch b/launch/autonomy.launch
index 26b856623..c4b6b3625 100644
--- a/launch/autonomy.launch
+++ b/launch/autonomy.launch
@@ -1,4 +1,6 @@
 <launch>
+  <rosparam command="load" file="$(find mrover)/config/esw.yaml"/>
+
   <arg name="tag_detector" default="true"/>
   <arg name="object_detector" default="false"/>
   <arg name="cost_map" default="true"/>
diff --git a/scripts/visualizer.py b/scripts/visualizer.py
index a849ec49c..2e16d6cba 100755
--- a/scripts/visualizer.py
+++ b/scripts/visualizer.py
@@ -11,14 +11,14 @@
 
 import rospy
 import sys
-from smach_msgs.msg import SmachContainerStatus, SmachContainerStructure
+from mrover.msg import StateMachineStructure, StateMachineStateUpdate
 from threading import Lock
 from dataclasses import dataclass
 from typing import Optional, List, Dict
 
 
-STRUCTURE_TOPIC = "/smach/container_structure"
-STATUS_TOPIC = "/smach/container_status"
+STRUCTURE_TOPIC = "nav_structure"
+STATUS_TOPIC = "nav_state"
 
 
 @dataclass
@@ -30,7 +30,7 @@ class State:
 class StateMachine:
     def __init__(self):
         self.states: Dict[str, State] = {}
-        self.structure: Optional[SmachContainerStructure] = None
+        self.structure: Optional[StateMachineStructure] = None
         self.mutex: Lock = Lock()
         self.cur_active: str = ""
         self.previous_state: str = ""
@@ -50,17 +50,18 @@ def set_active_state(self, active_state):
                     f"Current time: {now} Previous state: {self.previous_state} Current State: { self.cur_active}"
                 )
 
-    def _rebuild(self, structure: SmachContainerStructure):
+    def _rebuild(self, structure: StateMachineStructure):
         """
         rebuilds the state dictionary with a new structure message
         """
-        self.states = {child: State(child, []) for child in structure.children}
-        for start, end in zip(structure.outcomes_from, structure.outcomes_to):
-            if end != "None":
-                self.states[start].children.append(self.states[end])
+        self.states = {child.origin: State(child.origin, []) for child in structure.transitions}
+        for transition in structure.transitions:
+            origin = transition.origin
+            for to in transition.destinations:
+                self.states[origin].children.append(self.states[to])
         self.needs_redraw = True
 
-    def check_rebuild(self, structure: SmachContainerStructure):
+    def check_rebuild(self, structure: StateMachineStructure):
         """
         checks if the structure passed as input matches the structure already represented (thread safe)
         """
@@ -71,10 +72,10 @@ def check_rebuild(self, structure: SmachContainerStructure):
                 self._rebuild(structure)
                 self.structure = structure
 
-    def container_status_callback(self, status: SmachContainerStatus):
-        self.set_active_state(status.active_states[0])
+    def container_status_callback(self, status: StateMachineStateUpdate):
+        self.set_active_state(status.state)
 
-    def container_structure_callback(self, structure: SmachContainerStructure):
+    def container_structure_callback(self, structure: StateMachineStructure):
         self.check_rebuild(structure)
 
 
@@ -121,12 +122,12 @@ def update(self):
     rospy.init_node("smach_visualizer", anonymous=False, disable_signals=True, log_level=rospy.INFO)
     rospy.Subscriber(
         STRUCTURE_TOPIC,
-        SmachContainerStructure,
+        StateMachineStructure,
         state_machine.container_structure_callback,
     )
     rospy.Subscriber(
         STATUS_TOPIC,
-        SmachContainerStatus,
+        StateMachineStateUpdate,
         state_machine.container_status_callback,
     )
     app = QApplication([])  # type: ignore
diff --git a/src/navigation/navigation.py b/src/navigation/navigation.py
index 78088756b..c58d127da 100755
--- a/src/navigation/navigation.py
+++ b/src/navigation/navigation.py
@@ -45,7 +45,9 @@ def __init__(self, context: Context):
         )
         self.state_machine.add_transitions(
             SearchState(),
-            [ApproachTargetState(), LongRangeState(), WaypointState(), RecoveryState()],
+            [
+                ApproachTargetState(), LongRangeState(), WaypointState(), RecoveryState()
+            ],
         )
         self.state_machine.add_transitions(DoneState(), [WaypointState()])
         self.state_machine.add_transitions(