Make swerve_control_loops use continuous controller & zeroing

This updates the per-module control code to make use of the continuous
control loops and to actually zero the modules so that we can do simple
data collection.

Future TODOs:
* Update the control loop to use the static flatbuffers API.

Change-Id: I40c41fb207cbee9ddc3c31d4b5b023da468964d8
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 2d06837..1b8dcc3 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -67,7 +67,10 @@
         ":swerve_drivetrain_output_fbs",
         ":swerve_drivetrain_position_fbs",
         ":swerve_drivetrain_status_fbs",
+        ":swerve_zeroing_fbs",
         "//frc971/control_loops:control_loop",
+        "//frc971/control_loops:static_zeroing_single_dof_profiled_subsystem",
+        "//frc971/zeroing:continuous_absolute_encoder",
     ],
 )
 
@@ -76,6 +79,7 @@
     srcs = ["swerve_control_test.cc"],
     data = [
         ":aos_config",
+        "//frc971/control_loops/swerve/test_module:rotation_json",
     ],
     target_compatible_with = ["@platforms//os:linux"],
     deps = [
diff --git a/frc971/control_loops/swerve/swerve_control_loops.cc b/frc971/control_loops/swerve/swerve_control_loops.cc
index bd0d3bd..ec841c8 100644
--- a/frc971/control_loops/swerve/swerve_control_loops.cc
+++ b/frc971/control_loops/swerve/swerve_control_loops.cc
@@ -2,39 +2,114 @@
 
 namespace frc971::control_loops::swerve {
 
-SwerveControlLoops::SwerveControlLoops(::aos::EventLoop *event_loop,
-                                       const ::std::string &name)
+SwerveControlLoops::SwerveControlLoops(
+    ::aos::EventLoop *event_loop,
+    const frc971::control_loops::
+        StaticZeroingSingleDOFProfiledSubsystemCommonParams *rotation_params,
+    const SwerveZeroing *zeroing_params, const ::std::string &name)
     : frc971::controls::ControlLoop<Goal, Position, StatusStatic, OutputStatic>(
-          event_loop, name) {}
+          event_loop, name),
+      front_left_(rotation_params, zeroing_params->front_left()),
+      front_right_(rotation_params, zeroing_params->front_right()),
+      back_left_(rotation_params, zeroing_params->back_left()),
+      back_right_(rotation_params, zeroing_params->back_right()) {}
 
 void SwerveControlLoops::RunIteration(
     const Goal *goal, const Position *position,
     aos::Sender<OutputStatic>::StaticBuilder *output_builder,
     aos::Sender<StatusStatic>::StaticBuilder *status_builder) {
-  (void)goal, (void)position;
+  if (WasReset()) {
+    front_left_.Reset();
+    front_right_.Reset();
+    back_left_.Reset();
+    back_right_.Reset();
+  }
+  aos::FlatbufferFixedAllocatorArray<
+      frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal, 512>
+      front_left_goal_buffer, front_right_goal_buffer, back_left_goal_buffer,
+      back_right_goal_buffer;
 
-  if (output_builder != nullptr && goal != nullptr) {
+  front_left_goal_buffer.Finish(
+      frc971::control_loops::CreateStaticZeroingSingleDOFProfiledSubsystemGoal(
+          *front_left_goal_buffer.fbb(),
+          (goal != nullptr && goal->has_front_left_goal())
+              ? goal->front_left_goal()->rotation_angle()
+              : 0.0));
+  front_right_goal_buffer.Finish(
+      frc971::control_loops::CreateStaticZeroingSingleDOFProfiledSubsystemGoal(
+          *front_right_goal_buffer.fbb(),
+          (goal != nullptr && goal->has_front_right_goal())
+              ? goal->front_right_goal()->rotation_angle()
+              : 0.0));
+  back_left_goal_buffer.Finish(
+      frc971::control_loops::CreateStaticZeroingSingleDOFProfiledSubsystemGoal(
+          *back_left_goal_buffer.fbb(),
+          (goal != nullptr && goal->has_back_left_goal())
+              ? goal->back_left_goal()->rotation_angle()
+              : 0.0));
+  back_right_goal_buffer.Finish(
+      frc971::control_loops::CreateStaticZeroingSingleDOFProfiledSubsystemGoal(
+          *back_right_goal_buffer.fbb(),
+          (goal != nullptr && goal->has_back_right_goal())
+              ? goal->back_right_goal()->rotation_angle()
+              : 0.0));
+  const bool disabled = front_left_.Correct(
+      &front_left_goal_buffer.message(),
+      position->front_left()->rotation_position(), output_builder == nullptr);
+  front_right_.Correct(&front_right_goal_buffer.message(),
+                       position->front_right()->rotation_position(),
+                       output_builder == nullptr);
+  back_left_.Correct(&back_left_goal_buffer.message(),
+                     position->back_left()->rotation_position(),
+                     output_builder == nullptr);
+  back_right_.Correct(&back_right_goal_buffer.message(),
+                      position->back_right()->rotation_position(),
+                      output_builder == nullptr);
+
+  const double front_left_voltage = front_left_.UpdateController(disabled);
+  front_left_.UpdateObserver(front_left_voltage);
+  const double front_right_voltage = front_right_.UpdateController(disabled);
+  front_right_.UpdateObserver(front_right_voltage);
+  const double back_left_voltage = back_left_.UpdateController(disabled);
+  back_left_.UpdateObserver(back_left_voltage);
+  const double back_right_voltage = back_right_.UpdateController(disabled);
+  back_right_.UpdateObserver(back_right_voltage);
+  (void)goal, (void)position;
+  aos::FlatbufferFixedAllocatorArray<
+      frc971::control_loops::AbsoluteEncoderProfiledJointStatus, 512>
+      front_left_status_buffer, front_right_status_buffer,
+      back_left_status_buffer, back_right_status_buffer;
+  front_left_status_buffer.Finish(
+      front_left_.MakeStatus(front_left_status_buffer.fbb()));
+  front_right_status_buffer.Finish(
+      front_right_.MakeStatus(front_right_status_buffer.fbb()));
+  back_left_status_buffer.Finish(
+      back_left_.MakeStatus(back_left_status_buffer.fbb()));
+  back_right_status_buffer.Finish(
+      back_right_.MakeStatus(back_right_status_buffer.fbb()));
+
+  if (output_builder != nullptr) {
     OutputStatic *output = output_builder->get();
 
     auto front_left_output = output->add_front_left_output();
-    front_left_output->set_rotation_current(0);
+    front_left_output->set_rotation_current(front_left_voltage);
     front_left_output->set_translation_current(
-        goal->front_left_goal()->translation_current());
+        goal ? goal->front_left_goal()->translation_current() : 0.0);
 
     auto front_right_output = output->add_front_right_output();
-    front_right_output->set_rotation_current(0);
+    front_right_output->set_rotation_current(front_right_voltage);
     front_right_output->set_translation_current(
-        goal->front_right_goal()->translation_current());
+        goal ? goal->front_right_goal()->translation_current() : 0.0);
 
     auto back_left_output = output->add_back_left_output();
-    back_left_output->set_rotation_current(0);
+    back_left_output->set_rotation_current(back_left_voltage);
     back_left_output->set_translation_current(
-        goal->back_left_goal()->translation_current());
+        goal ? goal->back_left_goal()->translation_current() : 0.0);
 
     auto back_right_output = output->add_back_right_output();
-    back_right_output->set_rotation_current(0);
+    back_right_output->set_rotation_current(back_right_voltage);
     back_right_output->set_translation_current(
-        goal->back_right_goal()->translation_current());
+        goal ? goal->back_right_goal()->translation_current() : 0.0);
 
     // Ignore the return value of Send
     output_builder->CheckOk(output_builder->Send());
