Add ball sorting

Discard balls that are not our alliance color

Signed-off-by: Ravago Jones <ravagojones@gmail.com>
Change-Id: Ic9ff9bfc724c74d0c6de4364ddc3999945c062d2
diff --git a/frc971/control_loops/control_loop_test.h b/frc971/control_loops/control_loop_test.h
index a820215..576d476 100644
--- a/frc971/control_loops/control_loop_test.h
+++ b/frc971/control_loops/control_loop_test.h
@@ -61,6 +61,9 @@
   void set_team_id(uint16_t team_id) { team_id_ = team_id; }
   uint16_t team_id() const { return team_id_; }
 
+  void set_alliance(aos::Alliance alliance) { alliance_ = alliance; }
+  aos::Alliance alliance() const { return alliance_; }
+
   // Sets the enabled/disabled bit and (potentially) rebroadcasts the robot
   // state messages.
   void SetEnabled(bool enabled) {
@@ -123,9 +126,9 @@
       builder.add_enabled(enabled_);
       builder.add_autonomous(false);
       builder.add_team_id(team_id_);
+      builder.add_alliance(alliance_);
 
-      CHECK_EQ(new_state.Send(builder.Finish()),
-               aos::RawSender::Error::kOk);
+      CHECK_EQ(new_state.Send(builder.Finish()), aos::RawSender::Error::kOk);
 
       last_ds_time_ = monotonic_now();
       last_enabled_ = enabled_;
@@ -165,6 +168,7 @@
   const ::std::chrono::nanoseconds dt_;
 
   uint16_t team_id_ = 971;
+  aos::Alliance alliance_ = aos::Alliance::kInvalid;
   int32_t reader_pid_ = 1;
   double battery_voltage_ = 12.4;
 
diff --git a/y2022/control_loops/superstructure/BUILD b/y2022/control_loops/superstructure/BUILD
index ebd9ace..31b35fb 100644
--- a/y2022/control_loops/superstructure/BUILD
+++ b/y2022/control_loops/superstructure/BUILD
@@ -88,6 +88,7 @@
         "//y2022:constants",
         "//y2022/control_loops/superstructure/catapult",
         "//y2022/control_loops/superstructure/turret:aiming",
+        "//y2022/vision:ball_color_fbs",
     ],
 )
 
diff --git a/y2022/control_loops/superstructure/superstructure.cc b/y2022/control_loops/superstructure/superstructure.cc
index 15f9c2a..bc24fda 100644
--- a/y2022/control_loops/superstructure/superstructure.cc
+++ b/y2022/control_loops/superstructure/superstructure.cc
@@ -32,6 +32,10 @@
               "/drivetrain")),
       can_position_fetcher_(
           event_loop->MakeFetcher<CANPosition>("/superstructure")),
+      joystick_state_fetcher_(
+          event_loop->MakeFetcher<aos::JoystickState>("/aos")),
+      ball_color_fetcher_(event_loop->MakeFetcher<y2022::vision::BallColor>(
+          "/superstructure")),
       aimer_(values) {
   event_loop->SetRuntimeRealtimePriority(30);
 }
@@ -55,10 +59,43 @@
       frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal, 512>
       turret_loading_goal_buffer;
   aos::FlatbufferFixedAllocatorArray<CatapultGoal, 512> catapult_goal_buffer;
+  aos::FlatbufferFixedAllocatorArray<CatapultGoal, 512>
+      catapult_discarding_goal_buffer;
 
   const aos::monotonic_clock::time_point timestamp =
       event_loop()->context().monotonic_event_time;
 
+  if (joystick_state_fetcher_.Fetch() &&
+      joystick_state_fetcher_->has_alliance()) {
+    alliance_ = joystick_state_fetcher_->alliance();
+  }
+
+  if (ball_color_fetcher_.Fetch() && ball_color_fetcher_->has_ball_color()) {
+    ball_color_ = ball_color_fetcher_->ball_color();
+  }
+
+  if (alliance_ != aos::Alliance::kInvalid &&
+      ball_color_ != aos::Alliance::kInvalid && alliance_ != ball_color_) {
+    switch (state_) {
+      case SuperstructureState::IDLE:
+        break;
+      case SuperstructureState::TRANSFERRING:
+        break;
+      case SuperstructureState::LOADING:
+        break;
+      case SuperstructureState::LOADED:
+        discarding_ball_ = true;
+        break;
+      case SuperstructureState::SHOOTING:
+        if (!fire_) {
+          // we can still tell it not to shoot into the hub
+          // and change the turret and catapult goals
+          discarding_ball_ = true;
+        }
+        break;
+    }
+  }
+
   drivetrain_status_fetcher_.Fetch();
   const float velocity = robot_velocity();
 
