Merge "Add theta mode back to graph edit"
diff --git a/y2023/control_loops/python/graph_edit.py b/y2023/control_loops/python/graph_edit.py
index 204a9f9..011095b 100644
--- a/y2023/control_loops/python/graph_edit.py
+++ b/y2023/control_loops/python/graph_edit.py
@@ -227,10 +227,15 @@
         self.spline_edit = 0
         self.edit_control1 = True
 
-        self.roll_joint_thetas = None
-        self.roll_joint_point = None
+        self.joint_thetas = None
+        self.joint_points = None
         self.fig = plt.figure()
-        self.ax = self.fig.add_subplot(111)
+        self.axes = [
+            self.fig.add_subplot(3, 1, 1),
+            self.fig.add_subplot(3, 1, 2),
+            self.fig.add_subplot(3, 1, 3)
+        ]
+        self.fig.subplots_adjust(hspace=1.0)
         plt.show(block=False)
 
         self.index = 0
@@ -368,25 +373,34 @@
             self.outline.draw_theta(cr)
 
         set_color(cr, Color(0.0, 0.5, 1.0))
-        for segment in self.segments:
-            color = [0.2, random.random(), random.random()]
+        for i in range(len(self.segments)):
+            color = None
+            if i == self.index:
+                # Draw current spline in black
+                color = [0, 0, 0]
+            else:
+                color = [0.2, random.random(), random.random()]
             set_color(cr, Color(color[0], color[1], color[2]))
-            segment.DrawTo(cr, self.theta_version)
+            self.segments[i].DrawTo(cr, self.theta_version)
             with px(cr):
                 cr.stroke()
 
         set_color(cr, Color(0.0, 1.0, 0.5))
 
-        # Create the roll joint plot
-        if self.roll_joint_thetas:
-            self.ax.clear()
-            self.ax.plot(*self.roll_joint_thetas)
-            if self.roll_joint_point:
-                self.ax.scatter([self.roll_joint_point[0]],
-                                [self.roll_joint_point[1]],
-                                s=10,
-                                c="red")
-            plt.title("Roll Joint Angle")
+        # Create the plots
+        if self.joint_thetas:
+            if self.joint_points:
+                titles = ["Proximal", "Distal", "Roll joint"]
+                for i in range(len(self.joint_points)):
+                    self.axes[i].clear()
+                    self.axes[i].plot(self.joint_thetas[0],
+                                      self.joint_thetas[1][i])
+                    self.axes[i].scatter([self.joint_points[i][0]],
+                                         [self.joint_points[i][1]],
+                                         s=10,
+                                         c="red")
+                    self.axes[i].set_title(titles[i])
+            plt.title("Joint Angle")
             plt.xlabel("t (0 to 1)")
             plt.ylabel("theta (rad)")
 
@@ -408,23 +422,26 @@
         event.x = x / scale + self.center[0]
         event.y = y / scale + self.center[1]
 
-        for segment in self.segments:
-            self.roll_joint_thetas = segment.roll_joint_thetas()
+        segment = self.segments[self.index]
+        self.joint_thetas = segment.joint_thetas()
 
-            hovered_t = segment.intersection(event)
-            if hovered_t:
-                min_diff = np.inf
-                closest_t = None
-                closest_theta = None
-                for i in range(len(self.roll_joint_thetas[0])):
-                    t = self.roll_joint_thetas[0][i]
-                    diff = abs(t - hovered_t)
-                    if diff < min_diff:
-                        min_diff = diff
-                        closest_t = t
-                        closest_theta = self.roll_joint_thetas[1][i]
-                self.roll_joint_point = (closest_t, closest_theta)
-                break
+        hovered_t = segment.intersection(event)
+        if hovered_t:
+            min_diff = np.inf
+            closest_t = None
+            closest_thetas = None
+            for i in range(len(self.joint_thetas[0])):
+                t = self.joint_thetas[0][i]
+                diff = abs(t - hovered_t)
+                if diff < min_diff:
+                    min_diff = diff
+                    closest_t = t
+                    closest_thetas = [
+                        self.joint_thetas[1][0][i], self.joint_thetas[1][1][i],
+                        self.joint_thetas[1][2][i]
+                    ]
+            self.joint_points = [(closest_t, closest_theta)
+                                 for closest_theta in closest_thetas]
 
         event.x = o_x
         event.y = o_y
diff --git a/y2023/control_loops/python/graph_tools.py b/y2023/control_loops/python/graph_tools.py
index db703da..769c7b3 100644
--- a/y2023/control_loops/python/graph_tools.py
+++ b/y2023/control_loops/python/graph_tools.py
@@ -357,7 +357,7 @@
         pass
 
     @abc.abstractmethod
