Make ImuWatcher class for handling pico idiosyncracies

This provides a class to consistently handle the IMUValuesBatch coming
from the imu process on a pi.

Change-Id: Id2cdeb23fdd0a3e1c2bd01219678d957906112fa
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/imu_reader/BUILD b/frc971/imu_reader/BUILD
index 7d424c4..4825900 100644
--- a/frc971/imu_reader/BUILD
+++ b/frc971/imu_reader/BUILD
@@ -1,3 +1,6 @@
+load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+load("@com_github_google_flatbuffers//:typescript.bzl", "flatbuffer_ts_library")
+
 cc_library(
     name = "imu",
     srcs = [
@@ -14,8 +17,38 @@
         "//aos/util:crc32",
         "//frc971/wpilib:imu_batch_fbs",
         "//frc971/wpilib:imu_fbs",
-        "//y2022:constants",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/types:span",
     ],
 )
+
+flatbuffer_cc_library(
+    name = "imu_failures_fbs",
+    srcs = [
+        "imu_failures.fbs",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+flatbuffer_ts_library(
+    name = "imu_failures_ts_fbs",
+    srcs = [
+        "imu_failures.fbs",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "imu_watcher",
+    srcs = ["imu_watcher.cc"],
+    hdrs = ["imu_watcher.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":imu_failures_fbs",
+        "//aos/events:event_loop",
+        "//frc971/control_loops/drivetrain:drivetrain_config",
+        "//frc971/wpilib:imu_batch_fbs",
+        "//frc971/zeroing:imu_zeroer",
+        "//frc971/zeroing:wrap",
+    ],
+)
diff --git a/frc971/imu_reader/imu_failures.fbs b/frc971/imu_reader/imu_failures.fbs
new file mode 100644
index 0000000..e6d9d85
--- /dev/null
+++ b/frc971/imu_reader/imu_failures.fbs
@@ -0,0 +1,16 @@
+namespace frc971.controls;
+
+// Counters to track how many times different errors have occurred.
+table ImuFailures {
+  // Count of total number of checksum mismatches between the IMU and the
+  // pico itself.
+  imu_to_pico_checksum_mismatch:uint (id: 0);
+  // Count of total number of checksum mismatches between the pico board
+  // and the raspberry pi.
+  pico_to_pi_checksum_mismatch:uint (id: 1);
+  // Total number of dropped/missed messages.
+  missed_messages:uint (id: 2);
+  // Total number of messages dropped for any other conditions that can fault
+  // the zeroer (e.g., diagnostic failures in the IMU).
+  other_zeroing_faults:uint (id: 3);
+}
diff --git a/frc971/imu_reader/imu_watcher.cc b/frc971/imu_reader/imu_watcher.cc
new file mode 100644
index 0000000..0e0d656
--- /dev/null
+++ b/frc971/imu_reader/imu_watcher.cc
@@ -0,0 +1,117 @@
+#include "frc971/imu_reader/imu_watcher.h"
+
+#include "frc971/wpilib/imu_batch_generated.h"
+
+namespace frc971::controls {
+namespace {
+// Return the amount of distance that the drivetrain can travel before the
+// encoders will wrap. Necessary because the pico only sends over the encoders
+// in 16-bit counters, which will wrap relatively readily.
+double EncoderWrapDistance(double drivetrain_distance_per_encoder_tick) {
+  return drivetrain_distance_per_encoder_tick * (1 << 16);
+}
+}  // namespace
+ImuWatcher::ImuWatcher(
+    aos::EventLoop *event_loop,
+    const control_loops::drivetrain::DrivetrainConfig<double> &dt_config,
+    const double drivetrain_distance_per_encoder_tick,
+    std::function<
+        void(aos::monotonic_clock::time_point, aos::monotonic_clock::time_point,
+             std::optional<Eigen::Vector2d>, Eigen::Vector3d, Eigen::Vector3d)>
+        callback)
+    : dt_config_(dt_config),
+      callback_(std::move(callback)),
+      zeroer_(zeroing::ImuZeroer::FaultBehavior::kTemporary),
+      left_encoder_(
+          -EncoderWrapDistance(drivetrain_distance_per_encoder_tick) / 2.0,
+          EncoderWrapDistance(drivetrain_distance_per_encoder_tick)),
+      right_encoder_(
+          -EncoderWrapDistance(drivetrain_distance_per_encoder_tick) / 2.0,
+          EncoderWrapDistance(drivetrain_distance_per_encoder_tick)) {
+  event_loop->MakeWatcher("/localizer", [this, event_loop](
+                                            const IMUValuesBatch &values) {
+    CHECK(values.has_readings());
+    for (const IMUValues *value : *values.readings()) {
+      zeroer_.InsertAndProcessMeasurement(*value);
+      if (zeroer_.Faulted()) {
+        if (value->checksum_failed()) {
+          imu_fault_tracker_.pico_to_pi_checksum_mismatch++;
+        } else if (value->previous_reading_diag_stat()->checksum_mismatch()) {
+          imu_fault_tracker_.imu_to_pico_checksum_mismatch++;
+        } else {
+          imu_fault_tracker_.other_zeroing_faults++;
+        }
+      } else {
+        if (!first_valid_data_counter_.has_value()) {
+          first_valid_data_counter_ = value->data_counter();
+        }
+      }
+      if (first_valid_data_counter_.has_value()) {
+        total_imu_messages_received_++;
+        // Only update when we have good checksums, since the data counter
+        // could get corrupted.
+        if (!zeroer_.Faulted()) {
+          if (value->data_counter() < last_data_counter_) {
+            data_counter_offset_ += 1 << 16;
+          }
+          imu_fault_tracker_.missed_messages =
+              (1 + value->data_counter() + data_counter_offset_ -
+               first_valid_data_counter_.value()) -
+              total_imu_messages_received_;
+          last_data_counter_ = value->data_counter();
+        }
+      }
+      // Set encoders to nullopt if we are faulted at all (faults may include
+      // checksum mismatches).
+      const std::optional<Eigen::Vector2d> encoders =
+          zeroer_.Faulted()
+              ? std::nullopt
+              : std::make_optional(Eigen::Vector2d{
+                    left_encoder_.Unwrap(value->left_encoder()),
+                    right_encoder_.Unwrap(value->right_encoder())});
+      {
+        // If we can't trust the imu reading, just naively increment the
+        // pico timestamp.
+        const aos::monotonic_clock::time_point pico_timestamp =
+            zeroer_.Faulted()
+                ? (last_pico_timestamp_.has_value()
+                       ? last_pico_timestamp_.value() + kNominalDt
+                       : aos::monotonic_clock::epoch())
+                : aos::monotonic_clock::time_point(
+                      std::chrono::microseconds(value->pico_timestamp_us()));
+        // TODO(james): If we get large enough drift off of the pico,
+        // actually do something about it.
+        if (!pico_offset_.has_value()) {
+          pico_offset_ =
+              event_loop->context().monotonic_event_time - pico_timestamp;
+          last_pico_timestamp_ = pico_timestamp;
+        }
+        if (pico_timestamp < last_pico_timestamp_) {
+          pico_offset_.value() += std::chrono::microseconds(1ULL << 32);
+        }
+        const aos::monotonic_clock::time_point sample_timestamp =
+            pico_offset_.value() + pico_timestamp;
+        pico_offset_error_ =
+            event_loop->context().monotonic_event_time - sample_timestamp;
+        const bool zeroed = zeroer_.Zeroed();
+
+        // When not zeroed, we aim to approximate zero acceleration by doing a
+        // zero-order hold on the gyro and setting the accelerometer readings to
+        // gravity.
+        callback_(sample_timestamp,
+                  aos::monotonic_clock::time_point(std::chrono::nanoseconds(
+                      value->monotonic_timestamp_ns())),
+                  encoders, zeroed ? zeroer_.ZeroedGyro().value() : last_gyro_,
+                  zeroed ? zeroer_.ZeroedAccel().value()
+                         : dt_config_.imu_transform.transpose() *
+                               Eigen::Vector3d::UnitZ());
+
+        if (zeroed) {
+          last_gyro_ = zeroer_.ZeroedGyro().value();
+        }
+        last_pico_timestamp_ = pico_timestamp;
+      }
+    }
+  });
+}
+}  // namespace frc971::controls
diff --git a/frc971/imu_reader/imu_watcher.h b/frc971/imu_reader/imu_watcher.h
new file mode 100644
index 0000000..8867266
--- /dev/null
+++ b/frc971/imu_reader/imu_watcher.h
@@ -0,0 +1,104 @@
+#ifndef FRC971_IMU_READER_IMU_WATCHER_H_
+#define FRC971_IMU_READER_IMU_WATCHER_H_
+
+#include "aos/events/event_loop.h"
+#include "frc971/control_loops/drivetrain/drivetrain_config.h"
+#include "frc971/imu_reader/imu_failures_generated.h"
+#include "frc971/zeroing/imu_zeroer.h"
+#include "frc971/zeroing/wrap.h"
+
+namespace frc971::controls {
+// This class handles listening to an IMUValuesBatch channel sourced off of our
+// ADIS16505 pico board and calling the user-specified callback with the
+// relevant data. This intermediary is used to unwrap encoder readings, check
+// for checksum mismatches, zero the gyro/accelerometer, and translate
+// timestamps between devices.
+// TODO(james): Get unit tests for this class specifically written (we already
+// have tests for the code that exercises this).
+class ImuWatcher {
+ public:
+  // Expected frequency of messages from the pico-based IMU.
+  static constexpr std::chrono::microseconds kNominalDt{500};
+
+  // The callback specified by the user will take:
+  // sample_time_pico: The pico-based timestamp corresponding to the measurement
+  //   time. This will be offset by roughly pico_offset_error from the pi's
+  //   monotonic clock.
+  // sample_time_pi: Timestamp from the kernel for when the pi observed the
+  //   relevant measurement.
+  // encoders: Current encoder values, [left, right]. nullopt if we have faults.
+  // gyro: Current gyro readings, in the raw IMU axes (i.e., these must be
+  //   rotated by dt_config.imu_transform before being used). Suitable
+  //   for input to the down estimator.
+  // accel: Current accelerometer readings, in the raw IMU axes (i.e., these
+  //   must be rotated by dt_config.imu_transform before being used). Suitable
+  //   for input to the down estimator.
+  ImuWatcher(
+      aos::EventLoop *event_loop,
+      const control_loops::drivetrain::DrivetrainConfig<double> &dt_config,
+      double drivetrain_distance_per_encoder_tick,
+      std::function<void(
+          aos::monotonic_clock::time_point, aos::monotonic_clock::time_point,
+          std::optional<Eigen::Vector2d>, Eigen::Vector3d, Eigen::Vector3d)>
+          callback);
+
+  const zeroing::ImuZeroer &zeroer() const { return zeroer_; }
+
+  flatbuffers::Offset<ImuFailures> PopulateImuFailures(
+      flatbuffers::FlatBufferBuilder *fbb) const {
+    return ImuFailures::Pack(*fbb, &imu_fault_tracker_);
+  }
+
+  // t = pico_offset + pico_timestamp.
+  // Note that this can drift over sufficiently long time periods!
+  std::optional<std::chrono::nanoseconds> pico_offset() const {
+    return pico_offset_;
+  }
+  // pico_offset_error = actual_time - (pico_offset + pico_timestamp)
+  // If the pico clock and pi clock are exactly in sync, this will always be
+  // zero.
+  aos::monotonic_clock::duration pico_offset_error() const {
+    return pico_offset_error_;
+  }
+
+ private:
+  const control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
+  std::function<void(
+      aos::monotonic_clock::time_point, aos::monotonic_clock::time_point,
+      std::optional<Eigen::Vector2d>, Eigen::Vector3d, Eigen::Vector3d)>
+      callback_;
+
+  // Last observed pico measurement. Used to track IMU staleness.
+  std::optional<aos::monotonic_clock::time_point> last_pico_timestamp_;
+  // Estimate of the drift between the pi and pico clocks. See
+  // pico_offset_error() for definition.
+  aos::monotonic_clock::duration pico_offset_error_;
+  // Raw offset between the pico and pi clocks. Gets updated to compensate for
+  // wrapping in the pico timestamp.
+  std::optional<std::chrono::nanoseconds> pico_offset_;
+
+  zeroing::ImuZeroer zeroer_;
+
+  ImuFailuresT imu_fault_tracker_;
+  // The first observed data counter. This is used to help us track dropped
+  // messages.
+  std::optional<size_t> first_valid_data_counter_;
+  size_t total_imu_messages_received_ = 0;
+  // added to the current read data counter to allow the data counter to
+  // increase monotonically. Will be a multiple of 2 ** 16.
+  size_t data_counter_offset_ = 0;
+  // PRevious data counter value (data_counter_offset_ not included).
+  int last_data_counter_ = 0;
+
+  // Unwrappers for the left and right encoders (necessary because the pico only
+  // sends out 16-bit encoder counts).
+  zeroing::UnwrapSensor left_encoder_;
+  zeroing::UnwrapSensor right_encoder_;
+
+  // When we lose IMU readings (e.g., due to checksum mismatches), we perform a
+  // zero-order hold on gyro readings; in order to do this, store the most
+  // recent gyro readings.
+  Eigen::Vector3d last_gyro_ = Eigen::Vector3d::Zero();
+};
+}  // namespace frc971::controls
+#endif  // FRC971_IMU_READER_IMU_WATCHER_H_
diff --git a/y2022/localizer/BUILD b/y2022/localizer/BUILD
index 1d45516..e5c054b 100644
--- a/y2022/localizer/BUILD
+++ b/y2022/localizer/BUILD
@@ -38,12 +38,13 @@
         "localizer_status.fbs",
     ],
     gen_reflections = True,
-    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"],
+    deps = [
+        "//frc971/control_loops:control_loops_fbs",
+        "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
+        "//frc971/imu_reader:imu_failures_fbs",
+    ],
 )
 
 flatbuffer_ts_library(
@@ -53,6 +54,7 @@
     deps = [
         "//frc971/control_loops:control_loops_ts_fbs",
         "//frc971/control_loops/drivetrain:drivetrain_status_ts_fbs",
+        "//frc971/imu_reader:imu_failures_ts_fbs",
     ],
 )
 
@@ -106,12 +108,11 @@
         "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
         "//frc971/control_loops/drivetrain:improved_down_estimator",
         "//frc971/control_loops/drivetrain:localizer_fbs",
+        "//frc971/imu_reader:imu_watcher",
         "//frc971/input:joystick_state_fbs",
         "//frc971/vision:calibration_fbs",
         "//frc971/wpilib:imu_batch_fbs",
         "//frc971/wpilib:imu_fbs",
-        "//frc971/zeroing:imu_zeroer",
-        "//frc971/zeroing:wrap",
         "//y2022:constants",
         "//y2022/control_loops/superstructure:superstructure_status_fbs",
         "//y2022/vision:target_estimate_fbs",
diff --git a/y2022/localizer/localizer.cc b/y2022/localizer/localizer.cc
index 4b0588c..e7dfdf1 100644
--- a/y2022/localizer/localizer.cc
+++ b/y2022/localizer/localizer.cc
@@ -12,8 +12,7 @@
 
 namespace {
 constexpr double kG = 9.80665;
-constexpr std::chrono::microseconds kNominalDt(500);
-
+static constexpr std::chrono::microseconds kNominalDt = ImuWatcher::kNominalDt;
 // Field position of the target (the 2022 target is conveniently in the middle
 // of the field....).
 constexpr double kVisionTargetX = 0.0;
@@ -863,18 +862,10 @@
   }
 }
 
