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