001/************************* PROJECT RON *************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */
003/* Use of this source code is governed by an MIT-style license */
004/* that can be found in the repository LICENSE file.           */
005/***************************************************************/
006package com.stuypulse.robot.commands.swerve.PIDtoPose;
007
008import static edu.wpi.first.units.Units.*;
009
010import com.stuypulse.robot.Robot;
011import com.stuypulse.robot.constants.Field;
012import com.stuypulse.robot.constants.Gains.Swerve.Alignment;
013import com.stuypulse.robot.constants.Settings;
014import com.stuypulse.robot.subsystems.swerve.CommandSwerveDrivetrain;
015import com.stuypulse.robot.util.TranslationMotionProfile;
016
017import edu.wpi.first.math.controller.HolonomicDriveController;
018import edu.wpi.first.math.controller.PIDController;
019import edu.wpi.first.math.controller.ProfiledPIDController;
020import edu.wpi.first.math.filter.Debouncer;
021import edu.wpi.first.math.filter.Debouncer.DebounceType;
022import edu.wpi.first.math.filter.LinearFilter;
023
024import dev.doglog.DogLog;
025import edu.wpi.first.math.geometry.Pose2d;
026import edu.wpi.first.math.geometry.Rotation2d;
027import edu.wpi.first.math.geometry.Translation2d;
028import edu.wpi.first.math.kinematics.ChassisSpeeds;
029import edu.wpi.first.math.trajectory.TrapezoidProfile;
030import edu.wpi.first.wpilibj.smartdashboard.FieldObject2d;
031import edu.wpi.first.wpilibj2.command.Command;
032
033import java.util.function.BooleanSupplier;
034import java.util.function.Supplier;
035
036public class SwerveDrivePIDToPose extends Command {
037
038        private final CommandSwerveDrivetrain swerve;
039
040    private final HolonomicDriveController controller;
041
042        private final Supplier<Pose2d> targetPose;
043
044    // FILTERS
045    private final LinearFilter lowPass;
046    private final Debouncer debounceRC;
047
048    private double maxVelocity;
049
050        private double maxAcceleration;
051
052        private boolean isMotionProfiled;
053
054    private final BooleanSupplier isAligned;
055
056        private final FieldObject2d targetPose2d;
057
058        private Number xTolerance;
059
060        private Number yTolerance;
061
062        private Number thetaTolerance;
063
064        private Number maxVelocityWhenAligned;
065
066    private Supplier<Translation2d> translationSetpoint;
067
068        private Supplier<Boolean> canEnd;
069
070    public SwerveDrivePIDToPose(Pose2d targetPose) {
071        this(() -> targetPose);
072    }
073
074    public SwerveDrivePIDToPose(Supplier<Pose2d> targetPose) {
075        swerve = CommandSwerveDrivetrain.getInstance();
076        controller = new HolonomicDriveController(
077            new PIDController(Alignment.XY.kP, Alignment.XY.kI, Alignment.XY.kD),
078            new PIDController(Alignment.XY.kP, Alignment.XY.kI, Alignment.XY.kD),
079            new ProfiledPIDController(
080                Alignment.THETA.kP,
081                Alignment.THETA.kI,
082                Alignment.THETA.kD,
083                new TrapezoidProfile.Constraints(
084                    Settings.Swerve.Alignment.Constraints.DEFAULT_MAX_ANGULAR_VELOCITY,
085                    Settings.Swerve.Alignment.Constraints.DEFAULT_MAX_ANGULAR_ACCELERATION))
086        );
087        maxVelocity = Settings.Swerve.Alignment.Constraints.DEFAULT_MAX_VELOCITY;
088        maxAcceleration = Settings.Swerve.Alignment.Constraints.DEFAULT_MAX_ACCELERATION;
089        isMotionProfiled = true;
090        translationSetpoint = getNewTranslationSetpointGenerator();
091        this.targetPose = targetPose;
092        targetPose2d = Field.FIELD2D.getObject("Target Pose");
093        lowPass = LinearFilter.singlePoleIIR(0.05, 0.02);
094        isAligned = () -> isAlignedX()
095                && isAlignedY()
096                && isAlignedTheta()
097                && getVelocityError() < maxVelocityWhenAligned.doubleValue();
098        debounceRC =  new Debouncer(Settings.Swerve.Alignment.Tolerances.ALIGNMENT_DEBOUNCE.in(Seconds), DebounceType.kRising);
099        xTolerance = Settings.Swerve.Alignment.Tolerances.X_TOLERANCE.in(Meters);
100        yTolerance = Settings.Swerve.Alignment.Tolerances.Y_TOLERANCE.in(Meters);
101        thetaTolerance = Settings.Swerve.Alignment.Tolerances.THETA_TOLERANCE.getRadians();
102        maxVelocityWhenAligned = Settings.Swerve.Alignment.Tolerances.MAX_VELOCITY_WHEN_ALIGNED
103                .in(MetersPerSecond);
104        canEnd = () -> true;
105        addRequirements(swerve);
106    }
107
108    public SwerveDrivePIDToPose withTolerance(double x, double y, Rotation2d theta) {
109        xTolerance = x;
110        yTolerance = y;
111        thetaTolerance = theta.getRadians();
112        return this;
113    }
114
115    public SwerveDrivePIDToPose withTranslationalConstraints(
116        double maxVelocity, double maxAcceleration) {
117        this.maxVelocity = maxVelocity;
118        this.maxAcceleration = maxAcceleration;
119        return this;
120    }
121
122    public SwerveDrivePIDToPose withoutMotionProfile() {
123        this.isMotionProfiled = false;
124        return this;
125    }
126
127    public SwerveDrivePIDToPose withCanEnd(Supplier<Boolean> canEnd) {
128        this.canEnd = canEnd;
129        return this;
130    }
131
132    private Supplier<Translation2d> getNewTranslationSetpointGenerator() {
133        if (!isMotionProfiled) {
134            return () -> targetPose.get().getTranslation();
135        } else {
136            TranslationMotionProfile profile = new TranslationMotionProfile(
137                this.maxVelocity,
138                this.maxAcceleration,
139                swerve.getPose().getTranslation(),
140                new Translation2d());
141
142            return () -> profile.get(targetPose.get().getTranslation());
143        }
144    }
145
146    private double getVelocityError() {        
147        return Math.abs(lowPass.calculate(
148            new Translation2d(
149                controller.getXController().getError(),
150                controller.getYController().getError())
151                .getNorm()));
152    }
153
154    @Override
155    public void initialize() {
156        translationSetpoint = getNewTranslationSetpointGenerator();
157    }
158
159    public boolean isAlignedX() {
160        return Math.abs(targetPose.get().getX() - swerve.getPose().getX()) < xTolerance.doubleValue();
161    }
162
163    public boolean isAlignedY() {
164        return Math.abs(targetPose.get().getY() - swerve.getPose().getY()) < yTolerance.doubleValue();
165    }
166
167    public boolean isAlignedTheta() {
168        return Math.abs(
169                targetPose.get().getRotation().minus(swerve.getPose().getRotation())
170                        .getRadians()) < thetaTolerance.doubleValue();
171    }
172
173    @Override
174    public void execute() {
175        targetPose2d.setPose(
176                Robot.isBlue() ? targetPose.get()
177                        : Field.transformToOppositeAlliance(targetPose.get()));
178        final ChassisSpeeds output = controller.calculate(
179            swerve.getPose(),
180            new Pose2d(translationSetpoint.get(), targetPose.get().getRotation()),
181            0,
182            targetPose.get().getRotation());
183        swerve.setControl(
184                swerve
185                        .getRobotCentricSwerveRequest()
186                        .withVelocityX(output.vxMetersPerSecond)
187                        .withVelocityY(output.vyMetersPerSecond)
188                        .withRotationalRate(output.omegaRadiansPerSecond));
189        DogLog.log("Alignment/Target x", targetPose.get().getX());
190        DogLog.log("Alignment/Target y", targetPose.get().getY());
191        DogLog.log("Alignment/Target Angle", targetPose.get().getRotation().getDegrees());
192        DogLog.log(
193                "Alignment/Target Velocity Robot Relative X (m/s)",
194                output.vxMetersPerSecond);
195        DogLog.log(
196                "Alignment/Target Velocity Robot Relative Y (m/s)",
197                output.vyMetersPerSecond);
198        DogLog.log(
199                "Alignment/Target Angular Velocity (rad/s)",
200                output.omegaRadiansPerSecond);
201        DogLog.log("Alignment/Is Aligned", this.isAligned.getAsBoolean());
202        DogLog.log("Alignment/Is Aligned X", isAlignedX());
203        DogLog.log("Alignment/Is Aligned Y", isAlignedY());
204        DogLog.log("Alignment/Is Aligned Theta", isAlignedTheta());
205    }
206
207    @Override
208    public boolean isFinished() {
209        return debounceRC.calculate(this.isAligned.getAsBoolean()) && canEnd.get();
210    }
211
212    @Override
213    public void end(boolean interrupted) {
214        swerve.setControl(
215                swerve
216                        .getFieldCentricSwerveRequest()
217                        .withVelocityX(0)
218                        .withVelocityY(0)
219                        .withRotationalRate(0));
220        Field.clearFieldObject(targetPose2d);
221    }
222}