-namespace {
-// Period at which the encoder readings from the IMU board wrap.
-static double DrivetrainWrapPeriod() {
-  return y2022::constants::Values::DrivetrainEncoderToMeters(1 << 16);
-}
-}  // namespace
-
 EventLoopLocalizer::EventLoopLocalizer(
     aos::EventLoop *event_loop,
     const control_loops::drivetrain::DrivetrainConfig<double> &dt_config)
     : event_loop_(event_loop),
-      dt_config_(dt_config),
       model_based_(dt_config),
       status_sender_(event_loop_->MakeSender<LocalizerStatus>("/localizer")),
       output_sender_(event_loop_->MakeSender<LocalizerOutput>("/localizer")),
@@ -892,9 +883,15 @@
                   "/superstructure")),
       joystick_state_fetcher_(
           event_loop_->MakeFetcher<aos::JoystickState>("/roborio/aos")),
-      zeroer_(zeroing::ImuZeroer::FaultBehavior::kTemporary),
-      left_encoder_(-DrivetrainWrapPeriod() / 2.0, DrivetrainWrapPeriod()),
-      right_encoder_(-DrivetrainWrapPeriod() / 2.0, DrivetrainWrapPeriod()) {
+      imu_watcher_(event_loop, dt_config,
+                   y2022::constants::Values::DrivetrainEncoderToMeters(1),
+                   [this](aos::monotonic_clock::time_point sample_time_pico,
+                          aos::monotonic_clock::time_point sample_time_pi,
+                          std::optional<Eigen::Vector2d> encoders,
+                          Eigen::Vector3d gyro, Eigen::Vector3d accel) {
+                     HandleImu(sample_time_pico, sample_time_pi, encoders, gyro,
+                               accel);
+                   }) {
   event_loop->SetRuntimeRealtimePriority(10);
   event_loop_->MakeWatcher(
       "/drivetrain",
@@ -965,7 +962,7 @@
             model_based_.TallyRejection(RejectionReason::IMAGE_FROM_FUTURE);
             continue;
           }
-          capture_time -= pico_offset_error_;
+          capture_time -= imu_watcher_.pico_offset_error();
           model_based_.HandleImageMatch(
               capture_time, target_estimate_fetchers_[camera_index].get(),
               camera_index);
@@ -976,153 +973,83 @@
     estimate_timer->Setup(event_loop_->monotonic_now(),
                           std::chrono::milliseconds(100));
   });
