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.Rotations;
008
009import org.wpilib.math.filter.Debouncer;
010import org.wpilib.math.geometry.Rotation2d;
011import org.wpilib.math.util.Units;
012import org.wpilib.units.measure.Angle;
013import org.wpilib.units.measure.AngularVelocity;
014import org.wpilib.units.measure.Current;
015import org.wpilib.units.measure.Voltage;
016
017import com.ctre.phoenix6.BaseStatusSignal;
018import com.ctre.phoenix6.StatusCode;
019import com.ctre.phoenix6.StatusSignal;
020import com.ctre.phoenix6.configs.CANcoderConfiguration;
021import com.ctre.phoenix6.configs.TalonFXConfiguration;
022import com.ctre.phoenix6.controls.PositionTorqueCurrentFOC;
023import com.ctre.phoenix6.controls.PositionVoltage;
024import com.ctre.phoenix6.controls.TorqueCurrentFOC;
025import com.ctre.phoenix6.controls.VelocityTorqueCurrentFOC;
026import com.ctre.phoenix6.controls.VelocityVoltage;
027import com.ctre.phoenix6.controls.VoltageOut;
028import com.ctre.phoenix6.hardware.CANcoder;
029import com.ctre.phoenix6.hardware.ParentDevice;
030import com.ctre.phoenix6.hardware.TalonFX;
031import com.ctre.phoenix6.signals.FeedbackSensorSourceValue;
032import com.ctre.phoenix6.signals.InvertedValue;
033import com.ctre.phoenix6.signals.NeutralModeValue;
034import com.ctre.phoenix6.signals.SensorDirectionValue;
035import com.ctre.phoenix6.swerve.SwerveModuleConstants;
036
037import java.util.Queue;
038import java.util.function.Supplier;
039
040/**
041 * Module IO implementation for Talon FX drive motor controller, Talon FX turn motor controller, and
042 * CANcoder. Configured using a set of module constants from Phoenix.
043 *
044 * <p>Device configuration and other behaviors not exposed by TunerConstants can be customized here.
045 */
046public class ModuleIOTalonFX implements ModuleIO {
047    private static void tryUntilOk(int maxAttempts, Supplier<StatusCode> command) {
048        for (int i = 0; i < maxAttempts; i++) {
049            var error = command.get();
050            if (error.isOK()) break;
051        }
052    }
053
054    private final SwerveModuleConstants<
055                    TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>
056            constants;
057
058    // Hardware objects
059    private final TalonFX driveTalon;
060    private final TalonFX turnTalon;
061    private final CANcoder cancoder;
062
063    // Voltage control requests
064    private final VoltageOut voltageRequest = new VoltageOut(0);
065    private final PositionVoltage positionVoltageRequest = new PositionVoltage(0.0);
066    private final VelocityVoltage velocityVoltageRequest = new VelocityVoltage(0.0);
067
068    // Torque-current control requests
069    private final TorqueCurrentFOC torqueCurrentRequest = new TorqueCurrentFOC(0);
070    private final PositionTorqueCurrentFOC positionTorqueCurrentRequest =
071            new PositionTorqueCurrentFOC(0.0);
072    private final VelocityTorqueCurrentFOC velocityTorqueCurrentRequest =
073            new VelocityTorqueCurrentFOC(0.0);
074
075    // Timestamp inputs from Phoenix thread
076    private final Queue<Double> timestampQueue;
077
078    // Inputs from drive motor
079    private final StatusSignal<Angle> drivePosition;
080    private final Queue<Double> drivePositionQueue;
081    private final StatusSignal<AngularVelocity> driveVelocity;
082    private final StatusSignal<Voltage> driveAppliedVolts;
083    private final StatusSignal<Current> driveCurrent;
084
085    // Inputs from turn motor
086    private final StatusSignal<Angle> turnAbsolutePosition;
087    private final StatusSignal<Angle> turnPosition;
088    private final Queue<Double> turnPositionQueue;
089    private final StatusSignal<AngularVelocity> turnVelocity;
090    private final StatusSignal<Voltage> turnAppliedVolts;
091    private final StatusSignal<Current> turnCurrent;
092
093    // Connection debouncers
094    private final Debouncer driveConnectedDebounce =
095            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
096    private final Debouncer turnConnectedDebounce =
097            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
098    private final Debouncer turnEncoderConnectedDebounce =
099            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
100
101    public ModuleIOTalonFX(
102            SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>
103                    constants) {
104        this.constants = constants;
105        driveTalon = new TalonFX(constants.DriveMotorId, TunerConstants.kCANBus);
106        turnTalon = new TalonFX(constants.SteerMotorId, TunerConstants.kCANBus);
107        cancoder = new CANcoder(constants.EncoderId, TunerConstants.kCANBus);
108
109        // Configure drive motor
110        var driveConfig = constants.DriveMotorInitialConfigs;
111        driveConfig.MotorOutput.NeutralMode = NeutralModeValue.Brake;
112        driveConfig.Slot0 = constants.DriveMotorGains;
113        driveConfig.Feedback.SensorToMechanismRatio = constants.DriveMotorGearRatio;
114        driveConfig.TorqueCurrent.PeakForwardTorqueCurrent = constants.SlipCurrent;
115        driveConfig.TorqueCurrent.PeakReverseTorqueCurrent = -constants.SlipCurrent;
116        driveConfig.CurrentLimits.StatorCurrentLimit = constants.SlipCurrent;
117        driveConfig.CurrentLimits.StatorCurrentLimitEnable = true;
118        driveConfig.MotorOutput.Inverted =
119                constants.DriveMotorInverted
120                        ? InvertedValue.Clockwise_Positive
121                        : InvertedValue.CounterClockwise_Positive;
122        tryUntilOk(5, () -> driveTalon.getConfigurator().apply(driveConfig, 0.25));
123        tryUntilOk(5, () -> driveTalon.setPosition(0.0, 0.25));
124
125        // Configure turn motor
126        var turnConfig = new TalonFXConfiguration();
127        turnConfig.MotorOutput.NeutralMode = NeutralModeValue.Brake;
128        turnConfig.Slot0 = constants.SteerMotorGains;
129        turnConfig.Feedback.FeedbackRemoteSensorID = constants.EncoderId;
130        turnConfig.Feedback.FeedbackSensorSource =
131                switch (constants.FeedbackSource) {
132                    case RemoteCANcoder -> FeedbackSensorSourceValue.RemoteCANcoder;
133                    case FusedCANcoder -> FeedbackSensorSourceValue.FusedCANcoder;
134                    case SyncCANcoder -> FeedbackSensorSourceValue.SyncCANcoder;
135                    default ->
136                            throw new RuntimeException(
137                                    "You have selected a turn feedback source that is not supported by the default implementation of ModuleIOTalonFX. Please check the AdvantageKit documentation for more information on alternative configurations: https://docs.advantagekit.org/getting-started/template-projects/talonfx-swerve-template#custom-module-implementations");
138                };
139        turnConfig.Feedback.RotorToSensorRatio = constants.SteerMotorGearRatio;
140        turnConfig.MotionMagic.MotionMagicCruiseVelocity = 100.0 / constants.SteerMotorGearRatio;
141        turnConfig.MotionMagic.MotionMagicAcceleration =
142                turnConfig.MotionMagic.MotionMagicCruiseVelocity / 0.100;
143        turnConfig.MotionMagic.MotionMagicExpo_kV = 0.12 * constants.SteerMotorGearRatio;
144        turnConfig.MotionMagic.MotionMagicExpo_kA = 0.1;
145        turnConfig.ClosedLoopGeneral.ContinuousWrap = true;
146        turnConfig.MotorOutput.Inverted =
147                constants.SteerMotorInverted
148                        ? InvertedValue.Clockwise_Positive
149                        : InvertedValue.CounterClockwise_Positive;
150        tryUntilOk(5, () -> turnTalon.getConfigurator().apply(turnConfig, 0.25));
151
152        // Configure CANCoder
153        CANcoderConfiguration cancoderConfig = constants.EncoderInitialConfigs;
154        cancoderConfig.MagnetSensor.MagnetOffset = constants.EncoderOffset;
155        cancoderConfig.MagnetSensor.SensorDirection =
156                constants.EncoderInverted
157                        ? SensorDirectionValue.Clockwise_Positive
158                        : SensorDirectionValue.CounterClockwise_Positive;
159        cancoder.getConfigurator().apply(cancoderConfig);
160
161        // Create timestamp queue
162        timestampQueue = PhoenixOdometryThread.getInstance().makeTimestampQueue();
163
164        // Create drive status signals
165        drivePosition = driveTalon.getPosition();
166        drivePositionQueue =
167                PhoenixOdometryThread.getInstance().registerSignal(drivePosition.clone());
168        driveVelocity = driveTalon.getVelocity();
169        driveAppliedVolts = driveTalon.getMotorVoltage();
170        driveCurrent = driveTalon.getStatorCurrent();
171
172        // Create turn status signals
173        turnAbsolutePosition = cancoder.getAbsolutePosition();
174        turnPosition = turnTalon.getPosition();
175        turnPositionQueue =
176                PhoenixOdometryThread.getInstance().registerSignal(turnPosition.clone());
177        turnVelocity = turnTalon.getVelocity();
178        turnAppliedVolts = turnTalon.getMotorVoltage();
179        turnCurrent = turnTalon.getStatorCurrent();
180
181        // Configure periodic frames
182        BaseStatusSignal.setUpdateFrequencyForAll(
183                Swerve.ODOMETRY_FREQUENCY, drivePosition, turnPosition);
184        BaseStatusSignal.setUpdateFrequencyForAll(
185                50.0,
186                driveVelocity,
187                driveAppliedVolts,
188                driveCurrent,
189                turnAbsolutePosition,
190                turnVelocity,
191                turnAppliedVolts,
192                turnCurrent);
193        ParentDevice.optimizeBusUtilizationForAll(driveTalon, turnTalon);
194    }
195
196    @Override
197    public void updateInputs(ModuleIOInputs inputs) {
198        // Refresh all signals
199        var driveStatus =
200                BaseStatusSignal.refreshAll(
201                        drivePosition, driveVelocity, driveAppliedVolts, driveCurrent);
202        var turnStatus =
203                BaseStatusSignal.refreshAll(
204                        turnPosition, turnVelocity, turnAppliedVolts, turnCurrent);
205        var turnEncoderStatus = BaseStatusSignal.refreshAll(turnAbsolutePosition);
206
207        // Update drive inputs
208        inputs.driveConnected = driveConnectedDebounce.calculate(driveStatus.isOK());
209        inputs.drivePosition = drivePosition.getValue();
210        inputs.driveVelocity = driveVelocity.getValue();
211        inputs.driveAppliedVolts = driveAppliedVolts.getValue();
212        inputs.driveCurrentAmps = driveCurrent.getValue();
213
214        // Update turn inputs
215        inputs.turnConnected = turnConnectedDebounce.calculate(turnStatus.isOK());
216        inputs.turnEncoderConnected =
217                turnEncoderConnectedDebounce.calculate(turnEncoderStatus.isOK());
218        inputs.turnAbsolutePosition =
219                Rotation2d.fromRotations(turnAbsolutePosition.getValueAsDouble());
220        inputs.turnPosition = Rotation2d.fromRotations(turnPosition.getValueAsDouble());
221        inputs.turnVelocityRadPerSec = Units.rotationsToRadians(turnVelocity.getValueAsDouble());
222        inputs.turnAppliedVolts = turnAppliedVolts.getValueAsDouble();
223        inputs.turnCurrentAmps = turnCurrent.getValueAsDouble();
224
225        // Update odometry inputs
226        inputs.odometryTimestamps =
227                timestampQueue.stream().mapToDouble((Double value) -> value).toArray();
228        inputs.odometryDrivePositions =
229                drivePositionQueue.stream()
230                        .map((Double value) -> Rotations.of(value))
231                        .toArray(Angle[]::new);
232        inputs.odometryTurnPositions =
233                turnPositionQueue.stream()
234                        .map((Double value) -> Rotation2d.fromRotations(value))
235                        .toArray(Rotation2d[]::new);
236        timestampQueue.clear();
237        drivePositionQueue.clear();
238        turnPositionQueue.clear();
239    }
240
241    @Override
242    public void applyOutputs(ModuleIOOutputs outputs) {
243        switch (outputs.driveOutputMode) {
244            case VELOCITY -> {
245                driveTalon.setControl(
246                        switch (constants.DriveMotorClosedLoopOutput) {
247                            case Voltage ->
248                                    velocityVoltageRequest.withVelocity(
249                                            outputs.driveTargetVelocity);
250                            case TorqueCurrentFOC ->
251                                    velocityTorqueCurrentRequest.withVelocity(
252                                            outputs.driveTargetVelocity);
253                        });
254            }
255            case OPEN_LOOP -> {
256                driveTalon.setControl(
257                        switch (constants.DriveMotorClosedLoopOutput) {
258                            case Voltage -> voltageRequest.withOutput(outputs.driveTargetOpenLoop);
259                            case TorqueCurrentFOC ->
260                                    torqueCurrentRequest.withOutput(outputs.driveTargetOpenLoop);
261                        });
262            }
263        }
264        switch (outputs.turnOutputMode) {
265            case POSITION -> {
266                turnTalon.setControl(
267                        switch (constants.SteerMotorClosedLoopOutput) {
268                            case Voltage ->
269                                    positionVoltageRequest.withPosition(outputs.turnTargetPosition);
270                            case TorqueCurrentFOC ->
271                                    positionTorqueCurrentRequest.withPosition(
272                                            outputs.turnTargetPosition);
273                        });
274            }
275            case OPEN_LOOP -> {
276                turnTalon.setControl(
277                        switch (constants.SteerMotorClosedLoopOutput) {
278                            case Voltage -> voltageRequest.withOutput(outputs.turnTargetOpenLoop);
279                            case TorqueCurrentFOC ->
280                                    torqueCurrentRequest.withOutput(outputs.turnTargetOpenLoop);
281                        });
282            }
283        }
284    }
285}