@@ -92,8 +129,9 @@
 
     climber_servo = unsafe_goal->climber_servo();
 
-    turret_goal =
-        unsafe_goal->auto_aim() ? auto_aim_goal : unsafe_goal->turret();
+    turret_goal = unsafe_goal->auto_aim() && !discarding_ball_
+                      ? auto_aim_goal
+                      : unsafe_goal->turret();
 
     catapult_goal = unsafe_goal->catapult();
 
@@ -107,7 +145,7 @@
       std::optional<flatbuffers::Offset<
           frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal>>
           return_position_offset;
-      if (unsafe_goal != nullptr && unsafe_goal->has_catapult() &&
+      if (unsafe_goal->has_catapult() &&
           unsafe_goal->catapult()->has_return_position()) {
         return_position_offset = {aos::CopyFlatBuffer(
             unsafe_goal->catapult()->return_position(), catapult_goal_fbb)};
@@ -122,6 +160,27 @@
       catapult_goal = &catapult_goal_buffer.message();
     }
 
+    if (discarding_ball_) {
+      flatbuffers::FlatBufferBuilder *catapult_goal_fbb =
+          catapult_discarding_goal_buffer.fbb();
+      std::optional<flatbuffers::Offset<
+          frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal>>
+          return_position_offset;
+      if (unsafe_goal->has_catapult() &&
+          unsafe_goal->catapult()->has_return_position()) {
+        return_position_offset = {aos::CopyFlatBuffer(
+            unsafe_goal->catapult()->return_position(), catapult_goal_fbb)};
+      }
+      CatapultGoal::Builder builder(*catapult_goal_fbb);
+      builder.add_shot_position(kDiscardingPosition);
+      builder.add_shot_velocity(kDiscardingVelocity);
+      if (return_position_offset.has_value()) {
+        builder.add_return_position(return_position_offset.value());
+      }
+      catapult_discarding_goal_buffer.Finish(builder.Finish());
+      catapult_goal = &catapult_discarding_goal_buffer.message();
+    }
+
     if (unsafe_goal->has_turret_intake()) {
       have_active_intake_request = true;
     }
@@ -333,6 +392,8 @@
       } else if (timestamp >
                  loading_timer_ + constants::Values::kExtraLoadingTime()) {
         state_ = SuperstructureState::LOADED;
+        // reset color and wait for a new one once we know the ball is in place
+        ball_color_ = aos::Alliance::kInvalid;
         reseating_in_catapult_ = false;
       }
       break;
@@ -346,10 +407,11 @@
           }
           turret_goal = &turret_loading_goal_buffer.message();
         }
