incorporate arm into superstructure.cc by adding arm constants

Change-Id: Icfd6ba91044b7a8f29f10edf8c4c44ae1a443d2c
Signed-off-by: Yash Maheshwari <yashmahe2018@gmail.com>
Signed-off-by: James (Peilun) Li <jamespeilunli@gmail.com>
diff --git a/y2024_bot3/constants/9971.json b/y2024_bot3/constants/9971.json
index 5583b44..4eb1acf 100644
--- a/y2024_bot3/constants/9971.json
+++ b/y2024_bot3/constants/9971.json
@@ -1,3 +1,5 @@
+{% from 'y2024_bot3/constants/common.jinja2' import arm_zero %}
+
 {
   "cameras": [
     {
@@ -13,6 +15,21 @@
       "calibration": {% include 'y2024_bot3/constants/calib_files/calibration_orin1-9971-1_cam-24-11_2024-03-24_11-52-49.json' %}
     }
   ],
-  "robot": {},
+  "robot": {
+    "arm_constants": {
+      {% set _ = arm_zero.update(
+          {
+              "measured_absolute_position" : 0.0992895926495078
+          }
+      ) %}
+      "zeroing_constants": {{ arm_zero | tojson(indent=2)}},
+      "potentiometer_offset": {{ 0 }},
+      "arm_positions": {
+        "intake": 0,
+        "idle": 1,
+        "amp": 2
+      }
+    }
+  },
   {% include 'y2024_bot3/constants/common.json' %}
 }
diff --git a/y2024_bot3/constants/BUILD b/y2024_bot3/constants/BUILD
index 1ce69a0..323bd49 100644
--- a/y2024_bot3/constants/BUILD
+++ b/y2024_bot3/constants/BUILD
@@ -26,6 +26,7 @@
         "common.jinja2",
         "common.json",
         "//y2024_bot3/constants/calib_files",
+        "//y2024_bot3/control_loops/superstructure/arm:arm_json",
         "//y2024_bot3/vision/maps",
     ],
     parameters = {},
@@ -40,6 +41,7 @@
         "common.jinja2",
         "common.json",
         "//y2024_bot3/constants/calib_files",
+        "//y2024_bot3/control_loops/superstructure/arm:arm_json",
         "//y2024_bot3/vision/maps",
     ],
     parameters = {},
diff --git a/y2024_bot3/constants/common.jinja2 b/y2024_bot3/constants/common.jinja2
index 9b7af82..809f8c8 100644
--- a/y2024_bot3/constants/common.jinja2
+++ b/y2024_bot3/constants/common.jinja2
@@ -1,3 +1,16 @@
 {% set pi = 3.14159265 %}
 
 {%set zeroing_sample_size = 200 %}
+
+{# Arm #}
+{% set arm_encoder_ratio = (4.0 / 1.0) %}
+
+{%
+set arm_zero = {
+    "average_filter_size": zeroing_sample_size,
+    "one_revolution_distance": pi * 2.0 * arm_encoder_ratio,
+    "zeroing_threshold": 0.0005,
+    "moving_buffer_size": 20,
+    "allowable_encoder_error": 0.9
+}
+%}
diff --git a/y2024_bot3/constants/common.json b/y2024_bot3/constants/common.json
index e7b70b4..d5abec0 100644
--- a/y2024_bot3/constants/common.json
+++ b/y2024_bot3/constants/common.json
@@ -1,2 +1,21 @@
 "common": {
+  "arm": {
+    "zeroing_voltage": 3.0,
+    "operating_voltage": 12.0,
+    "zeroing_profile_params": {
+      "max_velocity": 0.5,
+      "max_acceleration": 3.0
+    },
+    "default_profile_params":{
+      "max_velocity": 2.0,
+      "max_acceleration": 10.0
+    },
+    "range": {
+        "lower_hard": -50,
+        "upper_hard": 50,
+        "lower": -50,
+        "upper": 50
+    },
+    "loop": {% include 'y2024_bot3/control_loops/superstructure/arm/integral_arm_plant.json' %}
+  }
 }
diff --git a/y2024_bot3/constants/constants.fbs b/y2024_bot3/constants/constants.fbs
index 0c7a792..99999bd 100644
--- a/y2024_bot3/constants/constants.fbs
+++ b/y2024_bot3/constants/constants.fbs
@@ -10,12 +10,26 @@
   calibration:frc971.vision.calibration.CameraCalibration (id: 0);
 }
 
