Add 2022 aimer to superstructure

Change-Id: I1ae24de0ed43d73578dabc63a9c9efc79c904552
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/y2022/control_loops/superstructure/BUILD b/y2022/control_loops/superstructure/BUILD
index 14fd403..e7b6559 100644
--- a/y2022/control_loops/superstructure/BUILD
+++ b/y2022/control_loops/superstructure/BUILD
@@ -87,6 +87,7 @@
         "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
         "//y2022:constants",
         "//y2022/control_loops/superstructure/catapult",
+        "//y2022/control_loops/superstructure/turret:aiming",
     ],
 )
 
diff --git a/y2022/control_loops/superstructure/superstructure.cc b/y2022/control_loops/superstructure/superstructure.cc
index affc25b..cc61851 100644
--- a/y2022/control_loops/superstructure/superstructure.cc
+++ b/y2022/control_loops/superstructure/superstructure.cc
@@ -55,6 +55,13 @@
   drivetrain_status_fetcher_.Fetch();
   const float velocity = robot_velocity();
 
+  const turret::Aimer::Goal *auto_aim_goal = nullptr;
+  if (drivetrain_status_fetcher_.get() != nullptr) {
+    aimer_.Update(drivetrain_status_fetcher_.get(),
+                  turret::Aimer::ShotMode::kShootOnTheFly);
+    auto_aim_goal = aimer_.TurretGoal();
+  }
+
   const frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal
       *turret_goal = nullptr;
   double roller_speed_compensated_front = 0.0;
@@ -63,8 +70,6 @@
   double flipper_arms_voltage = 0.0;
 
   if (unsafe_goal != nullptr) {
-    turret_goal = unsafe_goal->turret();
-
     roller_speed_compensated_front =
         unsafe_goal->roller_speed_front() +
         std::max(velocity * unsafe_goal->roller_speed_compensation(), 0.0);
@@ -74,9 +79,10 @@
         std::min(velocity * unsafe_goal->roller_speed_compensation(), 0.0);
 
     transfer_roller_speed = unsafe_goal->transfer_roller_speed();
-  }
 
-  // TODO: Aimer sets turret_goal here
+    turret_goal =
+        unsafe_goal->auto_aim() ? auto_aim_goal : unsafe_goal->turret();
+  }
 
   // Supersturcture state machine:
   // 1. IDLE: Wait until an intake beambreak is triggerred, meaning that a ball
@@ -337,6 +343,9 @@
   intake_back_.set_min_position(collision_avoidance_.min_intake_back_goal());
   intake_back_.set_max_position(collision_avoidance_.max_intake_back_goal());
 
+  const flatbuffers::Offset<AimerStatus> aimer_offset =
+      aimer_.PopulateStatus(status->fbb());
+
   // Disable the catapult if we want to restart to prevent damage with flippers
   const flatbuffers::Offset<PotAndAbsoluteEncoderProfiledJointStatus>
       catapult_status_offset =
@@ -410,6 +419,8 @@
   status_builder.add_state(state_);
   status_builder.add_intake_state(intake_state_);
 
+  status_builder.add_aimer(aimer_offset);
+
   (void)status->Send(status_builder.Finish());
 }
 
diff --git a/y2022/control_loops/superstructure/superstructure.h b/y2022/control_loops/superstructure/superstructure.h
index 13a790a..1ed63be 100644
--- a/y2022/control_loops/superstructure/superstructure.h
+++ b/y2022/control_loops/superstructure/superstructure.h
@@ -12,6 +12,7 @@
 #include "y2022/control_loops/superstructure/superstructure_output_generated.h"
 #include "y2022/control_loops/superstructure/superstructure_position_generated.h"
 #include "y2022/control_loops/superstructure/superstructure_status_generated.h"
+#include "y2022/control_loops/superstructure/turret/aiming.h"
 
 namespace y2022 {
 namespace control_loops {
@@ -74,6 +75,8 @@
 
   int prev_shot_count_ = 0;
 
+  turret::Aimer aimer_;
+
   bool flippers_open_ = false;
   bool reseating_in_catapult_ = false;
   bool fire_ = false;
diff --git a/y2022/control_loops/superstructure/superstructure_goal.fbs b/y2022/control_loops/superstructure/superstructure_goal.fbs
index 379f6ba..c86f338 100644
--- a/y2022/control_loops/superstructure/superstructure_goal.fbs
+++ b/y2022/control_loops/superstructure/superstructure_goal.fbs
@@ -45,6 +45,9 @@
   // Aborts the shooting process if the ball has been loaded into the catapult
   // and the superstructure is in the LOADED state.
   cancel_shot:bool (id: 10);
+
+  // If true, auto-track the turret to point at the goal.
+  auto_aim:bool (id: 11);
 }
 
 
diff --git a/y2022/control_loops/superstructure/superstructure_lib_test.cc b/y2022/control_loops/superstructure/superstructure_lib_test.cc
index 7425ac0..5dd5ddf 100644
--- a/y2022/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2022/control_loops/superstructure/superstructure_lib_test.cc
@@ -430,10 +430,20 @@
   }
 
   void SendRobotVelocity(double robot_velocity) {
+    SendDrivetrainStatus(robot_velocity, {0.0, 0.0}, 0.0);
+  }
+
+  void SendDrivetrainStatus(double robot_velocity, Eigen::Vector2d pos,
+                            double theta) {
     // Send a robot velocity to test compensation
     auto builder = drivetrain_status_sender_.MakeBuilder();
     auto drivetrain_status_builder = builder.MakeBuilder<DrivetrainStatus>();
     drivetrain_status_builder.add_robot_speed(robot_velocity);
+    drivetrain_status_builder.add_estimated_left_velocity(robot_velocity);
+    drivetrain_status_builder.add_estimated_right_velocity(robot_velocity);
+    drivetrain_status_builder.add_x(pos.x());
+    drivetrain_status_builder.add_y(pos.y());
+    drivetrain_status_builder.add_theta(theta);
     builder.CheckOk(builder.Send(drivetrain_status_builder.Finish()));
   }
 
