Add arm code on superstructure

Signed-off-by: Maxwell Henderson <mxwhenderson@gmail.com>
Change-Id: I20e40f8ccb8a17b7b3f18390fd2f428f5a1fff1e
Signed-off-by: Alexander Yee <xander.yee@gmail.com>
diff --git a/y2023/control_loops/superstructure/BUILD b/y2023/control_loops/superstructure/BUILD
index 718a575..f5eff79 100644
--- a/y2023/control_loops/superstructure/BUILD
+++ b/y2023/control_loops/superstructure/BUILD
@@ -68,7 +68,6 @@
         "superstructure.h",
     ],
     deps = [
-        # ":collision_avoidance_lib",
         ":superstructure_goal_fbs",
         ":superstructure_output_fbs",
         ":superstructure_position_fbs",
@@ -78,6 +77,7 @@
         "//frc971/control_loops:control_loop",
         "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
         "//y2023:constants",
+        "//y2023/control_loops/superstructure/arm",
     ],
 )
 
diff --git a/y2023/control_loops/superstructure/arm/arm.h b/y2023/control_loops/superstructure/arm/arm.h
index a595ace..27588f1 100644
--- a/y2023/control_loops/superstructure/arm/arm.h
+++ b/y2023/control_loops/superstructure/arm/arm.h
@@ -26,18 +26,16 @@
 
   // if true, tune down all the constants for testing.
   static constexpr bool kGrannyMode() { return false; }
+
   // the operating voltage.
   static constexpr double kOperatingVoltage() {
     return kGrannyMode() ? 5.0 : 12.0;
   }
-
   static constexpr double kDt() { return 0.00505; }
-
   static constexpr double kAlpha0Max() { return kGrannyMode() ? 5.0 : 15.0; }
   static constexpr double kAlpha1Max() { return kGrannyMode() ? 5.0 : 15.0; }
 
   static constexpr double kVMax() { return kGrannyMode() ? 5.0 : 11.5; }
-
   static constexpr double kPathlessVMax() { return 5.0; }
   static constexpr double kGotoPathVMax() { return 6.0; }
 
@@ -50,16 +48,19 @@
   void Reset();
 
   ArmState state() const { return state_; }
-  bool estopped() const { return state_ == ArmState::ESTOP; }
 
+  bool estopped() const { return state_ == ArmState::ESTOP; }
   bool zeroed() const {
     return (proximal_zeroing_estimator_.zeroed() &&
             distal_zeroing_estimator_.zeroed());
   }
+
   // Returns the maximum position for the intake.  This is used to override the
   // intake position to release the box when the state machine is lifting.
   double max_intake_override() const { return max_intake_override_; }
+
   uint32_t current_node() const { return current_node_; }
+
   double path_distance_to_go() { return follower_.path_distance_to_go(); }
 
  private:
@@ -71,6 +72,7 @@
   }
 
   std::shared_ptr<const constants::Values> values_;
+
   ArmState state_;
 
   ::frc971::zeroing::PotAndAbsoluteEncoderZeroingEstimator
@@ -80,15 +82,16 @@
 
   double proximal_offset_;
   double distal_offset_;
+
   double max_intake_override_;
 
   const ::Eigen::Matrix<double, 2, 2> alpha_unitizer_;
+
   double vmax_ = kVMax();
 
   frc971::control_loops::arm::Dynamics dynamics_;
 
   ::std::vector<TrajectoryAndParams> trajectories_;
-
   SearchGraph search_graph_;
 
   bool close_enough_for_full_power_;
@@ -96,7 +99,6 @@
   size_t brownout_count_;
 
   EKF arm_ekf_;
-
   TrajectoryFollower follower_;
 
   const ::std::vector<::Eigen::Matrix<double, 2, 1>> points_;
@@ -106,6 +108,7 @@
 
   EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
 };
+
 }  // namespace arm
 }  // namespace superstructure
 }  // namespace control_loops
