Pull some year-generic code out of y2022 localizer

Pull out some utilities that I want to re-use in 2023.

Change-Id: I5cd03e942708d240441b8356f37eed029e0d8710
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index 86c1ec0..7e0500e 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -831,3 +831,17 @@
         "//aos/network/www:proxy",
     ],
 )
+
+cc_library(
+    name = "localization_utils",
+    srcs = ["localization_utils.cc"],
+    hdrs = ["localization_utils.h"],
+    deps = [
+        ":drivetrain_output_fbs",
+        "//aos/events:event_loop",
+        "//aos/network:message_bridge_server_fbs",
+        "//frc971/input:joystick_state_fbs",
+        "//frc971/vision:calibration_fbs",
+        "@org_tuxfamily_eigen//:eigen",
+    ],
+)
diff --git a/frc971/control_loops/drivetrain/localization_utils.cc b/frc971/control_loops/drivetrain/localization_utils.cc
new file mode 100644
index 0000000..c7968e1
--- /dev/null
+++ b/frc971/control_loops/drivetrain/localization_utils.cc
@@ -0,0 +1,72 @@
+#include "frc971/control_loops/drivetrain/localization_utils.h"
+
+namespace frc971::control_loops::drivetrain {
+
+LocalizationUtils::LocalizationUtils(aos::EventLoop *event_loop)
+    : output_fetcher_(event_loop->MakeFetcher<Output>("/drivetrain")),
+      clock_offset_fetcher_(
+          event_loop->MakeFetcher<aos::message_bridge::ServerStatistics>(
+              "/aos")),
+      joystick_state_fetcher_(
+          event_loop->MakeFetcher<aos::JoystickState>("/roborio/aos")) {}
+
+Eigen::Vector2d LocalizationUtils::VoltageOrZero(
+    aos::monotonic_clock::time_point now) {
+  output_fetcher_.Fetch();
+  // Determine if the robot is likely to be disabled currently.
+  const bool disabled = (output_fetcher_.get() == nullptr) ||
+                        (output_fetcher_.context().monotonic_event_time +
+                             std::chrono::milliseconds(10) <
+                         now);
+  return disabled ? Eigen::Vector2d::Zero()
+                  : Eigen::Vector2d{output_fetcher_->left_voltage(),
+                                    output_fetcher_->right_voltage()};
+}
+
+bool LocalizationUtils::MaybeInAutonomous() {
+  joystick_state_fetcher_.Fetch();
+  return (joystick_state_fetcher_.get() != nullptr)
+             ? joystick_state_fetcher_->autonomous()
+             : true;
+}
+
+std::optional<aos::monotonic_clock::duration> LocalizationUtils::ClockOffset(
+    std::string_view node) {
+  std::optional<aos::monotonic_clock::duration> monotonic_offset;
+  clock_offset_fetcher_.Fetch();
+  if (clock_offset_fetcher_.get() != nullptr) {
+    for (const auto connection : *clock_offset_fetcher_->connections()) {
+      if (connection->has_node() && connection->node()->has_name() &&
+          connection->node()->name()->string_view() == node) {
+        if (connection->has_monotonic_offset()) {
+          monotonic_offset =
+              std::chrono::nanoseconds(connection->monotonic_offset());
+        } else {
+          // If we don't have a monotonic offset, that means we aren't
+          // connected.
+          return std::nullopt;
+        }
+        break;
+      }
+    }
+  }
+  CHECK(monotonic_offset.has_value());
+  return monotonic_offset;
+}
+
+// Technically, this should be able to do a single memcpy, but the extra
+// verbosity here seems appropriate.
+Eigen::Matrix<double, 4, 4> FlatbufferToTransformationMatrix(
+    const frc971::vision::calibration::TransformationMatrix &flatbuffer) {
+  CHECK_EQ(16u, CHECK_NOTNULL(flatbuffer.data())->size());
+  Eigen::Matrix<double, 4, 4> result;
+  result.setIdentity();
+  for (int row = 0; row < 4; ++row) {
+    for (int col = 0; col < 4; ++col) {
+      result(row, col) = (*flatbuffer.data())[row * 4 + col];
+    }
+  }
+  return result;
+}
+
+}  // namespace frc971::control_loops::drivetrain
diff --git a/frc971/control_loops/drivetrain/localization_utils.h b/frc971/control_loops/drivetrain/localization_utils.h
new file mode 100644
index 0000000..26242f9
--- /dev/null
+++ b/frc971/control_loops/drivetrain/localization_utils.h
@@ -0,0 +1,51 @@
+#ifndef FRC971_CONTROL_LOOPS_DRIVETRAIN_LOCALIZATION_UTILS_H_
+#define FRC971_CONTROL_LOOPS_DRIVETRAIN_LOCALIZATION_UTILS_H_
+#include <Eigen/Dense>
+
+#include "aos/events/event_loop.h"
+#include "aos/network/message_bridge_server_generated.h"
+#include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
+#include "frc971/input/joystick_state_generated.h"
+#include "frc971/vision/calibration_generated.h"
+
+namespace frc971::control_loops::drivetrain {
+// This class provides a variety of checks that have generally proved useful for
+// the localizer but which have no clear place to live otherwise.
+// Specifically, it tracks:
+// * Drivetrain voltages, including checks for whether the Output message
+//   has timed out.
+// * Offsets between monotonic clocks on different devices.
+// * Whether we are in autonomous mode.
+class LocalizationUtils {
+ public:
+  LocalizationUtils(aos::EventLoop *event_loop);
+
+  // Returns the latest drivetrain output voltage, or zero if no output is
+  // available (which happens when the robot is disabled; when the robot is
+  // disabled, the voltage is functionally zero). Return value will be
+  // [left_voltage, right_voltage]
+  Eigen::Vector2d VoltageOrZero(aos::monotonic_clock::time_point now);
+
+  // Returns true if either there is no JoystickState message available or if
+  // we are currently in autonomous mode.
+  bool MaybeInAutonomous();
+
+  // Returns the offset between our node and the specified node (or nullopt if
+  // no offset is available). The sign of this will be such that the time on
+  // the remote node = time on our node + ClockOffset().
+  std::optional<aos::monotonic_clock::duration> ClockOffset(
+      std::string_view node);
+
+ private:
+  aos::Fetcher<frc971::control_loops::drivetrain::Output> output_fetcher_;
+  aos::Fetcher<aos::message_bridge::ServerStatistics> clock_offset_fetcher_;
+  aos::Fetcher<aos::JoystickState> joystick_state_fetcher_;
+};
+
+// Converts a flatbuffer TransformationMatrix to an Eigen matrix.
+Eigen::Matrix<double, 4, 4> FlatbufferToTransformationMatrix(
+    const frc971::vision::calibration::TransformationMatrix &flatbuffer);
+
+}  // namespace frc971::control_loops::drivetrain
+
+#endif  // FRC971_CONTROL_LOOPS_DRIVETRAIN_LOCALIZATION_UTILS_H_
diff --git a/y2022/localizer/BUILD b/y2022/localizer/BUILD
index e5c054b..a3b1774 100644
--- a/y2022/localizer/BUILD
+++ b/y2022/localizer/BUILD
@@ -107,6 +107,7 @@
         "//frc971/control_loops/drivetrain:drivetrain_output_fbs",
         "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
         "//frc971/control_loops/drivetrain:improved_down_estimator",