@@ -1151,6 +1161,40 @@
   EXPECT_EQ(superstructure_status_fetcher_->state(), SuperstructureState::IDLE);
 }
 
+// Tests that the turret switches to auto-aiming when we set auto_aim to
+// true.
+TEST_F(SuperstructureTest, TurretAutoAim) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+
+  // Set ourselves up 5m from the target--the turret goal should be 90 deg (we
+  // need to shoot out the right of the robot, and we shoot out of the back of
+  // the turret).
+  SendDrivetrainStatus(0.0, {0.0, 5.0}, 0.0);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+
+    goal_builder.add_auto_aim(true);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()),
+              aos::RawSender::Error::kOk);
+  }
+
+  // Give it time to stabilize.
+  RunFor(chrono::seconds(2));
+
+  superstructure_status_fetcher_.Fetch();
+  EXPECT_NEAR(M_PI_2, superstructure_status_fetcher_->turret()->position(),
+              5e-4);
+  EXPECT_FLOAT_EQ(M_PI_2,
+                  superstructure_status_fetcher_->aimer()->turret_position());
+  EXPECT_FLOAT_EQ(0,
+                  superstructure_status_fetcher_->aimer()->turret_velocity());
+}
+
 }  // namespace testing
 }  // namespace superstructure
 }  // namespace control_loops
diff --git a/y2022/control_loops/superstructure/superstructure_status.fbs b/y2022/control_loops/superstructure/superstructure_status.fbs
index c74cbab..1005c46 100644
--- a/y2022/control_loops/superstructure/superstructure_status.fbs
+++ b/y2022/control_loops/superstructure/superstructure_status.fbs
@@ -25,6 +25,18 @@
   SHOOTING,
 }
 
+table AimerStatus {
+  // The current goal angle for the turret auto-tracking, in radians.
+  turret_position:double (id: 0);
+  // The current goal velocity for the turret, in radians / sec.
+  turret_velocity:double (id: 1);
+  // The current distance to the target, in meters.
+  target_distance:double (id: 2);
+  // The current "shot distance." When shooting on the fly, this may be
+  // different from the static distance to the target.
+  shot_distance:double (id: 3);
+}
+
 table Status {
   // All subsystems know their location.
   zeroed:bool (id: 0);
@@ -61,6 +73,8 @@
 
   // The number of shots we have taken.
   shot_count:int32 (id: 9);
+
+  aimer:AimerStatus (id: 15);
 }
 
 root_type Status;
diff --git a/y2022/control_loops/superstructure/turret/BUILD b/y2022/control_loops/superstructure/turret/BUILD
index c2948d7..2f7ad14 100644
--- a/y2022/control_loops/superstructure/turret/BUILD
+++ b/y2022/control_loops/superstructure/turret/BUILD
@@ -32,3 +32,21 @@
         "//frc971/control_loops:state_feedback_loop",
     ],
 )
