Add more debug information to 2022 localizer

Adds data in a similar format to 2020's LocalizerDebug message,
specifically intended for consumption by live debugging webpages
where attempting to send large 2 kHz messages (the normal
LocalizerStatus) may not work well, especially on the real field.

Change-Id: I747b66668567760a42c77653e2d8c05b0ebfa4f3
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/y2022/BUILD b/y2022/BUILD
index 4589c98..1a5f692 100644
--- a/y2022/BUILD
+++ b/y2022/BUILD
@@ -116,6 +116,7 @@
         "//aos/network:remote_message_fbs",
         "//y2022/localizer:localizer_status_fbs",
         "//y2022/localizer:localizer_output_fbs",
+        "//y2022/localizer:localizer_visualization_fbs",
     ],
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
diff --git a/y2022/localizer/BUILD b/y2022/localizer/BUILD
index 9deab1a..9e8c3d4 100644
--- a/y2022/localizer/BUILD
+++ b/y2022/localizer/BUILD
@@ -1,4 +1,4 @@
-load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_ts_library")
 load("//aos:flatbuffers.bzl", "cc_static_flatbuffer")
 load("@npm//@bazel/typescript:index.bzl", "ts_library")
 
@@ -39,6 +39,30 @@
     visibility = ["//visibility:public"],
 )
 
