From d7fabee877e1d5310f40935b0121065ed22e54d0 Mon Sep 17 00:00:00 2001
From: suryatho <suryathoppae@gmail.com>
Date: Tue, 21 Nov 2023 20:49:13 -0500
Subject: [PATCH] Code from thursday Added: Classes for creating Trajectories
 from Choreo and custom generators

New trajectory following command for those classes

New swerve controller based on the new trajectories
---
 .../frc2023/RobotContainer.java               |  12 ++
 .../frc2023/commands/DriveTrajectoryNew.java  | 158 ++++++++++++++++++
 .../trajectory/ChoreoTrajectoryFactory.java   | 148 ++++++++++++++++
 .../CustomSwerveDriveController.java          |  68 ++++++++
 .../trajectory/CustomTrajectoryFactory.java   |  57 +++++++
 .../trajectory/TrajectoryImplementation.java  |  18 ++
 6 files changed, 461 insertions(+)
 create mode 100644 src/main/java/org/littletonrobotics/frc2023/commands/DriveTrajectoryNew.java
 create mode 100644 src/main/java/org/littletonrobotics/frc2023/util/trajectory/ChoreoTrajectoryFactory.java
 create mode 100644 src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomSwerveDriveController.java
 create mode 100644 src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomTrajectoryFactory.java
 create mode 100644 src/main/java/org/littletonrobotics/frc2023/util/trajectory/TrajectoryImplementation.java

diff --git a/src/main/java/org/littletonrobotics/frc2023/RobotContainer.java b/src/main/java/org/littletonrobotics/frc2023/RobotContainer.java
index 2c02f14..af43b3b 100644
--- a/src/main/java/org/littletonrobotics/frc2023/RobotContainer.java
+++ b/src/main/java/org/littletonrobotics/frc2023/RobotContainer.java
@@ -17,7 +17,10 @@
 import edu.wpi.first.wpilibj2.command.Commands;
 import edu.wpi.first.wpilibj2.command.button.CommandXboxController;
 import edu.wpi.first.wpilibj2.command.button.Trigger;
+import java.io.File;
+import java.util.Optional;
 import org.littletonrobotics.frc2023.Constants.Mode;
+import org.littletonrobotics.frc2023.commands.DriveTrajectoryNew;
 import org.littletonrobotics.frc2023.commands.DriveWithJoysticks;
 import org.littletonrobotics.frc2023.commands.FeedForwardCharacterization;
 import org.littletonrobotics.frc2023.commands.autos.*;
@@ -35,6 +38,7 @@
 import org.littletonrobotics.frc2023.util.Alert.AlertType;
 import org.littletonrobotics.frc2023.util.AllianceFlipUtil;
 import org.littletonrobotics.frc2023.util.SparkMaxBurnManager;
+import org.littletonrobotics.frc2023.util.trajectory.ChoreoTrajectoryFactory;
 import org.littletonrobotics.junction.networktables.LoggedDashboardChooser;
 import org.littletonrobotics.junction.networktables.LoggedDashboardNumber;
 
@@ -148,6 +152,14 @@ public RobotContainer() {
             new FeedForwardCharacterization.FeedForwardCharacterizationData("drive"),
             (Double voltage) -> drive.runCharacterizationVolts(voltage),
             drive::getCharacterizationVelocity));
+    // only add to auto map if there
+    Optional<File> trajectoryFile = ChoreoTrajectoryFactory.getTrajectoryFile("OnePieceBalance");
+    trajectoryFile.ifPresent(
+        file ->
+            autoChooser.addOption(
+                "Test New DriveTrajectory Command",
+                new DriveTrajectoryNew(
+                    drive, ChoreoTrajectoryFactory.createTrajectoryFromFile(file))));
 
     // System.out.println("[Init] Instantiating auto routines (Drive Characterization)");
     // autoSelector.addRoutine(
