001/************************* PROJECT RON *************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */
003/* Use of this source code is governed by an MIT-style license */
004/* that can be found in the repository LICENSE file.           */
005/***************************************************************/
006package com.stuypulse.robot.util.simulation;
007
008import static edu.wpi.first.units.Units.*;
009
010import com.ctre.phoenix6.hardware.TalonFX;
011import com.ctre.phoenix6.sim.TalonFXSimState;
012import com.stuypulse.robot.constants.Motors.TalonFXConfig;
013import edu.wpi.first.units.measure.*;
014import edu.wpi.first.wpilibj.simulation.DCMotorSim;
015import edu.wpi.first.wpilibj.simulation.ElevatorSim;
016import edu.wpi.first.wpilibj.simulation.FlywheelSim;
017import edu.wpi.first.wpilibj.simulation.SingleJointedArmSim;
018
019public class TalonFXSimulation extends TalonFX {
020
021    private interface SystemSim {
022
023        public void setInputVoltage(Voltage voltage);
024
025        public void update(double dtSeconds);
026
027        public Angle getAngularPosition();
028
029        public AngularVelocity getAngularVelocity();
030
031        public static SystemSim of(DCMotorSim sim) {
032            return new SystemSim() {
033
034                @Override
035                public void setInputVoltage(Voltage voltage) {
036                    sim.setInputVoltage(voltage.in(Volts));
037                }
038
039                @Override
040                public void update(double dt) {
041                    sim.update(dt);
042                }
043
044                @Override
045                public Angle getAngularPosition() {
046                    return Radians.of(sim.getAngularPositionRad());
047                }
048
049                @Override
050                public AngularVelocity getAngularVelocity() {
051                    return RadiansPerSecond.of(sim.getAngularVelocityRadPerSec());
052                }
053            };
054        }
055
056        public static SystemSim of(FlywheelSim sim) {
057            return new SystemSim() {
058
059                @Override
060                public void setInputVoltage(Voltage voltage) {
061                    sim.setInputVoltage(voltage.in(Volts));
062                }
063
064                @Override
065                public void update(double dt) {
066                    sim.update(dt);
067                }
068
069                @Override
070                public Angle getAngularPosition() {
071                    // flywheels don't track position
072                    return Radians.of(0);
073                }
074
075                @Override
076                public AngularVelocity getAngularVelocity() {
077                    return RadiansPerSecond.of(sim.getAngularVelocityRadPerSec());
078                }
079            };
080        }
081
082        public static SystemSim of(ElevatorSim sim) {
083            return new SystemSim() {
084
085                @Override
086                public void setInputVoltage(Voltage voltage) {
087                    sim.setInputVoltage(voltage.in(Volts));
088                }
089
090                @Override
091                public void update(double dt) {
092                    sim.update(dt);
093                }
094
095                @Override
096                public Angle getAngularPosition() {
097                    return Radians.of(sim.getPositionMeters());
098                }
099
100                @Override
101                public AngularVelocity getAngularVelocity() {
102                    return RadiansPerSecond.of(sim.getVelocityMetersPerSecond());
103                }
104            };
105        }
106
107        public static SystemSim of(SingleJointedArmSim sim) {
108            return new SystemSim() {
109
110                @Override
111                public void setInputVoltage(Voltage voltage) {
112                    sim.setInputVoltage(voltage.in(Volts));
113                }
114
115                @Override
116                public void update(double dt) {
117                    sim.update(dt);
118                }
119
120                @Override
121                public Angle getAngularPosition() {
122                    return Radians.of(sim.getAngleRads());
123                }
124
125                @Override
126                public AngularVelocity getAngularVelocity() {
127                    return RadiansPerSecond.of(sim.getVelocityRadPerSec());
128                }
129            };
130        }
131    }
132
133    private final SystemSim simMotor;
134
135    private TalonFXSimulation(int port, SystemSim adapter) {
136        super(port);
137        this.simMotor = adapter;
138    }
139
140    public TalonFXSimulation(int port, DCMotorSim sim) {
141        this(port, SystemSim.of(sim));
142    }
143
144    public TalonFXSimulation(int port, FlywheelSim sim) {
145        this(port, SystemSim.of(sim));
146    }
147
148    public TalonFXSimulation(int port, ElevatorSim sim) {
149        this(port, SystemSim.of(sim));
150    }
151
152    public TalonFXSimulation(int port, SingleJointedArmSim sim) {
153        this(port, SystemSim.of(sim));
154    }
155
156    public void configure(TalonFXConfig config) {
157        config.configure(this);
158    }
159
160    public void update(double dtSeconds) {
161        final TalonFXSimState state = this.getSimState();
162        simMotor.setInputVoltage(Volts.of(state.getMotorVoltage()));
163        simMotor.update(dtSeconds);
164        state.setRawRotorPosition(simMotor.getAngularPosition().in(Rotations));
165        state.setRotorVelocity(simMotor.getAngularVelocity().in(RotationsPerSecond));
166    }
167
168    public void update(Time dtSeconds) {
169        this.update(dtSeconds.in(Seconds));
170    }
171}