@@ -44,16 +119,20 @@
     StatusStatic *status = status_builder->get();
 
     auto front_left_status = status->add_front_left_status();
-    PopulateSwerveModuleRotation(front_left_status);
+    PopulateSwerveModuleRotation(front_left_status,
+                                 &front_left_status_buffer.message());
 
     auto front_right_status = status->add_front_right_status();
-    PopulateSwerveModuleRotation(front_right_status);
+    PopulateSwerveModuleRotation(front_right_status,
+                                 &front_right_status_buffer.message());
 
     auto back_left_status = status->add_back_left_status();
-    PopulateSwerveModuleRotation(back_left_status);
+    PopulateSwerveModuleRotation(back_left_status,
+                                 &back_left_status_buffer.message());
 
     auto back_right_status = status->add_back_right_status();
-    PopulateSwerveModuleRotation(back_right_status);
+    PopulateSwerveModuleRotation(back_right_status,
+                                 &back_right_status_buffer.message());
 
     // Ignore the return value of Send
     status_builder->CheckOk(status_builder->Send());
diff --git a/frc971/control_loops/swerve/swerve_control_loops.h b/frc971/control_loops/swerve/swerve_control_loops.h
index 78f6ba0..68016ab 100644
--- a/frc971/control_loops/swerve/swerve_control_loops.h
+++ b/frc971/control_loops/swerve/swerve_control_loops.h
@@ -4,19 +4,24 @@
 #include "aos/time/time.h"
 #include "frc971/control_loops/control_loop.h"
 #include "frc971/control_loops/profiled_subsystem_generated.h"