+table ArmPositions {
+  intake:double (id: 0);
+  idle:double (id: 1);
+  amp:double (id: 2);
+}
+
+table PotAndAbsEncoderConstants {
+  zeroing_constants:frc971.zeroing.PotAndAbsoluteEncoderZeroingConstants (id: 0);
+  potentiometer_offset:double (id: 1);
+  arm_positions:ArmPositions (id: 2);
+}
+
 table RobotConstants {
+  arm_constants:PotAndAbsEncoderConstants (id: 0);
 }
 
 // Common table for constants unrelated to the robot
 table Common {
   target_map:frc971.vision.TargetMap (id: 0);
+  arm:frc971.control_loops.StaticZeroingSingleDOFProfiledSubsystemCommonParams (id: 1);
 }
 
 table Constants {
diff --git a/y2024_bot3/constants/test_data/test_team.json b/y2024_bot3/constants/test_data/test_team.json
index deacef4..06286e7 100644
--- a/y2024_bot3/constants/test_data/test_team.json
+++ b/y2024_bot3/constants/test_data/test_team.json
@@ -1,4 +1,21 @@
+{% from 'y2024_bot3/constants/common.jinja2' import arm_zero %}
+
 {
-  "robot": {},
+  "robot": {
+    "arm_constants": {
+      {% set _ = arm_zero.update(
+          {
+              "measured_absolute_position" : 0.0992895926495078
+          }
+      ) %}
+      "zeroing_constants": {{ arm_zero | tojson(indent=2)}},
+      "potentiometer_offset": {{ 0 }},
+      "arm_positions": {
+        "intake": 0,
+        "idle": 1,
+        "amp": 2
+      }
+    }
+  },
   {% include 'y2024_bot3/constants/common.json' %}
 }
diff --git a/y2024_bot3/control_loops/python/BUILD b/y2024_bot3/control_loops/python/BUILD
index 9ec6b57..853719d 100644
--- a/y2024_bot3/control_loops/python/BUILD
+++ b/y2024_bot3/control_loops/python/BUILD
@@ -6,3 +6,19 @@
     visibility = ["//visibility:public"],
     deps = ["//y2024_bot3/control_loops:python_init"],
 )
+
+py_binary(
+    name = "arm",
+    srcs = [
+        "arm.py",
+    ],
+    legacy_create_init = False,
+    target_compatible_with = ["@platforms//cpu:x86_64"],
+    deps = [
+        ":python_init",
+        "//frc971/control_loops/python:controls",
+        "//frc971/control_loops/python:linear_system",
+        "@pip//glog",
+        "@pip//python_gflags",
+    ],
+)
diff --git a/y2024_bot3/control_loops/python/arm.py b/y2024_bot3/control_loops/python/arm.py
new file mode 100644
index 0000000..ef02be8
--- /dev/null
+++ b/y2024_bot3/control_loops/python/arm.py
@@ -0,0 +1,61 @@
+#!/usr/bin/python3
+
+from aos.util.trapezoid_profile import TrapezoidProfile
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import linear_system
+from frc971.control_loops.python import controls
+import numpy
+import sys
+from matplotlib import pylab
+import gflags
+import glog
+
+FLAGS = gflags.FLAGS
+
+try:
+    gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
+except gflags.DuplicateFlagError:
+    pass
+
+gflags.DEFINE_bool('hybrid', False, 'If true, make it hybrid.')
+
+kArm = linear_system.LinearSystemParams(
+    name='Arm',
+    motor=control_loop.KrakenFOC(),
+    G=(14. / 60.) * (32. / 48.),
+    radius=36 * 0.005 / numpy.pi / 2.0,
+    mass=5.0,
+    q_pos=0.20,
+    q_vel=80.0,
+    kalman_q_pos=0.12,
+    kalman_q_vel=2.0,
+    kalman_q_voltage=8.0,
+    kalman_r_position=0.05,
+    dt=0.005,
+)
+
+
+def main(argv):
+    if FLAGS.plot:
+        R = numpy.matrix([[0.4], [0.0]])
+        linear_system.PlotKick(kArm, R)
+        linear_system.PlotMotion(kArm,
+                                 R,
+                                 max_velocity=2.0,
+                                 max_acceleration=15.0)
+        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 arm pivot and integral arm pivot.'
+        )
+    else:
+        namespaces = ['y2024_bot3', 'control_loops', 'superstructure', 'arm']
+        linear_system.WriteLinearSystem(kArm, argv[1:4], argv[4:7], namespaces)
+
+
+if __name__ == '__main__':
+    argv = FLAGS(sys.argv)
+    glog.init()
+    sys.exit(main(argv))
diff --git a/y2024_bot3/control_loops/superstructure/BUILD b/y2024_bot3/control_loops/superstructure/BUILD
index 8d909bf..bcac518 100644
--- a/y2024_bot3/control_loops/superstructure/BUILD
+++ b/y2024_bot3/control_loops/superstructure/BUILD
@@ -142,6 +142,7 @@
         "//frc971/control_loops:position_sensor_sim",
         "//frc971/control_loops:subsystem_simulator",
         "//frc971/control_loops:team_number_test_environment",