diff --git a/src/main/java/org/littletonrobotics/frc2023/commands/DriveTrajectoryNew.java b/src/main/java/org/littletonrobotics/frc2023/commands/DriveTrajectoryNew.java
new file mode 100644
index 0000000..4839a39
--- /dev/null
+++ b/src/main/java/org/littletonrobotics/frc2023/commands/DriveTrajectoryNew.java
@@ -0,0 +1,158 @@
+// Copyright (c) 2023 FRC 6328
+// http://github.com/Mechanical-Advantage
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file at
+// the root directory of this project.
+
+package org.littletonrobotics.frc2023.commands;
+
+import static org.littletonrobotics.frc2023.util.trajectory.CustomSwerveDriveController.DriveDynamicState;
+
+import edu.wpi.first.math.controller.PIDController;
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.kinematics.ChassisSpeeds;
+import edu.wpi.first.math.util.Units;
+import edu.wpi.first.wpilibj.Timer;
+import edu.wpi.first.wpilibj2.command.CommandBase;
+import org.littletonrobotics.frc2023.Constants;
+import org.littletonrobotics.frc2023.subsystems.drive.Drive;
+import org.littletonrobotics.frc2023.util.Alert;
+import org.littletonrobotics.frc2023.util.LoggedTunableNumber;
+import org.littletonrobotics.frc2023.util.trajectory.CustomSwerveDriveController;
+import org.littletonrobotics.frc2023.util.trajectory.TrajectoryImplementation;
+import org.littletonrobotics.junction.Logger;
+
+public class DriveTrajectoryNew extends CommandBase {
+  private static final Alert generatorAlert =
+      new Alert("Failed to generate trajectory, check constants.", Alert.AlertType.ERROR);
+
+  private static boolean supportedRobot = true;
+  private static double maxVelocityMetersPerSec;
+  private static double maxAccelerationMetersPerSec2;
+  private static double maxCentripetalAccelerationMetersPerSec2;
+
+  private static final LoggedTunableNumber driveKp =
+      new LoggedTunableNumber("DriveTrajectory/DriveKp");
+  private static final LoggedTunableNumber driveKd =
+      new LoggedTunableNumber("DriveTrajectory/DriveKd");
+  private static final LoggedTunableNumber turnKp =
+      new LoggedTunableNumber("DriveTrajectory/TurnKp");
+  private static final LoggedTunableNumber turnKd =
+      new LoggedTunableNumber("DriveTrajectory/TurnKd");
+
+  private final PIDController xController = new PIDController(0.0, 0.0, 0.0);
+  private final PIDController yController = new PIDController(0.0, 0.0, 0.0);
+  private final PIDController thetaController = new PIDController(0.0, 0.0, 0.0);
+
+  private CustomSwerveDriveController driveController =
+      new CustomSwerveDriveController(xController, yController, thetaController);
+
+  private final Drive drive;
+  private TrajectoryImplementation trajectory;
+  private final Timer timer = new Timer();
+
+  static {
+    switch (Constants.getRobot()) {
+      case ROBOT_2023C:
+      case ROBOT_2023P:
+        maxVelocityMetersPerSec = Units.inchesToMeters(140.0);
+        maxAccelerationMetersPerSec2 = Units.inchesToMeters(80.0);
+        maxCentripetalAccelerationMetersPerSec2 = Units.inchesToMeters(130.0);
+
+        driveKp.initDefault(6.0);
+        driveKd.initDefault(0.0);
+        turnKp.initDefault(8.0);
+        turnKd.initDefault(0.0);
+        break;
+      case ROBOT_SIMBOT:
+        maxVelocityMetersPerSec = Units.inchesToMeters(140.0);
+        maxAccelerationMetersPerSec2 = Units.inchesToMeters(80.0);
+        maxCentripetalAccelerationMetersPerSec2 = Units.inchesToMeters(130.0);
+
+        driveKp.initDefault(2.5);
+        driveKd.initDefault(0.0);
+        turnKp.initDefault(7.0);
+        turnKd.initDefault(0.0);
+        break;
+      default:
+        supportedRobot = false;
+        break;
+    }
+  }
+
+  public DriveTrajectoryNew(Drive drive, TrajectoryImplementation trajectory) {
+    this.drive = drive;
+    this.trajectory = trajectory;
+
+    addRequirements(drive);
+  }
+
+  @Override
+  public void initialize() {
+    // Log trajectory
+    Logger.getInstance().recordOutput("Odometry/Trajectory", trajectory.getTrajectoryPoses());
+
+    // Reset all controllers
+    timer.reset();
+    timer.start();
+    xController.reset();
+    yController.reset();
+    thetaController.reset();
+
+    // Reset PID gains
+    xController.setP(driveKp.get());
+    xController.setD(driveKd.get());
+    yController.setP(driveKp.get());
+    yController.setD(driveKd.get());
+    thetaController.setP(turnKp.get());
+    thetaController.setD(turnKd.get());
+  }
+
+  @Override
+  public void execute() {
+    // Update from tunable numbers
+    if (driveKd.hasChanged(hashCode())
+        || driveKp.hasChanged(hashCode())
+        || turnKd.hasChanged(hashCode())
+        || turnKp.hasChanged(hashCode())) {
+      xController.setP(driveKp.get());
+      xController.setD(driveKd.get());
+      yController.setP(driveKp.get());
+      yController.setD(driveKd.get());
+      thetaController.setP(turnKp.get());
+      thetaController.setD(turnKd.get());
+    }
+
+    DriveDynamicState state = trajectory.sample(timer.get()).maybeFlip();
+    Pose2d statePose = state.pose();
+    Logger.getInstance().recordOutput("Odometry/TrajectorySetpoint", statePose);
+    ChassisSpeeds nextDriveState = driveController.calculate(drive.getPose(), state);
+
+    Logger.getInstance()
+        .recordOutput(
+            "Drive/ffSpeeds",
+            new double[] {state.velocityX(), state.velocityY(), state.angularVelocity()});
+    Logger.getInstance()
+        .recordOutput(
+            "Drive/autoSpeeds",
+            new double[] {
+              nextDriveState.vxMetersPerSecond,
+              nextDriveState.vyMetersPerSecond,
+              nextDriveState.omegaRadiansPerSecond
+            });
+    drive.runVelocity(nextDriveState);
+  }
+
+  @Override
+  public boolean isFinished() {
+    return timer.hasElapsed(trajectory.getDuration());
+  }
+
+  @Override
+  public void end(boolean interrupted) {
+    drive.stop();
+    Logger.getInstance().recordOutput("Odometry/Trajectory", new double[] {});
+    Logger.getInstance().recordOutput("Odometry/TrajectorySetpoint", new double[] {});
+  }
+}
diff --git a/src/main/java/org/littletonrobotics/frc2023/util/trajectory/ChoreoTrajectoryFactory.java b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/ChoreoTrajectoryFactory.java
new file mode 100644
index 0000000..97fa3ec
--- /dev/null
+++ b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/ChoreoTrajectoryFactory.java
@@ -0,0 +1,148 @@
+// Copyright (c) 2023 FRC 6328
+// http://github.com/Mechanical-Advantage
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file at
+// the root directory of this project.
+
+package org.littletonrobotics.frc2023.util.trajectory;
+
+import static org.littletonrobotics.frc2023.util.trajectory.CustomSwerveDriveController.DriveDynamicState;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import edu.wpi.first.math.MathUtil;
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.wpilibj.DriverStation;
+import edu.wpi.first.wpilibj.Filesystem;
+import java.io.File;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Optional;
+
+public class ChoreoTrajectoryFactory {
+  private static final ObjectMapper mapper = new ObjectMapper();
+  private static final HashMap<String, File> FILE_HASH_MAP = new HashMap<>();
+  // load trajectory files into map
+  static {
+    File trajectory_dir = Filesystem.getDeployDirectory();
+    if (trajectory_dir.isDirectory() && trajectory_dir.listFiles() != null) {
+      for (File file : trajectory_dir.listFiles()) {
+        String fileName = file.getName();
+        // check if json
+        int pointIndex = fileName.lastIndexOf(".");
+        String type = fileName.substring(pointIndex);
+        if (!type.equals(".json")) continue;
+        // use name of file without ".json"
+        String trimmedName = fileName.substring(0, pointIndex - 1);
+        FILE_HASH_MAP.put(trimmedName, file);
+      }
+    }
+  }
+
+  // TODO: does this add another unnecessary purpose to this file?
+
+  /**
+   * @param trajectoryName Name of trajectory file without ".json"
+   * @return Optional of trajectory file
+   */
+  public static Optional<File> getTrajectoryFile(String trajectoryName) {
+    if (!FILE_HASH_MAP.containsKey(trajectoryName)) {
+      DriverStation.reportWarning("No trajectory with name: " + trajectoryName, false);
+      return Optional.empty();
+    }
+    return Optional.of(FILE_HASH_MAP.get(trajectoryName));
+  }
+
+  /**
+   * @param trajectoryFile File to make trajectory from
+   * @return
+   */
+  public static TrajectoryImplementation createTrajectoryFromFile(File trajectoryFile) {
+    try {
+      return createTrajectoryImplementation(
+          mapper.readValue(trajectoryFile, new TypeReference<>() {}));
+    } catch (IOException e) {
+      e.printStackTrace();
+      return createTrajectoryImplementation(Collections.emptyList());
+    }
+  }
+
+  private static TrajectoryImplementation createTrajectoryImplementation(
+      List<ChoreoTrajectoryState> states) {
+    return new TrajectoryImplementation() {
+      @Override
+      public double getDuration() {
+        return states.get(states.size() - 1).timestamp();
+      }
+
+      @Override
+      public Pose2d[] getTrajectoryPoses() {
+        return states.stream()
+            .map(state -> new Pose2d(state.x(), state.y(), new Rotation2d(state.heading())))
+            .toArray(Pose2d[]::new);
+      }
+
+      @Override
+      public DriveDynamicState sample(double time) {
+        ChoreoTrajectoryState before = null;
+        ChoreoTrajectoryState after = null;
+
+        for (ChoreoTrajectoryState state : states) {
+          if (time == state.timestamp()) {
+            return ChoreoTrajectoryState.toDynamicState(state);
+          }
+
+          if (state.timestamp() < time) {
+            before = state;
+          } else {
+            after = state;
+            break;
+          }
+        }
+
+        if (before == null) return ChoreoTrajectoryState.toDynamicState(states.get(0));
+
+        if (after == null)
+          return ChoreoTrajectoryState.toDynamicState(states.get(states.size() - 1));
+
+        double s = 1 - ((after.timestamp() - time) / (after.timestamp() - before.timestamp()));
+
+        DriveDynamicState beforeState = ChoreoTrajectoryState.toDynamicState(before);
+        DriveDynamicState afterState = ChoreoTrajectoryState.toDynamicState(after);
+
+        Pose2d interpolatedPose = beforeState.pose().interpolate(afterState.pose(), s);
+        double interpolatedVelocityX =
+            MathUtil.interpolate(beforeState.velocityX(), afterState.velocityX(), s);
+        double interpolatedVelocityY =
+            MathUtil.interpolate(beforeState.velocityY(), afterState.velocityY(), s);
+        double interpolatedAngularVelocity =
+            MathUtil.interpolate(beforeState.angularVelocity(), afterState.angularVelocity(), s);
+
+        return new DriveDynamicState(
+            interpolatedPose,
+            interpolatedVelocityX,
+            interpolatedVelocityY,
+            interpolatedAngularVelocity);
+      }
+    };
+  }
+
+  private record ChoreoTrajectoryState(
+      double timestamp,
+      double x,
+      double y,
+      double heading,
+      double velocityX,
+      double velocityY,
+      double angularVelocity) {
+    public static DriveDynamicState toDynamicState(ChoreoTrajectoryState state) {
+      Pose2d pose = new Pose2d(state.x(), state.y(), new Rotation2d(state.heading()));
+      return new DriveDynamicState(
+          pose, state.velocityX(), state.velocityY(), state.angularVelocity());
+    }
+  }
+}
diff --git a/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomSwerveDriveController.java b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomSwerveDriveController.java
new file mode 100644
index 0000000..9da51db
--- /dev/null
+++ b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomSwerveDriveController.java
@@ -0,0 +1,68 @@
+// Copyright (c) 2023 FRC 6328
+// http://github.com/Mechanical-Advantage
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file at
+// the root directory of this project.
+
+package org.littletonrobotics.frc2023.util.trajectory;
+
+import static org.littletonrobotics.frc2023.util.AllianceFlipUtil.*;
+
+import edu.wpi.first.math.controller.PIDController;
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.kinematics.ChassisSpeeds;
+
+public class CustomSwerveDriveController {
+  private PIDController xController, yController, thetaController;
+
+  private Pose2d m_poseError;
+  private Rotation2d m_rotationError;
+
+  public CustomSwerveDriveController(
+      PIDController xController, PIDController yController, PIDController thetaController) {
+    this.xController = xController;
+    this.yController = yController;
+    this.thetaController = thetaController;
+
+    this.thetaController.enableContinuousInput(-Math.PI, Math.PI);
+  }
+
+  public ChassisSpeeds calculate(Pose2d currentPose, DriveDynamicState state) {
+    // Calculate feedforward velocities (field-relative).
+    double xFF = state.velocityX();
+    double yFF = state.velocityY();
+    double thetaFF = state.angularVelocity();
+
+    Pose2d poseRef = state.pose();
+
+    m_poseError = poseRef.relativeTo(currentPose);
+    m_rotationError = poseRef.getRotation().minus(currentPose.getRotation());
+
+    // Calculate feedback velocities (based on position error).
+    double xFeedback = xController.calculate(currentPose.getX(), poseRef.getX());
+    double yFeedback = yController.calculate(currentPose.getY(), poseRef.getY());
+    double thetaFeedback =
+        thetaController.calculate(
+            currentPose.getRotation().getRadians(), poseRef.getRotation().getRadians());
+
+    // Return next output.
+    return ChassisSpeeds.fromFieldRelativeSpeeds(
+        xFF + xFeedback, yFF + yFeedback, thetaFF + thetaFeedback, currentPose.getRotation());
+  }
+
+  public record DriveDynamicState(
+      Pose2d pose, double velocityX, double velocityY, double angularVelocity) {
+    public DriveDynamicState maybeFlip() {
+      double x = apply(pose.getX());
+      double y = pose.getY();
+      Rotation2d heading = apply(pose.getRotation());
+      return new DriveDynamicState(
+          new Pose2d(x, y, heading),
+          shouldFlip() ? -velocityX : velocityX,
+          velocityY,
+          shouldFlip() ? -angularVelocity : angularVelocity);
+    }
+  }
+}
diff --git a/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomTrajectoryFactory.java b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomTrajectoryFactory.java
new file mode 100644
index 0000000..d2084a7
--- /dev/null
+++ b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/CustomTrajectoryFactory.java
@@ -0,0 +1,57 @@
+// Copyright (c) 2023 FRC 6328
+// http://github.com/Mechanical-Advantage
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file at
+// the root directory of this project.
+
+package org.littletonrobotics.frc2023.util.trajectory;
+
+import static org.littletonrobotics.frc2023.util.trajectory.CustomSwerveDriveController.DriveDynamicState;
+
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.trajectory.Trajectory;
+import edu.wpi.first.math.trajectory.TrajectoryConfig;
+import java.util.List;
+import org.littletonrobotics.frc2023.util.AllianceFlipUtil;
+
+public class CustomTrajectoryFactory {
+  public static TrajectoryImplementation createFromWaypoints(
+      TrajectoryConfig config, List<Waypoint> waypoints) {
+    CustomTrajectoryGenerator trajectoryGenerator = new CustomTrajectoryGenerator();
+    trajectoryGenerator.generate(config, waypoints);
+
+    return new TrajectoryImplementation() {
+      @Override
+      public double getDuration() {
+        return trajectoryGenerator.getDriveTrajectory().getTotalTimeSeconds();
+      }
+
+      @Override
+      public Pose2d[] getTrajectoryPoses() {
+        return trajectoryGenerator.getDriveTrajectory().getStates().stream()
+            .map(state -> AllianceFlipUtil.apply(state.poseMeters))
+            .toArray(Pose2d[]::new);
+      }
+
+      @Override
+      public DriveDynamicState sample(double time) {
+        Trajectory.State driveState = trajectoryGenerator.getDriveTrajectory().sample(time);
+        RotationSequence.State rotationState =
+            trajectoryGenerator.getHolonomicRotationSequence().sample(time);
+
+        // TODO: Check if this is actually what I need to do
+        double velocityX =
+            driveState.poseMeters.getRotation().getCos() * driveState.velocityMetersPerSecond;
+        double velocityY =
+            driveState.poseMeters.getRotation().getSin() * driveState.velocityMetersPerSecond;
+
+        return new DriveDynamicState(
+            new Pose2d(driveState.poseMeters.getTranslation(), rotationState.position),
+            velocityX,
+            velocityY,
+            rotationState.velocityRadiansPerSec);
+      }
+    };
+  }
+}
diff --git a/src/main/java/org/littletonrobotics/frc2023/util/trajectory/TrajectoryImplementation.java b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/TrajectoryImplementation.java
new file mode 100644
index 0000000..589c622
--- /dev/null
+++ b/src/main/java/org/littletonrobotics/frc2023/util/trajectory/TrajectoryImplementation.java
@@ -0,0 +1,18 @@
+// Copyright (c) 2023 FRC 6328
+// http://github.com/Mechanical-Advantage
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file at
+// the root directory of this project.
+
+package org.littletonrobotics.frc2023.util.trajectory;
+
+import edu.wpi.first.math.geometry.Pose2d;
+
+public interface TrajectoryImplementation {
+  double getDuration();
+
+  Pose2d[] getTrajectoryPoses();
+
+  CustomSwerveDriveController.DriveDynamicState sample(double time);
+}