Process LocalizerOutput messages in 2022 drivetrain

Make it so that the roborio will actually pay attention to the pi's
localizer.

Change-Id: I8db97cbba65beef8c2ec863a792e57960ead996a
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/y2022/control_loops/drivetrain/BUILD b/y2022/control_loops/drivetrain/BUILD
index 2858dc7..a1f1e47 100644
--- a/y2022/control_loops/drivetrain/BUILD
+++ b/y2022/control_loops/drivetrain/BUILD
@@ -1,3 +1,5 @@
+load("//aos:config.bzl", "aos_config")
+
 genrule(
     name = "genrule_drivetrain",
     outs = [
@@ -69,6 +71,19 @@
     ],
 )
 
+cc_library(
+    name = "localizer",
+    srcs = ["localizer.cc"],
+    hdrs = ["localizer.h"],
+    deps = [
+        "//aos/events:event_loop",
+        "//aos/network:message_bridge_server_fbs",
+        "//frc971/control_loops/drivetrain:hybrid_ekf",
+        "//frc971/control_loops/drivetrain:localizer",
+        "//y2022/control_loops/localizer:localizer_output_fbs",
+    ],
+)
+
 cc_binary(
     name = "drivetrain",
     srcs = [
@@ -78,8 +93,40 @@
     visibility = ["//visibility:public"],
     deps = [
         ":drivetrain_base",
+        ":localizer",
         "//aos:init",
         "//aos/events:shm_event_loop",
         "//frc971/control_loops/drivetrain:drivetrain_lib",
     ],
 )
+
+aos_config(
+    name = "simulation_config",
+    src = "drivetrain_simulation_config.json",
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//frc971/control_loops/drivetrain:simulation_channels",
+        "//y2022:config",
+    ],
+)
+
+cc_test(
+    name = "localizer_test",
+    srcs = ["localizer_test.cc"],
+    data = [":simulation_config"],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        ":drivetrain_base",
+        ":localizer",
+        "//aos/events:simulated_event_loop",
+        "//aos/events/logging:log_writer",
+        "//aos/network:team_number",
+        "//aos/network:testing_time_converter",
+        "//frc971/control_loops:control_loop_test",
+        "//frc971/control_loops:team_number_test_environment",
+        "//frc971/control_loops/drivetrain:drivetrain_lib",
+        "//frc971/control_loops/drivetrain:drivetrain_test_lib",
+        "//y2022/control_loops/localizer:localizer_output_fbs",
+    ],
+)
diff --git a/y2022/control_loops/drivetrain/drivetrain_main.cc b/y2022/control_loops/drivetrain/drivetrain_main.cc
index e29f950..ecdcdbb 100644
--- a/y2022/control_loops/drivetrain/drivetrain_main.cc
+++ b/y2022/control_loops/drivetrain/drivetrain_main.cc
@@ -4,6 +4,7 @@
 #include "aos/init.h"
 #include "frc971/control_loops/drivetrain/drivetrain.h"
 #include "y2022/control_loops/drivetrain/drivetrain_base.h"
+#include "y2022/control_loops/drivetrain/localizer.h"
 
 using ::frc971::control_loops::drivetrain::DrivetrainLoop;
 
@@ -14,7 +15,7 @@
       aos::configuration::ReadConfig("config.json");
 
   ::aos::ShmEventLoop event_loop(&config.message());
