From a2dbab96f06b8433827b8e27a46e6c3fccfd2a30 Mon Sep 17 00:00:00 2001
From: yapplejack <jacklf@umich.edu>
Date: Wed, 2 Oct 2024 15:17:00 -0400
Subject: [PATCH] changed input to desaturateWheelSpeeds from
 getRobotVelocity() to desiredChaassisSpeeds

---
 src/main/java/swervelib/SwerveDrive.java     | 22 ++++++++++++++------
 src/main/java/swervelib/math/SwerveMath.java |  2 +-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/main/java/swervelib/SwerveDrive.java b/src/main/java/swervelib/SwerveDrive.java
index c6460e98..09097143 100644
--- a/src/main/java/swervelib/SwerveDrive.java
+++ b/src/main/java/swervelib/SwerveDrive.java
@@ -518,7 +518,7 @@ public void drive(ChassisSpeeds velocity, boolean isOpenLoop, Translation2d cent
     // Calculate required module states via kinematics
     SwerveModuleState[] swerveModuleStates = kinematics.toSwerveModuleStates(velocity, centerOfRotationMeters);
 
-    setRawModuleStates(swerveModuleStates, isOpenLoop);
+    setRawModuleStates(swerveModuleStates, velocity, isOpenLoop);
   }
 
   /**
@@ -568,14 +568,15 @@ public double getMaximumAngularVelocity()
    * Set the module states (azimuth and velocity) directly.
    *
    * @param desiredStates A list of SwerveModuleStates to send to the modules.
+   * @param desiredChassisSpeed The desired chassis speeds to set the robot to achieve.
    * @param isOpenLoop    Whether to use closed-loop velocity control. Set to true to disable closed-loop.
    */
-  private void setRawModuleStates(SwerveModuleState[] desiredStates, boolean isOpenLoop)
+  private void setRawModuleStates(SwerveModuleState[] desiredStates, ChassisSpeeds desiredChassisSpeed, boolean isOpenLoop)
   {
     // Desaturates wheel speeds
     if (attainableMaxTranslationalSpeedMetersPerSecond != 0 || attainableMaxRotationalVelocityRadiansPerSecond != 0)
     {
-      SwerveDriveKinematics.desaturateWheelSpeeds(desiredStates, getRobotVelocity(),
+      SwerveDriveKinematics.desaturateWheelSpeeds(desiredStates, desiredChassisSpeed,
                                                   maxSpeedMPS,
                                                   attainableMaxTranslationalSpeedMetersPerSecond,
                                                   attainableMaxRotationalVelocityRadiansPerSecond);
@@ -593,14 +594,23 @@ private void setRawModuleStates(SwerveModuleState[] desiredStates, boolean isOpe
 
   /**
    * Set the module states (azimuth and velocity) directly. Used primarily for auto paths.
+   * Does not allow for usage of desaturateWheelSpeeds(SwerveModuleState[] moduleStates, 
+   * ChassisSpeeds desiredChassisSpeed, double attainableMaxModuleSpeedMetersPerSecond, 
+   * double attainableMaxTranslationalSpeedMetersPerSecond, double attainableMaxRotationalVelocityRadiansPerSecond)
    *
    * @param desiredStates A list of SwerveModuleStates to send to the modules.
    * @param isOpenLoop    Whether to use closed-loop velocity control. Set to true to disable closed-loop.
    */
   public void setModuleStates(SwerveModuleState[] desiredStates, boolean isOpenLoop)
   {
-    setRawModuleStates(kinematics.toSwerveModuleStates(kinematics.toChassisSpeeds(desiredStates)),
-                       isOpenLoop);
+    desiredStates = kinematics.toSwerveModuleStates(kinematics.toChassisSpeeds(desiredStates));
+    SwerveDriveKinematics.desaturateWheelSpeeds(desiredStates, maxSpeedMPS);
+
+    // Sets states
+    for (SwerveModule module : swerveModules)
+    {
+      module.setDesiredState(desiredStates[module.moduleNumber], isOpenLoop, false);
+    }
   }
 
   /**
@@ -614,7 +624,7 @@ public void setChassisSpeeds(ChassisSpeeds chassisSpeeds)
     SwerveDriveTelemetry.desiredChassisSpeeds[0] = chassisSpeeds.vxMetersPerSecond;
     SwerveDriveTelemetry.desiredChassisSpeeds[2] = Math.toDegrees(chassisSpeeds.omegaRadiansPerSecond);
 
-    setRawModuleStates(kinematics.toSwerveModuleStates(chassisSpeeds), false);
+    setRawModuleStates(kinematics.toSwerveModuleStates(chassisSpeeds), chassisSpeeds, false);
   }
 
   /**
diff --git a/src/main/java/swervelib/math/SwerveMath.java b/src/main/java/swervelib/math/SwerveMath.java
index b77838c0..ca278a9e 100644
--- a/src/main/java/swervelib/math/SwerveMath.java
+++ b/src/main/java/swervelib/math/SwerveMath.java
@@ -136,7 +136,7 @@ public static double calculateDegreesPerSteeringRotation(
   public static double calculateMaxAngularVelocity(
       double maxSpeed, double furthestModuleX, double furthestModuleY)
   {
-    return maxSpeed / Math.hypot(furthestModuleX, furthestModuleY);
+    return maxSpeed / new Rotation2d(furthestModuleX, furthestModuleY).getRadians();
   }
 
   /**