From c97113a621c6bcf160ccf7a9a9c0bddd811c45d9 Mon Sep 17 00:00:00 2001
From: David Chen <dychen2523@gmail.com>
Date: Mon, 26 Feb 2024 17:43:16 -0500
Subject: [PATCH] working on climb state machine

---
 src/main/java/frc/robot/RobotContainer.java   | 22 +++++++
 .../statemachines/ClimbStateMachine.java      | 63 +++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 src/main/java/frc/robot/statemachines/ClimbStateMachine.java

diff --git a/src/main/java/frc/robot/RobotContainer.java b/src/main/java/frc/robot/RobotContainer.java
index 2563157..6e67558 100644
--- a/src/main/java/frc/robot/RobotContainer.java
+++ b/src/main/java/frc/robot/RobotContainer.java
@@ -21,8 +21,11 @@
 import edu.wpi.first.wpilibj2.command.Command;
 import edu.wpi.first.wpilibj2.command.Commands;
 import edu.wpi.first.wpilibj2.command.InstantCommand;
+import edu.wpi.first.wpilibj2.command.SelectCommand;
 import edu.wpi.first.wpilibj2.command.button.CommandXboxController;
 import frc.robot.commands.DriveCommands;
+import frc.robot.statemachines.ClimbStateMachine;
+import frc.robot.statemachines.ClimbStateMachine.CLIMB_STATES;
 import frc.robot.subsystems.drive.Drive;
 import frc.robot.subsystems.drive.GyroIO;
 import frc.robot.subsystems.drive.GyroIOPigeon2;
@@ -48,6 +51,9 @@
 import frc.robot.subsystems.shooter.FlywheelIOSim;
 import frc.robot.subsystems.shooter.FlywheelIOTalonFX;
 import frc.robot.subsystems.shooter.Shooter;
+
+import java.util.Map;
+
 import org.littletonrobotics.junction.networktables.LoggedDashboardChooser;
 import org.littletonrobotics.junction.networktables.LoggedDashboardNumber;
 
@@ -75,6 +81,20 @@ public class RobotContainer {
   private final LoggedDashboardNumber elevatorDistance =
       new LoggedDashboardNumber("elevatorDistance");
 
+  private final ClimbStateMachine climbStateMachine;
+
+  private CLIMB_STATES climbSelect() {
+    return climbStateMachine.getTargetState();
+  }
+
+  private final Command climbCommand = 
+      new SelectCommand<>(
+          Map.ofEntries(
+            Map.entry(CLIMB_STATES.NONE, null),
+            Map.entry(CLIMB_STATES.EXTEND_CLIMB, null),
+            Map.entry(CLIMB_STATES.RETRACT_CLIMB, null),
+            Map.entry(CLIMB_STATES.SCORE_TRAP, null)), this::climbSelect);
+
   /** The container for the robot. Contains subsystems, OI devices, and commands. */
   public RobotContainer() {
     switch (Constants.getMode()) {
@@ -153,6 +173,8 @@ public RobotContainer() {
         break;
     }
 
+    climbStateMachine = new ClimbStateMachine(elevator, shooter, pivot);
+
     // Set up auto routines
     autoChooser = new LoggedDashboardChooser<>("Auto Choices", AutoBuilder.buildAutoChooser());
 
diff --git a/src/main/java/frc/robot/statemachines/ClimbStateMachine.java b/src/main/java/frc/robot/statemachines/ClimbStateMachine.java
new file mode 100644
index 0000000..bf2e520
--- /dev/null
+++ b/src/main/java/frc/robot/statemachines/ClimbStateMachine.java
@@ -0,0 +1,63 @@
+package frc.robot.statemachines;
+
+import java.lang.invoke.MethodHandles.Lookup.ClassOption;
+
+import org.littletonrobotics.junction.Logger;
+
+import frc.robot.subsystems.elevator.Elevator;
+import frc.robot.subsystems.pivot.Pivot;
+import frc.robot.subsystems.shooter.Shooter;
+
+public class ClimbStateMachine {
+    private final Elevator elevator;
+    private final Pivot pivot;
+    private final Shooter shooter;
+
+    public ClimbStateMachine(Elevator elevator, Shooter shooter, Pivot pivot) {
+        this.elevator = elevator;
+        this.pivot = pivot;
+        this.shooter = shooter;
+    }
+    public enum CLIMB_STATES {
+        NONE,
+        EXTEND_CLIMB,
+        RETRACT_CLIMB,
+        SCORE_TRAP,
+    }
+
+    private CLIMB_STATES targetState = CLIMB_STATES.NONE;
+    int counter = 0;
+
+    public CLIMB_STATES getTargetState() {
+        return targetState;
+    }
+
+    public void advanceTargetState() {
+        if (atTargetState()) counter = Math.min(counter+1, 3);
+
+        switch (counter) {
+            case 0:
+                targetState = CLIMB_STATES.NONE;
+                break;
+            case 1:
+                targetState = CLIMB_STATES.EXTEND_CLIMB;
+                break;
+            case 2:
+                targetState = CLIMB_STATES.RETRACT_CLIMB;
+                break;
+            case 3:
+                targetState = CLIMB_STATES.SCORE_TRAP;
+            default:
+                targetState = CLIMB_STATES.NONE;
+                break;
+        }
+        Logger.recordOutput("Climb Target State", targetState);
+    }
+
+    public boolean atTargetState() {
+        if (elevator.extenderAtSetpoint() && pivot.pivotAtSetpoint() && shooter.atFlywheelSetpoints()) {
+            return true;
+        }
+        return false;
+    }
+}