-  ::frc971::control_loops::drivetrain::DeadReckonEkf localizer(
+  ::y2022::control_loops::drivetrain::Localizer localizer(
       &event_loop,
       ::y2022::control_loops::drivetrain::GetDrivetrainConfig());
   std::unique_ptr<DrivetrainLoop> drivetrain = std::make_unique<DrivetrainLoop>(
diff --git a/y2022/control_loops/drivetrain/drivetrain_simulation_config.json b/y2022/control_loops/drivetrain/drivetrain_simulation_config.json
new file mode 100644
index 0000000..37c1bb7
--- /dev/null
+++ b/y2022/control_loops/drivetrain/drivetrain_simulation_config.json
@@ -0,0 +1,6 @@
+{
+  "imports": [
+    "../../y2022.json",
+    "../../../frc971/control_loops/drivetrain/drivetrain_simulation_channels.json"
+  ]
+}
diff --git a/y2022/control_loops/drivetrain/localizer.cc b/y2022/control_loops/drivetrain/localizer.cc
new file mode 100644
index 0000000..3d02312
--- /dev/null
+++ b/y2022/control_loops/drivetrain/localizer.cc
@@ -0,0 +1,102 @@
+#include "y2022/control_loops/drivetrain/localizer.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace drivetrain {
+
+Localizer::Localizer(
+    aos::EventLoop *event_loop,
+    const frc971::control_loops::drivetrain::DrivetrainConfig<double>
+        &dt_config)
+    : event_loop_(event_loop),
+      dt_config_(dt_config),
+      ekf_(dt_config),
+      localizer_output_fetcher_(
+          event_loop_->MakeFetcher<frc971::controls::LocalizerOutput>(
+              "/localizer")),
+      clock_offset_fetcher_(
+          event_loop_->MakeFetcher<aos::message_bridge::ServerStatistics>(
+              "/aos")) {
+  ekf_.set_ignore_accel(true);
+
+  event_loop->OnRun([this, event_loop]() {
+    ekf_.ResetInitialState(event_loop->monotonic_now(),
+                           HybridEkf::State::Zero(), ekf_.P());
+  });
+
+  target_selector_.set_has_target(false);
+}
+
+void Localizer::Reset(
+    aos::monotonic_clock::time_point t,
+    const frc971::control_loops::drivetrain::HybridEkf<double>::State &state) {
+  // Go through and clear out all of the fetchers so that we don't get behind.
+  localizer_output_fetcher_.Fetch();
+  ekf_.ResetInitialState(t, state.cast<float>(), ekf_.P());
+}
+
+void Localizer::Update(const Eigen::Matrix<double, 2, 1> &U,
+                       aos::monotonic_clock::time_point now,
+                       double left_encoder, double right_encoder,
+                       double gyro_rate, const Eigen::Vector3d &accel) {
+  ekf_.UpdateEncodersAndGyro(left_encoder, right_encoder, gyro_rate,
+                             U.cast<float>(), accel.cast<float>(), now);
+  if (localizer_output_fetcher_.Fetch()) {
+    clock_offset_fetcher_.Fetch();
+    bool message_bridge_connected = true;
+    std::chrono::nanoseconds monotonic_offset{0};
+    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() == "imu") {
+          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, in which case we should break the loop but shouldn't
+            // populate the offset.
+            message_bridge_connected = false;
+          }
+          break;
+        }
+      }
+    }
+    if (!message_bridge_connected) {
+      return;
+    }
+    aos::monotonic_clock::time_point capture_time(
+        std::chrono::nanoseconds(
+            localizer_output_fetcher_->monotonic_timestamp_ns()) -
+        monotonic_offset);
+    // TODO: Finish implementing simple x/y/theta updater with state_at_capture.
+    // TODO: Implement turret/camera processing logic on pi side.
+    const std::optional<State> state_at_capture =
+        ekf_.LastStateBeforeTime(capture_time);
+    Eigen::Matrix<float, HybridEkf::kNOutputs, HybridEkf::kNStates> H;
+    H.setZero();
+    H(0, StateIdx::kX) = 1;
+    H(1, StateIdx::kY) = 1;
+    H(2, StateIdx::kTheta) = 1;
+    const Eigen::Vector3f Z{
+        static_cast<float>(localizer_output_fetcher_->x()),
+        static_cast<float>(localizer_output_fetcher_->y()),
+        static_cast<float>(localizer_output_fetcher_->theta())};
+    Eigen::Matrix3f R = Eigen::Matrix3f::Zero();
+    R.diagonal() << 0.01, 0.01, 1e-4;
+    const Input U_correct = ekf_.MostRecentInput();
+    ekf_.Correct(
+        Eigen::Vector3f::Zero(), &U_correct, {},
+        [H, state_at_capture, Z](const State &,
+                                 const Input &) -> Eigen::Vector3f {
+          Eigen::Vector3f error = H * state_at_capture.value() - Z;
+          error(2) = aos::math::NormalizeAngle(error(2));
+          return error;
+        },
+        [H](const State &) { return H; }, R, now);
+  }
+}
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace y2022
diff --git a/y2022/control_loops/drivetrain/localizer.h b/y2022/control_loops/drivetrain/localizer.h
new file mode 100644
index 0000000..50d083a
--- /dev/null
+++ b/y2022/control_loops/drivetrain/localizer.h
@@ -0,0 +1,80 @@
+#ifndef Y2022_CONTROL_LOOPS_DRIVETRAIN_LOCALIZER_H_
+#define Y2022_CONTROL_LOOPS_DRIVETRAIN_LOCALIZER_H_
+
+#include <string_view>
+
+#include "aos/events/event_loop.h"
+#include "frc971/control_loops/drivetrain/hybrid_ekf.h"
+#include "frc971/control_loops/drivetrain/localizer.h"
+#include "y2022/control_loops/localizer/localizer_output_generated.h"
+#include "aos/network/message_bridge_server_generated.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace drivetrain {
+
+// This class handles the localization for the 2022 robot. Rather than actually
+// doing any work on the roborio, we farm all the localization out to a
+// raspberry pi and it then sends out LocalizerOutput messages that we treat as
+// measurement updates. See //y2022/control_loops/localizer.
+// TODO(james): Needs tests. Should refactor out some of the code from the 2020
+// localizer test.
+class Localizer : public frc971::control_loops::drivetrain::LocalizerInterface {
+ public:
+  typedef frc971::control_loops::TypedPose<float> Pose;
+  typedef frc971::control_loops::drivetrain::HybridEkf<float> HybridEkf;
+  typedef typename HybridEkf::State State;
+  typedef typename HybridEkf::StateIdx StateIdx;
+  typedef typename HybridEkf::StateSquare StateSquare;
+  typedef typename HybridEkf::Input Input;
+  typedef typename HybridEkf::Output Output;
+  Localizer(aos::EventLoop *event_loop,
+            const frc971::control_loops::drivetrain::DrivetrainConfig<double>
+                &dt_config);
+  frc971::control_loops::drivetrain::HybridEkf<double>::State Xhat()
+      const override {
+    return ekf_.X_hat().cast<double>();
+  }
+  frc971::control_loops::drivetrain::TrivialTargetSelector *target_selector()
+      override {
+    return &target_selector_;
+  }
+
+  void Update(const ::Eigen::Matrix<double, 2, 1> &U,
+              aos::monotonic_clock::time_point now, double left_encoder,
+              double right_encoder, double gyro_rate,
+              const Eigen::Vector3d &accel) override;
+
+  void Reset(aos::monotonic_clock::time_point t,
+             const frc971::control_loops::drivetrain::HybridEkf<double>::State
+                 &state) override;
+
+  void ResetPosition(aos::monotonic_clock::time_point t, double x, double y,
+                     double theta, double /*theta_override*/,
+                     bool /*reset_theta*/) override {
+    const double left_encoder = ekf_.X_hat(StateIdx::kLeftEncoder);
+    const double right_encoder = ekf_.X_hat(StateIdx::kRightEncoder);
+    ekf_.ResetInitialState(t,
+                           (HybridEkf::State() << x, y, theta, left_encoder, 0,
+                            right_encoder, 0, 0, 0, 0, 0, 0)
+                               .finished(),
+                           ekf_.P());
+  }
+
+ private:
+  aos::EventLoop *const event_loop_;
+  const frc971::control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
+  HybridEkf ekf_;
+
+  aos::Fetcher<frc971::controls::LocalizerOutput> localizer_output_fetcher_;
+  aos::Fetcher<aos::message_bridge::ServerStatistics> clock_offset_fetcher_;
+
+  // Target selector to allow us to satisfy the LocalizerInterface requirements.
+  frc971::control_loops::drivetrain::TrivialTargetSelector target_selector_;
+};
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace y2022
+
+#endif  // Y2022_CONTROL_LOOPS_DRIVETRAIN_LOCALIZER_H_
diff --git a/y2022/control_loops/drivetrain/localizer_test.cc b/y2022/control_loops/drivetrain/localizer_test.cc
new file mode 100644
index 0000000..87fc3fd
--- /dev/null
+++ b/y2022/control_loops/drivetrain/localizer_test.cc
@@ -0,0 +1,208 @@
+#include "y2022/control_loops/drivetrain/localizer.h"
+
+#include <queue>
+
+#include "aos/events/logging/log_writer.h"
+#include "aos/network/message_bridge_server_generated.h"
+#include "aos/network/team_number.h"
+#include "aos/network/testing_time_converter.h"
+#include "frc971/control_loops/control_loop_test.h"
+#include "frc971/control_loops/drivetrain/drivetrain.h"
+#include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
+#include "frc971/control_loops/team_number_test_environment.h"
+#include "y2022/control_loops/localizer/localizer_output_generated.h"
+#include "gtest/gtest.h"
+#include "y2022/control_loops/drivetrain/drivetrain_base.h"
+
+DEFINE_string(output_folder, "",
+              "If set, logs all channels to the provided logfile.");
+
+namespace y2022 {
+namespace control_loops {
+namespace drivetrain {
+namespace testing {
+
+using frc971::control_loops::drivetrain::DrivetrainConfig;
+using frc971::control_loops::drivetrain::Goal;
+using frc971::control_loops::drivetrain::LocalizerControl;
+
+namespace {
+DrivetrainConfig<double> GetTest2022DrivetrainConfig() {
+  DrivetrainConfig<double> config = GetDrivetrainConfig();
+  return config;
+}
+}
+
+namespace chrono = std::chrono;
+using aos::monotonic_clock;
+using frc971::control_loops::drivetrain::DrivetrainLoop;
+using frc971::control_loops::drivetrain::testing::DrivetrainSimulation;
+
+// TODO(james): Make it so this actually tests the full system of the localizer.
+class LocalizedDrivetrainTest : public frc971::testing::ControlLoopTest {
+ protected:
+  // We must use the 2022 drivetrain config so that we don't have to deal
+  // with shifting:
+  LocalizedDrivetrainTest()
+      : frc971::testing::ControlLoopTest(
+            aos::configuration::ReadConfig(
+                "y2022/control_loops/drivetrain/simulation_config.json"),
+            GetTest2022DrivetrainConfig().dt),
+        roborio_(aos::configuration::GetNode(configuration(), "roborio")),
+        imu_(aos::configuration::GetNode(configuration(), "imu")),
+        test_event_loop_(MakeEventLoop("test", roborio_)),
+        imu_test_event_loop_(MakeEventLoop("test", imu_)),
+        drivetrain_goal_sender_(
+            test_event_loop_->MakeSender<Goal>("/drivetrain")),
+        localizer_output_sender_(
+            imu_test_event_loop_->MakeSender<frc971::controls::LocalizerOutput>(
+                "/localizer")),
+        drivetrain_goal_fetcher_(
+            test_event_loop_->MakeFetcher<Goal>("/drivetrain")),
+        drivetrain_status_fetcher_(
+            test_event_loop_
+                ->MakeFetcher<frc971::control_loops::drivetrain::Status>(
+                    "/drivetrain")),
+        localizer_control_sender_(
+            test_event_loop_->MakeSender<LocalizerControl>("/drivetrain")),
+        drivetrain_event_loop_(MakeEventLoop("drivetrain", roborio_)),
+        dt_config_(GetTest2022DrivetrainConfig()),
+        localizer_(drivetrain_event_loop_.get(), dt_config_),
+        drivetrain_(dt_config_, drivetrain_event_loop_.get(), &localizer_),
+        drivetrain_plant_event_loop_(MakeEventLoop("plant", roborio_)),
+        drivetrain_plant_imu_event_loop_(MakeEventLoop("plant", imu_)),
+        drivetrain_plant_(drivetrain_plant_event_loop_.get(),
+                          drivetrain_plant_imu_event_loop_.get(), dt_config_,
+                          std::chrono::microseconds(500)) {
+    set_team_id(frc971::control_loops::testing::kTeamNumber);
+    set_battery_voltage(12.0);
+
+    if (!FLAGS_output_folder.empty()) {
+      logger_event_loop_ = MakeEventLoop("logger", roborio_);
+      logger_ = std::make_unique<aos::logger::Logger>(logger_event_loop_.get());
+      logger_->StartLoggingOnRun(FLAGS_output_folder);
+    }
+
+    test_event_loop_->OnRun([this]() { SetStartingPosition({3.0, 2.0, 0.0}); });
+
+    imu_test_event_loop_
+        ->AddTimer([this]() {
+          auto builder = localizer_output_sender_.MakeBuilder();
+          frc971::controls::LocalizerOutput::Builder output_builder =
+              builder.MakeBuilder<frc971::controls::LocalizerOutput>();
+          output_builder.add_monotonic_timestamp_ns(
+              imu_test_event_loop_->monotonic_now().time_since_epoch().count());
+          output_builder.add_x(drivetrain_plant_.state()(0));
+          output_builder.add_y(drivetrain_plant_.state()(1));
+          output_builder.add_theta(drivetrain_plant_.state()(2));
+        })
+        ->Setup(imu_test_event_loop_->monotonic_now(),
+                std::chrono::milliseconds(5));
+  }
+
+  virtual ~LocalizedDrivetrainTest() override {}
+
+  void SetStartingPosition(const Eigen::Matrix<double, 3, 1> &xytheta) {
+    *drivetrain_plant_.mutable_state() << xytheta.x(), xytheta.y(),
+        xytheta(2, 0), 0.0, 0.0;
+    Eigen::Matrix<double, Localizer::HybridEkf::kNStates, 1> localizer_state;
+    localizer_state.setZero();
+    localizer_state.block<3, 1>(0, 0) = xytheta;
+    localizer_.Reset(monotonic_now(), localizer_state);
+  }
+
+  void VerifyNearGoal(double eps = 1e-2) {
+    drivetrain_goal_fetcher_.Fetch();
+    EXPECT_NEAR(drivetrain_goal_fetcher_->left_goal(),
+                drivetrain_plant_.GetLeftPosition(), eps);
+    EXPECT_NEAR(drivetrain_goal_fetcher_->right_goal(),
+                drivetrain_plant_.GetRightPosition(), eps);
+  }
+
+  ::testing::AssertionResult IsNear(double expected, double actual,
+                                    double epsilon) {
+    if (std::abs(expected - actual) < epsilon) {
+      return ::testing::AssertionSuccess();
+    } else {
+      return ::testing::AssertionFailure()
+             << "Expected " << expected << " but got " << actual
+             << " with a max difference of " << epsilon
+             << " and an actual difference of " << std::abs(expected - actual);
+    }
+  }
+  ::testing::AssertionResult VerifyEstimatorAccurate(double eps) {
+    const Eigen::Matrix<double, 5, 1> true_state = drivetrain_plant_.state();
+    ::testing::AssertionResult result(true);
+    if (!(result = IsNear(localizer_.x(), true_state(0), eps))) {
+      return result;
+    }
+    if (!(result = IsNear(localizer_.y(), true_state(1), eps))) {
+      return result;
+    }
+    if (!(result = IsNear(localizer_.theta(), true_state(2), eps))) {
+      return result;
+    }
+    if (!(result = IsNear(localizer_.left_velocity(), true_state(3), eps))) {
+      return result;
+    }
+    if (!(result = IsNear(localizer_.right_velocity(), true_state(4), eps))) {
+      return result;
+    }
+    return result;
+  }
+
+  const aos::Node *const roborio_;
+  const aos::Node *const imu_;
+
+  std::unique_ptr<aos::EventLoop> test_event_loop_;
+  std::unique_ptr<aos::EventLoop> imu_test_event_loop_;
+  aos::Sender<Goal> drivetrain_goal_sender_;
+  aos::Sender<frc971::controls::LocalizerOutput> localizer_output_sender_;
+  aos::Fetcher<Goal> drivetrain_goal_fetcher_;
+  aos::Fetcher<frc971::control_loops::drivetrain::Status>
+      drivetrain_status_fetcher_;
+  aos::Sender<LocalizerControl> localizer_control_sender_;
+
+  std::unique_ptr<aos::EventLoop> drivetrain_event_loop_;
+  const frc971::control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
+
+  Localizer localizer_;
+  DrivetrainLoop drivetrain_;
+
+  std::unique_ptr<aos::EventLoop> drivetrain_plant_event_loop_;
+  std::unique_ptr<aos::EventLoop> drivetrain_plant_imu_event_loop_;
+  DrivetrainSimulation drivetrain_plant_;
+
+  void SendGoal(double left, double right) {
+    auto builder = drivetrain_goal_sender_.MakeBuilder();
+
+    Goal::Builder drivetrain_builder = builder.MakeBuilder<Goal>();
+    drivetrain_builder.add_controller_type(
+        frc971::control_loops::drivetrain::ControllerType::MOTION_PROFILE);
+    drivetrain_builder.add_left_goal(left);
+    drivetrain_builder.add_right_goal(right);
+
+    EXPECT_EQ(builder.Send(drivetrain_builder.Finish()),
+              aos::RawSender::Error::kOk);
+  }
+
+ private:
+  std::unique_ptr<aos::EventLoop> logger_event_loop_;
+  std::unique_ptr<aos::logger::Logger> logger_;
+};
+
+TEST_F(LocalizedDrivetrainTest, Nominal) {
+  SetEnabled(true);
+  EXPECT_TRUE(VerifyEstimatorAccurate(1e-7));
+
+  SendGoal(-1.0, 1.0);
+
+  RunFor(chrono::seconds(10));
+  VerifyNearGoal();
+  EXPECT_TRUE(VerifyEstimatorAccurate(5e-3));
+}
+
+}  // namespace testing
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace y2022
diff --git a/y2022/control_loops/localizer/localizer.cc b/y2022/control_loops/localizer/localizer.cc
index d6afd56..dd2e67d 100644
--- a/y2022/control_loops/localizer/localizer.cc
+++ b/y2022/control_loops/localizer/localizer.cc
@@ -558,6 +558,10 @@
             auto builder = output_sender_.MakeBuilder();
             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));