Add boilerplate localizer code for LED control

Will use this to turn LEDs facing away from the target off

Signed-off-by: Milind Upadhyay <milind.upadhyay@gmail.com>
Change-Id: I4e03778f6b06db87b26ab3986dfebab7840d6de1
diff --git a/y2022/BUILD b/y2022/BUILD
index b16067d..4420894 100644
--- a/y2022/BUILD
+++ b/y2022/BUILD
@@ -92,6 +92,7 @@
             "//aos/network:timestamp_fbs",
             "//aos/network:remote_message_fbs",
             "//frc971/vision:vision_fbs",
+            "//y2022/localizer:localizer_output_fbs",
             "//y2022/vision:calibration_fbs",
             "//y2022/vision:target_estimate_fbs",
         ],
diff --git a/y2022/localizer/localizer.cc b/y2022/localizer/localizer.cc
index 8f9ffa1..1aa2458 100644
--- a/y2022/localizer/localizer.cc
+++ b/y2022/localizer/localizer.cc
@@ -102,6 +102,8 @@
   DiscretizeQAFast(Q_continuous_accel_, A_continuous_accel_, kNominalDt,
                    &Q_discrete_accel_, &A_discrete_accel_);
   P_accel_ = Q_discrete_accel_;
+
+  led_outputs_.fill(LedOutput::ON);
 }
 
 Eigen::Matrix<double, ModelBasedLocalizer::kNModelStates,
@@ -487,7 +489,8 @@
 }
 
 // Node names of the pis to listen for cameras from.
-const std::array<std::string_view, 4> kPisToUse{"pi1", "pi2", "pi3", "pi4"};
+constexpr std::array<std::string_view, ModelBasedLocalizer::kNumPis> kPisToUse{
+    "pi1", "pi2", "pi3", "pi4"};
 }  // namespace
 
 const Eigen::Matrix<double, 4, 4> ModelBasedLocalizer::CameraTransform(
@@ -667,6 +670,9 @@
   const double camera_yaw =
       std::atan2(camera_z_in_field.y(), camera_z_in_field.x());
 
+  // TODO(milind): actually control this
+  led_outputs_[camera_index] = LedOutput::ON;
+
   TargetEstimateDebugT debug;
   debug.camera = static_cast<uint8_t>(camera_index);
   debug.camera_x = H_field_camera_measured(0, 3);
@@ -1010,6 +1016,11 @@
           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());
+
             LocalizerOutput::Builder output_builder =
                 builder.MakeBuilder<LocalizerOutput>();
             // TODO(james): Should we bother to try to estimate time offsets for
@@ -1027,6 +1038,7 @@
             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();
           }
diff --git a/y2022/localizer/localizer.h b/y2022/localizer/localizer.h
index c19bf05..e9a61e1 100644
--- a/y2022/localizer/localizer.h
+++ b/y2022/localizer/localizer.h
@@ -63,6 +63,8 @@
 // * Tune for ADIS16505/real robot.
 class ModelBasedLocalizer {
  public:
+  static constexpr size_t kNumPis = 4;
+
   static constexpr size_t kX = 0;
   static constexpr size_t kY = 1;
   static constexpr size_t kTheta = 2;
@@ -141,6 +143,8 @@
 
   size_t NumQueuedImageDebugs() const { return image_debugs_.size(); }
 
+  std::array<LedOutput, kNumPis> led_outputs() const { return led_outputs_; }
+
  private:
   struct CombinedState {
     AccelState accel_state = AccelState::Zero();
@@ -298,6 +302,7 @@
   Statistics statistics_;
 
   aos::SizedArray<TargetEstimateDebugT, kDebugBufferSize> image_debugs_;
+  std::array<LedOutput, kNumPis> led_outputs_;
 
   friend class testing::LocalizerTest;
 };
@@ -322,7 +327,8 @@
   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>, 4>
+  std::array<aos::Fetcher<y2022::vision::TargetEstimate>,
+             ModelBasedLocalizer::kNumPis>
       target_estimate_fetchers_;
   aos::Fetcher<y2022::control_loops::superstructure::Status>
       superstructure_fetcher_;
diff --git a/y2022/localizer/localizer_output.fbs b/y2022/localizer/localizer_output.fbs
index b07278b..ec3302a 100644
--- a/y2022/localizer/localizer_output.fbs
+++ b/y2022/localizer/localizer_output.fbs
@@ -10,6 +10,12 @@
   z:double (id: 3);
 }
 
+// Used to tell different LEDs to be on or off
+enum LedOutput : byte {
+  ON,
+  OFF
+}
+
 table LocalizerOutput {
   // Timestamp (on the source node) that this sample corresponds with. This
   // may be older than the sent time to account for delays in sensor readings.
@@ -24,6 +30,10 @@
   // Whether we have zeroed the IMU (may go false if we observe a fault with the
   // IMU).
   zeroed:bool (id: 5);
+
+  // Whether each led should be on.
+  // Indices correspond to pi number.
+  led_outputs:[LedOutput] (id: 6);
 }
 
 root_type LocalizerOutput;
