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.sim;
006
007import static org.wpilib.units.Units.*;
008
009import org.wpilib.simulation.DCMotorSim;
010import org.wpilib.simulation.ElevatorSim;
011import org.wpilib.simulation.FlywheelSim;
012import org.wpilib.simulation.SingleJointedArmSim;
013import org.wpilib.units.measure.*;
014
015/**
016 *
017 *
018 * <h2>System Simulation</h2>
019 *
020 * <p>This interface defines the contract for a system simulation, which can be used to abstract the
021 * behavior of all WPILib linear system sims in one interface.
022 */
023public interface SystemSim<T> {
024    /**
025     * Sets the input voltage for the simulation.
026     *
027     * @param voltage The input voltage to set for the simulation.
028     */
029    void setInputVoltage(Voltage voltage);
030
031    /**
032     * Updates the simulation
033     *
034     * @param dt The time step to advance the simulation
035     */
036    void update(Time dt);
037
038    /**
039     * Gets the current position of the mechanism in the simulation.
040     *
041     * @return The current position of the mechanism.
042     */
043    Angle getMechanismPosition();
044
045    /**
046     * Gets the current velocity of the mechanism in the simulation.
047     *
048     * @return The current velocity of the mechanism.
049     */
050    AngularVelocity getMechanismVelocity();
051
052    /**
053     * Gets the wrapped WPILib linear system simulation instance.
054     *
055     * @return The wrapped WPILib linear system simulation instance.
056     */
057    T getLinearSystemSim();
058
059    static SystemSim<DCMotorSim> of(DCMotorSim dcMotorSim) {
060        return new SystemSim<DCMotorSim>() {
061            @Override
062            public void setInputVoltage(Voltage voltage) {
063                dcMotorSim.setInputVoltage(voltage.in(Volts));
064            }
065
066            @Override
067            public void update(Time dt) {
068                dcMotorSim.update(dt.in(Seconds));
069            }
070
071            @Override
072            public Angle getMechanismPosition() {
073                return Radians.of(dcMotorSim.getAngularPosition());
074            }
075
076            @Override
077            public AngularVelocity getMechanismVelocity() {
078                return RadiansPerSecond.of(dcMotorSim.getAngularVelocity());
079            }
080
081            @Override
082            public DCMotorSim getLinearSystemSim() {
083                return dcMotorSim;
084            }
085        };
086    }
087
088    static SystemSim<FlywheelSim> of(FlywheelSim flywheelSim) {
089        return new SystemSim<FlywheelSim>() {
090            private Angle position = Rotations.zero();
091
092            @Override
093            public void setInputVoltage(Voltage voltage) {
094                flywheelSim.setInputVoltage(voltage.in(Volts));
095            }
096
097            @Override
098            public void update(Time dt) {
099                flywheelSim.update(dt.in(Seconds));
100                position =
101                        position.plus(
102                                RadiansPerSecond.of(flywheelSim.getAngularVelocity()).times(dt));
103            }
104
105            @Override
106            public Angle getMechanismPosition() {
107                return position;
108            }
109
110            @Override
111            public AngularVelocity getMechanismVelocity() {
112                return RadiansPerSecond.of(flywheelSim.getAngularVelocity());
113            }
114
115            @Override
116            public FlywheelSim getLinearSystemSim() {
117                return flywheelSim;
118            }
119        };
120    }
121
122    static SystemSim<SingleJointedArmSim> of(SingleJointedArmSim armSim) {
123        return new SystemSim<SingleJointedArmSim>() {
124            @Override
125            public void setInputVoltage(Voltage voltage) {
126                armSim.setInputVoltage(voltage.in(Volts));
127            }
128
129            @Override
130            public void update(Time dt) {
131                armSim.update(dt.in(Seconds));
132            }
133
134            @Override
135            public Angle getMechanismPosition() {
136                return Radians.of(armSim.getAngle());
137            }
138
139            @Override
140            public AngularVelocity getMechanismVelocity() {
141                return RadiansPerSecond.of(armSim.getVelocity());
142            }
143
144            @Override
145            public SingleJointedArmSim getLinearSystemSim() {
146                return armSim;
147            }
148        };
149    }
150
151    static SystemSim<ElevatorSim> of(ElevatorSim elevatorSim, Distance drumRadius) {
152        return new SystemSim<ElevatorSim>() {
153            @Override
154            public void setInputVoltage(Voltage voltage) {
155                elevatorSim.setInputVoltage(voltage.in(Volts));
156            }
157
158            @Override
159            public void update(Time dt) {
160                elevatorSim.update(dt.in(Seconds));
161            }
162
163            @Override
164            public Angle getMechanismPosition() {
165                return Radians.of(elevatorSim.getPosition() / drumRadius.in(Meters));
166            }
167
168            @Override
169            public AngularVelocity getMechanismVelocity() {
170                return RadiansPerSecond.of(elevatorSim.getVelocity() / drumRadius.in(Meters));
171            }
172
173            @Override
174            public ElevatorSim getLinearSystemSim() {
175                return elevatorSim;
176            }
177        };
178    }
179}