Move shared flywheel code into frc971

Signed-off-by: Maxwell Henderson <mxwhenderson@gmail.com>
Change-Id: Ieac317a3e5bc8243e63473f485a2467b74aed348
diff --git a/frc971/control_loops/flywheel/BUILD b/frc971/control_loops/flywheel/BUILD
new file mode 100644
index 0000000..6399abe
--- /dev/null
+++ b/frc971/control_loops/flywheel/BUILD
@@ -0,0 +1,49 @@
+load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+load("@com_github_google_flatbuffers//:typescript.bzl", "flatbuffer_ts_library")
+
+flatbuffer_cc_library(
+    name = "flywheel_controller_status_fbs",
+    srcs = [
+        "flywheel_controller_status.fbs",
+    ],
+    gen_reflections = 1,
+    visibility = ["//visibility:public"],
+)
+
+flatbuffer_ts_library(
+    name = "flywheel_controller_status_ts_fbs",
+    srcs = [
+        "flywheel_controller_status.fbs",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "flywheel_test_plant",
+    hdrs = [
+        "flywheel_test_plant.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":flywheel_controller",
+    ],
+)
+
+cc_library(
+    name = "flywheel_controller",
+    srcs = [
+        "flywheel_controller.cc",
+    ],
+    hdrs = [
+        "flywheel_controller.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":flywheel_controller_status_fbs",
+        "//frc971/control_loops:control_loop",
+        "//frc971/control_loops:hybrid_state_feedback_loop",
+        "//frc971/control_loops:profiled_subsystem",
+    ],
+)
diff --git a/frc971/control_loops/flywheel/flywheel_controller.cc b/frc971/control_loops/flywheel/flywheel_controller.cc
new file mode 100644
index 0000000..adea0e8
--- /dev/null
+++ b/frc971/control_loops/flywheel/flywheel_controller.cc
@@ -0,0 +1,170 @@
+#include "frc971/control_loops/flywheel/flywheel_controller.h"
+
+#include <chrono>
+
+#include "aos/logging/logging.h"
+namespace frc971 {
+namespace control_loops {
+namespace flywheel {
+
+// Class to current limit battery current for a flywheel controller.
+class CurrentLimitedStateFeedbackController
+    : public StateFeedbackLoop<3, 1, 1, double,
+                               StateFeedbackHybridPlant<3, 1, 1>,
+                               HybridKalman<3, 1, 1>> {
+ public:
+  // Builds a CurrentLimitedStateFeedbackController given the coefficients, bemf
+  // coefficient (units of radians/sec / volt), and motor resistance in ohms.
+  CurrentLimitedStateFeedbackController(
+      StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>,
+                        HybridKalman<3, 1, 1>> &&other,
+      double bemf, double resistance)
+      : StateFeedbackLoop(std::move(other)),
+        bemf_(bemf),
+        resistance_(resistance) {}
+
+  double resistance() const { return resistance_; }
+  double bemf() const { return bemf_; }
+
+  void CapU() override {
+    const double bemf_voltage = X_hat(1) / bemf_;
+    // Solve the system of equations:
+    //
+    //   motor_current = (u - bemf_voltage) / resistance
+    //   battery_current = ((u - bemf_voltage) / resistance) * u / 12.0
+    //   0.0 = u * u - u * bemf_voltage - max_current * 12.0 * resistance
+    //
+    // And we have a quadratic!
+    const double a = 1;
+    const double b = -bemf_voltage;
+    const double c_positive = -70.0 * 12.0 * resistance_;
+    const double c_negative = -25.0 * 12.0 * resistance_;
+
+    // Root is always positive.
+    const double root_positive = std::sqrt(b * b - 4.0 * a * c_positive);
+    const double root_negative = std::sqrt(b * b - 4.0 * a * c_negative);
+    const double upper_limit = (-b + root_positive) / (2.0 * a);
+    const double lower_limit = (-b - root_negative) / (2.0 * a);
+
+    // Limit to the battery voltage and the current limit voltage.
+    mutable_U(0, 0) = std::clamp(U(0, 0), lower_limit, upper_limit);
+    if (R(1) > 50.0) {
+      mutable_U(0, 0) = std::clamp(U(0, 0), 1.0, 12.0);
+    } else {
+      mutable_U(0, 0) = std::clamp(U(0, 0), 0.0, 12.0);
+    }
+  }
+
+ private:
+  double bemf_ = 0.0;
+  double resistance_ = 0.0;
+};
+
+FlywheelController::FlywheelController(
+    StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>,
+                      HybridKalman<3, 1, 1>> &&loop,
+    double bemf, double resistance)
+    : loop_(new CurrentLimitedStateFeedbackController(std::move(loop), bemf,
+                                                      resistance)) {
+  history_.fill(std::pair<double, ::aos::monotonic_clock::time_point>(
+      0, ::aos::monotonic_clock::epoch()));
+  Y_.setZero();
+}
+
+void FlywheelController::set_goal(double angular_velocity_goal) {
+  loop_->mutable_next_R() << 0.0, angular_velocity_goal, 0.0;
+  last_goal_ = angular_velocity_goal;
+}
+
+void FlywheelController::set_position(
+    double current_position,
+    const aos::monotonic_clock::time_point position_timestamp) {
+  // Project time forwards.
+  const int newest_history_position =
+      ((history_position_ == 0) ? kHistoryLength : history_position_) - 1;
+
+  if (!first_) {
+    loop_->UpdateObserver(
+        loop_->U(),
+        position_timestamp - std::get<1>(history_[newest_history_position]));
+  } else {
+    first_ = false;
+  }
+
+  // Update position in the model.
+  Y_ << current_position;
+
+  // Add the position to the history.
+  history_[history_position_] =
+      std::pair<double, ::aos::monotonic_clock::time_point>(current_position,
+                                                            position_timestamp);
+  history_position_ = (history_position_ + 1) % kHistoryLength;
+
+  loop_->Correct(Y_);
+}
+
+double FlywheelController::voltage() const { return loop_->U(0, 0); }
+
+double FlywheelController::current() const {
+  return ((voltage() - (velocity() / loop_->bemf())) / (loop_->resistance())) *
+         voltage() / 12.0;
+}
+
+void FlywheelController::Update(bool disabled) {
+  loop_->mutable_R() = loop_->next_R();
+  if (loop_->R(1, 0) < 1.0) {
+    // Kill power at low angular velocities.
+    disabled = true;
+  }
+
+  loop_->UpdateController(disabled);
+}
+
+flatbuffers::Offset<FlywheelControllerStatus> FlywheelController::SetStatus(
+    flatbuffers::FlatBufferBuilder *fbb) {
+  // Compute the oldest point in the history.
+  const int oldest_history_position = history_position_;
+  const int newest_history_position =
+      ((history_position_ == 0) ? kHistoryLength : history_position_) - 1;
+  const int second_newest_history_position =
+      ((newest_history_position == 0) ? kHistoryLength
+                                      : newest_history_position) -
+      1;
+
+  const double total_loop_time = ::aos::time::DurationInSeconds(
+      std::get<1>(history_[newest_history_position]) -
+      std::get<1>(history_[oldest_history_position]));
+
+  const double distance_traveled =
+      std::get<0>(history_[newest_history_position]) -
+      std::get<0>(history_[oldest_history_position]);
+
+  const double last_loop_time = ::aos::time::DurationInSeconds(
+      std::get<1>(history_[newest_history_position]) -
+      std::get<1>(history_[second_newest_history_position]));
+
+  const double last_distance_traveled =
+      std::get<0>(history_[newest_history_position]) -
+      std::get<0>(history_[second_newest_history_position]);
+
+  // Compute the distance moved over that time period.
+  avg_angular_velocity_ = (distance_traveled) / (total_loop_time);
+
+  FlywheelControllerStatusBuilder builder(*fbb);
+
+  builder.add_avg_angular_velocity(avg_angular_velocity_);
+  builder.add_dt_angular_velocity(last_distance_traveled / last_loop_time);
+  builder.add_angular_velocity(loop_->X_hat(1, 0));
+  builder.add_voltage_error(loop_->X_hat(2, 0));
+  builder.add_commanded_current(current());
+  builder.add_angular_velocity_goal(last_goal_);
+  return builder.Finish();
+}
+
+FlywheelController::~FlywheelController() {}
+
+double FlywheelController::velocity() const { return loop_->X_hat(1, 0); }
+
+}  // namespace flywheel
+}  // namespace control_loops
+}  // namespace frc971
diff --git a/frc971/control_loops/flywheel/flywheel_controller.h b/frc971/control_loops/flywheel/flywheel_controller.h
new file mode 100644
index 0000000..db8baeb
--- /dev/null
+++ b/frc971/control_loops/flywheel/flywheel_controller.h
@@ -0,0 +1,80 @@
+#ifndef FRC971_CONTROL_LOOPS_SHOOTER_FLYWHEEL_CONTROLLER_H_
+#define FRC971_CONTROL_LOOPS_SHOOTER_FLYWHEEL_CONTROLLER_H_
+
+#include <memory>
+
+#include "aos/time/time.h"
+#include "frc971/control_loops/control_loop.h"
+#include "frc971/control_loops/flywheel/flywheel_controller_status_generated.h"
+#include "frc971/control_loops/hybrid_state_feedback_loop.h"
+#include "frc971/control_loops/state_feedback_loop.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace flywheel {
+
+class CurrentLimitedStateFeedbackController;
+
+// Handles the velocity control of each flywheel.
+class FlywheelController {
+ public:
+  FlywheelController(
+      StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>,
+                        HybridKalman<3, 1, 1>> &&loop,
+      double bemf, double resistance);
+
+  ~FlywheelController();
+
+  // Sets the velocity goal in radians/sec
+  void set_goal(double angular_velocity_goal);
+  double goal() const { return last_goal_; }
+  // Sets the current encoder position in radians
+  void set_position(double current_position,
+                    const aos::monotonic_clock::time_point position_timestamp);
+
+  // Populates the status structure.
+  flatbuffers::Offset<FlywheelControllerStatus> SetStatus(
+      flatbuffers::FlatBufferBuilder *fbb);
+
+  // Returns the control loop calculated voltage.
+  double voltage() const;
+
+  // Returns the expected battery current for the last U.
+  double current() const;
+
+  // Returns the instantaneous velocity.
+  double velocity() const;
+
+  // Executes the control loop for a cycle.
+  void Update(bool disabled);
+
+  double avg_angular_velocity() { return avg_angular_velocity_; }
+
+ private:
+  // The current sensor measurement.
+  Eigen::Matrix<double, 1, 1> Y_;
+  // The control loop.
+  ::std::unique_ptr<CurrentLimitedStateFeedbackController> loop_;
+
+  // History array for calculating a filtered angular velocity.
+  static constexpr int kHistoryLength = 10;
+  ::std::array<std::pair<double, ::aos::monotonic_clock::time_point>,
+               kHistoryLength>
+      history_;
+  ptrdiff_t history_position_ = 0;
+
+  // Average velocity logging.
+  double avg_angular_velocity_ = 0;
+
+  double last_goal_ = 0;
+
+  bool first_ = true;
+
+  DISALLOW_COPY_AND_ASSIGN(FlywheelController);
+};
+
+}  // namespace flywheel
+}  // namespace control_loops
+}  // namespace frc971
+
+#endif  // FRC971_CONTROL_LOOPS_SHOOTER_FLYWHEEL_CONTROLLER_H_
diff --git a/frc971/control_loops/flywheel/flywheel_controller_status.fbs b/frc971/control_loops/flywheel/flywheel_controller_status.fbs
new file mode 100644
index 0000000..0f95818
--- /dev/null
+++ b/frc971/control_loops/flywheel/flywheel_controller_status.fbs
@@ -0,0 +1,22 @@
+namespace frc971.control_loops.flywheel;
+
+table FlywheelControllerStatus {
+  // The current average velocity in radians/second
+  avg_angular_velocity:double (id: 0);
+
+  // The current instantaneous filtered velocity in radians/second.
+  angular_velocity:double (id: 1);
+
+  // The target speed selected by the lookup table or from manual override
+  // Can be compared to velocity to determine if ready.
+  angular_velocity_goal:double (id: 2);
+
+  // Voltage error.
+  voltage_error:double (id: 3);
+
+  // The commanded battery current.
+  commanded_current:double (id: 4);
+
+  // The angular velocity of the flywheel computed using delta x / delta t
+  dt_angular_velocity:double (id: 5);
+}
diff --git a/frc971/control_loops/flywheel/flywheel_test_plant.h b/frc971/control_loops/flywheel/flywheel_test_plant.h
new file mode 100644
index 0000000..01e347d
--- /dev/null
+++ b/frc971/control_loops/flywheel/flywheel_test_plant.h
@@ -0,0 +1,47 @@
+#ifndef FRC971_CONTROL_LOOPS_FLYWHEEL_FLYWHEEL_TEST_PLANT_H_
+#define FRC971_CONTROL_LOOPS_FLYWHEEL_FLYWHEEL_TEST_PLANT_H_
+
+#include "frc971/control_loops/flywheel/flywheel_controller.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace flywheel {
+
+class FlywheelPlant : public StateFeedbackPlant<2, 1, 1> {
+ public:
+  explicit FlywheelPlant(StateFeedbackPlant<2, 1, 1> &&other, double bemf,
+                         double resistance)
+      : StateFeedbackPlant<2, 1, 1>(::std::move(other)),
+        bemf_(bemf),
+        resistance_(resistance) {}
+
+  void CheckU(const Eigen::Matrix<double, 1, 1> &U) override {
+    EXPECT_LE(U(0, 0), U_max(0, 0) + 0.00001 + voltage_offset_);
+    EXPECT_GE(U(0, 0), U_min(0, 0) - 0.00001 + voltage_offset_);
+  }
+
+  double motor_current(const Eigen::Matrix<double, 1, 1> U) const {
+    return (U(0) - X(1) / bemf_) / resistance_;
+  }
+
+  double battery_current(const Eigen::Matrix<double, 1, 1> U) const {
+    return motor_current(U) * U(0) / 12.0;
+  }
+
+  double voltage_offset() const { return voltage_offset_; }
+  void set_voltage_offset(double voltage_offset) {
+    voltage_offset_ = voltage_offset;
+  }
+
+ private:
+  double voltage_offset_ = 0.0;
+
+  double bemf_;
+  double resistance_;
+};
+
+}  // namespace flywheel
+}  // namespace control_loops
+}  // namespace frc971
+
+#endif  // FRC971_CONTROL_LOOPS_FLYWHEEL_FLYWHEEL_TEST_PLANT_H_