Source code

001/************************ PROJECT ZIXI ************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved.*/
003/* This work is licensed under the terms of the MIT license.  */
004/**************************************************************/
005package com.stuypulse.robot.subsystems.swerve;
006
007import static org.wpilib.units.Units.*;
008
009import com.stuypulse.robot.constants.GlobalSettings;
010import com.stuypulse.robot.constants.GlobalSettings.RobotMode;
011import com.stuypulse.robot.subsystems.vision.Vision.VisionConsumer;
012import com.stuypulse.robot.util.PeriodicSubsystem;
013
014import org.wpilib.driverstation.RobotState;
015import org.wpilib.math.estimator.SwerveDrivePoseEstimator;
016import org.wpilib.math.geometry.Pose2d;
017import org.wpilib.math.geometry.Rotation2d;
018import org.wpilib.math.geometry.Translation2d;
019import org.wpilib.math.geometry.Twist2d;
020import org.wpilib.math.kinematics.ChassisVelocities;
021import org.wpilib.math.kinematics.SwerveDriveKinematics;
022import org.wpilib.math.kinematics.SwerveModulePosition;
023import org.wpilib.math.kinematics.SwerveModuleVelocity;
024import org.wpilib.util.Alert;
025import org.wpilib.util.UsageReporting;
026
027import java.util.Optional;
028import java.util.concurrent.locks.Lock;
029import java.util.concurrent.locks.ReentrantLock;
030import org.littletonrobotics.junction.AutoLogOutput;
031import org.littletonrobotics.junction.Logger;
032
033// import org.wpilib.wpilibj2.command.sysid.SysIdRoutine;
034
035public class Swerve extends PeriodicSubsystem {
036    // TunerConstants doesn't include these constants, so they are declared locally
037    static final double ODOMETRY_FREQUENCY = TunerConstants.kCANBus.isNetworkFD() ? 250.0 : 100.0;
038    public static final double DRIVE_BASE_RADIUS =
039            Math.max(
040                    Math.max(
041                            Math.hypot(
042                                    TunerConstants.FrontLeft.LocationX,
043                                    TunerConstants.FrontLeft.LocationY),
044                            Math.hypot(
045                                    TunerConstants.FrontRight.LocationX,
046                                    TunerConstants.FrontRight.LocationY)),
047                    Math.max(
048                            Math.hypot(
049                                    TunerConstants.BackLeft.LocationX,
050                                    TunerConstants.BackLeft.LocationY),
051                            Math.hypot(
052                                    TunerConstants.BackRight.LocationX,
053                                    TunerConstants.BackRight.LocationY)));
054
055    static final Lock odometryLock = new ReentrantLock();
056    private final GyroIO gyroIO;
057    private final GyroIOInputsAutoLogged gyroInputs = new GyroIOInputsAutoLogged();
058    private final Module[] modules = new Module[4]; // FL, FR, BL, BR
059    // private final SysIdRoutine sysId;
060    private final Alert gyroDisconnectedAlert =
061            new Alert(
062                    "SwerveGyroDisconnect",
063                    "Disconnected gyro, using kinematics as fallback.",
064                    Alert.Level.HIGH);
065
066    private SwerveDriveKinematics kinematics = new SwerveDriveKinematics(getModuleTranslations());
067    private Rotation2d rawGyroRotation = Rotation2d.ZERO;
068    private SwerveModulePosition[] lastModulePositions = // For delta tracking
069            new SwerveModulePosition[] {
070                new SwerveModulePosition(),
071                new SwerveModulePosition(),
072                new SwerveModulePosition(),
073                new SwerveModulePosition()
074            };
075    private SwerveDrivePoseEstimator poseEstimator =
076            new SwerveDrivePoseEstimator(
077                    kinematics, rawGyroRotation, lastModulePositions, Pose2d.ZERO);
078
079    public Swerve(
080            GyroIO gyroIO,
081            ModuleIO flModuleIO,
082            ModuleIO frModuleIO,
083            ModuleIO blModuleIO,
084            ModuleIO brModuleIO) {
085        this.gyroIO = gyroIO;
086        modules[0] = new Module(flModuleIO, 0, TunerConstants.FrontLeft);
087        modules[1] = new Module(frModuleIO, 1, TunerConstants.FrontRight);
088        modules[2] = new Module(blModuleIO, 2, TunerConstants.BackLeft);
089        modules[3] = new Module(brModuleIO, 3, TunerConstants.BackRight);
090
091        // Usage reporting for swerve template
092        UsageReporting.reportUsage("RobotDrive", "RobotDriveSwerve_AdvantageKit");
093
094        // Start odometry thread
095        PhoenixOdometryThread.getInstance().start();
096
097        // Configure AutoBuilder for PathPlanner
098        // AutoBuilder.configure(
099        // this::getPose,
100        // this::setPose,
101        // this::getChassisSpeeds,
102        // this::runVelocity,
103        // new PPHolonomicDriveController(
104        // new PIDConstants(5.0, 0.0, 0.0), new PIDConstants(5.0, 0.0, 0.0)),
105        // PP_CONFIG,
106        // () -> DriverStation.getAlliance().orElse(Alliance.Blue) == Alliance.Red,
107        // this);
108        // Pathfinding.setPathfinder(new LocalADStarAK());
109        // PathPlannerLogging.setLogActivePathCallback(
110        // (activePath) -> {
111        // Logger.recordOutput("Odometry/Trajectory", activePath.toArray(new
112        // Pose2d[0]));
113        // });
114        // PathPlannerLogging.setLogTargetPoseCallback(
115        // (targetPose) -> {
116        // Logger.recordOutput("Odometry/TrajectorySetpoint", targetPose);
117        // });
118
119        // Configure SysId
120        // sysId =
121        // new SysIdRoutine(
122        // new SysIdRoutine.Config(
123        // null,
124        // null,
125        // null,
126        // (state) -> Logger.recordOutput("Drive/SysIdState", state.toString())),
127        // new SysIdRoutine.Mechanism(
128        // (voltage) -> runCharacterization(voltage.in(Volts)), null, this));
129    }
130
131    @Override
132    public void periodic() {
133        odometryLock.lock(); // Prevents odometry updates while reading data
134        gyroIO.updateInputs(gyroInputs);
135        Logger.processInputs("Drive/Gyro", gyroInputs);
136        for (var module : modules) {
137            module.periodic();
138        }
139        odometryLock.unlock();
140
141        // Stop moving when disabled
142        if (RobotState.isDisabled()) {
143            for (var module : modules) {
144                module.stop();
145            }
146        }
147
148        // Log empty setpoint states when disabled
149        if (RobotState.isDisabled()) {
150            Logger.recordOutput("SwerveStates/Setpoints", new SwerveModuleVelocity[] {});
151            Logger.recordOutput("SwerveStates/SetpointsOptimized", new SwerveModuleVelocity[] {});
152        }
153
154        // Update odometry
155        double[] sampleTimestamps =
156                modules[0].getOdometryTimestamps(); // All signals are sampled together
157        int sampleCount = sampleTimestamps.length;
158        for (int i = 0; i < sampleCount; i++) {
159            // Read wheel positions and deltas from each module
160            SwerveModulePosition[] modulePositions = new SwerveModulePosition[4];
161            SwerveModulePosition[] moduleDeltas = new SwerveModulePosition[4];
162            for (int moduleIndex = 0; moduleIndex < 4; moduleIndex++) {
163                modulePositions[moduleIndex] = modules[moduleIndex].getOdometryPositions()[i];
164                moduleDeltas[moduleIndex] =
165                        new SwerveModulePosition(
166                                modulePositions[moduleIndex].distance
167                                        - lastModulePositions[moduleIndex].distance,
168                                modulePositions[moduleIndex].angle);
169                lastModulePositions[moduleIndex] = modulePositions[moduleIndex];
170            }
171
172            // Update gyro angle
173            if (gyroInputs.connected) {
174                // Use the real gyro angle
175                rawGyroRotation = gyroInputs.odometryYawPositions[i];
176            } else {
177                // Use the angle delta from the kinematics and module deltas
178                Twist2d twist = kinematics.toTwist2d(moduleDeltas);
179                rawGyroRotation = rawGyroRotation.plus(new Rotation2d(twist.dtheta));
180            }
181
182            // Apply update
183            poseEstimator.updateWithTime(sampleTimestamps[i], rawGyroRotation, modulePositions);
184        }
185
186        // Update gyro alert
187        gyroDisconnectedAlert.set(
188                !gyroInputs.connected && GlobalSettings.CURRENT_MODE != RobotMode.SIM);
189    }
190
191    @Override
192    public void periodicAfterScheduler() {
193        for (Module module : modules) {
194            module.periodicAfterScheduler();
195        }
196    }
197
198    /**
199     * Runs the drive at the desired velocity.
200     *
201     * @param speeds Speeds in meters/sec
202     */
203    public void runVelocity(ChassisVelocities speeds) {
204        // Calculate module setpoints
205        ChassisVelocities discreteSpeeds = speeds.discretize(0.02);
206        SwerveModuleVelocity[] setpointStates = kinematics.toSwerveModuleVelocities(discreteSpeeds);
207        setpointStates =
208                SwerveDriveKinematics.desaturateWheelVelocities(
209                        setpointStates, TunerConstants.kSpeedAt12Volts);
210
211        // Log unoptimized setpoints and setpoint speeds
212        Logger.recordOutput("SwerveStates/Setpoints", setpointStates);
213        Logger.recordOutput("SwerveChassisSpeeds/Setpoints", discreteSpeeds);
214
215        // Send setpoints to modules
216        for (int i = 0; i < 4; i++) {
217            modules[i].runSetpoint(setpointStates[i]);
218        }
219
220        // Log optimized setpoints (runSetpoint mutates each state)
221        Logger.recordOutput("SwerveStates/SetpointsOptimized", setpointStates);
222    }
223
224    /** Runs the drive in a straight line with the specified drive output. */
225    public void runCharacterization(double output) {
226        for (int i = 0; i < 4; i++) {
227            modules[i].runCharacterization(output);
228        }
229    }
230
231    /** Stops the drive. */
232    public void stop() {
233        runVelocity(new ChassisVelocities());
234    }
235
236    /**
237     * Stops the drive and turns the modules to an X arrangement to resist movement. The modules
238     * will return to their normal orientations the next time a nonzero velocity is requested.
239     */
240    public void stopWithX() {
241        Rotation2d[] headings = new Rotation2d[4];
242        for (int i = 0; i < 4; i++) {
243            Optional<Rotation2d> rotation = getModuleTranslations()[i].getAngle();
244            headings[i] = rotation.isPresent() ? rotation.get() : headings[i];
245        }
246        kinematics.resetHeadings(headings);
247        stop();
248    }
249
250    /** Returns a command to run a quasistatic test in the specified direction. */
251    // public Command sysIdQuasistatic(SysIdRoutine.Direction direction) {
252    //     return run(() -> runCharacterization(0.0))
253    //             .withTimeout(1.0)
254    //             .andThen(sysId.quasistatic(direction));
255    // }
256
257    // /** Returns a command to run a dynamic test in the specified direction. */
258    // public Command sysIdDynamic(SysIdRoutine.Direction direction) {
259    //     return run(() ->
260    // runCharacterization(0.0)).withTimeout(1.0).andThen(sysId.dynamic(direction));
261    // }
262
263    /** Returns the module states (turn angles and drive velocities) for all of the modules. */
264    @AutoLogOutput(key = "SwerveStates/Measured")
265    private SwerveModuleVelocity[] getModuleVelocities() {
266        SwerveModuleVelocity[] states = new SwerveModuleVelocity[4];
267        for (int i = 0; i < 4; i++) {
268            states[i] = modules[i].getVelocity();
269        }
270        return states;
271    }
272
273    /** Returns the module positions (turn angles and drive positions) for all of the modules. */
274    private SwerveModulePosition[] getModulePositions() {
275        SwerveModulePosition[] states = new SwerveModulePosition[4];
276        for (int i = 0; i < 4; i++) {
277            states[i] = modules[i].getPosition();
278        }
279        return states;
280    }
281
282    /** Returns the measured chassis speeds of the robot. */
283    @AutoLogOutput(key = "SwerveChassisSpeeds/Measured")
284    private ChassisVelocities getChassisSpeeds() {
285        return kinematics.toChassisVelocities(getModuleVelocities());
286    }
287
288    /** Returns the position of each module in radians. */
289    public double[] getWheelRadiusCharacterizationPositions() {
290        double[] values = new double[4];
291        for (int i = 0; i < 4; i++) {
292            values[i] = modules[i].getWheelRadiusCharacterizationPosition();
293        }
294        return values;
295    }
296
297    /** Returns the average velocity of the modules in rotations/sec (Phoenix native units). */
298    public double getFFCharacterizationVelocity() {
299        double output = 0.0;
300        for (int i = 0; i < 4; i++) {
301            output += modules[i].getFFCharacterizationVelocity() / 4.0;
302        }
303        return output;
304    }
305
306    /** Returns the current odometry pose. */
307    @AutoLogOutput(key = "Odometry/Robot")
308    public Pose2d getPose() {
309        return poseEstimator.getEstimatedPosition();
310    }
311
312    /** Returns the current odometry rotation. */
313    public Rotation2d getRotation() {
314        return getPose().getRotation();
315    }
316
317    /** Resets the current odometry pose. */
318    public void resetPose(Pose2d pose) {
319        poseEstimator.resetPosition(rawGyroRotation, getModulePositions(), pose);
320    }
321
322    public void resetRotation(Rotation2d rotation) {
323        poseEstimator.resetRotation(rotation);
324    }
325
326    public VisionConsumer visionConsumer = poseEstimator::addVisionMeasurement;
327
328    /** Returns the maximum linear speed in meters per sec. */
329    public double getMaxLinearSpeedMetersPerSec() {
330        return TunerConstants.kSpeedAt12Volts.in(MetersPerSecond);
331    }
332
333    /** Returns the maximum angular speed in radians per sec. */
334    public double getMaxAngularSpeedRadPerSec() {
335        return getMaxLinearSpeedMetersPerSec() / DRIVE_BASE_RADIUS;
336    }
337
338    /** Returns an array of module translations. */
339    public static Translation2d[] getModuleTranslations() {
340        return new Translation2d[] {
341            new Translation2d(
342                    TunerConstants.FrontLeft.LocationX, TunerConstants.FrontLeft.LocationY),
343            new Translation2d(
344                    TunerConstants.FrontRight.LocationX, TunerConstants.FrontRight.LocationY),
345            new Translation2d(TunerConstants.BackLeft.LocationX, TunerConstants.BackLeft.LocationY),
346            new Translation2d(
347                    TunerConstants.BackRight.LocationX, TunerConstants.BackRight.LocationY)
348        };
349    }
350}