-  event_loop_->MakeWatcher(
-      "/localizer", [this](const frc971::IMUValuesBatch &values) {
-        CHECK(values.has_readings());
-        output_fetcher_.Fetch();
-        for (const IMUValues *value : *values.readings()) {
-          zeroer_.InsertAndProcessMeasurement(*value);
-          if (zeroer_.Faulted()) {
-            if (value->checksum_failed()) {
-              imu_fault_tracker_.pico_to_pi_checksum_mismatch++;
-            } else if (value->previous_reading_diag_stat()->checksum_mismatch()) {
-              imu_fault_tracker_.imu_to_pico_checksum_mismatch++;
-            } else {
-              imu_fault_tracker_.other_zeroing_faults++;
-            }
-          } else {
-            if (!first_valid_data_counter_.has_value()) {
-              first_valid_data_counter_ = value->data_counter();
-            }
-          }
-          if (first_valid_data_counter_.has_value()) {
-            total_imu_messages_received_++;
-            // Only update when we have good checksums, since the data counter
-            // could get corrupted.
-            if (!zeroer_.Faulted()) {
-              if (value->data_counter() < last_data_counter_) {
-                data_counter_offset_ += 1 << 16;
-              }
-              imu_fault_tracker_.missed_messages =
-                  (1 + value->data_counter() + data_counter_offset_ -
-                   first_valid_data_counter_.value()) -
-                  total_imu_messages_received_;
-              last_data_counter_ = value->data_counter();
-            }
-          }
-          const std::optional<Eigen::Vector2d> encoders =
-              zeroer_.Faulted()
-                  ? std::nullopt
-                  : std::make_optional(Eigen::Vector2d{
-                        left_encoder_.Unwrap(value->left_encoder()),
-                        right_encoder_.Unwrap(value->right_encoder())});
-          {
-            // If we can't trust the imu reading, just naively increment the
-            // pico timestamp.
-            const aos::monotonic_clock::time_point pico_timestamp =
-                zeroer_.Faulted()
-                    ? (last_pico_timestamp_.has_value()
-                           ? last_pico_timestamp_.value() + kNominalDt
-                           : aos::monotonic_clock::epoch())
-                    : aos::monotonic_clock::time_point(
-                          std::chrono::microseconds(
-                              value->pico_timestamp_us()));
-            // TODO(james): If we get large enough drift off of the pico,
-            // actually do something about it.
-            if (!pico_offset_.has_value()) {
-              pico_offset_ =
-                  event_loop_->context().monotonic_event_time - pico_timestamp;
-              last_pico_timestamp_ = pico_timestamp;
-            }
-            if (pico_timestamp < last_pico_timestamp_) {
-              pico_offset_.value() += std::chrono::microseconds(1ULL << 32);
-            }
-            const aos::monotonic_clock::time_point sample_timestamp =
-                pico_offset_.value() + pico_timestamp;
-            pico_offset_error_ =
-                event_loop_->context().monotonic_event_time - sample_timestamp;
-            const bool disabled =
-                (output_fetcher_.get() == nullptr) ||
-                (output_fetcher_.context().monotonic_event_time +
-                     std::chrono::milliseconds(10) <
-                 event_loop_->context().monotonic_event_time);
-            const bool zeroed = zeroer_.Zeroed();
-            // For gyros, use the most recent gyro reading if we aren't zeroed,
-            // to avoid missing integration cycles.
-            model_based_.HandleImu(
-                sample_timestamp,
-                zeroed ? zeroer_.ZeroedGyro().value() : last_gyro_,
-                zeroed ? zeroer_.ZeroedAccel().value()
-                       : dt_config_.imu_transform.transpose() *
-                             Eigen::Vector3d::UnitZ(),
-                encoders,
-                disabled ? Eigen::Vector2d::Zero()
-                         : Eigen::Vector2d{output_fetcher_->left_voltage(),
-                                           output_fetcher_->right_voltage()});
-            last_pico_timestamp_ = pico_timestamp;
+}
 
-            if (zeroed) {
-              last_gyro_ = zeroer_.ZeroedGyro().value();
-            }
-          }
-          {
-            auto builder = status_sender_.MakeBuilder();
-            const flatbuffers::Offset<ModelBasedStatus> model_based_status =
-                model_based_.PopulateStatus(builder.fbb());
-            const flatbuffers::Offset<control_loops::drivetrain::ImuZeroerState>
-                zeroer_status = zeroer_.PopulateStatus(builder.fbb());
-            const flatbuffers::Offset<ImuFailures> imu_failures =
-                ImuFailures::Pack(*builder.fbb(), &imu_fault_tracker_);
-            LocalizerStatus::Builder status_builder =
-                builder.MakeBuilder<LocalizerStatus>();
-            status_builder.add_model_based(model_based_status);
-            status_builder.add_zeroed(zeroer_.Zeroed());
-            status_builder.add_faulted_zero(zeroer_.Faulted());
-            status_builder.add_zeroing(zeroer_status);
-            status_builder.add_imu_failures(imu_failures);
-            if (encoders.has_value()) {
-              status_builder.add_left_encoder(encoders.value()(0));
-              status_builder.add_right_encoder(encoders.value()(1));
-            }
-            if (pico_offset_.has_value()) {
-              status_builder.add_pico_offset_ns(pico_offset_.value().count());
-              status_builder.add_pico_offset_error_ns(
-                  pico_offset_error_.count());
-            }
-            builder.CheckOk(builder.Send(status_builder.Finish()));
-          }
-          if (last_output_send_ + std::chrono::milliseconds(5) <
-              event_loop_->context().monotonic_event_time) {
-            auto builder = output_sender_.MakeBuilder();
+void EventLoopLocalizer::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,
+    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()});
+  }
+  {
+    auto builder = status_sender_.MakeBuilder();
+    const flatbuffers::Offset<ModelBasedStatus> model_based_status =
+        model_based_.PopulateStatus(builder.fbb());
+    const flatbuffers::Offset<control_loops::drivetrain::ImuZeroerState>
+        zeroer_status = imu_watcher_.zeroer().PopulateStatus(builder.fbb());
+    const flatbuffers::Offset<ImuFailures> imu_failures =
+        imu_watcher_.PopulateImuFailures(builder.fbb());
+    LocalizerStatus::Builder status_builder =
+        builder.MakeBuilder<LocalizerStatus>();
+    status_builder.add_model_based(model_based_status);
+    status_builder.add_zeroed(imu_watcher_.zeroer().Zeroed());
+    status_builder.add_faulted_zero(imu_watcher_.zeroer().Faulted());
+    status_builder.add_zeroing(zeroer_status);
+    status_builder.add_imu_failures(imu_failures);
+    if (encoders.has_value()) {
+      status_builder.add_left_encoder(encoders.value()(0));
+      status_builder.add_right_encoder(encoders.value()(1));
+    }
+    if (imu_watcher_.pico_offset().has_value()) {
+      status_builder.add_pico_offset_ns(
+          imu_watcher_.pico_offset().value().count());
+      status_builder.add_pico_offset_error_ns(
+          imu_watcher_.pico_offset_error().count());
+    }
+    builder.CheckOk(builder.Send(status_builder.Finish()));
+  }
+  if (last_output_send_ + std::chrono::milliseconds(5) <
+      event_loop_->context().monotonic_event_time) {
+    auto builder = output_sender_.MakeBuilder();
 
-            const auto led_outputs_offset =
-                builder.fbb()->CreateVector(model_based_.led_outputs().data(),
-                                            model_based_.led_outputs().size());
+    const auto led_outputs_offset = builder.fbb()->CreateVector(
+        model_based_.led_outputs().data(), model_based_.led_outputs().size());
 
-            LocalizerOutput::Builder output_builder =
-                builder.MakeBuilder<LocalizerOutput>();
-            // TODO(james): Should we bother to try to estimate time offsets for
-            // the pico?
-            output_builder.add_monotonic_timestamp_ns(
-                value->monotonic_timestamp_ns());
-            output_builder.add_x(model_based_.xytheta()(0));
-            output_builder.add_y(model_based_.xytheta()(1));
-            output_builder.add_theta(model_based_.xytheta()(2));
-            output_builder.add_zeroed(zeroer_.Zeroed());
-            output_builder.add_image_accepted_count(model_based_.total_accepted());
-            const Eigen::Quaterniond &orientation = model_based_.orientation();
-            Quaternion quaternion;
-            quaternion.mutate_x(orientation.x());
-            quaternion.mutate_y(orientation.y());
-            quaternion.mutate_z(orientation.z());
-            quaternion.mutate_w(orientation.w());
-            output_builder.add_orientation(&quaternion);
-            output_builder.add_led_outputs(led_outputs_offset);
-            builder.CheckOk(builder.Send(output_builder.Finish()));
-            last_output_send_ = event_loop_->monotonic_now();
-          }
-        }
-      });
+    LocalizerOutput::Builder output_builder =
+        builder.MakeBuilder<LocalizerOutput>();
+    output_builder.add_monotonic_timestamp_ns(
+        std::chrono::duration_cast<std::chrono::nanoseconds>(
+            sample_time_pi.time_since_epoch())
+            .count());
+    output_builder.add_x(model_based_.xytheta()(0));
+    output_builder.add_y(model_based_.xytheta()(1));
+    output_builder.add_theta(model_based_.xytheta()(2));
+    output_builder.add_zeroed(imu_watcher_.zeroer().Zeroed());
+    output_builder.add_image_accepted_count(model_based_.total_accepted());
+    const Eigen::Quaterniond &orientation = model_based_.orientation();
+    Quaternion quaternion;
+    quaternion.mutate_x(orientation.x());
+    quaternion.mutate_y(orientation.y());
+    quaternion.mutate_z(orientation.z());
+    quaternion.mutate_w(orientation.w());
+    output_builder.add_orientation(&quaternion);
+    output_builder.add_led_outputs(led_outputs_offset);
+    builder.CheckOk(builder.Send(output_builder.Finish()));
+    last_output_send_ = event_loop_->monotonic_now();
+  }
 }
 
 std::optional<aos::monotonic_clock::duration> EventLoopLocalizer::ClockOffset(
diff --git a/y2022/localizer/localizer.h b/y2022/localizer/localizer.h
index fc15e9f..917f131 100644
--- a/y2022/localizer/localizer.h
+++ b/y2022/localizer/localizer.h
@@ -19,6 +19,7 @@
 #include "y2022/localizer/localizer_status_generated.h"
 #include "y2022/localizer/localizer_visualization_generated.h"
 #include "y2022/vision/target_estimate_generated.h"
+#include "frc971/imu_reader/imu_watcher.h"
 
 namespace frc971::controls {
 
@@ -58,10 +59,6 @@
 // until the branches stop diverging--this will indicate that the model
 // matches the accelerometer readings again, and so we will swap back to
 // the model-based state.
-//
-// TODO:
-// * Implement paying attention to camera readings.
-// * Tune for ADIS16505/real robot.
 class ModelBasedLocalizer {
  public:
   static constexpr size_t kNumPis = 4;
@@ -331,8 +328,11 @@
  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,
+                 Eigen::Vector3d accel);
   aos::EventLoop *event_loop_;
-  const control_loops::drivetrain::DrivetrainConfig<double> &dt_config_;
   ModelBasedLocalizer model_based_;
   aos::Sender<LocalizerStatus> status_sender_;
   aos::Sender<LocalizerOutput> output_sender_;
@@ -350,22 +350,8 @@
       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.
-  // Note that this can drift over sufficiently long time periods!
-  std::optional<std::chrono::nanoseconds> pico_offset_;
 
-  ImuFailuresT imu_fault_tracker_;
-  std::optional<size_t> first_valid_data_counter_;
-  size_t total_imu_messages_received_ = 0;
-  size_t data_counter_offset_ = 0;
-  int last_data_counter_ = 0;
-
-  Eigen::Vector3d last_gyro_ = Eigen::Vector3d::Zero();
-
-  zeroing::UnwrapSensor left_encoder_;
-  zeroing::UnwrapSensor right_encoder_;
+  ImuWatcher imu_watcher_;
 };
 }  // namespace frc971::controls
 #endif  // Y2022_LOCALIZER_LOCALIZER_H_
diff --git a/y2022/localizer/localizer_status.fbs b/y2022/localizer/localizer_status.fbs
index 05cfc73..9e5081d 100644
--- a/y2022/localizer/localizer_status.fbs
+++ b/y2022/localizer/localizer_status.fbs
@@ -1,4 +1,5 @@
 include "frc971/control_loops/drivetrain/drivetrain_status.fbs";
+include "frc971/imu_reader/imu_failures.fbs";
 
 namespace frc971.controls;
 
@@ -92,13 +93,6 @@
   statistics:CumulativeStatistics (id: 18);
 }
 
-table ImuFailures {
-  imu_to_pico_checksum_mismatch:uint (id: 0);
-  pico_to_pi_checksum_mismatch:uint (id: 1);
-  missed_messages:uint (id: 2);
-  other_zeroing_faults:uint (id: 3);
-}
-
 table LocalizerStatus {
   model_based:ModelBasedStatus (id: 0);
   // Whether the IMU is zeroed or not.