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.CANdiConfiguration;
021import com.ctre.phoenix6.configs.TalonFXSConfiguration;
022import com.ctre.phoenix6.controls.PositionVoltage;
023import com.ctre.phoenix6.controls.VelocityVoltage;
024import com.ctre.phoenix6.controls.VoltageOut;
025import com.ctre.phoenix6.hardware.CANdi;
026import com.ctre.phoenix6.hardware.ParentDevice;
027import com.ctre.phoenix6.hardware.TalonFXS;
028import com.ctre.phoenix6.signals.BrushedMotorWiringValue;
029import com.ctre.phoenix6.signals.ExternalFeedbackSensorSourceValue;
030import com.ctre.phoenix6.signals.InvertedValue;
031import com.ctre.phoenix6.signals.MotorArrangementValue;
032import com.ctre.phoenix6.signals.NeutralModeValue;
033import com.ctre.phoenix6.swerve.SwerveModuleConstants;
034
035import java.util.Queue;
036import java.util.function.Supplier;
037
038/**
039 * Module IO implementation for Talon FXS drive motor controller, Talon FXS turn motor controller,
040 * and CANdi (PWM 1). Configured using a set of module constants from Phoenix.
041 *
042 * <p>Device configuration and other behaviors not exposed by TunerConstants can be customized here.
043 */
044public class ModuleIOTalonFXS implements ModuleIO {
045    private static void tryUntilOk(int maxAttempts, Supplier<StatusCode> command) {
046        for (int i = 0; i < maxAttempts; i++) {
047            var error = command.get();
048            if (error.isOK()) break;
049        }
050    }
051
052    // Hardware objects
053    private final TalonFXS driveTalon;
054    private final TalonFXS turnTalon;
055    private final CANdi candi;
056
057    // Voltage control requests
058    private final VoltageOut voltageRequest = new VoltageOut(0);
059    private final PositionVoltage positionVoltageRequest = new PositionVoltage(0.0);
060    private final VelocityVoltage velocityVoltageRequest = new VelocityVoltage(0.0);
061
062    // Timestamp inputs from Phoenix thread
063    private final Queue<Double> timestampQueue;
064
065    // Inputs from drive motor
066    private final StatusSignal<Angle> drivePosition;
067    private final Queue<Double> drivePositionQueue;
068    private final StatusSignal<AngularVelocity> driveVelocity;
069    private final StatusSignal<Voltage> driveAppliedVolts;
070    private final StatusSignal<Current> driveCurrent;
071
072    // Inputs from turn motor
073    private final StatusSignal<Angle> turnAbsolutePosition;
074    private final StatusSignal<Angle> turnPosition;
075    private final Queue<Double> turnPositionQueue;
076    private final StatusSignal<AngularVelocity> turnVelocity;
077    private final StatusSignal<Voltage> turnAppliedVolts;
078    private final StatusSignal<Current> turnCurrent;
079
080    // Connection debouncers
081    private final Debouncer driveConnectedDebounce =
082            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
083    private final Debouncer turnConnectedDebounce =
084            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
085    private final Debouncer turnEncoderConnectedDebounce =
086            new Debouncer(0.5, Debouncer.DebounceType.FALLING);
087
088    public ModuleIOTalonFXS(
089            SwerveModuleConstants<TalonFXSConfiguration, TalonFXSConfiguration, CANdiConfiguration>
090                    constants) {
091        driveTalon = new TalonFXS(constants.DriveMotorId, TunerConstants.kCANBus);
092        turnTalon = new TalonFXS(constants.SteerMotorId, TunerConstants.kCANBus);
093        candi = new CANdi(constants.EncoderId, TunerConstants.kCANBus);
094
095        // Configure drive motor
096        var driveConfig = constants.DriveMotorInitialConfigs;
097        driveConfig.Commutation.MotorArrangement =
098                switch (constants.DriveMotorType) {
099                    case TalonFXS_NEO_JST -> MotorArrangementValue.NEO_JST;
100                    case TalonFXS_VORTEX_JST -> MotorArrangementValue.VORTEX_JST;
101                    default -> MotorArrangementValue.Disabled;
102                };
103        driveConfig.MotorOutput.NeutralMode = NeutralModeValue.Brake;
104        driveConfig.Slot0 = constants.DriveMotorGains;
105        driveConfig.ExternalFeedback.SensorToMechanismRatio = constants.DriveMotorGearRatio;
106        driveConfig.CurrentLimits.StatorCurrentLimit = constants.SlipCurrent;
107        driveConfig.CurrentLimits.StatorCurrentLimitEnable = true;
108        driveConfig.MotorOutput.Inverted =
109                constants.DriveMotorInverted
110                        ? InvertedValue.Clockwise_Positive
111                        : InvertedValue.CounterClockwise_Positive;
112        tryUntilOk(5, () -> driveTalon.getConfigurator().apply(driveConfig, 0.25));
113        tryUntilOk(5, () -> driveTalon.setPosition(0.0, 0.25));
114
115        // Configure turn motor
116        var turnConfig = new TalonFXSConfiguration();
117        turnConfig.Commutation.MotorArrangement =
118                switch (constants.SteerMotorType) {
119                    case TalonFXS_Minion_JST -> MotorArrangementValue.Minion_JST;
120                    case TalonFXS_NEO_JST -> MotorArrangementValue.NEO_JST;
121                    case TalonFXS_VORTEX_JST -> MotorArrangementValue.VORTEX_JST;
122                    case TalonFXS_NEO550_JST -> MotorArrangementValue.NEO550_JST;
123                    case TalonFXS_Brushed_AB, TalonFXS_Brushed_AC, TalonFXS_Brushed_BC ->
124                            MotorArrangementValue.Brushed_DC;
125                    default -> MotorArrangementValue.Disabled;
126                };
127        turnConfig.Commutation.BrushedMotorWiring =
128                switch (constants.SteerMotorType) {
129                    case TalonFXS_Brushed_AC -> BrushedMotorWiringValue.Leads_A_and_C;
130                    case TalonFXS_Brushed_BC -> BrushedMotorWiringValue.Leads_B_and_C;
131                    default -> BrushedMotorWiringValue.Leads_A_and_B;
132                };
133        turnConfig.MotorOutput.NeutralMode = NeutralModeValue.Brake;
134        turnConfig.Slot0 = constants.SteerMotorGains;
135        turnConfig.ExternalFeedback.FeedbackRemoteSensorID = constants.EncoderId;
136        turnConfig.ExternalFeedback.ExternalFeedbackSensorSource =
137                switch (constants.FeedbackSource) {
138                    case RemoteCANdiPWM1 -> ExternalFeedbackSensorSourceValue.RemoteCANdiPWM1;
139                    case FusedCANdiPWM1 -> ExternalFeedbackSensorSourceValue.FusedCANdiPWM1;
140                    case SyncCANdiPWM1 -> ExternalFeedbackSensorSourceValue.SyncCANdiPWM1;
141                    default ->
142                            throw new RuntimeException(
143                                    "You have selected a turn feedback source that is not supported by the default implementation of ModuleIOTalonFXS (CANdi PWM 1). 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");
144                };
145        turnConfig.ExternalFeedback.RotorToSensorRatio = constants.SteerMotorGearRatio;
146        turnConfig.MotionMagic.MotionMagicCruiseVelocity = 100.0 / constants.SteerMotorGearRatio;
147        turnConfig.MotionMagic.MotionMagicAcceleration =
148                turnConfig.MotionMagic.MotionMagicCruiseVelocity / 0.100;
149        turnConfig.MotionMagic.MotionMagicExpo_kV = 0.12 * constants.SteerMotorGearRatio;
150        turnConfig.MotionMagic.MotionMagicExpo_kA = 0.1;
151        turnConfig.ClosedLoopGeneral.ContinuousWrap = true;
152        turnConfig.MotorOutput.Inverted =
153                constants.SteerMotorInverted
154                        ? InvertedValue.Clockwise_Positive
155                        : InvertedValue.CounterClockwise_Positive;
156        tryUntilOk(5, () -> turnTalon.getConfigurator().apply(turnConfig, 0.25));
157
158        // Configure CANdi
159        CANdiConfiguration candiConfig = constants.EncoderInitialConfigs;
160        candiConfig.PWM1.AbsoluteSensorOffset = constants.EncoderOffset;
161        candiConfig.PWM1.SensorDirection = constants.EncoderInverted;
162        candi.getConfigurator().apply(candiConfig);
163
164        // Create timestamp queue
165        timestampQueue = PhoenixOdometryThread.getInstance().makeTimestampQueue();
166
167        // Create drive status signals
168        drivePosition = driveTalon.getPosition();
169        drivePositionQueue =
170                PhoenixOdometryThread.getInstance().registerSignal(drivePosition.clone());
171        driveVelocity = driveTalon.getVelocity();
172        driveAppliedVolts = driveTalon.getMotorVoltage();
173        driveCurrent = driveTalon.getStatorCurrent();
174
175        // Create turn status signals
176        turnAbsolutePosition = candi.getPWM1Position();
177        turnPosition = turnTalon.getPosition();
178        turnPositionQueue =
179                PhoenixOdometryThread.getInstance().registerSignal(turnPosition.clone());
180        turnVelocity = turnTalon.getVelocity();
181        turnAppliedVolts = turnTalon.getMotorVoltage();
182        turnCurrent = turnTalon.getStatorCurrent();
183
184        // Configure periodic frames
185        BaseStatusSignal.setUpdateFrequencyForAll(
186                Swerve.ODOMETRY_FREQUENCY, drivePosition, turnPosition);
187        BaseStatusSignal.setUpdateFrequencyForAll(
188                50.0,
189                driveVelocity,
190                driveAppliedVolts,
191                driveCurrent,
192                turnAbsolutePosition,
193                turnVelocity,
194                turnAppliedVolts,
195                turnCurrent);
196        ParentDevice.optimizeBusUtilizationForAll(driveTalon, turnTalon);
197    }
198
199    @Override
200    public void updateInputs(ModuleIOInputs inputs) {
201        // Refresh all signals
202        var driveStatus =
203                BaseStatusSignal.refreshAll(
204                        drivePosition, driveVelocity, driveAppliedVolts, driveCurrent);
205        var turnStatus =
206                BaseStatusSignal.refreshAll(
207                        turnPosition, turnVelocity, turnAppliedVolts, turnCurrent);
208        var turnEncoderStatus = BaseStatusSignal.refreshAll(turnAbsolutePosition);
209
210        // Update drive inputs
211        inputs.driveConnected = driveConnectedDebounce.calculate(driveStatus.isOK());
212        inputs.drivePosition = drivePosition.getValue();
213        inputs.driveVelocity = driveVelocity.getValue();
214        inputs.driveAppliedVolts = driveAppliedVolts.getValue();
215        inputs.driveCurrentAmps = driveCurrent.getValue();
216
217        // Update turn inputs
218        inputs.turnConnected = turnConnectedDebounce.calculate(turnStatus.isOK());
219        inputs.turnEncoderConnected =
220                turnEncoderConnectedDebounce.calculate(turnEncoderStatus.isOK());
221        inputs.turnAbsolutePosition =
222                Rotation2d.fromRotations(turnAbsolutePosition.getValueAsDouble());
223        inputs.turnPosition = Rotation2d.fromRotations(turnPosition.getValueAsDouble());
224        inputs.turnVelocityRadPerSec = Units.rotationsToRadians(turnVelocity.getValueAsDouble());
225        inputs.turnAppliedVolts = turnAppliedVolts.getValueAsDouble();
226        inputs.turnCurrentAmps = turnCurrent.getValueAsDouble();
227
228        // Update odometry inputs
229        inputs.odometryTimestamps =
230                timestampQueue.stream().mapToDouble((Double value) -> value).toArray();
231        inputs.odometryDrivePositions =
232                drivePositionQueue.stream()
233                        .map((Double value) -> Rotations.of(value))
234                        .toArray(Angle[]::new);
235        inputs.odometryTurnPositions =
236                turnPositionQueue.stream()
237                        .map((Double value) -> Rotation2d.fromRotations(value))
238                        .toArray(Rotation2d[]::new);
239        timestampQueue.clear();
240        drivePositionQueue.clear();
241        turnPositionQueue.clear();
242    }
243
244    @Override
245    public void applyOutputs(ModuleIOOutputs outputs) {
246        switch (outputs.driveOutputMode) {
247            case VELOCITY ->
248                    driveTalon.setControl(
249                            velocityVoltageRequest.withVelocity(outputs.driveTargetVelocity));
250            case OPEN_LOOP ->
251                    driveTalon.setControl(voltageRequest.withOutput(outputs.driveTargetOpenLoop));
252        }
253        switch (outputs.turnOutputMode) {
254            case POSITION ->
255                    turnTalon.setControl(
256                            positionVoltageRequest.withPosition(
257                                    outputs.turnTargetPosition.in(Rotations)));
258            case OPEN_LOOP ->
259                    turnTalon.setControl(voltageRequest.withOutput(outputs.turnTargetOpenLoop));
260        }
261    }
262}