diff --git a/y2023/control_loops/superstructure/superstructure.cc b/y2023/control_loops/superstructure/superstructure.cc
index 1125e21..fa83d20 100644
--- a/y2023/control_loops/superstructure/superstructure.cc
+++ b/y2023/control_loops/superstructure/superstructure.cc
@@ -12,6 +12,8 @@
 namespace control_loops {
 namespace superstructure {
 
+using ::aos::monotonic_clock;
+
 using frc971::control_loops::AbsoluteEncoderProfiledJointStatus;
 using frc971::control_loops::PotAndAbsoluteEncoderProfiledJointStatus;
 using frc971::control_loops::RelativeEncoderProfiledJointStatus;
@@ -26,19 +28,19 @@
           event_loop->MakeFetcher<frc971::control_loops::drivetrain::Status>(
               "/drivetrain")),
       joystick_state_fetcher_(
-          event_loop->MakeFetcher<aos::JoystickState>("/aos")) {
-  (void)values;
-}
+          event_loop->MakeFetcher<aos::JoystickState>("/aos")),
+      arm_(values_) {}
 
 void Superstructure::RunIteration(const Goal *unsafe_goal,
                                   const Position *position,
                                   aos::Sender<Output>::Builder *output,
                                   aos::Sender<Status>::Builder *status) {
-  (void)unsafe_goal;
-  (void)position;
+  const monotonic_clock::time_point timestamp =
+      event_loop()->context().monotonic_event_time;
 
   if (WasReset()) {
     AOS_LOG(ERROR, "WPILib reset, restarting\n");
+    arm_.Reset();
   }
 
   OutputT output_struct;
@@ -47,11 +49,31 @@
     alliance_ = joystick_state_fetcher_->alliance();
   }
   drivetrain_status_fetcher_.Fetch();
-  output->CheckOk(output->Send(Output::Pack(*output->fbb(), &output_struct)));
-  
+
+  const uint32_t arm_goal_position =
+      unsafe_goal != nullptr ? unsafe_goal->arm_goal_position() : 0u;
+
+  flatbuffers::Offset<superstructure::ArmStatus> arm_status_offset =
+      arm_.Iterate(
+          timestamp, unsafe_goal != nullptr ? &(arm_goal_position) : nullptr,
+          position->arm(),
+          unsafe_goal != nullptr ? unsafe_goal->trajectory_override() : false,
+          output != nullptr ? &output_struct.proximal_voltage : nullptr,
+          output != nullptr ? &output_struct.distal_voltage : nullptr,
+          unsafe_goal != nullptr ? unsafe_goal->intake() : false,
+          unsafe_goal != nullptr ? unsafe_goal->spit() : false,
+
+          status->fbb());
+
+  if (output) {
+    output->CheckOk(output->Send(Output::Pack(*output->fbb(), &output_struct)));
+  }
+
   Status::Builder status_builder = status->MakeBuilder<Status>();
   status_builder.add_zeroed(true);
   status_builder.add_estopped(false);
+  status_builder.add_arm(arm_status_offset);
+
   (void)status->Send(status_builder.Finish());
 }
 
diff --git a/y2023/control_loops/superstructure/superstructure.h b/y2023/control_loops/superstructure/superstructure.h
index 0740de3..2d0fc43 100644
--- a/y2023/control_loops/superstructure/superstructure.h
+++ b/y2023/control_loops/superstructure/superstructure.h
@@ -5,6 +5,7 @@
 #include "frc971/control_loops/control_loop.h"
 #include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
 #include "y2023/constants.h"
+#include "y2023/control_loops/superstructure/arm/arm.h"
 #include "y2023/control_loops/superstructure/superstructure_goal_generated.h"
 #include "y2023/control_loops/superstructure/superstructure_output_generated.h"
 #include "y2023/control_loops/superstructure/superstructure_position_generated.h"
@@ -45,6 +46,8 @@
       drivetrain_status_fetcher_;
   aos::Fetcher<aos::JoystickState> joystick_state_fetcher_;
 
+  arm::Arm arm_;
+
   aos::Alliance alliance_ = aos::Alliance::kInvalid;
 
   DISALLOW_COPY_AND_ASSIGN(Superstructure);
diff --git a/y2023/control_loops/superstructure/superstructure_goal.fbs b/y2023/control_loops/superstructure/superstructure_goal.fbs
index 47ce7b2..936cbee 100644
--- a/y2023/control_loops/superstructure/superstructure_goal.fbs
+++ b/y2023/control_loops/superstructure/superstructure_goal.fbs
@@ -8,13 +8,16 @@
     // Controls distal, proximal, and roll joints
     arm_goal_position:uint32 (id: 0);
 
-    wrist:frc971.control_loops.StaticZeroingSingleDOFProfiledSubsystemGoal (id: 1);
+    // Overrides the current path to go to the next path
+    trajectory_override:bool (id: 1);
+
+    wrist:frc971.control_loops.StaticZeroingSingleDOFProfiledSubsystemGoal (id: 2);
 
     // If this is true, the rollers should intake.
-    intake:bool (id: 2);
+    intake:bool (id: 3);
 
     // If this is true, the rollers should spit.
-    spit:bool (id: 3);
+    spit:bool (id: 4);
 }
 
 