+        "//y2024_bot3/control_loops/superstructure/arm:arm_plants",
     ],
 )
 
diff --git a/y2024_bot3/control_loops/superstructure/arm/BUILD b/y2024_bot3/control_loops/superstructure/arm/BUILD
new file mode 100644
index 0000000..d6d2fec
--- /dev/null
+++ b/y2024_bot3/control_loops/superstructure/arm/BUILD
@@ -0,0 +1,42 @@
+package(default_visibility = ["//y2024_bot3:__subpackages__"])
+
+genrule(
+    name = "genrule_arm",
+    outs = [
+        "arm_plant.h",
+        "arm_plant.cc",
+        "arm_plant.json",
+        "integral_arm_plant.h",
+        "integral_arm_plant.cc",
+        "integral_arm_plant.json",
+    ],
+    cmd = "$(location //y2024_bot3/control_loops/python:arm) $(OUTS)",
+    target_compatible_with = ["@platforms//os:linux"],
+    tools = [
+        "//y2024_bot3/control_loops/python:arm",
+    ],
+)
+
+cc_library(
+    name = "arm_plants",
+    srcs = [
+        "arm_plant.cc",
+        "integral_arm_plant.cc",
+    ],
+    hdrs = [
+        "arm_plant.h",
+        "integral_arm_plant.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//frc971/control_loops:hybrid_state_feedback_loop",
+        "//frc971/control_loops:state_feedback_loop",
+    ],
+)
+
+filegroup(
+    name = "arm_json",
+    srcs = ["integral_arm_plant.json"],
+    visibility = ["//visibility:public"],
+)
diff --git a/y2024_bot3/control_loops/superstructure/superstructure.cc b/y2024_bot3/control_loops/superstructure/superstructure.cc
index c18da15..2c510ad 100644
--- a/y2024_bot3/control_loops/superstructure/superstructure.cc
+++ b/y2024_bot3/control_loops/superstructure/superstructure.cc
@@ -26,7 +26,9 @@
       constants_fetcher_(event_loop),
       robot_constants_(&constants_fetcher_.constants()),
       joystick_state_fetcher_(
-          event_loop->MakeFetcher<aos::JoystickState>("/aos")) {
+          event_loop->MakeFetcher<aos::JoystickState>("/aos")),
+      arm_(robot_constants_->common()->arm(),
+           robot_constants_->robot()->arm_constants()->zeroing_constants()) {
   event_loop->SetRuntimeRealtimePriority(30);
 }
 