-    def roll_joint_thetas(self):
+    def joint_thetas(self):
         pass
 
     @abc.abstractmethod
@@ -490,12 +490,16 @@
     def get_controls_theta(self):
         return (self.start, self.control1, self.control2, self.end)
 
-    def roll_joint_thetas(self):
+    def joint_thetas(self):
         ts = []
-        thetas = []
+        thetas = [[], [], []]
         for alpha in subdivide_multistep():
+            proximal, distal = spline_eval(self.start, self.control1,
+                                           self.control2, self.end, alpha)
             roll_joint = get_roll_joint_theta_multistep(
                 self.alpha_rolls, alpha)
-            thetas.append(roll_joint)
+            thetas[0].append(proximal)
+            thetas[1].append(distal)
+            thetas[2].append(roll_joint)
             ts.append(alpha)
         return ts, thetas
diff --git a/y2023/control_loops/superstructure/arm/arm.cc b/y2023/control_loops/superstructure/arm/arm.cc
index 07d54ec..fd265c4 100644
--- a/y2023/control_loops/superstructure/arm/arm.cc
+++ b/y2023/control_loops/superstructure/arm/arm.cc
@@ -51,6 +51,31 @@
 
 void Arm::Reset() { state_ = ArmState::UNINITIALIZED; }
 
+namespace {
+
+// Proximal joint center in xy space
+constexpr std::pair<double, double> kJointCenter = {-0.203, 0.787};
+
+std::tuple<double, double, int> ArmThetasToXY(double theta_proximal,
+                                              double theta_distal) {
+  double theta_proximal_shifted = M_PI / 2.0 - theta_proximal;
+  double theta_distal_shifted = M_PI / 2.0 - theta_distal;
+
+  double x = std::cos(theta_proximal_shifted) * kArmConstants.l0 +
+             std::cos(theta_distal_shifted) * kArmConstants.l1 +
+             kJointCenter.first;
+  double y = std::sin(theta_proximal_shifted) * kArmConstants.l0 +
+             std::sin(theta_distal_shifted) * kArmConstants.l1 +
+             kJointCenter.second;
+
+  int circular_index =
+      std::floor((theta_distal_shifted - theta_proximal_shifted) / M_PI);
+
+  return std::make_tuple(x, y, circular_index);
+}
+
+}  // namespace
+
 flatbuffers::Offset<superstructure::ArmStatus> Arm::Iterate(
     const ::aos::monotonic_clock::time_point /*monotonic_now*/,
     const uint32_t *unsafe_goal, const superstructure::ArmPosition *position,
@@ -273,6 +298,9 @@
       roll_joint_estimator_state_offset =
           roll_joint_zeroing_estimator_.GetEstimatorState(fbb);
 
+  const auto [arm_x, arm_y, arm_circular_index] =
+      ArmThetasToXY(arm_ekf_.X_hat(0), arm_ekf_.X_hat(2));
+
   superstructure::ArmStatus::Builder status_builder(*fbb);
   status_builder.add_proximal_estimator_state(proximal_estimator_state_offset);
   status_builder.add_distal_estimator_state(distal_estimator_state_offset);
@@ -296,6 +324,10 @@
   status_builder.add_voltage_error1(arm_ekf_.X_hat(5));
   status_builder.add_voltage_error2(roll_joint_loop_.X_hat(2));
 
+  status_builder.add_arm_x(arm_x);
+  status_builder.add_arm_y(arm_y);
+  status_builder.add_arm_circular_index(arm_circular_index);
+
   if (!disable) {
     *proximal_output = ::std::max(
         -kOperatingVoltage(), ::std::min(kOperatingVoltage(), follower_.U(0)));
diff --git a/y2023/control_loops/superstructure/superstructure_status.fbs b/y2023/control_loops/superstructure/superstructure_status.fbs
index 2555134..2966d75 100644
--- a/y2023/control_loops/superstructure/superstructure_status.fbs
+++ b/y2023/control_loops/superstructure/superstructure_status.fbs
@@ -44,6 +44,12 @@
   voltage_error1:float (id: 13);
   voltage_error2:float (id: 23);
 
+  // Current arm position in meters for use with UI
+  arm_x:float (id: 24);
+  arm_y:float (id: 25);
+  // Circular index to handle theta wrapping
+  arm_circular_index:int (id: 26);
+
   // True if we are zeroed.
   zeroed:bool (id: 14);