+
         if (unsafe_goal->cancel_shot()) {
           // Cancel the shot process
           state_ = SuperstructureState::IDLE;
-        } else if (unsafe_goal->fire()) {
+        } else if (unsafe_goal->fire() || discarding_ball_) {
           // Start if we were asked to and the turret is at goal
           state_ = SuperstructureState::SHOOTING;
           prev_shot_count_ = catapult_.shot_count();
@@ -435,6 +497,7 @@
       if (catapult_.shot_count() > prev_shot_count_ && near_return_position) {
         prev_shot_count_ = catapult_.shot_count();
         fire_ = false;
+        discarding_ball_ = false;
         state_ = SuperstructureState::IDLE;
       }
 
@@ -542,6 +605,7 @@
   status_builder.add_reseating_in_catapult(reseating_in_catapult_);
   status_builder.add_fire(fire_);
   status_builder.add_moving_too_fast(moving_too_fast);
+  status_builder.add_discarding_ball(discarding_ball_);
   status_builder.add_ready_to_fire(state_ == SuperstructureState::LOADED &&
                                    turret_near_goal && !collided);
   status_builder.add_state(state_);
diff --git a/y2022/control_loops/superstructure/superstructure.h b/y2022/control_loops/superstructure/superstructure.h
index 14fa8ab..48d07f1 100644
--- a/y2022/control_loops/superstructure/superstructure.h
+++ b/y2022/control_loops/superstructure/superstructure.h
@@ -13,6 +13,7 @@
 #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"
+#include "y2022/vision/ball_color_generated.h"
 
 namespace y2022 {
 namespace control_loops {
@@ -35,6 +36,8 @@
   static constexpr double kCatapultGoalThreshold = 0.05;
   // potentiometer will be more noisy
   static constexpr double kFlipperGoalThreshold = 0.05;
+  static constexpr double kDiscardingPosition = 0.35;
+  static constexpr double kDiscardingVelocity = 6.0;
 
   explicit Superstructure(::aos::EventLoop *event_loop,
                           std::shared_ptr<const constants::Values> values,
@@ -72,6 +75,8 @@
   aos::Fetcher<frc971::control_loops::drivetrain::Status>
       drivetrain_status_fetcher_;
   aos::Fetcher<CANPosition> can_position_fetcher_;
+  aos::Fetcher<aos::JoystickState> joystick_state_fetcher_;
+  aos::Fetcher<y2022::vision::BallColor> ball_color_fetcher_;
 
   int prev_shot_count_ = 0;
 
@@ -80,6 +85,9 @@
   bool flippers_open_ = false;
   bool reseating_in_catapult_ = false;
   bool fire_ = false;
+  bool discarding_ball_ = false;
+  aos::Alliance alliance_ = aos::Alliance::kInvalid;
+  aos::Alliance ball_color_ = aos::Alliance::kInvalid;
 
   aos::monotonic_clock::time_point front_intake_beambreak_timer_ =
       aos::monotonic_clock::min_time;
diff --git a/y2022/control_loops/superstructure/superstructure_lib_test.cc b/y2022/control_loops/superstructure/superstructure_lib_test.cc
index 7123a47..e18c2d9 100644
--- a/y2022/control_loops/superstructure/superstructure_lib_test.cc
+++ b/y2022/control_loops/superstructure/superstructure_lib_test.cc
@@ -328,9 +328,11 @@
         values_(std::make_shared<constants::Values>(constants::MakeValues(
             frc971::control_loops::testing::kTeamNumber))),
         roborio_(aos::configuration::GetNode(configuration(), "roborio")),
+        logger_pi_(aos::configuration::GetNode(configuration(), "logger")),
         superstructure_event_loop(MakeEventLoop("Superstructure", roborio_)),
         superstructure_(superstructure_event_loop.get(), values_),
         test_event_loop_(MakeEventLoop("test", roborio_)),
+        ball_color_event_loop_(MakeEventLoop("ball color test", logger_pi_)),
         superstructure_goal_fetcher_(
             test_event_loop_->MakeFetcher<Goal>("/superstructure")),
         superstructure_goal_sender_(
@@ -345,6 +347,9 @@
             test_event_loop_->MakeSender<Position>("/superstructure")),
         drivetrain_status_sender_(
             test_event_loop_->MakeSender<DrivetrainStatus>("/drivetrain")),
+        ball_color_sender_(
+            ball_color_event_loop_->MakeSender<y2022::vision::BallColor>(
+                "/superstructure")),
         superstructure_plant_event_loop_(MakeEventLoop("plant", roborio_)),
         superstructure_plant_(superstructure_plant_event_loop_.get(), values_,
                               dt()) {
@@ -476,10 +481,12 @@
   std::shared_ptr<const constants::Values> values_;
 
   const aos::Node *const roborio_;
+  const aos::Node *const logger_pi_;
 
   ::std::unique_ptr<::aos::EventLoop> superstructure_event_loop;
   ::y2022::control_loops::superstructure::Superstructure superstructure_;
   ::std::unique_ptr<::aos::EventLoop> test_event_loop_;
+  ::std::unique_ptr<aos::EventLoop> ball_color_event_loop_;
   ::aos::PhasedLoopHandler *phased_loop_handle_ = nullptr;
 
   ::aos::Fetcher<Goal> superstructure_goal_fetcher_;
@@ -489,6 +496,7 @@
   ::aos::Fetcher<Position> superstructure_position_fetcher_;
   ::aos::Sender<Position> superstructure_position_sender_;
   ::aos::Sender<DrivetrainStatus> drivetrain_status_sender_;
+  ::aos::Sender<y2022::vision::BallColor> ball_color_sender_;
 
   ::std::unique_ptr<::aos::EventLoop> superstructure_plant_event_loop_;
   SuperstructureSimulation superstructure_plant_;
@@ -1285,6 +1293,80 @@
             shot_params.shot_angle);
 }
 
+// Tests that balls get discarded when they are the wrong color.
+TEST_F(SuperstructureTest, BallDiscarding) {
+  set_alliance(aos::Alliance::kInvalid);
+  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);
+
+  RunFor(chrono::milliseconds(500));
+  set_alliance(aos::Alliance::kBlue);
+
+  {
+    auto builder = superstructure_goal_sender_.MakeBuilder();
+
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+
+    goal_builder.add_auto_aim(true);
+    goal_builder.add_preloaded(true);
+    goal_builder.add_turret_intake(RequestedIntake::kFront);
+
+    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);
+
+  {
+    auto builder = ball_color_sender_.MakeBuilder();
+
+    y2022::vision::BallColor::Builder ball_color_builder =
+        builder.MakeBuilder<y2022::vision::BallColor>();
+
+    ball_color_builder.add_ball_color(aos::Alliance::kBlue);
+
+    ASSERT_EQ(builder.Send(ball_color_builder.Finish()),
+              aos::RawSender::Error::kOk);
+  }
+
+  RunFor(chrono::milliseconds(100));
+  superstructure_status_fetcher_.Fetch();
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::LOADED);
+
+  {
+    auto builder = ball_color_sender_.MakeBuilder();
+
+    y2022::vision::BallColor::Builder ball_color_builder =
+        builder.MakeBuilder<y2022::vision::BallColor>();
+
+    ball_color_builder.add_ball_color(aos::Alliance::kRed);
+
+    ASSERT_EQ(builder.Send(ball_color_builder.Finish()),
+              aos::RawSender::Error::kOk);
+  }
+
+  RunFor(dt());
+
+  superstructure_status_fetcher_.Fetch();
+  EXPECT_EQ(superstructure_status_fetcher_->state(),
+            SuperstructureState::SHOOTING);
+
+  RunFor(chrono::milliseconds(2000));
+  superstructure_status_fetcher_.Fetch();
+  EXPECT_NEAR(constants::Values::kTurretFrontIntakePos(),
+              superstructure_status_fetcher_->turret()->position(), 5e-4);
+}
+
 }  // 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 9cf9a5f..ac60ce3 100644
--- a/y2022/control_loops/superstructure/superstructure_status.fbs
+++ b/y2022/control_loops/superstructure/superstructure_status.fbs
@@ -59,6 +59,9 @@
   ready_to_fire:bool (id: 20);
   // Whether the robot is moving too fast to shoot
   moving_too_fast:bool (id: 21);
+  // True if the robot has detected that it is holding
+  // the wrong color ball and is now discarding it.
+  discarding_ball:bool (id: 22);
   // Whether the catapult was told to fire,
   // meaning that the turret and flippers are ready for firing
   // and we were asked to fire. Different from fire flag in goal.
diff --git a/y2022/y2022_logger.json b/y2022/y2022_logger.json
index f54ccd7..2054452 100644
--- a/y2022/y2022_logger.json
+++ b/y2022/y2022_logger.json
@@ -42,7 +42,7 @@
       ]
     },
     {
-      "name": "/aos/remote_timestamps/roborio/superstructure/y2022-vision-BallColor",
+      "name": "/logger/aos/remote_timestamps/roborio/superstructure/y2022-vision-BallColor",
       "type": "aos.message_bridge.RemoteMessage",
       "source_node": "logger",
       "logger": "NOT_LOGGED",