+#include "frc971/control_loops/profiled_subsystem_static.h"
+#include "frc971/control_loops/static_zeroing_single_dof_profiled_subsystem.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_can_position_generated.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_goal_generated.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_output_static.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_position_generated.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_status_static.h"
+#include "frc971/control_loops/swerve/swerve_zeroing_static.h"
+#include "frc971/zeroing/continuous_absolute_encoder.h"
 
 namespace frc971::control_loops::swerve {
 
 inline void PopulateSwerveModuleRotation(
-    SwerveModuleStatusStatic *swerve_module_table) {
+    SwerveModuleStatusStatic *swerve_module_table,
+    const frc971::control_loops::AbsoluteEncoderProfiledJointStatus
+        *rotation_status) {
   auto rotation = swerve_module_table->add_rotation();
-  auto estimator_state = rotation->add_estimator_state();
-  (void)estimator_state;
+  CHECK(rotation->FromFlatbuffer(rotation_status));
 }
 
 // Handles the translation and rotation current for each swerve module
@@ -24,14 +29,24 @@
     : public ::frc971::controls::ControlLoop<Goal, Position, StatusStatic,
                                              OutputStatic> {
  public:
-  explicit SwerveControlLoops(::aos::EventLoop *event_loop,
-                              const ::std::string &name = "/swerve");
+  using AbsoluteEncoderSubsystem =
+      ::frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystem<
+          ::frc971::zeroing::ContinuousAbsoluteEncoderZeroingEstimator,
+          ::frc971::control_loops::AbsoluteEncoderProfiledJointStatus>;
+
+  explicit SwerveControlLoops(
+      ::aos::EventLoop *event_loop,
+      const frc971::control_loops::
+          StaticZeroingSingleDOFProfiledSubsystemCommonParams *rotation_params,
+      const SwerveZeroing *zeroing_params,
+      const ::std::string &name = "/swerve");
 
  protected:
   void RunIteration(
       const Goal *goal, const Position *position,
       aos::Sender<OutputStatic>::StaticBuilder *output_builder,
       aos::Sender<StatusStatic>::StaticBuilder *status_builder) override;
+  AbsoluteEncoderSubsystem front_left_, front_right_, back_left_, back_right_;
 };
 
 }  // namespace frc971::control_loops::swerve
diff --git a/frc971/control_loops/swerve/swerve_control_test.cc b/frc971/control_loops/swerve/swerve_control_test.cc
index 9378f1c..051d1a7 100644
--- a/frc971/control_loops/swerve/swerve_control_test.cc
+++ b/frc971/control_loops/swerve/swerve_control_test.cc
@@ -4,15 +4,16 @@
 #include <memory>
 #include <vector>
 