+flatbuffer_cc_library(
+    name = "localizer_visualization_fbs",
+    srcs = ["localizer_visualization.fbs"],
+    gen_reflections = 1,
+    includes = [
+        ":localizer_status_fbs_includes",
+        "//frc971/control_loops:control_loops_fbs_includes",
+        "//frc971/control_loops/drivetrain:drivetrain_status_fbs_includes",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+flatbuffer_ts_library(
+    name = "localizer_visualization_ts_fbs",
+    srcs = ["localizer_visualization.fbs"],
+    includes = [
+        ":localizer_status_fbs_includes",
+        "//frc971/control_loops:control_loops_fbs_includes",
+        "//frc971/control_loops/drivetrain:drivetrain_status_fbs_includes",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+)
+
 cc_static_flatbuffer(
     name = "localizer_schema",
     function = "frc971::controls::LocalizerStatusSchema",
@@ -54,7 +78,9 @@
     deps = [
         ":localizer_output_fbs",
         ":localizer_status_fbs",
+        ":localizer_visualization_fbs",
         "//aos/containers:ring_buffer",
+        "//aos/containers:sized_array",
         "//aos/events:event_loop",
         "//aos/network:message_bridge_server_fbs",
         "//aos/time",
diff --git a/y2022/localizer/localizer.cc b/y2022/localizer/localizer.cc
index d0e9704..ea80e73 100644
--- a/y2022/localizer/localizer.cc
+++ b/y2022/localizer/localizer.cc
@@ -483,7 +483,7 @@
 }
 
 // Node names of the pis to listen for cameras from.
-const std::array<std::string, 4> kPisToUse{"pi1", "pi2", "pi3", "pi4"};
+const std::array<std::string_view, 4> kPisToUse{"pi1", "pi2", "pi3", "pi4"};
 }
 
 const Eigen::Matrix<double, 4, 4> ModelBasedLocalizer::CameraTransform(
@@ -524,7 +524,8 @@
 const std::optional<Eigen::Vector2d>
 ModelBasedLocalizer::CameraMeasuredRobotPosition(
     const OldPosition &state, const y2022::vision::TargetEstimate *target,
-    std::optional<RejectionReason> *rejection_reason) const {
+    std::optional<RejectionReason> *rejection_reason,
+    Eigen::Matrix<double, 4, 4> *H_field_camera_measured) const {
   if (!target->has_camera_calibration()) {
     *rejection_reason = RejectionReason::NO_CALIBRATION;
     return std::nullopt;
@@ -567,26 +568,28 @@
       aos::math::NormalizeAngle(M_PI + camera_yaw + target->angle_to_target());
   const double distance_from_target = target->distance();
   // Extract the implied camera position on the field.
-  Eigen::Matrix<double, 4, 4> H_field_camera_measured = H_field_camera;
+  *H_field_camera_measured = H_field_camera;
   // TODO(james): Are we going to need to evict the roll/pitch components of the
   // camera extrinsics this year as well?
-  H_field_camera_measured(0, 3) =
+  (*H_field_camera_measured)(0, 3) =
       distance_from_target * std::cos(heading_from_target) + kVisionTargetX;
-  H_field_camera_measured(1, 3) =
+  (*H_field_camera_measured)(1, 3) =
       distance_from_target * std::sin(heading_from_target) + kVisionTargetY;
   const Eigen::Matrix<double, 4, 4> H_field_robot_measured =
-      H_field_camera_measured * H_robot_camera.inverse();
+      *H_field_camera_measured * H_robot_camera.inverse();
   return H_field_robot_measured.block<2, 1>(0, 3);
 }
 
 void ModelBasedLocalizer::HandleImageMatch(
     aos::monotonic_clock::time_point sample_time,
-    const y2022::vision::TargetEstimate *target) {
+    const y2022::vision::TargetEstimate *target, int camera_index) {
   std::optional<RejectionReason> rejection_reason;
 
   const OldPosition &state = GetStateForTime(sample_time);
+  Eigen::Matrix<double, 4, 4> H_field_camera_measured;
   const std::optional<Eigen::Vector2d> measured_robot_position =
-      CameraMeasuredRobotPosition(state, target, &rejection_reason);
+      CameraMeasuredRobotPosition(state, target, &rejection_reason,
+                                  &H_field_camera_measured);
   // Technically, rejection_reason should always be set if
   // measured_robot_position is nullopt, but in the future we may have more
   // recoverable rejection reasons that we wish to allow to propagate further
@@ -648,6 +651,25 @@
 
   statistics_.total_accepted++;
   statistics_.total_candidates++;
+
+  const Eigen::Vector3d camera_z_in_field =
+      H_field_camera_measured.block<3, 3>(0, 0) * Eigen::Vector3d::UnitZ();
+  const double camera_yaw =
+      std::atan2(camera_z_in_field.y(), camera_z_in_field.x());
+
+  TargetEstimateDebugT debug;
+  debug.camera = static_cast<uint8_t>(camera_index);
+  debug.camera_x = H_field_camera_measured(0, 3);
+  debug.camera_y = H_field_camera_measured(1, 3);
+  debug.camera_theta = camera_yaw;
+  debug.implied_robot_x = measured_robot_position.value().x();
+  debug.implied_robot_y = measured_robot_position.value().y();
+  debug.implied_robot_theta = xytheta()(2);
+  debug.implied_turret_goal =
+      aos::math::NormalizeAngle(camera_yaw + target->angle_to_target());
+  debug.accepted = true;
+  debug.image_age_sec = aos::time::DurationInSeconds(t_ - sample_time);
+  image_debugs_.push_back(debug);
 }
 
 void ModelBasedLocalizer::HandleTurret(
@@ -700,8 +722,8 @@
   return model_state_builder.Finish();
 }
 
-flatbuffers::Offset<ModelBasedStatus> ModelBasedLocalizer::PopulateStatus(
-    flatbuffers::FlatBufferBuilder *fbb) {
+flatbuffers::Offset<CumulativeStatistics>
+ModelBasedLocalizer::PopulateStatistics(flatbuffers::FlatBufferBuilder *fbb) {
   const auto rejections_offset = fbb->CreateVector(
       statistics_.rejection_counts.data(), statistics_.rejection_counts.size());
 
@@ -709,8 +731,13 @@
   stats_builder.add_total_accepted(statistics_.total_accepted);
   stats_builder.add_total_candidates(statistics_.total_candidates);
   stats_builder.add_rejection_reason_count(rejections_offset);
+  return stats_builder.Finish();
+}
+
+flatbuffers::Offset<ModelBasedStatus> ModelBasedLocalizer::PopulateStatus(
+    flatbuffers::FlatBufferBuilder *fbb) {
   const flatbuffers::Offset<CumulativeStatistics> stats_offset =
-      stats_builder.Finish();
+      PopulateStatistics(fbb);
 
   const flatbuffers::Offset<control_loops::drivetrain::DownEstimatorState>
       down_estimator_offset = down_estimator_.PopulateStatus(fbb, t_);
@@ -754,6 +781,58 @@
   return builder.Finish();
 }
 
+flatbuffers::Offset<LocalizerVisualization>
+ModelBasedLocalizer::PopulateVisualization(
+    flatbuffers::FlatBufferBuilder *fbb) {
+  const flatbuffers::Offset<CumulativeStatistics> stats_offset =
+      PopulateStatistics(fbb);
+
+  aos::SizedArray<flatbuffers::Offset<TargetEstimateDebug>, kDebugBufferSize>
+      debug_offsets;
+
+  for (const TargetEstimateDebugT& debug : image_debugs_) {
+    debug_offsets.push_back(PackTargetEstimateDebug(debug, fbb));
+  }
+
+  image_debugs_.clear();
+
+  const flatbuffers::Offset<
+      flatbuffers::Vector<flatbuffers::Offset<TargetEstimateDebug>>>
+      debug_offset =
+          fbb->CreateVector(debug_offsets.data(), debug_offsets.size());
+
+  LocalizerVisualization::Builder builder(*fbb);
+  builder.add_statistics(stats_offset);
+  builder.add_targets(debug_offset);
+  return builder.Finish();
+}
+
+void ModelBasedLocalizer::TallyRejection(const RejectionReason reason) {
+  statistics_.total_candidates++;
+  statistics_.rejection_counts[static_cast<size_t>(reason)]++;
+  TargetEstimateDebugT debug;
+  debug.accepted = false;
+  debug.rejection_reason = reason;
+  image_debugs_.push_back(debug);
+}
+
+flatbuffers::Offset<TargetEstimateDebug>
+ModelBasedLocalizer::PackTargetEstimateDebug(
+    const TargetEstimateDebugT &debug, flatbuffers::FlatBufferBuilder *fbb) {
+  if (!debug.accepted) {
+    TargetEstimateDebug::Builder builder(*fbb);
+    builder.add_accepted(debug.accepted);
+    builder.add_rejection_reason(debug.rejection_reason);
+    return builder.Finish();
+  } else {
+    flatbuffers::Offset<TargetEstimateDebug> offset =
+        TargetEstimateDebug::Pack(*fbb, &debug);
+    flatbuffers::GetMutableTemporaryPointer(*fbb, offset)
+        ->clear_rejection_reason();
+    return offset;
+  }
+}
+
 namespace {
 // Period at which the encoder readings from the IMU board wrap.
 static double DrivetrainWrapPeriod() {
@@ -768,6 +847,8 @@
       model_based_(dt_config),
       status_sender_(event_loop_->MakeSender<LocalizerStatus>("/localizer")),
       output_sender_(event_loop_->MakeSender<LocalizerOutput>("/localizer")),
+      visualization_sender_(
+          event_loop_->MakeSender<LocalizerVisualization>("/localizer")),
       output_fetcher_(
           event_loop_->MakeFetcher<frc971::control_loops::drivetrain::Output>(
               "/drivetrain")),
@@ -798,12 +879,12 @@
                                   status.turret()->velocity());
       });
 
-  for (const auto &pi : kPisToUse) {
+  for (size_t camera_index = 0; camera_index < kPisToUse.size(); ++camera_index) {
     event_loop_->MakeWatcher(
-        "/" + pi + "/camera",
-        [this, pi](const y2022::vision::TargetEstimate &target) {
+        absl::StrCat("/", kPisToUse[camera_index], "/camera"),
+        [this, camera_index](const y2022::vision::TargetEstimate &target) {
           const std::optional<aos::monotonic_clock::duration> monotonic_offset =
-              ClockOffset(pi);
+              ClockOffset(kPisToUse[camera_index]);
           if (!monotonic_offset.has_value()) {
             return;
           }
@@ -814,7 +895,15 @@
             model_based_.TallyRejection(RejectionReason::IMAGE_FROM_FUTURE);
             return;
           }
-          model_based_.HandleImageMatch(capture_time, &target);
+          model_based_.HandleImageMatch(capture_time, &target, camera_index);
+          if (model_based_.NumQueuedImageDebugs() ==
+                  ModelBasedLocalizer::kDebugBufferSize ||
+              (last_visualization_send_ + kMinVisualizationPeriod <
+               event_loop_->monotonic_now())) {
+            auto builder = visualization_sender_.MakeBuilder();
+            visualization_sender_.CheckOk(
+                builder.Send(model_based_.PopulateVisualization(builder.fbb())));
+          }
         });
   }
   event_loop_->MakeWatcher(
diff --git a/y2022/localizer/localizer.h b/y2022/localizer/localizer.h
index 978eb15..4b7eff0 100644
--- a/y2022/localizer/localizer.h
+++ b/y2022/localizer/localizer.h
@@ -4,6 +4,7 @@
 #include "Eigen/Dense"
 #include "Eigen/Geometry"
 #include "aos/containers/ring_buffer.h"
+#include "aos/containers/sized_array.h"
 #include "aos/events/event_loop.h"
 #include "aos/network/message_bridge_server_generated.h"
 #include "aos/time/time.h"
@@ -15,6 +16,7 @@
 #include "y2022/control_loops/superstructure/superstructure_status_generated.h"
 #include "y2022/localizer/localizer_output_generated.h"
 #include "y2022/localizer/localizer_status_generated.h"
+#include "y2022/localizer/localizer_visualization_generated.h"
 #include "y2022/vision/target_estimate_generated.h"
 
 namespace frc971::controls {
@@ -95,6 +97,8 @@
   // ~20 gives ~55-60% CPU.
   static constexpr int kBranchPeriod = 100;
 
+  static constexpr size_t kDebugBufferSize = 10;
+
   typedef Eigen::Matrix<double, kNModelStates, 1> ModelState;
   typedef Eigen::Matrix<double, kNAccelStates, 1> AccelState;
   typedef Eigen::Matrix<double, kNModelInputs, 1> ModelInput;
@@ -108,7 +112,8 @@
   void HandleTurret(aos::monotonic_clock::time_point sample_time,
                     double turret_position, double turret_velocity);
   void HandleImageMatch(aos::monotonic_clock::time_point sample_time,
-                        const y2022::vision::TargetEstimate *target);
+                        const y2022::vision::TargetEstimate *target,
+                        int camera_index);
   void HandleReset(aos::monotonic_clock::time_point,
                    const Eigen::Vector3d &xytheta);
 
@@ -129,10 +134,12 @@
 
   void set_longitudinal_offset(double offset) { long_offset_ = offset; }
 
-  void TallyRejection(const RejectionReason reason) {
-    statistics_.total_candidates++;
-    statistics_.rejection_counts[static_cast<size_t>(reason)]++;
-  }
+  void TallyRejection(const RejectionReason reason);
+
+  flatbuffers::Offset<LocalizerVisualization> PopulateVisualization(
+      flatbuffers::FlatBufferBuilder *fbb);
+
+  size_t NumQueuedImageDebugs() const { return image_debugs_.size(); }
 
  private:
   struct CombinedState {
@@ -197,9 +204,17 @@
 
   // Returns the robot x/y position implied by the specified camera data and
   // estimated state from that time.
+  // H_field_camera is passed in so that it can be used as a debugging output.
   const std::optional<Eigen::Vector2d> CameraMeasuredRobotPosition(
       const OldPosition &state, const y2022::vision::TargetEstimate *target,
-      std::optional<RejectionReason> *rejection_reason) const;
+      std::optional<RejectionReason> *rejection_reason,
+      Eigen::Matrix<double, 4, 4> *H_field_camera_measured) const;
+
+  flatbuffers::Offset<TargetEstimateDebug> PackTargetEstimateDebug(
+      const TargetEstimateDebugT &debug, flatbuffers::FlatBufferBuilder *fbb);
+
+  flatbuffers::Offset<CumulativeStatistics> PopulateStatistics(
+      flatbuffers::FlatBufferBuilder *fbb);
 
   const control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
   const StateFeedbackHybridPlantCoefficients<2, 2, 2, double>
@@ -282,11 +297,15 @@
   };
   Statistics statistics_;
 
+  aos::SizedArray<TargetEstimateDebugT, kDebugBufferSize> image_debugs_;
+
   friend class testing::LocalizerTest;
 };
 
 class EventLoopLocalizer {
  public:
+  static constexpr std::chrono::milliseconds kMinVisualizationPeriod{100};
+
   EventLoopLocalizer(
       aos::EventLoop *event_loop,
       const control_loops::drivetrain::DrivetrainConfig<double> &dt_config);
@@ -300,11 +319,14 @@
   ModelBasedLocalizer model_based_;
   aos::Sender<LocalizerStatus> status_sender_;
   aos::Sender<LocalizerOutput> output_sender_;
+  aos::Sender<LocalizerVisualization> visualization_sender_;
   aos::Fetcher<frc971::control_loops::drivetrain::Output> output_fetcher_;
   aos::Fetcher<aos::message_bridge::ServerStatistics> clock_offset_fetcher_;
   zeroing::ImuZeroer zeroer_;
   aos::monotonic_clock::time_point last_output_send_ =
       aos::monotonic_clock::min_time;
+  aos::monotonic_clock::time_point last_visualization_send_ =
+      aos::monotonic_clock::min_time;
   std::optional<aos::monotonic_clock::time_point> last_pico_timestamp_;
   aos::monotonic_clock::duration pico_offset_error_;
   // t = pico_offset_ + pico_timestamp.
diff --git a/y2022/localizer/localizer_visualization.fbs b/y2022/localizer/localizer_visualization.fbs
new file mode 100644
index 0000000..6cca9e8
--- /dev/null
+++ b/y2022/localizer/localizer_visualization.fbs
@@ -0,0 +1,26 @@
+include "y2022/localizer/localizer_status.fbs";
+
+namespace frc971.controls;
+
+table TargetEstimateDebug {
+  camera:uint8 (id: 0);
+  camera_x:double (id: 1);
+  camera_y:double (id: 2);
+  camera_theta:double (id: 3);
+  implied_robot_x:double (id: 4);
+  implied_robot_y:double (id: 5);
+  implied_robot_theta:double (id: 6);
+  implied_turret_goal:double (id: 7);
+  accepted:bool (id: 8);
+  rejection_reason:RejectionReason  (id: 9);
+  // Image age (more human-readable than trying to interpret raw nanosecond
+  // values).
+  image_age_sec:double (id: 10);
+}
+
+table LocalizerVisualization {
+  targets:[TargetEstimateDebug] (id: 0);
+  statistics:CumulativeStatistics (id: 1);
+}
+
+root_type LocalizerVisualization;
diff --git a/y2022/y2022_imu.json b/y2022/y2022_imu.json
index 0044ccc..3d97c3c 100644
--- a/y2022/y2022_imu.json
+++ b/y2022/y2022_imu.json
@@ -251,6 +251,37 @@
     },
     {
       "name": "/localizer",
+      "type": "frc971.controls.LocalizerVisualization",
+      "source_node": "imu",
+      "frequency": 200,
+      "max_size": 2000,
+      "logger": "LOCAL_AND_REMOTE_LOGGER",
+      "logger_nodes": [
+        "logger"
+      ],
+      "destination_nodes": [
+        {
+          "name": "logger",
+          "priority": 5,
+          "timestamp_logger": "LOCAL_AND_REMOTE_LOGGER",
+          "timestamp_logger_nodes": [
+            "imu"
+          ],
+          "time_to_live": 5000000
+        }
+      ]
+    },
+    {
+      "name": "/imu/aos/remote_timestamps/logger/localizer/frc971-controls-LocalizerVisualization",
+      "type": "aos.message_bridge.RemoteMessage",
+      "source_node": "imu",
+      "logger": "NOT_LOGGED",
+      "frequency": 200,
+      "num_senders": 2,
+      "max_size": 200
+    },
+    {
+      "name": "/localizer",
       "type": "frc971.controls.LocalizerOutput",
       "source_node": "imu",
       "frequency": 200,