Create y2022 localizer

TODO:
* Update Roborio's localizer to listen to LocalizerOutput messages.
* Actually listen to camera messages.
* Correct the interface between this and the pico-reading process for
  encoder values.

Change-Id: If1c7b0cca60581780fbe39fe2f78fa6b4a6933e3
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/y2022/control_loops/localizer/localizer.h b/y2022/control_loops/localizer/localizer.h
new file mode 100644
index 0000000..f2d76ef
--- /dev/null
+++ b/y2022/control_loops/localizer/localizer.h
@@ -0,0 +1,218 @@
+#ifndef Y2022_CONTROL_LOOPS_LOCALIZER_LOCALIZER_H_
+#define Y2022_CONTROL_LOOPS_LOCALIZER_LOCALIZER_H_
+
+#include "Eigen/Dense"
+#include "Eigen/Geometry"
+
+#include "aos/events/event_loop.h"
+#include "aos/containers/ring_buffer.h"
+#include "aos/time/time.h"
+#include "y2020/vision/sift/sift_generated.h"
+#include "y2022/control_loops/localizer/localizer_status_generated.h"
+#include "y2022/control_loops/localizer/localizer_output_generated.h"
+#include "frc971/control_loops/drivetrain/improved_down_estimator.h"
+#include "frc971/control_loops/drivetrain/drivetrain_position_generated.h"
+#include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
+#include "frc971/control_loops/drivetrain/localizer_generated.h"
+#include "frc971/zeroing/imu_zeroer.h"
+
+namespace frc971::controls {
+
+namespace testing {
+class LocalizerTest;
+}
+
+// Localizer implementation that makes use of a 6-axis IMU, encoder readings,
+// drivetrain voltages, and camera returns to localize the robot. Meant to
+// be run on a raspberry pi.
+//
+// This operates on the principle that the drivetrain can be in one of two
+// modes:
+// 1) A "normal" mode where it is obeying the regular drivetrain model, with
+//    minimal lateral motion and no major external disturbances. This is
+//    referred to as the "model" mode in the code/variable names.
+// 2) An non-holonomic mode where the robot is just flying around on a 2-D
+//    plane with no meaningful constraints (referred to as an "accel" model
+//    in the code, because we rely primarily on the accelerometer readings).
+//
+// In order to determine which mode to be in, we attempt to track whether the
+// two models are diverging substantially. In order to do this, we maintain a
+// 1-second long queue of "branches". A branch is generated every X iterations
+// and contains a model state and an accel state. When the branch starts, the
+// two will have identical states. For the remaining 1 second, the model state
+// will evolve purely according to the drivetrian model, and the accel state
+// will evolve purely using IMU readings.
+//
+// When the branch reaches 1 second in age, we calculate a residual associated
+// with how much the accel and model based states diverged. If they have
+// diverged substantially, that implies that the model is a poor match for
+// whatever has been happening to the robot in the past second, so if we were
+// previously relying on the model, we will override the current "actual"
+// state with the branched accel state, and then continue to update the accel
+// state based on IMU readings.
+// If we are currently in the accel state, we will continue generating branches
+// 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.
+// * Tune down CPU usage to run on a pi.
+class ModelBasedLocalizer {
+ public:
+  static constexpr size_t kX = 0;
+  static constexpr size_t kY = 1;
+  static constexpr size_t kTheta = 2;
+
+  static constexpr size_t kVelocityX = 3;
+  static constexpr size_t kVelocityY = 4;
+  static constexpr size_t kNAccelStates = 5;
+
+  static constexpr size_t kLeftEncoder = 3;
+  static constexpr size_t kLeftVelocity = 4;
+  static constexpr size_t kLeftVoltageError = 5;
+  static constexpr size_t kRightEncoder = 6;
+  static constexpr size_t kRightVelocity = 7;
+  static constexpr size_t kRightVoltageError = 8;
+  static constexpr size_t kNModelStates = 9;
+
+  static constexpr size_t kNModelOutputs = 3;
+
+  static constexpr size_t kNAccelOuputs = 1;
+
+  static constexpr size_t kAccelX = 0;
+  static constexpr size_t kAccelY = 1;
+  static constexpr size_t kThetaRate = 2;
+  static constexpr size_t kNAccelInputs = 3;
+
+  static constexpr size_t kLeftVoltage = 0;
+  static constexpr size_t kRightVoltage = 1;
+  static constexpr size_t kNModelInputs = 2;
+
+  typedef Eigen::Matrix<double, kNModelStates, 1> ModelState;
+  typedef Eigen::Matrix<double, kNAccelStates, 1> AccelState;
+  typedef Eigen::Matrix<double, kNModelInputs, 1> ModelInput;
+  typedef Eigen::Matrix<double, kNAccelInputs, 1> AccelInput;
+
+  ModelBasedLocalizer(
+      const control_loops::drivetrain::DrivetrainConfig<double> &dt_config);
+  void HandleImu(aos::monotonic_clock::time_point t,
+                 const Eigen::Vector3d &gyro, const Eigen::Vector3d &accel,
+                 const Eigen::Vector2d encoders, const Eigen::Vector2d voltage);
+  void HandleImageMatch(aos::monotonic_clock::time_point,
+                        const vision::sift::ImageMatchResult *) {
+    LOG(FATAL) << "Unimplemented.";
+  }
+  void HandleReset(aos::monotonic_clock::time_point,
+                   const Eigen::Vector3d &xytheta);
+
+  flatbuffers::Offset<ModelBasedStatus> PopulateStatus(
+      flatbuffers::FlatBufferBuilder *fbb);
+
+  Eigen::Vector3d xytheta() const {
+    if (using_model_) {
+      return current_state_.model_state.block<3, 1>(0, 0);
+    } else {
+      return current_state_.accel_state.block<3, 1>(0, 0);
+    }
+  }
+
+  AccelState accel_state() const { return current_state_.accel_state; };
+
+ private:
+  struct CombinedState {
+    AccelState accel_state;
+    ModelState model_state;
+    aos::monotonic_clock::time_point branch_time;
+    double accumulated_divergence;
+  };
+
+  static flatbuffers::Offset<AccelBasedState> BuildAccelState(
+      flatbuffers::FlatBufferBuilder *fbb, const AccelState &state);
+
+  static flatbuffers::Offset<ModelBasedState> BuildModelState(
+      flatbuffers::FlatBufferBuilder *fbb, const ModelState &state);
+
+  Eigen::Matrix<double, kNModelStates, kNModelStates> AModel(
+      const ModelState &state) const;
+  Eigen::Matrix<double, kNAccelStates, kNAccelStates> AAccel() const;
+  ModelState DiffModel(const ModelState &state, const ModelInput &U) const;
+  AccelState DiffAccel(const AccelState &state, const AccelInput &U) const;
+
+  ModelState UpdateModel(const ModelState &model, const ModelInput &input,
+                         aos::monotonic_clock::duration dt) const;
+  AccelState UpdateAccel(const AccelState &accel, const AccelInput &input,
+                         aos::monotonic_clock::duration dt) const;
+
+  AccelState AccelStateForModelState(const ModelState &state) const;
+  ModelState ModelStateForAccelState(const AccelState &state,
+                                     const Eigen::Vector2d &encoders,
+                                     const double yaw_rate) const;
+  double ModelDivergence(const CombinedState &state,
+                         const AccelInput &accel_inputs,
+                         const Eigen::Vector2d &filtered_accel,
+                         const ModelInput &model_inputs);
+
+  const control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
+  const StateFeedbackHybridPlantCoefficients<2, 2, 2, double>
+      velocity_drivetrain_coefficients_;
+  Eigen::Matrix<double, kNModelStates, kNModelStates> A_continuous_model_;
+  Eigen::Matrix<double, kNAccelStates, kNAccelStates> A_continuous_accel_;
+  Eigen::Matrix<double, kNModelStates, kNModelInputs> B_continuous_model_;
+  Eigen::Matrix<double, kNAccelStates, kNAccelInputs> B_continuous_accel_;
+
+  Eigen::Matrix<double, kNModelStates, kNModelStates> Q_continuous_model_;
+
+  Eigen::Matrix<double, kNModelStates, kNModelStates> P_model_;
+  // When we are following the model, we will, on each iteration:
+  // 1) Perform a model-based update of a single state.
+  // 2) Add a hypothetical non-model-based entry based on the current state.
+  // 3) Evict too-old non-model-based entries.
+  control_loops::drivetrain::DrivetrainUkf down_estimator_;
+
+  // Buffer of old branches these are all created by initializing a new
+  // model-based state based on the current state, and then initializing a new
+  // accel-based state on top of that new model-based state (to eliminate the
+  // impact of any lateral motion).
+  // We then integrate up all of these states and observe how much the model and
+  // accel based states of each branch compare to one another.
+  aos::RingBuffer<CombinedState, 2000> branches_;
+
+  CombinedState current_state_;
+  aos::monotonic_clock::time_point t_ = aos::monotonic_clock::min_time;
+  bool using_model_;
+
+  double last_residual_ = 0.0;
+  double filtered_residual_ = 0.0;
+  Eigen::Vector2d filtered_residual_accel_ = Eigen::Vector2d::Zero();
+  Eigen::Vector3d abs_accel_ = Eigen::Vector3d::Zero();
+  double velocity_residual_ = 0.0;
+  double accel_residual_ = 0.0;
+  double theta_rate_residual_ = 0.0;
+  int hysteresis_count_ = 0;
+
+  int clock_resets_ = 0;
+
+  friend class testing::LocalizerTest;
+};
+
+class EventLoopLocalizer {
+ public:
+  EventLoopLocalizer(
+      aos::EventLoop *event_loop,
+      const control_loops::drivetrain::DrivetrainConfig<double> &dt_config);
+
+ private:
+  aos::EventLoop *event_loop_;
+  ModelBasedLocalizer model_based_;
+  aos::Sender<LocalizerStatus> status_sender_;
+  aos::Sender<LocalizerOutput> output_sender_;
+  aos::Fetcher<frc971::control_loops::drivetrain::Position> position_fetcher_;
+  aos::Fetcher<frc971::control_loops::drivetrain::Output> output_fetcher_;
+  zeroing::ImuZeroer zeroer_;
+  aos::monotonic_clock::time_point last_output_send_ =
+      aos::monotonic_clock::min_time;
+};
+}  // namespace frc971::controls
+#endif // Y2022_CONTROL_LOOPS_LOCALIZER_LOCALIZER_H_