+#include "absl/strings/str_format.h"
 #include "gtest/gtest.h"
 
 #include "aos/events/shm_event_loop.h"
 #include "frc971/control_loops/control_loop_test.h"
 #include "frc971/control_loops/swerve/swerve_control_loops.h"
-#include "frc971/control_loops/swerve/swerve_drivetrain_can_position_generated.h"
-#include "frc971/control_loops/swerve/swerve_drivetrain_goal_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_can_position_static.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_goal_static.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_output_generated.h"
-#include "frc971/control_loops/swerve/swerve_drivetrain_position_generated.h"
+#include "frc971/control_loops/swerve/swerve_drivetrain_position_static.h"
 #include "frc971/control_loops/swerve/swerve_drivetrain_status_generated.h"
 #include "frc971/control_loops/team_number_test_environment.h"
 
@@ -27,7 +28,7 @@
  public:
   SwerveControlSimulation(::aos::EventLoop *event_loop, chrono::nanoseconds dt)
       : event_loop_(event_loop),
-        position_sender_(event_loop_->MakeSender<Position>("/swerve")),
+        position_sender_(event_loop_->MakeSender<PositionStatic>("/swerve")),
         can_position_sender_(event_loop_->MakeSender<CanPosition>("/swerve")),
         goal_fetcher_(event_loop_->MakeFetcher<Goal>("/swerve")),
         status_fetcher_(event_loop_->MakeFetcher<Status>("/swerve")),
@@ -45,12 +46,14 @@
 
   // Sends a queue message with the position data.
   void SendPositionMessage() {
-    auto builder = position_sender_.MakeBuilder();
+    auto builder = position_sender_.MakeStaticBuilder();
 
-    Position::Builder position_builder = builder.MakeBuilder<Position>();
+    builder->add_front_left()->add_rotation_position();
+    builder->add_front_right()->add_rotation_position();
+    builder->add_back_left()->add_rotation_position();
+    builder->add_back_right()->add_rotation_position();
 
-    EXPECT_EQ(builder.Send(position_builder.Finish()),
-              aos::RawSender::Error::kOk);
+    EXPECT_EQ(builder.Send(), aos::RawSender::Error::kOk);
   }
 
   void VerifyNearGoal() {
@@ -88,7 +91,7 @@
  private:
   ::aos::EventLoop *event_loop_;
 
-  ::aos::Sender<Position> position_sender_;
+  ::aos::Sender<PositionStatic> position_sender_;
   ::aos::Sender<CanPosition> can_position_sender_;
   ::aos::Sender<Goal> goal_sender_;
 
@@ -112,7 +115,57 @@
         goal_sender_(swerve_test_event_loop_->MakeSender<Goal>("/swerve")),
 
         swerve_control_event_loop_(MakeEventLoop("swerve_control")),
-        swerve_control_loops_(swerve_control_event_loop_.get(), "/swerve"),
+        subsystem_params_(
+            aos::JsonToFlatbuffer<
+                frc971::control_loops::
+                    StaticZeroingSingleDOFProfiledSubsystemCommonParams>(
+                absl::StrFormat(R"json({
+                    "zeroing_voltage": 3.0,
+                    "operating_voltage": 12.0,
+                    "zeroing_profile_params": {
+                      "max_velocity": 0.5,
+                      "max_acceleration": 3.0
+                    },
+                    "default_profile_params":{
+                      "max_velocity": 12.0,
+                      "max_acceleration": 55.0
+                    },
+                    "range": {
+                        "lower_hard": -inf,
+                        "upper_hard": inf,
+                        "lower": -inf,
+                        "upper": inf
+                    },
+                    "loop": %s
+                    })json",
+                                aos::util::ReadFileToStringOrDie(
+                                    "frc971/control_loops/swerve/test_module/"
+                                    "integral_rotation_plant.json")))),
+        zeroing_params_(aos::JsonToFlatbuffer<SwerveZeroing>(R"json({
+        "front_left": {
+          "average_filter_size": 10,
+          "one_revolution_distance": 6,
+          "moving_buffer_size": 10
+        },
+        "front_right": {
+          "average_filter_size": 10,
+          "one_revolution_distance": 6,
+          "moving_buffer_size": 10
+        },
+        "back_left": {
+          "average_filter_size": 10,
+          "one_revolution_distance": 6,
+          "moving_buffer_size": 10
+        },
+        "back_right": {
+          "average_filter_size": 10,
+          "one_revolution_distance": 6,
+          "moving_buffer_size": 10
+        }
+        })json")),
+        swerve_control_loops_(swerve_control_event_loop_.get(),
+                              &subsystem_params_.message(),
+                              &zeroing_params_.message(), "/swerve"),
 
         swerve_control_simulation_event_loop_(MakeEventLoop("simulation")),
         swerve_control_simulation_(swerve_control_simulation_event_loop_.get(),
@@ -125,6 +178,11 @@
   ::aos::Sender<Goal> goal_sender_;
 
   ::std::unique_ptr<::aos::EventLoop> swerve_control_event_loop_;
+  aos::FlatbufferDetachedBuffer<
+      frc971::control_loops::
+          StaticZeroingSingleDOFProfiledSubsystemCommonParams>
+      subsystem_params_;
+  aos::FlatbufferDetachedBuffer<SwerveZeroing> zeroing_params_;
   SwerveControlLoops swerve_control_loops_;
 
   ::std::unique_ptr<::aos::EventLoop> swerve_control_simulation_event_loop_;
@@ -158,4 +216,4 @@
   swerve_control_simulation_.VerifyNearGoal();
 }
 