@@ -38,33 +40,71 @@
                                   const Position *position,
                                   aos::Sender<Output>::Builder *output,
                                   aos::Sender<Status>::Builder *status) {
-  (void)position;
-
   if (WasReset()) {
     AOS_LOG(ERROR, "WPILib reset, restarting\n");
   }
 
   OutputT output_struct;
 
-  if (unsafe_goal != nullptr) {
-  }
-
   if (joystick_state_fetcher_.Fetch() &&
       joystick_state_fetcher_->has_alliance()) {
     alliance_ = joystick_state_fetcher_->alliance();
   }
 
+  aos::FlatbufferFixedAllocatorArray<
+      frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal, 512>
+      arm_goal_buffer;
+
+  ArmStatus arm_status = ArmStatus::IDLE;
+  double arm_position =
+      robot_constants_->robot()->arm_constants()->arm_positions()->idle();
+  if (unsafe_goal != nullptr) {
+    switch (unsafe_goal->arm_position()) {
+      case PivotGoal::INTAKE:
+        arm_status = ArmStatus::INTAKE;
+        arm_position = robot_constants_->robot()
+                           ->arm_constants()
+                           ->arm_positions()
+                           ->intake();
+        break;
+      case PivotGoal::AMP:
+        arm_status = ArmStatus::AMP;
+        arm_position =
+            robot_constants_->robot()->arm_constants()->arm_positions()->amp();
+        break;
+      default:
+        arm_position =
+            robot_constants_->robot()->arm_constants()->arm_positions()->idle();
+    }
+  }
+
+  arm_goal_buffer.Finish(
+      frc971::control_loops::CreateStaticZeroingSingleDOFProfiledSubsystemGoal(
+          *arm_goal_buffer.fbb(), arm_position));
+
+  const frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal
+      *arm_goal = &arm_goal_buffer.message();
+
+  // static_zeroing_single_dof_profiled_subsystem.h
+  const flatbuffers::Offset<PotAndAbsoluteEncoderProfiledJointStatus>
+      arm_offset =
+          arm_.Iterate(arm_goal, position->arm(),
+                       output != nullptr ? &output_struct.arm_voltage : nullptr,
+                       status->fbb());
+
   if (output) {
     output->CheckOk(output->Send(Output::Pack(*output->fbb(), &output_struct)));
   }
 
   Status::Builder status_builder = status->MakeBuilder<Status>();
 
-  const bool zeroed = true;
-  const bool estopped = false;
+  const bool zeroed = arm_.zeroed();
+  const bool estopped = arm_.estopped();
 
   status_builder.add_zeroed(zeroed);
   status_builder.add_estopped(estopped);
+  status_builder.add_arm(arm_offset);
+  status_builder.add_arm_status(arm_status);
 
   (void)status->Send(status_builder.Finish());
 }
diff --git a/y2024_bot3/control_loops/superstructure/superstructure.h b/y2024_bot3/control_loops/superstructure/superstructure.h
index 6362bc0..b7499c2 100644
--- a/y2024_bot3/control_loops/superstructure/superstructure.h
+++ b/y2024_bot3/control_loops/superstructure/superstructure.h
@@ -35,6 +35,8 @@
 
   double robot_velocity() const;
 
+  inline const PotAndAbsoluteEncoderSubsystem &arm() const { return arm_; }
+
  protected:
   virtual void RunIteration(const Goal *unsafe_goal, const Position *position,
                             aos::Sender<Output>::Builder *output,
@@ -47,6 +49,8 @@
 
   aos::Alliance alliance_ = aos::Alliance::kInvalid;
 
+  PotAndAbsoluteEncoderSubsystem arm_;
+
   DISALLOW_COPY_AND_ASSIGN(Superstructure);
 };
 
diff --git a/y2024_bot3/control_loops/superstructure/superstructure_lib_test.cc b/y2024_bot3/control_loops/superstructure/superstructure_lib_test.cc
index fe26af2..d9f684c 100644
--- a/y2024_bot3/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2024_bot3/control_loops/superstructure/superstructure_lib_test.cc
@@ -12,6 +12,7 @@
 #include "frc971/control_loops/team_number_test_environment.h"
 #include "frc971/zeroing/absolute_encoder.h"
 #include "y2024_bot3/constants/simulated_constants_sender.h"
+#include "y2024_bot3/control_loops/superstructure/arm/arm_plant.h"
 #include "y2024_bot3/control_loops/superstructure/superstructure.h"
 #include "y2024_bot3/control_loops/superstructure/superstructure_can_position_generated.h"
 