+
+cc_library(
+    name = "aiming",
+    srcs = ["aiming.cc"],
+    hdrs = ["aiming.h"],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        "//aos:flatbuffers",
+        "//frc971/control_loops:control_loops_fbs",
+        "//frc971/control_loops:pose",
+        "//frc971/control_loops:profiled_subsystem_fbs",
+        "//frc971/control_loops/aiming",
+        "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
+        "//y2022:constants",
+        "//y2022/control_loops/drivetrain:drivetrain_base",
+        "//y2022/control_loops/superstructure:superstructure_status_fbs",
+    ],
+)
diff --git a/y2022/control_loops/superstructure/turret/aiming.cc b/y2022/control_loops/superstructure/turret/aiming.cc
new file mode 100644
index 0000000..4c0309c
--- /dev/null
+++ b/y2022/control_loops/superstructure/turret/aiming.cc
@@ -0,0 +1,78 @@
+#include "y2022/control_loops/superstructure/turret/aiming.h"
+
+#include "y2022/constants.h"
+#include "y2022/control_loops/drivetrain/drivetrain_base.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace superstructure {
+namespace turret {
+
+using frc971::control_loops::Pose;
+using frc971::control_loops::aiming::ShotConfig;
+using frc971::control_loops::aiming::RobotState;
+
+namespace {
+// Average speed-over-ground of the ball on its way to the target. Our current
+// model assumes constant ball velocity regardless of shot distance.
+constexpr double kBallSpeedOverGround = 12.0;  // m/s
+
+// If the turret is at zero, then it will be at this angle at which the shot
+// will leave the robot. I.e., if the turret is at zero, then the shot will go
+// straight out the back of the robot.
+constexpr double kTurretZeroOffset = M_PI;
+
+flatbuffers::DetachedBuffer MakePrefilledGoal() {
+  flatbuffers::FlatBufferBuilder fbb;
+  fbb.ForceDefaults(true);
+  Aimer::Goal::Builder builder(fbb);
+  builder.add_unsafe_goal(0);
+  builder.add_goal_velocity(0);
+  builder.add_ignore_profile(true);
+  fbb.Finish(builder.Finish());
+  return fbb.Release();
+}
+}  // namespace
+
+Aimer::Aimer() : goal_(MakePrefilledGoal()) {}
+
+void Aimer::Update(const Status *status, ShotMode shot_mode) {
+  const Pose robot_pose({status->x(), status->y(), 0}, status->theta());
+  const Pose goal({0.0, 0.0, 0.0}, 0.0);
+
+  const Eigen::Vector2d linear_angular =
+      drivetrain::GetDrivetrainConfig().Tlr_to_la() *
+      Eigen::Vector2d(status->estimated_left_velocity(),
+                      status->estimated_right_velocity());
+  const double xdot = linear_angular(0) * std::cos(status->theta());
+  const double ydot = linear_angular(0) * std::sin(status->theta());
+
+  current_goal_ =
+      frc971::control_loops::aiming::AimerGoal(
+          ShotConfig{goal, shot_mode, constants::Values::kTurretRange(),
+                     kBallSpeedOverGround,
+                     /*wrap_mode=*/0.0, kTurretZeroOffset},
+          RobotState{robot_pose,
+                     {xdot, ydot},
+                     linear_angular(1),
+                     goal_.message().unsafe_goal()});
+
+  goal_.mutable_message()->mutate_unsafe_goal(current_goal_.position);
+  goal_.mutable_message()->mutate_goal_velocity(
+      std::clamp(current_goal_.velocity, -2.0, 2.0));
+}
+
+flatbuffers::Offset<AimerStatus> Aimer::PopulateStatus(
+    flatbuffers::FlatBufferBuilder *fbb) const {
+  AimerStatus::Builder builder(*fbb);
+  builder.add_turret_position(current_goal_.position);
+  builder.add_turret_velocity(current_goal_.velocity);
+  builder.add_target_distance(current_goal_.target_distance);
+  builder.add_shot_distance(DistanceToGoal());
+  return builder.Finish();
+}
+
+}  // namespace turret
+}  // namespace superstructure
+}  // namespace control_loops
+}  // namespace y2022
diff --git a/y2022/control_loops/superstructure/turret/aiming.h b/y2022/control_loops/superstructure/turret/aiming.h
new file mode 100644
index 0000000..9494103
--- /dev/null
+++ b/y2022/control_loops/superstructure/turret/aiming.h
@@ -0,0 +1,40 @@
+#ifndef Y2022_CONTROL_LOOPS_SUPERSTRUCTURE_TURRET_AIMING_H_
+#define Y2022_CONTROL_LOOPS_SUPERSTRUCTURE_TURRET_AIMING_H_
+
+#include "aos/flatbuffers.h"
+#include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
+#include "frc971/control_loops/pose.h"
+#include "frc971/control_loops/profiled_subsystem_generated.h"
+#include "frc971/control_loops/aiming/aiming.h"
+#include "y2022/control_loops/superstructure/superstructure_status_generated.h"
+
+namespace y2022::control_loops::superstructure::turret {
+
+// This class manages taking in drivetrain status messages and generating turret
+// goals so that it gets aimed at the goal.
+class Aimer {
+ public:
+  typedef frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal
+      Goal;
+  typedef frc971::control_loops::drivetrain::Status Status;
+  typedef frc971::control_loops::aiming::ShotMode ShotMode;
+
+  Aimer();
+
+  void Update(const Status *status, ShotMode shot_mode);
+
+  const Goal *TurretGoal() const { return &goal_.message(); }
+
+  // Returns the distance to the goal, in meters.
+  double DistanceToGoal() const { return current_goal_.virtual_shot_distance; }
+
+  flatbuffers::Offset<AimerStatus> PopulateStatus(
+      flatbuffers::FlatBufferBuilder *fbb) const;
+
+ private:
+  aos::FlatbufferDetachedBuffer<Goal> goal_;
+  frc971::control_loops::aiming::TurretGoal current_goal_;
+};
+
+}  // namespace y2022::control_loops::superstructure::turret
+#endif  // Y2020_CONTROL_LOOPS_SUPERSTRUCTURE_TURRET_AIMING_H_