-}  // namespace frc971::control_loops::swerve::testing
\ No newline at end of file
+}  // namespace frc971::control_loops::swerve::testing
diff --git a/frc971/control_loops/swerve/test_module/BUILD b/frc971/control_loops/swerve/test_module/BUILD
new file mode 100644
index 0000000..972bddf
--- /dev/null
+++ b/frc971/control_loops/swerve/test_module/BUILD
@@ -0,0 +1,36 @@
+py_binary(
+    name = "rotation",
+    srcs = [
+        "rotation.py",
+    ],
+    target_compatible_with = ["@platforms//cpu:x86_64"],
+    deps = [
+        "//frc971/control_loops/python:angular_system_current",
+        "//frc971/control_loops/python:controls",
+        "@pip//glog",
+        "@pip//python_gflags",
+    ],
+)
+
+genrule(
+    name = "genrule_rotation",
+    outs = [
+        "rotation_plant.h",
+        "rotation_plant.cc",
+        "rotation_plant.json",
+        "integral_rotation_plant.h",
+        "integral_rotation_plant.cc",
+        "integral_rotation_plant.json",
+    ],
+    cmd = "$(location :rotation) $(OUTS)",
+    target_compatible_with = ["@platforms//os:linux"],
+    tools = [
+        ":rotation",
+    ],
+)
+
+filegroup(
+    name = "rotation_json",
+    srcs = ["integral_rotation_plant.json"],
+    visibility = ["//visibility:public"],
+)
diff --git a/frc971/control_loops/swerve/test_module/rotation.py b/frc971/control_loops/swerve/test_module/rotation.py
new file mode 100644
index 0000000..3f94e53
--- /dev/null
+++ b/frc971/control_loops/swerve/test_module/rotation.py
@@ -0,0 +1,57 @@
+#!/usr/bin/python3
+
+from aos.util.trapezoid_profile import TrapezoidProfile
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import angular_system_current
+from frc971.control_loops.python import controls
+import numpy
+import sys
+from matplotlib import pylab
+import gflags
+import glog
+import matplotlib
+
+FLAGS = gflags.FLAGS
+
+try:
+    gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
+except gflags.DuplicateFlagError:
+    pass
+
+kRotation = angular_system_current.AngularSystemCurrentParams(
+    name='Rotation',
+    motor=control_loop.KrakenFOC(),
+    G=9.0 / 24.0 * 14.0 / 72.0,
+    J=3.1 / 1000.0,
+    q_pos=0.05,
+    q_vel=20.0,
+    kalman_q_pos=0.12,
+    kalman_q_vel=2.0,
+    kalman_q_voltage=4.0,
+    kalman_r_position=0.05,
+    radius=25 * 0.0254,
+    wrap_point=2.0 * numpy.pi)
+
+
+def main(argv):
+    if FLAGS.plot:
+        R = numpy.matrix([[numpy.pi / 2.0], [0.0]])
+        angular_system_current.PlotKick(kRotation, R)
+        angular_system_current.PlotMotion(kRotation, R)
+        return
+
+    # Write the generated constants out to a file.
+    if len(argv) != 7:
+        glog.fatal(
+            'Expected .h file name and .cc file name for the wrist and integral wrist.'
+        )
+    else:
+        namespaces = ['frc971', 'control_loops', 'swerve', 'test_module']
+        angular_system_current.WriteAngularSystemCurrent(
+            kRotation, argv[1:4], argv[4:7], namespaces)
+
+
+if __name__ == '__main__':
+    argv = FLAGS(sys.argv)
+    glog.init()
+    sys.exit(main(argv))
diff --git a/y2024_swerve/BUILD b/y2024_swerve/BUILD
index e432358..2357c6e 100644
--- a/y2024_swerve/BUILD
+++ b/y2024_swerve/BUILD
@@ -32,6 +32,7 @@
         "//aos/starter:irq_affinity",
         ":wpilib_interface",
         ":swerve_publisher",