+        "//frc971/control_loops/drivetrain:localization_utils",
         "//frc971/control_loops/drivetrain:localizer_fbs",
         "//frc971/imu_reader:imu_watcher",
         "//frc971/input:joystick_state_fbs",
diff --git a/y2022/localizer/localizer.cc b/y2022/localizer/localizer.cc
index e7dfdf1..46eff94 100644
--- a/y2022/localizer/localizer.cc
+++ b/y2022/localizer/localizer.cc
@@ -871,18 +871,10 @@
       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")),
-      clock_offset_fetcher_(
-          event_loop_->MakeFetcher<aos::message_bridge::ServerStatistics>(
-              "/aos")),
       superstructure_fetcher_(
           event_loop_
               ->MakeFetcher<y2022::control_loops::superstructure::Status>(
                   "/superstructure")),
-      joystick_state_fetcher_(
-          event_loop_->MakeFetcher<aos::JoystickState>("/roborio/aos")),
       imu_watcher_(event_loop, dt_config,
                    y2022::constants::Values::DrivetrainEncoderToMeters(1),
                    [this](aos::monotonic_clock::time_point sample_time_pico,
@@ -891,7 +883,8 @@
                           Eigen::Vector3d gyro, Eigen::Vector3d accel) {
                      HandleImu(sample_time_pico, sample_time_pi, encoders, gyro,
                                accel);
-                   }) {
+                   }),
+      utils_(event_loop) {
   event_loop->SetRuntimeRealtimePriority(10);
   event_loop_->MakeWatcher(
       "/drivetrain",
@@ -929,10 +922,7 @@
             absl::StrCat("/", kPisToUse[camera_index], "/camera"));
   }
   aos::TimerHandler *estimate_timer = event_loop_->AddTimer([this]() {
-      joystick_state_fetcher_.Fetch();
-      const bool maybe_in_auto = (joystick_state_fetcher_.get() != nullptr)
-                                     ? joystick_state_fetcher_->autonomous()
-                                     : true;
+      const bool maybe_in_auto = utils_.MaybeInAutonomous();
       model_based_.set_use_aggressive_image_corrections(!maybe_in_auto);
       for (size_t camera_index = 0; camera_index < kPisToUse.size();
            ++camera_index) {
@@ -946,8 +936,10 @@
         }
         if (target_estimate_fetchers_[camera_index].Fetch()) {
           const std::optional<aos::monotonic_clock::duration> monotonic_offset =
-              ClockOffset(kPisToUse[camera_index]);
+              utils_.ClockOffset(kPisToUse[camera_index]);
           if (!monotonic_offset.has_value()) {
+            model_based_.TallyRejection(
+                RejectionReason::MESSAGE_BRIDGE_DISCONNECTED);
             continue;
           }
           // TODO(james): Get timestamp from message contents.
@@ -980,20 +972,11 @@
     aos::monotonic_clock::time_point sample_time_pi,
     std::optional<Eigen::Vector2d> encoders, Eigen::Vector3d gyro,
     Eigen::Vector3d accel) {
-  output_fetcher_.Fetch();
-  {
-    const bool disabled = (output_fetcher_.get() == nullptr) ||
-                          (output_fetcher_.context().monotonic_event_time +
-                               std::chrono::milliseconds(10) <
-                           event_loop_->context().monotonic_event_time);
-    // For gyros, use the most recent gyro reading if we aren't zeroed,
-    // to avoid missing integration cycles.
-    model_based_.HandleImu(
-        sample_time_pico, gyro, accel, encoders,
-        disabled ? Eigen::Vector2d::Zero()
-                 : Eigen::Vector2d{output_fetcher_->left_voltage(),
-                                   output_fetcher_->right_voltage()});
-  }
+
+  model_based_.HandleImu(
+      sample_time_pico, gyro, accel, encoders,
+      utils_.VoltageOrZero(event_loop_->context().monotonic_event_time));
+
   {
     auto builder = status_sender_.MakeBuilder();
     const flatbuffers::Offset<ModelBasedStatus> model_based_status =
@@ -1052,30 +1035,4 @@
   }
 }
 
-std::optional<aos::monotonic_clock::duration> EventLoopLocalizer::ClockOffset(
-    std::string_view pi) {
-  std::optional<aos::monotonic_clock::duration> monotonic_offset;
-  clock_offset_fetcher_.Fetch();
-  if (clock_offset_fetcher_.get() != nullptr) {
-    for (const auto connection : *clock_offset_fetcher_->connections()) {
-      if (connection->has_node() && connection->node()->has_name() &&
-          connection->node()->name()->string_view() == pi) {
-        if (connection->has_monotonic_offset()) {
-          monotonic_offset =
-              std::chrono::nanoseconds(connection->monotonic_offset());
-        } else {
-          // If we don't have a monotonic offset, that means we aren't
-          // connected.
-          model_based_.TallyRejection(
-              RejectionReason::MESSAGE_BRIDGE_DISCONNECTED);
-          return std::nullopt;
-        }
-        break;
-      }
-    }
-  }
-  CHECK(monotonic_offset.has_value());
-  return monotonic_offset;
-}
-
 }  // namespace frc971::controls
diff --git a/y2022/localizer/localizer.h b/y2022/localizer/localizer.h
index 917f131..a403ca8 100644
--- a/y2022/localizer/localizer.h
+++ b/y2022/localizer/localizer.h
@@ -12,6 +12,7 @@
 #include "frc971/input/joystick_state_generated.h"
 #include "frc971/control_loops/drivetrain/improved_down_estimator.h"
 #include "frc971/control_loops/drivetrain/localizer_generated.h"
+#include "frc971/control_loops/drivetrain/localization_utils.h"
 #include "frc971/zeroing/imu_zeroer.h"
 #include "frc971/zeroing/wrap.h"
 #include "y2022/control_loops/superstructure/superstructure_status_generated.h"
@@ -326,8 +327,6 @@
   ModelBasedLocalizer *localizer() { return &model_based_; }
 
  private:
-  std::optional<aos::monotonic_clock::duration> ClockOffset(
-      std::string_view pi);
   void HandleImu(aos::monotonic_clock::time_point sample_time_pico,
                  aos::monotonic_clock::time_point sample_time_pi,
                  std::optional<Eigen::Vector2d> encoders, Eigen::Vector3d gyro,
@@ -337,21 +336,18 @@
   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_;
   std::array<aos::Fetcher<y2022::vision::TargetEstimate>,
              ModelBasedLocalizer::kNumPis>
       target_estimate_fetchers_;
   aos::Fetcher<y2022::control_loops::superstructure::Status>
       superstructure_fetcher_;
-  aos::Fetcher<aos::JoystickState> joystick_state_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;
 
   ImuWatcher imu_watcher_;
+  control_loops::drivetrain::LocalizationUtils utils_;
 };
 }  // namespace frc971::controls
 #endif  // Y2022_LOCALIZER_LOCALIZER_H_