diff --git a/y2022/localizer/localizer_test.cc b/y2022/localizer/localizer_test.cc
index bf75446..6e9cd1e 100644
--- a/y2022/localizer/localizer_test.cc
+++ b/y2022/localizer/localizer_test.cc
@@ -2,12 +2,12 @@
 
 #include "aos/events/logging/log_writer.h"
 #include "aos/events/simulated_event_loop.h"
-#include "gtest/gtest.h"
 #include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
 #include "frc971/control_loops/pose.h"
-#include "y2022/vision/target_estimate_generated.h"
-#include "y2022/control_loops/superstructure/superstructure_status_generated.h"
+#include "gtest/gtest.h"
 #include "y2022/control_loops/drivetrain/drivetrain_base.h"
+#include "y2022/control_loops/superstructure/superstructure_status_generated.h"
+#include "y2022/vision/target_estimate_generated.h"
 
 DEFINE_string(output_folder, "",
               "If set, logs all channels to the provided logfile.");
@@ -18,11 +18,11 @@
 typedef ModelBasedLocalizer::ModelInput ModelInput;
 typedef ModelBasedLocalizer::AccelInput AccelInput;
 
-using frc971::vision::calibration::CameraCalibrationT;
-using frc971::vision::calibration::TransformationMatrixT;
+using frc971::control_loops::Pose;
 using frc971::control_loops::drivetrain::DrivetrainConfig;
 using frc971::control_loops::drivetrain::LocalizerControl;
-using frc971::control_loops::Pose;
+using frc971::vision::calibration::CameraCalibrationT;
+using frc971::vision::calibration::TransformationMatrixT;
 using y2022::vision::TargetEstimate;
 using y2022::vision::TargetEstimateT;
 
@@ -82,14 +82,12 @@
   config.is_simulated = true;
   return config;
 }
-}
+}  // namespace
 
 class LocalizerTest : public ::testing::Test {
  protected:
   LocalizerTest()
-      : dt_config_(
-            GetTest2022DrivetrainConfig()),
-        localizer_(dt_config_) {
+      : dt_config_(GetTest2022DrivetrainConfig()), localizer_(dt_config_) {
     localizer_.set_longitudinal_offset(0.0);
   }
   ModelState CallDiffModel(const ModelState &state, const ModelInput &U) {
@@ -102,7 +100,6 @@
 
   const control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
   ModelBasedLocalizer localizer_;
-
 };
 
 TEST_F(LocalizerTest, AccelIntegrationTest) {
@@ -233,7 +230,8 @@
   const Eigen::Vector2d encoders{0.0, 0.0};
   const Eigen::Vector2d voltages{0.0, 0.0};
   Eigen::Vector3d accel{5.0, 2.0, 9.80665};
-  Eigen::Vector3d accel_gs = dt_config_.imu_transform.inverse() * accel / 9.80665;
+  Eigen::Vector3d accel_gs =
+      dt_config_.imu_transform.inverse() * accel / 9.80665;
   while (t < start) {
     // Spin to fill up the buffer.
     localizer_.HandleImu(t, gyro, Eigen::Vector3d::UnitZ(), encoders, voltages);
@@ -744,10 +742,10 @@
   ASSERT_FALSE(status_fetcher_->model_based()->using_model());
   ASSERT_LT(0.1, status_fetcher_->model_based()->residual())
       << aos::FlatbufferToJson(status_fetcher_.get(), {.multi_line = true});
-  ASSERT_NEAR(drivetrain_plant_.state()(0),
-              status_fetcher_->model_based()->x(), 1.0);
-  ASSERT_NEAR(drivetrain_plant_.state()(1),
-              status_fetcher_->model_based()->y(), 1e-6);
+  ASSERT_NEAR(drivetrain_plant_.state()(0), status_fetcher_->model_based()->x(),
+              1.0);
+  ASSERT_NEAR(drivetrain_plant_.state()(1), status_fetcher_->model_based()->y(),
+              1e-6);
 }
 
 // Tests that image corrections in the nominal case (no errors) causes no
@@ -864,7 +862,8 @@
   EXPECT_TRUE(VerifyEstimatorAccurate(3.0));
   // y should be noticeably more accurate than x, since we are just driving
   // straight.
-  ASSERT_NEAR(drivetrain_plant_.state()(1), status_fetcher_->model_based()->y(), 0.1);
+  ASSERT_NEAR(drivetrain_plant_.state()(1), status_fetcher_->model_based()->y(),
+              0.1);
   ASSERT_TRUE(status_fetcher_->model_based()->has_statistics());
   ASSERT_LT(10,
             status_fetcher_->model_based()->statistics()->total_candidates());
@@ -872,4 +871,16 @@
             status_fetcher_->model_based()->statistics()->total_accepted());
 }
 
-}  // namespace frc91::controls::testing
+TEST_F(EventLoopLocalizerTest, LedOutputs) {
+  send_targets_ = true;
+
+  event_loop_factory_.RunFor(std::chrono::milliseconds(10));
+  CHECK(output_fetcher_.Fetch());
+  EXPECT_EQ(output_fetcher_->led_outputs()->size(),
+            ModelBasedLocalizer::kNumPis);
+  for (LedOutput led_output : *output_fetcher_->led_outputs()) {
+    EXPECT_EQ(led_output, LedOutput::ON);
+  }
+}
+
+}  // namespace frc971::controls::testing
diff --git a/y2022/vision/BUILD b/y2022/vision/BUILD
index bede03a..e9543cf 100644
--- a/y2022/vision/BUILD
+++ b/y2022/vision/BUILD
@@ -120,6 +120,7 @@
         "//frc971/vision:v4l2_reader",
         "//frc971/vision:vision_fbs",
         "//third_party:opencv",