+        "//y2024_swerve/control_loops:swerve_control_loops",
         "//frc971/can_logger",
         "//aos/network:message_bridge_client",
         "//aos/network:message_bridge_server",
diff --git a/y2024_swerve/control_loops/BUILD b/y2024_swerve/control_loops/BUILD
index 67fa349..0417af9 100644
--- a/y2024_swerve/control_loops/BUILD
+++ b/y2024_swerve/control_loops/BUILD
@@ -1,5 +1,21 @@
 load("//tools/build_rules:js.bzl", "ts_project")
 
+cc_binary(
+    name = "swerve_control_loops",
+    srcs = [
+        "swerve_control_loops_main.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//aos:init",
+        "//aos/events:shm_event_loop",
+        "//aos/time",
+        "//frc971/constants:constants_sender_lib",
+        "//frc971/control_loops/swerve:swerve_control_loops",
+        "//y2024_swerve/constants:constants_fbs",
+    ],
+)
+
 ts_project(
     name = "swerve_plotter",
     srcs = ["swerve_plotter.ts"],
diff --git a/y2024_swerve/control_loops/swerve_control_loops_main.cc b/y2024_swerve/control_loops/swerve_control_loops_main.cc
new file mode 100644
index 0000000..813744a
--- /dev/null
+++ b/y2024_swerve/control_loops/swerve_control_loops_main.cc
@@ -0,0 +1,30 @@
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "frc971/constants/constants_sender_lib.h"
+#include "frc971/control_loops/swerve/swerve_control_loops.h"
+#include "y2024_swerve/constants/constants_generated.h"
+
+using frc971::control_loops::swerve::SwerveControlLoops;
+
+int main(int argc, char **argv) {
+  ::aos::InitGoogle(&argc, &argv);
+
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig("aos_config.json");
+
+  frc971::constants::WaitForConstants<y2024_swerve::Constants>(
+      &config.message());
+
+  ::aos::ShmEventLoop event_loop(&config.message());
+
+  frc971::constants::ConstantsFetcher<y2024_swerve::Constants> constants(
+      &event_loop);
+
+  SwerveControlLoops swerve_control_loops(
+      &event_loop, constants.constants().common()->rotation(),
+      constants.constants().robot()->swerve_zeroing(), "/drivetrain");
+
+  event_loop.Run();
+
+  return 0;
+}