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.Amps;
008import static org.wpilib.units.Units.Radians;
009import static org.wpilib.units.Units.RadiansPerSecond;
010import static org.wpilib.units.Units.Volts;
011
012import org.wpilib.math.controller.PIDController;
013import org.wpilib.math.geometry.Rotation2d;
014import org.wpilib.math.system.DCMotor;
015import org.wpilib.math.system.Models;
016import org.wpilib.math.util.Units;
017import org.wpilib.simulation.DCMotorSim;
018import org.wpilib.system.Timer;
019import org.wpilib.units.measure.Angle;
020
021import com.ctre.phoenix6.configs.CANcoderConfiguration;
022import com.ctre.phoenix6.configs.TalonFXConfiguration;
023import com.ctre.phoenix6.swerve.SwerveModuleConstants;
024
025/**
026 * Physics sim implementation of module IO. The sim models are configured using a set of module
027 * constants from Phoenix. Simulation is always based on voltage control.
028 */
029public class ModuleIOSim implements ModuleIO {
030    // TunerConstants doesn't support separate sim constants, so they are declared
031    // locally
032    private static final double DRIVE_KP = 0.05;
033    private static final double DRIVE_KD = 0.0;
034    private static final double DRIVE_KS = 0.0;
035    private static final double DRIVE_KV_ROT =
036            0.91035; // Same units as TunerConstants: (volt * secs) / rotation
037    private static final double DRIVE_KV = 1.0 / Units.rotationsToRadians(1.0 / DRIVE_KV_ROT);
038    private static final double TURN_KP = 8.0;
039    private static final double TURN_KD = 0.0;
040    private static final DCMotor DRIVE_GEARBOX = DCMotor.getKrakenX60Foc(1);
041    private static final DCMotor TURN_GEARBOX = DCMotor.getKrakenX60Foc(1);
042
043    private final DCMotorSim driveSim;
044    private final DCMotorSim turnSim;
045
046    private boolean driveClosedLoop = false;
047    private boolean turnClosedLoop = false;
048    private PIDController driveController = new PIDController(DRIVE_KP, 0, DRIVE_KD);
049    private PIDController turnController = new PIDController(TURN_KP, 0, TURN_KD);
050    private double driveFFVolts = 0.0;
051    private double driveAppliedVolts = 0.0;
052    private double turnAppliedVolts = 0.0;
053
054    public ModuleIOSim(
055            SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>
056                    constants) {
057        // Create drive and turn sim models
058        driveSim =
059                new DCMotorSim(
060                        Models.singleJointedArmFromPhysicalConstants( // Recommended model by
061                                // documentation
062                                DRIVE_GEARBOX,
063                                constants.DriveInertia,
064                                constants.DriveMotorGearRatio),
065                        DRIVE_GEARBOX);
066        turnSim =
067                new DCMotorSim(
068                        Models.singleJointedArmFromPhysicalConstants(
069                                TURN_GEARBOX,
070                                constants.SteerInertia,
071                                constants.SteerMotorGearRatio),
072                        TURN_GEARBOX);
073
074        // Enable wrapping for turn PID
075        turnController.enableContinuousInput(-Math.PI, Math.PI);
076    }
077
078    @Override
079    public void updateInputs(ModuleIOInputs inputs) {
080        // Run closed-loop control
081        if (driveClosedLoop) {
082            driveAppliedVolts =
083                    driveFFVolts + driveController.calculate(driveSim.getAngularVelocity());
084        } else {
085            driveController.reset();
086        }
087        if (turnClosedLoop) {
088            turnAppliedVolts = turnController.calculate(turnSim.getAngularPosition());
089        } else {
090            turnController.reset();
091        }
092
093        // Update simulation state
094        driveSim.setInputVoltage(Math.clamp(driveAppliedVolts, -12.0, 12.0));
095        turnSim.setInputVoltage(Math.clamp(turnAppliedVolts, -12.0, 12.0));
096        driveSim.update(0.02);
097        turnSim.update(0.02);
098
099        // Update drive inputs
100        inputs.driveConnected = true;
101        inputs.drivePosition = Radians.of(driveSim.getAngularPosition());
102        inputs.driveVelocity = RadiansPerSecond.of(driveSim.getAngularVelocity());
103        inputs.driveAppliedVolts = Volts.of(driveAppliedVolts);
104        inputs.driveCurrentAmps = Amps.of(Math.abs(driveSim.getCurrentDraw()));
105
106        // Update turn inputs
107        inputs.turnConnected = true;
108        inputs.turnEncoderConnected = true;
109        inputs.turnAbsolutePosition = new Rotation2d(turnSim.getAngularPosition());
110        inputs.turnPosition = new Rotation2d(turnSim.getAngularPosition());
111        inputs.turnVelocityRadPerSec = turnSim.getAngularVelocity();
112        inputs.turnAppliedVolts = turnAppliedVolts;
113        inputs.turnCurrentAmps = Math.abs(turnSim.getCurrentDraw());
114
115        // Update odometry inputs (50Hz because high-frequency odometry in sim doesn't
116        // matter)
117        inputs.odometryTimestamps = new double[] {Timer.getMonotonicTimestamp()};
118        inputs.odometryDrivePositions = new Angle[] {inputs.drivePosition};
119        inputs.odometryTurnPositions = new Rotation2d[] {inputs.turnPosition};
120    }
121
122    @Override
123    public void applyOutputs(ModuleIOOutputs outputs) {
124        switch (outputs.driveOutputMode) {
125            case VELOCITY -> {
126                driveClosedLoop = true;
127                double velocityRadPerSec = outputs.driveTargetVelocity.in(RadiansPerSecond);
128                driveFFVolts =
129                        DRIVE_KS * Math.signum(velocityRadPerSec) + DRIVE_KV * velocityRadPerSec;
130                driveController.setSetpoint(velocityRadPerSec);
131            }
132            case OPEN_LOOP -> {
133                driveClosedLoop = false;
134                driveAppliedVolts = outputs.driveTargetOpenLoop;
135            }
136        }
137        switch (outputs.turnOutputMode) {
138            case POSITION -> {
139                turnClosedLoop = true;
140                turnController.setSetpoint(outputs.turnTargetPosition.in(Radians));
141            }
142            case OPEN_LOOP -> {
143                turnClosedLoop = false;
144                turnAppliedVolts = outputs.turnTargetOpenLoop;
145            }
146        }
147    }
148}