+        "//y2022/localizer:localizer_output_fbs",
     ],
 )
 
diff --git a/y2022/vision/camera_reader.cc b/y2022/vision/camera_reader.cc
index f1cb4e8..94b2d96 100644
--- a/y2022/vision/camera_reader.cc
+++ b/y2022/vision/camera_reader.cc
@@ -211,6 +211,26 @@
 
   reader_->SendLatestImage();
   read_image_timer_->Setup(event_loop_->monotonic_now());
+
+  // Disable the LEDs based on localizer output
+  if (localizer_output_fetcher_.Fetch()) {
+    const auto node_name = event_loop_->node()->name()->string_view();
+    const size_t pi_number =
+        std::atol(node_name.substr(node_name.size() - 1).data());
+
+    CHECK(localizer_output_fetcher_->has_led_outputs() &&
+          localizer_output_fetcher_->led_outputs()->size() > pi_number);
+
+    const LedOutput led_output =
+        localizer_output_fetcher_->led_outputs()->Get(pi_number);
+
+    if (led_output != prev_led_output_) {
+      gpio_disable_control_.GPIOWrite(led_output == LedOutput::OFF ? kGPIOHigh
+                                                                   : kGPIOLow);
+
+      prev_led_output_ = led_output;
+    }
+  }
 }
 
 }  // namespace vision
diff --git a/y2022/vision/camera_reader.h b/y2022/vision/camera_reader.h
index edc3a1a..707e04e 100644
--- a/y2022/vision/camera_reader.h
+++ b/y2022/vision/camera_reader.h
@@ -13,6 +13,7 @@
 #include "aos/network/team_number.h"
 #include "frc971/vision/v4l2_reader.h"
 #include "frc971/vision/vision_generated.h"
+#include "y2022/localizer/localizer_output_generated.h"
 #include "y2022/vision/calibration_data.h"
 #include "y2022/vision/calibration_generated.h"
 #include "y2022/vision/gpio.h"
@@ -23,6 +24,7 @@
 namespace vision {
 
 using namespace frc971::vision;
+using frc971::controls::LedOutput;
 
 // TODO<jim>: Probably need to break out LED control to separate process
 class CameraReader {
@@ -38,6 +40,9 @@
         target_estimator_(CameraIntrinsics(), CameraExtrinsics()),
         target_estimate_sender_(
             event_loop->MakeSender<TargetEstimate>("/camera")),
+        localizer_output_fetcher_(
+            event_loop->MakeFetcher<frc971::controls::LocalizerOutput>(
+                "/localizer")),
         read_image_timer_(event_loop->AddTimer([this]() { ReadImage(); })),
         gpio_imu_pin_(GPIOControl(GPIO_PIN_SCLK_IMU, kGPIOIn)),
         gpio_pwm_control_(GPIOPWMControl(GPIO_PIN_SCK_PWM, duty_cycle_)),
@@ -99,6 +104,9 @@
   TargetEstimator target_estimator_;
   aos::Sender<TargetEstimate> target_estimate_sender_;
 
+  LedOutput prev_led_output_ = LedOutput::ON;
+  aos::Fetcher<frc971::controls::LocalizerOutput> localizer_output_fetcher_;
+
   // We schedule this immediately to read an image. Having it on a timer
   // means other things can run on the event loop in between.
   aos::TimerHandler *const read_image_timer_;
diff --git a/y2022/y2022_pi_template.json b/y2022/y2022_pi_template.json
index 27bc251..3dff81e 100644
--- a/y2022/y2022_pi_template.json
+++ b/y2022/y2022_pi_template.json
@@ -208,6 +208,18 @@
       "max_size": 208
     },
     {
+      "name": "/localizer",
+      "type": "frc971.controls.LocalizerOutput",
+      "source_node": "imu",
+      "destination_nodes": [
+        {
+          "name": "pi{{ NUM }}",
+          "priority": 5,
+          "time_to_live": 5000000
+        }
+      ]
+    },
+    {
       "name": "/logger/aos",
       "type": "aos.starter.StarterRpc",
       "source_node": "logger",