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.util.talonfx.config;
006
007import static org.wpilib.units.Units.*;
008
009import org.wpilib.units.measure.*;
010
011import com.ctre.phoenix6.configs.*;
012import com.ctre.phoenix6.hardware.TalonFX;
013import com.ctre.phoenix6.signals.*;
014
015/** Wrapper class for configuring TalonFX motors */
016public class TalonFXConfig {
017    private final TalonFXConfiguration configuration = new TalonFXConfiguration();
018    private final Slot0Configs slot0Configs = new Slot0Configs();
019    private final Slot1Configs slot1Configs = new Slot1Configs();
020    private final Slot2Configs slot2Configs = new Slot2Configs();
021    private final MotorOutputConfigs motorOutputConfigs = new MotorOutputConfigs();
022    private final ClosedLoopRampsConfigs closedLoopRampsConfigs = new ClosedLoopRampsConfigs();
023    private final OpenLoopRampsConfigs openLoopRampsConfigs = new OpenLoopRampsConfigs();
024    private final CurrentLimitsConfigs currentLimitsConfigs = new CurrentLimitsConfigs();
025    private final FeedbackConfigs feedbackConfigs = new FeedbackConfigs();
026    private final MotionMagicConfigs motionMagicConfigs = new MotionMagicConfigs();
027    private final SoftwareLimitSwitchConfigs softwareLimitSwitchConfigs =
028            new SoftwareLimitSwitchConfigs();
029    private final ClosedLoopGeneralConfigs closedLoopGeneralConfigs =
030            new ClosedLoopGeneralConfigs();
031    private final VoltageConfigs voltageConfigs = new VoltageConfigs();
032    private final TorqueCurrentConfigs torqueCurrentConfigs = new TorqueCurrentConfigs();
033
034    private final double[] lastKP = new double[3];
035    private final double[] lastKI = new double[3];
036    private final double[] lastKD = new double[3];
037    private final double[] lastKS = new double[3];
038    private final double[] lastKV = new double[3];
039    private final double[] lastKA = new double[3];
040
041    public void configure(TalonFX motor) {
042        TalonFXConfiguration defaultConfig = new TalonFXConfiguration();
043        motor.getConfigurator().apply(defaultConfig);
044
045        motor.getConfigurator().apply(configuration);
046    }
047
048    public TalonFXConfiguration getConfiguration() {
049        return this.configuration;
050    }
051
052    public void updateGainsConfig(
053            TalonFX motor,
054            int slot,
055            double kP,
056            double kI,
057            double kD,
058            double kS,
059            double kV,
060            double kA) {
061        if (slot != 0 && slot != 1 && slot != 2) {
062            return;
063        }
064
065        final boolean changed =
066                kP != lastKP[slot]
067                        || kI != lastKI[slot]
068                        || kD != lastKD[slot]
069                        || kS != lastKS[slot]
070                        || kV != lastKV[slot]
071                        || kA != lastKA[slot];
072
073        if (!changed) {
074            return;
075        }
076
077        final SlotConfigs gainConfig =
078                new SlotConfigs().withKP(kP).withKI(kI).withKD(kD).withKS(kS).withKV(kV).withKA(kA);
079
080        gainConfig.SlotNumber = slot;
081
082        motor.getConfigurator().apply(gainConfig);
083
084        lastKP[slot] = kP;
085        lastKI[slot] = kI;
086        lastKD[slot] = kD;
087        lastKS[slot] = kS;
088        lastKV[slot] = kV;
089        lastKA[slot] = kA;
090
091        switch (slot) {
092            case 0 -> motor.getConfigurator().refresh(this.getConfiguration().Slot0);
093            case 1 -> motor.getConfigurator().refresh(this.getConfiguration().Slot1);
094            case 2 -> motor.getConfigurator().refresh(this.getConfiguration().Slot2);
095        }
096    }
097
098    // SLOT CONFIGS
099
100    public TalonFXConfig withPIDConstants(double kP, double kI, double kD, int slot) {
101        switch (slot) {
102            case 0:
103                slot0Configs.withKP(kP).withKI(kI).withKD(kD);
104                configuration.withSlot0(slot0Configs);
105                break;
106            case 1:
107                slot1Configs.withKP(kP).withKI(kI).withKD(kD);
108                configuration.withSlot1(slot1Configs);
109                break;
110            case 2:
111                slot2Configs.withKP(kP).withKI(kI).withKD(kD);
112                configuration.withSlot2(slot2Configs);
113                break;
114        }
115        return this;
116    }
117
118    public TalonFXConfig withFFConstants(double kS, double kV, double kA, int slot) {
119        return withFFConstants(kS, kV, kA, 0.0, slot);
120    }
121
122    public TalonFXConfig withFFConstants(double kS, double kV, double kA, double kG, int slot) {
123        switch (slot) {
124            case 0 -> {
125                slot0Configs.withKS(kS).withKV(kV).withKA(kA).withKG(kG);
126                configuration.withSlot0(slot0Configs);
127            }
128            case 1 -> {
129                slot1Configs.withKS(kS).withKV(kV).withKA(kA).withKG(kG);
130                configuration.withSlot1(slot1Configs);
131            }
132            case 2 -> {
133                slot2Configs.withKS(kS).withKV(kV).withKA(kA).withKG(kG);
134                configuration.withSlot2(slot2Configs);
135            }
136        }
137        return this;
138    }
139
140    public TalonFXConfig withStaticFeedforwardSign(
141            StaticFeedforwardSignValue staticFeedforwardSign, int slot) {
142        switch (slot) {
143            case 0 -> {
144                slot0Configs.withStaticFeedforwardSign(staticFeedforwardSign);
145                configuration.withSlot0(slot0Configs);
146            }
147            case 1 -> {
148                slot1Configs.withStaticFeedforwardSign(staticFeedforwardSign);
149                configuration.withSlot1(slot1Configs);
150            }
151            case 2 -> {
152                slot2Configs.withStaticFeedforwardSign(staticFeedforwardSign);
153                configuration.withSlot2(slot2Configs);
154            }
155        }
156
157        return this;
158    }
159
160    public TalonFXConfig withGravityType(GravityTypeValue gravityType) {
161        slot0Configs.withGravityType(gravityType);
162        slot1Configs.withGravityType(gravityType);
163        slot2Configs.withGravityType(gravityType);
164
165        configuration.withSlot0(slot0Configs);
166        configuration.withSlot1(slot1Configs);
167        configuration.withSlot2(slot2Configs);
168
169        return this;
170    }
171
172    /**
173     * Modifies this configuration's GainSchedErrorThreshold parameter, which is the position
174     * closed-loop error threshold to apply gains scheduling.
175     *
176     * @param value Gains scheduling behavior, see {@link
177     *     com.ctre.phoenix6.signals.GainSchedBehaviorValue}.
178     * @param thresholdRotations The position closed-loop error threshold for gain scheduling in
179     *     rotations.
180     * @param slot Slot to apply gains scheduling behavior to.
181     * @return Itself for method-chaining.
182     */
183    public TalonFXConfig withGainSchedBehavior(
184            GainSchedBehaviorValue value, double thresholdRotations, int slot) {
185        return this.withGainSchedBehavior(value, Rotations.of(thresholdRotations), slot);
186    }
187
188    /**
189     * Modifies this configuration's GainSchedErrorThreshold parameter, which is the position
190     * closed-loop error threshold to apply gains scheduling.
191     *
192     * @param value Gains scheduling behavior, see {@link
193     *     com.ctre.phoenix6.signals.GainSchedBehaviorValue}.
194     * @param threshold The position closed-loop error threshold for gain scheduling.
195     * @param slot Slot to apply gains scheduling behavior to.
196     * @return Itself for method-chaining.
197     */
198    public TalonFXConfig withGainSchedBehavior(
199            GainSchedBehaviorValue value, Angle threshold, int slot) {
200        closedLoopGeneralConfigs.withGainSchedErrorThreshold(threshold);
201        configuration.withClosedLoopGeneral(closedLoopGeneralConfigs);
202
203        switch (slot) {
204            case 0 -> {
205                slot0Configs.withGainSchedBehavior(value);
206                configuration.withSlot0(slot0Configs);
207            }
208            case 1 -> {
209                slot1Configs.withGainSchedBehavior(value);
210                configuration.withSlot1(slot1Configs);
211            }
212            case 2 -> {
213                slot2Configs.withGainSchedBehavior(value);
214                configuration.withSlot2(slot2Configs);
215            }
216        }
217
218        return this;
219    }
220
221    // MOTOR OUTPUT CONFIGS
222
223    public TalonFXConfig withInvertedValue(InvertedValue invertedValue) {
224        motorOutputConfigs.withInverted(invertedValue);
225
226        configuration.withMotorOutput(motorOutputConfigs);
227
228        return this;
229    }
230
231    public TalonFXConfig withNeutralMode(NeutralModeValue neutralMode) {
232        motorOutputConfigs.withNeutralMode(neutralMode);
233
234        configuration.withMotorOutput(motorOutputConfigs);
235
236        return this;
237    }
238
239    /**
240     * Modifies this configuration's velocity filter's time constant.
241     *
242     * @param filterTimeSeconds The configurable time constant in seconds of the Kalman velocity
243     *     filter.
244     * @return Itself for easier method-chaining
245     */
246    public TalonFXConfig withVelocityTimeFilter(double filterTimeSeconds) {
247        return this.withVelocityTimeFilter(Seconds.of(filterTimeSeconds));
248    }
249
250    /**
251     * Modifies this configuration's velocity filter's time constant.
252     *
253     * @param filterTime The configurable time constant of the Kalman velocity filter.
254     * @return Itself for easier method-chaining
255     */
256    public TalonFXConfig withVelocityTimeFilter(Time filterTime) {
257        feedbackConfigs.withVelocityFilterTimeConstant(filterTime);
258
259        configuration.withFeedback(feedbackConfigs);
260
261        return this;
262    }
263
264    // RAMP RATE CONFIGS
265
266    /**
267     * Modifies this configuration's open and closed loop ramp periods for the following control
268     * types:
269     *
270     * <ul>
271     *   <li>DutyCycle
272     *   <li>TorqueCurrent
273     *   <li>Voltage
274     * </ul>
275     *
276     * <br>
277     *
278     * @param rampRateSeconds Seconds for the motor to reach peak output from 0
279     * @return Itself for method-chaining
280     */
281    public TalonFXConfig withRampRate(double rampRateSeconds) {
282        return this.withRampRate(Seconds.of(rampRateSeconds));
283    }
284
285    /**
286     * Modifies this configuration's open and closed loop ramp periods for the following control
287     * types:
288     *
289     * <ul>
290     *   <li>DutyCycle
291     *   <li>TorqueCurrent
292     *   <li>Voltage
293     * </ul>
294     *
295     * <br>
296     *
297     * @param rampRate Time for the motor to reach peak output from 0
298     * @return Itself for easy method chaining
299     */
300    public TalonFXConfig withRampRate(Time rampRate) {
301        closedLoopRampsConfigs.withDutyCycleClosedLoopRampPeriod(rampRate);
302        closedLoopRampsConfigs.withTorqueClosedLoopRampPeriod(rampRate);
303        closedLoopRampsConfigs.withVoltageClosedLoopRampPeriod(rampRate);
304
305        openLoopRampsConfigs.withDutyCycleOpenLoopRampPeriod(rampRate);
306        openLoopRampsConfigs.withTorqueOpenLoopRampPeriod(rampRate);
307        openLoopRampsConfigs.withVoltageOpenLoopRampPeriod(rampRate);
308
309        configuration.withClosedLoopRamps(closedLoopRampsConfigs);
310        configuration.withOpenLoopRamps(openLoopRampsConfigs);
311
312        return this;
313    }
314
315    // CURRENT LIMIT CONFIGS
316    /**
317     * Modifies this configuration's supply current lower limit.
318     *
319     * @param currentLowerLimitAmps The amount of supply current in amperes allowed after the
320     *     regular {@link CurrentLimitsConfigs#SupplyCurrentLimit} is active for longer than the
321     *     time parameter.
322     * @param timeSeconds The amount of time in seconds that the regular {@link
323     *     CurrentLimitsConfigs#SupplyCurrentLimit} can be active before lowring it to the
324     *     currentLowerLimit parameter.
325     * @return
326     */
327    public TalonFXConfig withLowerLimitSupplyCurrent(
328            double currentLowerLimitAmps, double timeSeconds) {
329        return this.withLowerLimitSupplyCurrent(
330                Amps.of(currentLowerLimitAmps), Seconds.of(timeSeconds));
331    }
332
333    /**
334     * Modifies this configuration's supply current lower limit.
335     *
336     * @param currentLowerLimit The amount of supply current allowed after the regular {@link
337     *     CurrentLimitsConfigs#SupplyCurrentLimit} is active for longer than the time parameter.
338     * @param time The amount of time that the regular {@link
339     *     CurrentLimitsConfigs#SupplyCurrentLimit} can be active before lowring it to the
340     *     currentLowerLimit parameter.
341     * @return
342     */
343    public TalonFXConfig withLowerLimitSupplyCurrent(Current currentLowerLimit, Time time) {
344        currentLimitsConfigs.withSupplyCurrentLowerLimit(currentLowerLimit);
345        currentLimitsConfigs.withSupplyCurrentLowerTime(time);
346
347        configuration.withCurrentLimits(currentLimitsConfigs);
348
349        return this;
350    }
351
352    /**
353     * Modifies this configuration's supply current limit and enables it.
354     *
355     * @param currentLimitAmps Maximum allowed current in amperes drawn from the battery.
356     * @return Itself for easier method-chaining.
357     */
358    public TalonFXConfig withSupplyCurrentLimit(double currentLimitAmps) {
359        return this.withSupplyCurrentLimit(Amps.of(currentLimitAmps));
360    }
361
362    /**
363     * Modifies this configuration's supply current limit and enables it.
364     *
365     * @param currentLimit Maximum allowed current drawn from the battery.
366     * @return Itself for easier method-chaining.
367     */
368    public TalonFXConfig withSupplyCurrentLimit(Current currentLimit) {
369        currentLimitsConfigs.withSupplyCurrentLimit(currentLimit);
370        currentLimitsConfigs.withSupplyCurrentLimitEnable(true);
371
372        configuration.withCurrentLimits(currentLimitsConfigs);
373
374        return this;
375    }
376
377    public TalonFXConfig withSupplyCurrentLimitEnabled(boolean enabled) {
378        currentLimitsConfigs.withSupplyCurrentLimitEnable(enabled);
379
380        configuration.withCurrentLimits(currentLimitsConfigs);
381
382        return this;
383    }
384
385    /**
386     * Modifies this configuration's stator current limit and enables it.
387     *
388     * @param currentLimitAmps Amount of current in amperes allowed in the motor (motoring and regen
389     *     current).
390     * @return Itself for easier method chaining.
391     */
392    public TalonFXConfig withStatorCurrentLimit(double currentLimitAmps) {
393        return this.withStatorCurrentLimit(Amps.of(currentLimitAmps));
394    }
395
396    /**
397     * Modifies this configuration's stator current limit and enables it.
398     *
399     * @param currentLimit Amount of current allowed in the motor (motoring and regen current).
400     * @return Itself for easier method chaining.
401     */
402    public TalonFXConfig withStatorCurrentLimit(Current currentLimit) {
403        currentLimitsConfigs.withStatorCurrentLimit(currentLimit);
404        currentLimitsConfigs.withStatorCurrentLimitEnable(true);
405
406        configuration.withCurrentLimits(currentLimitsConfigs);
407
408        return this;
409    }
410
411    public TalonFXConfig withStatorCurrentLimitEnabled(boolean enabled) {
412        currentLimitsConfigs.withStatorCurrentLimitEnable(enabled);
413
414        configuration.withCurrentLimits(currentLimitsConfigs);
415
416        return this;
417    }
418
419    /**
420     * Modifies this configuration's torque current limits. Limits apply during torque current
421     * control modes.
422     *
423     * @param peakForwardTorqueCurrentAmps Maximum forward torque current output in amperes.
424     * @param peakReverseTorqueCurrentAmps Maximum reverse torque current output in amperes.
425     * @param neutralToleranceAmps Current range in amperes where requested torque current will
426     *     result in zero bridge output.
427     * @return Itself for easier method-chaining.
428     */
429    public TalonFXConfig withTorqueCurrentLimits(
430            double peakForwardTorqueCurrentAmps,
431            double peakReverseTorqueCurrentAmps,
432            double neutralToleranceAmps) {
433        return this.withTorqueCurrentLimits(
434                Amps.of(peakForwardTorqueCurrentAmps),
435                Amps.of(peakReverseTorqueCurrentAmps),
436                Amps.of(neutralToleranceAmps));
437    }
438
439    /**
440     * Modifies this configuration's torque current limits. Limits apply during torque current
441     * control modes.
442     *
443     * @param peakForwardTorqueCurrent Maximum forward torque current output.
444     * @param peakReverseTorqueCurrent Maximum reverse torque current output.
445     * @param neutralTolerance Current range where requested torque current will result in zero
446     *     bridge output.
447     * @return Itself for easier method-chaining.
448     */
449    public TalonFXConfig withTorqueCurrentLimits(
450            Current peakForwardTorqueCurrent,
451            Current peakReverseTorqueCurrent,
452            Current neutralTolerance) {
453        torqueCurrentConfigs.withPeakForwardTorqueCurrent(peakForwardTorqueCurrent);
454        torqueCurrentConfigs.withPeakReverseTorqueCurrent(peakReverseTorqueCurrent);
455        torqueCurrentConfigs.withTorqueNeutralDeadband(neutralTolerance);
456
457        configuration.withTorqueCurrent(torqueCurrentConfigs);
458
459        return this;
460    }
461
462    // VOLTAGE LIMIT CONFIGS
463
464    /**
465     * Modifies this configuration's voltage limits. Limits apply during voltage control modes.
466     *
467     * @param peakForwardVolts Maximum forward voltage output in volts.
468     * @param peakReverseVolts Maximum reverse voltage output in volts.
469     * @return Itself for easier method chaining.
470     */
471    public TalonFXConfig withVoltageLimits(double peakForwardVolts, double peakReverseVolts) {
472        return this.withVoltageLimits(Volts.of(peakForwardVolts), Volts.of(peakReverseVolts));
473    }
474
475    /**
476     * Modifies this configuration's voltage limits. Limits apply during voltage control modes.
477     *
478     * @param peakForwardVoltage Maximum forward voltage output.
479     * @param peakReverseVoltage Maximum reverse voltage output.
480     * @return Itself for easier method chaining.
481     */
482    public TalonFXConfig withVoltageLimits(Voltage peakForwardVoltage, Voltage peakReverseVoltage) {
483        voltageConfigs.withPeakForwardVoltage(peakForwardVoltage);
484        voltageConfigs.withPeakReverseVoltage(peakReverseVoltage);
485
486        configuration.withVoltage(voltageConfigs);
487
488        return this;
489    }
490
491    // SOFTWARE LIMIT CONFIGS
492    /**
493     * Modifies this configuration's software limit switch.
494     *
495     * @param forwardEnable Whether to set the motor output to neutral if its position exceeds the
496     *     forward threshold.
497     * @param reverseEnable Whether to set the motor output to neutral if its position exceeds the
498     *     reverse threshold.
499     * @param forwardThresholdRotations The threshold angle in rotations for the application of the
500     *     forward limit.
501     * @param reverseThresholdRotations The threshold angle in rotations for the application of the
502     *     reverse limit.
503     * @return Itself for easier method-chaining.
504     */
505    public TalonFXConfig withSoftLimits(
506            boolean forwardEnable,
507            boolean reverseEnable,
508            double forwardThresholdRotations,
509            double reverseThresholdRotations) {
510        return this.withSoftLimits(
511                forwardEnable,
512                reverseEnable,
513                Rotations.of(forwardThresholdRotations),
514                Rotations.of(reverseThresholdRotations));
515    }
516
517    /**
518     * Modifies this configuration's software limit switch.
519     *
520     * @param forwardEnable Whether to set the motor output to neutral if its position exceeds the
521     *     forward threshold.
522     * @param reverseEnable Whether to set the motor output to neutral if its position exceeds the
523     *     reverse threshold.
524     * @param forwardThreshold The threshold angle for the application of the forward limit.
525     * @param reverseThreshold The threshold angle for the application of the reverse limit.
526     * @return Itself for easier method-chaining.
527     */
528    public TalonFXConfig withSoftLimits(
529            boolean forwardEnable,
530            boolean reverseEnable,
531            Angle forwardThreshold,
532            Angle reverseThreshold) {
533        softwareLimitSwitchConfigs.withForwardSoftLimitEnable(forwardEnable);
534        softwareLimitSwitchConfigs.withReverseSoftLimitEnable(reverseEnable);
535        softwareLimitSwitchConfigs.withForwardSoftLimitThreshold(forwardThreshold);
536        softwareLimitSwitchConfigs.withReverseSoftLimitThreshold(reverseThreshold);
537
538        configuration.withSoftwareLimitSwitch(softwareLimitSwitchConfigs);
539
540        return this;
541    }
542
543    // MOTION MAGIC CONFIGS
544    /**
545     * Modifies this configuration's motion magic profile.
546     *
547     * @param maxVelocityRotPerSec Maximum/cruise velocity of the motion profile in rotations per
548     *     second.
549     * @param maxAccelerationRotPerSecSquared Maximum acceleration of the motion profile in
550     *     rotations per second squared.
551     * @return Itself for method-chaining.
552     */
553    public TalonFXConfig withMotionProfile(
554            double maxVelocityRotPerSec, double maxAccelerationRotPerSecSquared) {
555        return this.withMotionProfile(
556                RotationsPerSecond.of(maxVelocityRotPerSec),
557                RotationsPerSecondPerSecond.of(maxAccelerationRotPerSecSquared));
558    }
559
560    /**
561     * Modifies this configuration's motion magic profile.
562     *
563     * @param maxVelocity Maximum/cruise velocity of the motion profile.
564     * @param maxAcceleration Maximum acceleration of the motion profile.
565     * @return Itself for method-chaining.
566     */
567    public TalonFXConfig withMotionProfile(
568            AngularVelocity maxVelocity, AngularAcceleration maxAcceleration) {
569        motionMagicConfigs.withMotionMagicCruiseVelocity(maxVelocity);
570        motionMagicConfigs.withMotionMagicAcceleration(maxAcceleration);
571
572        configuration.withMotionMagic(motionMagicConfigs);
573
574        return this;
575    }
576
577    // FEEDBACK CONFIGS
578
579    public TalonFXConfig withRemoteSensor(
580            int ID, FeedbackSensorSourceValue source, double rotorToSensorRatio) {
581        feedbackConfigs.withFeedbackRemoteSensorID(ID);
582        feedbackConfigs.withFeedbackSensorSource(source);
583        feedbackConfigs.withRotorToSensorRatio(rotorToSensorRatio);
584
585        configuration.withFeedback(feedbackConfigs);
586
587        return this;
588    }
589
590    public TalonFXConfig withSensorToMechanismRatio(double sensorToMechanismRatio) {
591        feedbackConfigs.withSensorToMechanismRatio(sensorToMechanismRatio);
592
593        configuration.withFeedback(feedbackConfigs);
594
595        return this;
596    }
597}