diff --git a/y2023/control_loops/superstructure/superstructure_lib_test.cc b/y2023/control_loops/superstructure/superstructure_lib_test.cc
index 0b9fb7d..ae01f0b 100644
--- a/y2023/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2023/control_loops/superstructure/superstructure_lib_test.cc
@@ -18,6 +18,9 @@
 namespace control_loops {
 namespace superstructure {
 namespace testing {
+namespace {
+constexpr double kNoiseScalar = 0.01;
+}  // namespace
 namespace chrono = std::chrono;
 
 using ::aos::monotonic_clock;
@@ -40,6 +43,89 @@
     frc971::control_loops::RelativeEncoderProfiledJointStatus,
     RelativeEncoderSubsystem::State, constants::Values::PotConstants>;
 
+class ArmSimulation {
+ public:
+  explicit ArmSimulation(
+      const ::frc971::constants::PotAndAbsoluteEncoderZeroingConstants
+          &proximal_zeroing_constants,
+      const ::frc971::constants::PotAndAbsoluteEncoderZeroingConstants
+          &distal_zeroing_constants,
+      std::chrono::nanoseconds dt)
+      : proximal_zeroing_constants_(proximal_zeroing_constants),
+        proximal_pot_encoder_(M_PI * 2.0 *
+                              constants::Values::kProximalEncoderRatio()),
+        distal_zeroing_constants_(distal_zeroing_constants),
+        distal_pot_encoder_(M_PI * 2.0 *
+                            constants::Values::kDistalEncoderRatio()),
+        dynamics_(arm::kArmConstants),
+        dt_(dt) {
+    X_.setZero();
+  }
+
+  void InitializePosition(::Eigen::Matrix<double, 2, 1> position) {
+    X_ << position(0), 0.0, position(1), 0.0;
+
+    proximal_pot_encoder_.Initialize(
+        X_(0), kNoiseScalar, 0.0,
+        proximal_zeroing_constants_.measured_absolute_position);
+    distal_pot_encoder_.Initialize(
+        X_(2), kNoiseScalar, 0.0,
+        distal_zeroing_constants_.measured_absolute_position);
+  }
+
+  flatbuffers::Offset<ArmPosition> GetSensorValues(
+      flatbuffers::FlatBufferBuilder *fbb) {
+    frc971::PotAndAbsolutePosition::Builder proximal_builder(*fbb);
+    flatbuffers::Offset<frc971::PotAndAbsolutePosition> proximal_offset =
+        proximal_pot_encoder_.GetSensorValues(&proximal_builder);
+
+    frc971::PotAndAbsolutePosition::Builder distal_builder(*fbb);
+    flatbuffers::Offset<frc971::PotAndAbsolutePosition> distal_offset =
+        distal_pot_encoder_.GetSensorValues(&distal_builder);
+
+    ArmPosition::Builder arm_position_builder(*fbb);
+    arm_position_builder.add_proximal(proximal_offset);
+    arm_position_builder.add_distal(distal_offset);
+
+    return arm_position_builder.Finish();
+  }
+
+  double proximal_position() const { return X_(0, 0); }
+  double proximal_velocity() const { return X_(1, 0); }
+  double distal_position() const { return X_(2, 0); }
+  double distal_velocity() const { return X_(3, 0); }
+
+  void Simulate(::Eigen::Matrix<double, 2, 1> U) {
+    constexpr double voltage_check =
+        superstructure::arm::Arm::kOperatingVoltage();
+
+    AOS_CHECK_LE(::std::abs(U(0)), voltage_check);
+    AOS_CHECK_LE(::std::abs(U(1)), voltage_check);
+
+    X_ = dynamics_.UnboundedDiscreteDynamics(
+        X_, U,
+        std::chrono::duration_cast<std::chrono::duration<double>>(dt_).count());
+
+    // TODO(austin): Estop on grose out of bounds.
+    proximal_pot_encoder_.MoveTo(X_(0));
+    distal_pot_encoder_.MoveTo(X_(2));
+  }
+
+ private:
+  ::Eigen::Matrix<double, 4, 1> X_;
+
+  const ::frc971::constants::PotAndAbsoluteEncoderZeroingConstants
+      proximal_zeroing_constants_;
+  PositionSensorSimulator proximal_pot_encoder_;
+
+  const ::frc971::constants::PotAndAbsoluteEncoderZeroingConstants
+      distal_zeroing_constants_;
+  PositionSensorSimulator distal_pot_encoder_;
+
+  ::frc971::control_loops::arm::Dynamics dynamics_;
+
+  std::chrono::nanoseconds dt_;
+};
 // Class which simulates the superstructure and sends out queue messages with
 // the position.
 class SuperstructureSimulation {
@@ -49,21 +135,26 @@
                            chrono::nanoseconds dt)
       : event_loop_(event_loop),
         dt_(dt),
+        arm_(values->arm_proximal.zeroing, values->arm_distal.zeroing, dt_),
         superstructure_position_sender_(
             event_loop_->MakeSender<Position>("/superstructure")),
         superstructure_status_fetcher_(
             event_loop_->MakeFetcher<Status>("/superstructure")),
         superstructure_output_fetcher_(
             event_loop_->MakeFetcher<Output>("/superstructure")) {
-    (void)values;
-    (void)dt_;
-
+    InitializeArmPosition(arm::NeutralPosPoint());
     phased_loop_handle_ = event_loop_->AddPhasedLoop(
         [this](int) {
           // Skip this the first time.
           if (!first_) {
             EXPECT_TRUE(superstructure_output_fetcher_.Fetch());
             EXPECT_TRUE(superstructure_status_fetcher_.Fetch());
+
+            arm_.Simulate(
+                (::Eigen::Matrix<double, 2, 1>()
+                     << superstructure_output_fetcher_->proximal_voltage(),
+                 superstructure_output_fetcher_->distal_voltage())
+                    .finished());
           }
           first_ = false;
           SendPositionMessage();
@@ -71,13 +162,22 @@
         dt);
   }
 
+  void InitializeArmPosition(::Eigen::Matrix<double, 2, 1> position) {
+    arm_.InitializePosition(position);
+  }
+
+  const ArmSimulation &arm() const { return arm_; }
+
   // Sends a queue message with the position of the superstructure.
   void SendPositionMessage() {
     ::aos::Sender<Position>::Builder builder =
         superstructure_position_sender_.MakeBuilder();
 
-    Position::Builder position_builder = builder.MakeBuilder<Position>();
+    flatbuffers::Offset<ArmPosition> arm_offset =
+        arm_.GetSensorValues(builder.fbb());
 
+    Position::Builder position_builder = builder.MakeBuilder<Position>();
+    position_builder.add_arm(arm_offset);
     CHECK_EQ(builder.Send(position_builder.Finish()),
              aos::RawSender::Error::kOk);
   }
@@ -87,6 +187,8 @@
   const chrono::nanoseconds dt_;
   ::aos::PhasedLoopHandler *phased_loop_handle_ = nullptr;
 
+  ArmSimulation arm_;
+
   ::aos::Sender<Position> superstructure_position_sender_;
   ::aos::Fetcher<Status> superstructure_status_fetcher_;
   ::aos::Fetcher<Output> superstructure_output_fetcher_;
@@ -139,6 +241,10 @@
   void VerifyNearGoal() {
     superstructure_goal_fetcher_.Fetch();
     superstructure_status_fetcher_.Fetch();
+
+    ASSERT_TRUE(superstructure_goal_fetcher_.get() != nullptr) << ": No goal";
+    ASSERT_TRUE(superstructure_status_fetcher_.get() != nullptr)
+        << ": No status";
   }
 
   void CheckIfZeroed() {
@@ -233,6 +339,8 @@
     auto builder = superstructure_goal_sender_.MakeBuilder();
     Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
 
+    goal_builder.add_arm_goal_position(arm::NeutralPosIndex());
+
     ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
   }
 
@@ -289,6 +397,73 @@
   CheckIfZeroed();
 }
 
+// Tests that we don't freak out without a goal.
+TEST_F(SuperstructureTest, ArmSimpleGoal) {
+  SetEnabled(true);
+  RunFor(chrono::seconds(20));
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_goal_position(arm::NeutralToConePos1Index());
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+  EXPECT_EQ(ArmState::RUNNING, superstructure_status_fetcher_->arm()->state());
+}
+
+// Tests that we can can execute a move.
+TEST_F(SuperstructureTest, ArmMoveAndMoveBack) {
+  SetEnabled(true);
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_goal_position(arm::NeutralToConePos1Index());
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(chrono::seconds(10));
+
+  VerifyNearGoal();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_goal_position(arm::NeutralToConePos2Index());
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(chrono::seconds(10));
+  VerifyNearGoal();
+}
+
+// Tests that we can can execute a move which moves through multiple nodes.
+TEST_F(SuperstructureTest, ArmMultistepMove) {
+  SetEnabled(true);
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_goal_position(arm::NeutralPosIndex());
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(chrono::seconds(10));
+
+  VerifyNearGoal();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_goal_position(arm::ConePosIndex());
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+
+  RunFor(chrono::seconds(10));
+  VerifyNearGoal();
+}
+
 }  // namespace testing
 }  // namespace superstructure
 }  // namespace control_loops