@@ -56,10 +57,28 @@
         superstructure_status_fetcher_(
             event_loop_->MakeFetcher<Status>("/superstructure")),
         superstructure_output_fetcher_(
-            event_loop_->MakeFetcher<Output>("/superstructure")) {
-    (void)dt_;
-    (void)simulated_robot_constants;
+            event_loop_->MakeFetcher<Output>("/superstructure")),
+        arm_(new CappedTestPlant(arm::MakeArmPlant()),
+             PositionSensorSimulator(simulated_robot_constants->robot()
+                                         ->arm_constants()
+                                         ->zeroing_constants()
+                                         ->one_revolution_distance()),
+             {.subsystem_params = {simulated_robot_constants->common()->arm(),
+                                   simulated_robot_constants->robot()
+                                       ->arm_constants()
+                                       ->zeroing_constants()},
+              .potentiometer_offset = simulated_robot_constants->robot()
+                                          ->arm_constants()
+                                          ->potentiometer_offset()},
+             frc971::constants::Range::FromFlatbuffer(
+                 simulated_robot_constants->common()->arm()->range()),
+             simulated_robot_constants->robot()
+                 ->arm_constants()
+                 ->zeroing_constants()
+                 ->measured_absolute_position(),
+             dt_)
 
+  {
     phased_loop_handle_ = event_loop_->AddPhasedLoop(
         [this](int) {
           // Skip this the first time.
@@ -78,8 +97,15 @@
     ::aos::Sender<Position>::Builder builder =
         superstructure_position_sender_.MakeBuilder();
 
+    frc971::PotAndAbsolutePosition::Builder arm_builder =
+        builder.MakeBuilder<frc971::PotAndAbsolutePosition>();
+    flatbuffers::Offset<frc971::PotAndAbsolutePosition> arm_offset =
+        arm_.encoder()->GetSensorValues(&arm_builder);
+
     Position::Builder position_builder = builder.MakeBuilder<Position>();
 
+    position_builder.add_arm(arm_offset);
+
     CHECK_EQ(builder.Send(position_builder.Finish()),
              aos::RawSender::Error::kOk);
   }
@@ -94,6 +120,8 @@
   ::aos::Fetcher<Status> superstructure_status_fetcher_;
   ::aos::Fetcher<Output> superstructure_output_fetcher_;
 
+  PotAndAbsoluteEncoderSimulator arm_;
+
   bool first_ = true;
 };
 
@@ -154,6 +182,29 @@
         << ": No status";
     ASSERT_TRUE(superstructure_output_fetcher_.get() != nullptr)
         << ": No output";
+
+    double set_point;
+    auto arm_positions =
+        simulated_robot_constants_->robot()->arm_constants()->arm_positions();
+
+    switch (superstructure_goal_fetcher_->arm_position()) {
+      case PivotGoal::IDLE:
+        set_point = arm_positions->idle();
+        break;
+      case PivotGoal::INTAKE:
+        set_point = arm_positions->intake();
+        break;
+      case PivotGoal::AMP:
+        set_point = arm_positions->amp();
+        break;
+      default:
+        LOG(FATAL) << "PivotGoal is not IDLE, INTAKE, or AMP";
+    }
+
+    EXPECT_NEAR(
+        set_point,
+        superstructure_status_fetcher_->arm()->unprofiled_goal_position(),
+        0.03);
   }
 
   void CheckIfZeroed() {
@@ -211,7 +262,6 @@
 
 // Tests that the superstructure does nothing when the goal is to remain
 // still.
-
 TEST_F(SuperstructureTest, DoesNothing) {
   SetEnabled(true);
   WaitUntilZeroed();
@@ -240,4 +290,50 @@
   CheckIfZeroed();
 }
 
+TEST_F(SuperstructureTest, ArmPivotMovement) {
+  SetEnabled(true);
+  WaitUntilZeroed();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_position(PivotGoal::INTAKE);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+  RunFor(chrono::seconds(20));
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+  EXPECT_EQ(superstructure_status_fetcher_->arm_status(), ArmStatus::INTAKE);
+
+  VerifyNearGoal();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_position(PivotGoal::AMP);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+  RunFor(chrono::seconds(10));
+
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+  EXPECT_EQ(superstructure_status_fetcher_->arm_status(), ArmStatus::AMP);
+
+  VerifyNearGoal();
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_arm_position(PivotGoal::IDLE);
+
+    ASSERT_EQ(builder.Send(goal_builder.Finish()), aos::RawSender::Error::kOk);
+  }
+  RunFor(chrono::seconds(10));
+
+  ASSERT_TRUE(superstructure_status_fetcher_.Fetch());
+  EXPECT_EQ(superstructure_status_fetcher_->arm_status(), ArmStatus::IDLE);
+
+  VerifyNearGoal();
+}
+
 }  // namespace y2024_bot3::control_loops::superstructure::testing