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.Radians;
008import static org.wpilib.units.Units.RadiansPerSecond;
009import static org.wpilib.units.Units.RotationsPerSecond;
010
011import com.stuypulse.robot.subsystems.swerve.ModuleIO.DriveOutputMode;
012import com.stuypulse.robot.subsystems.swerve.ModuleIO.ModuleIOOutputs;
013import com.stuypulse.robot.subsystems.swerve.ModuleIO.TurnOutputMode;
014
015import org.wpilib.math.geometry.Rotation2d;
016import org.wpilib.math.kinematics.SwerveModulePosition;
017import org.wpilib.math.kinematics.SwerveModuleVelocity;
018import org.wpilib.util.Alert;
019
020import com.ctre.phoenix6.configs.CANcoderConfiguration;
021import com.ctre.phoenix6.configs.TalonFXConfiguration;
022import com.ctre.phoenix6.swerve.SwerveModuleConstants;
023
024import org.littletonrobotics.junction.Logger;
025
026public class Module {
027    private final ModuleIO io;
028    private final ModuleIOInputsAutoLogged inputs = new ModuleIOInputsAutoLogged();
029    private final ModuleIOOutputs outputs = new ModuleIOOutputs();
030    private final int index;
031    private final SwerveModuleConstants<
032                    TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>
033            constants;
034
035    private final Alert driveDisconnectedAlert;
036    private final Alert turnDisconnectedAlert;
037    private final Alert turnEncoderDisconnectedAlert;
038    private SwerveModulePosition[] odometryPositions = new SwerveModulePosition[] {};
039
040    public Module(
041            ModuleIO io,
042            int index,
043            SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>
044                    constants) {
045        this.io = io;
046        this.index = index;
047        this.constants = constants;
048        driveDisconnectedAlert =
049                new Alert(
050                        "ModuleDriveDisconnect_" + Integer.toString(index),
051                        "Disconnected drive motor on module " + Integer.toString(index) + ".",
052                        Alert.Level.HIGH);
053        turnDisconnectedAlert =
054                new Alert(
055                        "ModuleTurnDisconnect_" + Integer.toString(index),
056                        "Disconnected turn motor on module " + Integer.toString(index) + ".",
057                        Alert.Level.HIGH);
058        turnEncoderDisconnectedAlert =
059                new Alert(
060                        "ModuleEncoderDisconnect_" + Integer.toString(index),
061                        "Disconnected turn encoder on module " + Integer.toString(index) + ".",
062                        Alert.Level.HIGH);
063    }
064
065    public void periodic() {
066        io.updateInputs(inputs);
067        Logger.processInputs("Drive/Module" + Integer.toString(index), inputs);
068
069        // Calculate positions for odometry
070        int sampleCount = inputs.odometryTimestamps.length; // All signals are sampled together
071        odometryPositions = new SwerveModulePosition[sampleCount];
072        for (int i = 0; i < sampleCount; i++) {
073            double positionMeters =
074                    inputs.odometryDrivePositions[i].in(Radians) * constants.WheelRadius;
075            Rotation2d angle = inputs.odometryTurnPositions[i];
076            odometryPositions[i] = new SwerveModulePosition(positionMeters, angle);
077        }
078
079        // Update alerts
080        driveDisconnectedAlert.set(!inputs.driveConnected);
081        turnDisconnectedAlert.set(!inputs.turnConnected);
082        turnEncoderDisconnectedAlert.set(!inputs.turnEncoderConnected);
083    }
084
085    /** Runs the module with the specified setpoint state. Mutates the state to optimize it. */
086    public void runSetpoint(SwerveModuleVelocity state) {
087        // Optimize velocity setpoint
088        state = state.optimize(getAngle());
089        state = state.cosineScale(inputs.turnPosition);
090
091        // Apply setpoints
092        outputs.driveOutputMode = DriveOutputMode.VELOCITY;
093        outputs.driveTargetVelocity = RadiansPerSecond.of(state.velocity / constants.WheelRadius);
094        outputs.turnOutputMode = TurnOutputMode.POSITION;
095        outputs.turnTargetPosition = state.angle.getMeasure();
096    }
097
098    /** Runs the module with the specified output while controlling to zero degrees. */
099    public void runCharacterization(double output) {
100        outputs.driveOutputMode = DriveOutputMode.OPEN_LOOP;
101        outputs.driveTargetOpenLoop = output;
102        outputs.turnOutputMode = TurnOutputMode.POSITION;
103        outputs.turnTargetPosition = Radians.zero();
104    }
105
106    /** Disables all outputs to motors. */
107    public void stop() {
108        outputs.driveOutputMode = DriveOutputMode.OPEN_LOOP;
109        outputs.driveTargetOpenLoop = 0.0;
110        outputs.turnOutputMode = TurnOutputMode.OPEN_LOOP;
111        outputs.turnTargetOpenLoop = 0.0;
112    }
113
114    /** Returns the current turn angle of the module. */
115    public Rotation2d getAngle() {
116        return inputs.turnPosition;
117    }
118
119    /** Returns the current drive position of the module in meters. */
120    public double getPositionMeters() {
121        return inputs.drivePosition.in(Radians) * constants.WheelRadius;
122    }
123
124    /** Returns the current drive velocity of the module in meters per second. */
125    public double getVelocityMetersPerSec() {
126        return inputs.driveVelocity.in(RadiansPerSecond) * constants.WheelRadius;
127    }
128
129    /** Returns the module position (turn angle and drive position). */
130    public SwerveModulePosition getPosition() {
131        return new SwerveModulePosition(getPositionMeters(), getAngle());
132    }
133
134    /** Returns the module velocity (turn angle and drive velocity). */
135    public SwerveModuleVelocity getVelocity() {
136        return new SwerveModuleVelocity(getVelocityMetersPerSec(), getAngle());
137    }
138
139    /** Returns the module positions received this cycle. */
140    public SwerveModulePosition[] getOdometryPositions() {
141        return odometryPositions;
142    }
143
144    /** Returns the timestamps of the samples received this cycle. */
145    public double[] getOdometryTimestamps() {
146        return inputs.odometryTimestamps;
147    }
148
149    /** Returns the module position in radians. */
150    public double getWheelRadiusCharacterizationPosition() {
151        return inputs.drivePosition.in(Radians);
152    }
153
154    /** Returns the module velocity in rotations/sec (Phoenix native units). */
155    public double getFFCharacterizationVelocity() {
156        return inputs.driveVelocity.in(RotationsPerSecond);
157    }
158
159    public void periodicAfterScheduler() {
160        io.applyOutputs(outputs);
161    }
162}