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.util.swerve.swerveinput;
007
008import com.stuypulse.robot.constants.Settings;
009
010import edu.wpi.first.math.MathUtil;
011import edu.wpi.first.math.filter.LinearFilter;
012import edu.wpi.first.math.filter.SlewRateLimiter;
013import edu.wpi.first.math.geometry.Translation2d;
014import edu.wpi.first.wpilibj2.command.button.CommandXboxController;
015
016import static edu.wpi.first.units.Units.*;
017
018/**
019 * <h2>DriveInputProcessor</h2>
020 * <p>Class for processing driver input for the drivetrain. It's intended to replace StuyLib's VStream 
021 * and Filters and use pure WPILib and Java.</p>
022 * <p>This will be used to be as close as the original StuyLib implementation as possible</p>
023 * 
024 * It will be used for
025 * <ul>
026 *   <li>Applying deadbands</li>
027 *   <li>Applying power curves</li>
028 *   <li>Applying rate limits</li>
029 *   <li>Applying a low pass filter</li>
030 * </ul>
031 * 
032 * <p>To use within a command, create an instance of this class and call the {@link #update()} method in the {@link edu.wpi.first.wpilibj2.command.Command#execute()} method of the command.</p>
033 * <p>To get the processed speed, call the {@link #get()} method.</p>
034 */
035public class DriveInputProcessor {
036    private final CommandXboxController controller;
037    private final double deadband;
038    private final double clamp;
039    private final double power;
040    /** The max velocity in meters per second/ */
041    private final double maxVelocity;
042
043    private final SlewRateLimiter xRateLimiter;
044    private final SlewRateLimiter yRateLimiter;
045
046    private final LinearFilter xLowPassFilter;
047    private final LinearFilter yLowPassFilter;
048
049    /**
050     * The speed vector after the full processing and filtering.
051     */
052    private Translation2d processedSpeed;
053
054    /**
055     * <h4>Constructor for the DriveInputProcessor</h4>
056     * <p>Creates a new DriveInputProcessor with the specified parameters and processes the input periodically.</p>
057     * @param controller The CommandXboxController to get driver input from
058     * @param deadband The deadband to apply to the input (0-1)
059     * @param clamp The maximum magnitude of the input vector (0-1)
060     * @param power The power to apply to the input
061     * @param maxVelocity The maximum velocity to scale the input to (m/s)
062     * @param maxAcceleration The maximum acceleration to apply to the input (m/s^2)
063     * @param rc The time constant for the low pass filter (seconds)
064     */
065    public DriveInputProcessor(CommandXboxController controller, double deadband, double clamp, double power, double maxVelocity, double maxAcceleration, double rc) {
066        this.controller = controller;
067        this.deadband = deadband;
068        this.clamp = clamp;
069        this.power = power;
070        this.maxVelocity = maxVelocity;
071
072        this.xRateLimiter = new SlewRateLimiter(maxAcceleration);
073        this.yRateLimiter = new SlewRateLimiter(maxAcceleration);
074        this.xLowPassFilter = LinearFilter.singlePoleIIR(rc, Settings.DT.in(Seconds));
075        this.yLowPassFilter = LinearFilter.singlePoleIIR(rc, Settings.DT.in(Seconds));
076
077        this.processedSpeed = Translation2d.kZero;
078    }
079
080    /**
081     * Read the raw joystick axes and store them as a velocity vector in {@link #processedSpeed}
082     *
083     * @return This instance of the class
084     */
085    private DriveInputProcessor getDriverInputAsVelocity() {
086        this.processedSpeed = new Translation2d(-controller.getLeftY(), -controller.getLeftX());
087        return this;
088    }
089
090    /**
091     * Apply a deadband on the X and Y axes to the current {@link #processedSpeed} vector
092     * using {@link MathUtil#applyDeadband}
093     *
094     * @return This instance of the class
095     */
096    private DriveInputProcessor applyDeadband() {
097        double deadbandX = MathUtil.applyDeadband(this.processedSpeed.getX(), deadband);
098        double deadbandY = MathUtil.applyDeadband(this.processedSpeed.getY(), deadband);
099        this.processedSpeed = new Translation2d(deadbandX, deadbandY);
100        return this;
101    }
102
103    /** 
104     * Limits the magnitude of the current {@link #processedSpeed} vector to a maximum
105     * 
106     * @return This instance of the class
107    */
108    private DriveInputProcessor applyClamp() {
109        double magnitude = this.processedSpeed.getNorm();
110
111        if (magnitude > clamp) {
112            this.processedSpeed = this.processedSpeed.times(clamp / magnitude);
113        }
114
115        return this;
116    }
117
118    /**
119     * Apply a power curve to the current {@link #processedSpeed} vector.
120     * The vector's magnitude is raised according to the configured power
121     * while preserving direction.
122     *
123     * @return This instance of the class
124     */
125    private DriveInputProcessor applyPowerCurve() {
126        this.processedSpeed = this.processedSpeed.times(Math.pow(this.processedSpeed.getNorm(), power - 1));
127        return this;
128    }
129
130    /**
131     * Scale the {@link #processedSpeed} vector to the robot's maximum velocity.
132     * Multiplies the current vector by maxVelocity (m/s).
133     *
134     * @return This instance of the class
135     */
136    private DriveInputProcessor applyScalingToMaxVelocity() {
137        this.processedSpeed = this.processedSpeed.times(maxVelocity);
138        return this;
139    }
140
141    /**
142     * Applies a rate limit to {@link #processedSpeed} using 
143     * a {@link edu.wpi.first.math.filter.SlewRateLimiter} for each axis. 
144     * This limits the rate of change of the velocity vector to the 
145     * max acceleration (m/s^2).
146     *
147     * @return This instance of the class
148     */
149    private DriveInputProcessor applyRateLimit() {
150        double limitedX = xRateLimiter.calculate(this.processedSpeed.getX());
151        double limitedY = yRateLimiter.calculate(this.processedSpeed.getY());
152        this.processedSpeed = new Translation2d(limitedX, limitedY);
153        return this;
154    }
155    
156    /**
157     * Applies a low pass filter to {@link #processedSpeed} using
158     * a {@link edu.wpi.first.math.filter.LinearFilter} for each axis. 
159     * This smooths out the velocity vector and reduces noise.
160     *
161     * @return This instance of the class
162     */
163    private DriveInputProcessor applyLowPassFilter() {
164        double filteredX = xLowPassFilter.calculate(this.processedSpeed.getX());
165        double filteredY = yLowPassFilter.calculate(this.processedSpeed.getY());
166        this.processedSpeed = new Translation2d(filteredX, filteredY);
167        return this;
168    }
169    
170
171    /**
172     * Update method that processes all filters and updates the filtered speed.
173     * This should be called within the start of the {@link edu.wpi.first.wpilibj2.command.Command#execute()} 
174     * method of the command using DriveInputProcessor.
175     */
176    public void update() {
177        getDriverInputAsVelocity()
178            .applyDeadband()
179            .applyPowerCurve()
180            .applyClamp()
181            .applyScalingToMaxVelocity()
182            .applyRateLimit()
183            .applyLowPassFilter();
184    }
185
186    /**
187     * Get the processed speed.
188     * @return A Translation2d representing the processed speed
189     */
190    public Translation2d get() {
191        return processedSpeed;
192    }
193}