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.subsystems.shooter;
007
008import static edu.wpi.first.units.Units.*;
009
010import com.ctre.phoenix6.controls.ControlRequest;
011import com.ctre.phoenix6.controls.Follower;
012import com.ctre.phoenix6.controls.VelocityTorqueCurrentFOC;
013import com.ctre.phoenix6.controls.VoltageOut;
014import com.ctre.phoenix6.signals.MotorAlignmentValue;
015import com.stuypulse.robot.constants.Motors;
016import com.stuypulse.robot.constants.Ports;
017import com.stuypulse.robot.constants.Settings;
018import com.stuypulse.robot.util.SysId;
019import com.stuypulse.robot.util.simulation.RobotVisualizer;
020import com.stuypulse.robot.util.simulation.TalonFXSimulation;
021import edu.wpi.first.math.system.plant.DCMotor;
022import edu.wpi.first.math.system.plant.LinearSystemId;
023import edu.wpi.first.units.measure.AngularVelocity;
024import edu.wpi.first.units.measure.Voltage;
025import edu.wpi.first.wpilibj.simulation.FlywheelSim;
026import edu.wpi.first.wpilibj2.command.sysid.SysIdRoutine;
027import java.util.Optional;
028
029public class ShooterSim extends Shooter {
030
031    private final FlywheelSim shooterSim;
032
033    private final TalonFXSimulation shooterMotorLeft;
034
035    private final TalonFXSimulation shooterMotorCenter;
036
037    private final TalonFXSimulation shooterMotorRight;
038
039    private final VelocityTorqueCurrentFOC shooterController;
040
041    private final Follower shooterFollowerController;
042
043    private Optional<Voltage> voltageOverride;
044
045    public ShooterSim() {
046        shooterSim = new FlywheelSim(
047                LinearSystemId.createFlywheelSystem(
048                        DCMotor.getKrakenX60(3),
049                        Settings.Shooter.J.in(KilogramSquareMeters),
050                        Settings.Shooter.GEAR_RATIO),
051                DCMotor.getKrakenX60(3));
052        shooterMotorRight = new TalonFXSimulation(Ports.Shooter.SHOOTER_MOTOR_RIGHT, shooterSim);
053        shooterMotorCenter = new TalonFXSimulation(Ports.Shooter.SHOOTER_MOTOR_CENTER, shooterSim);
054        shooterMotorLeft = new TalonFXSimulation(Ports.Shooter.SHOOTER_MOTOR_LEFT, shooterSim);
055        // leader
056        Motors.Shooter.SHOOTER_MOTOR_RIGHT.configure(shooterMotorRight);
057        Motors.Shooter.SHOOTER_MOTOR_CENTER.configure(shooterMotorCenter);
058        Motors.Shooter.SHOOTER_MOTOR_LEFT.configure(shooterMotorLeft);
059        shooterController = new VelocityTorqueCurrentFOC(getState().getTargetAngularVelocity());
060        shooterFollowerController = new Follower(shooterMotorRight.getDeviceID(), MotorAlignmentValue.Opposed);
061        shooterMotorCenter.setControl(shooterFollowerController);
062        shooterMotorLeft.setControl(shooterFollowerController);
063        voltageOverride = Optional.empty();
064    }
065
066    @Override
067    public AngularVelocity getCurrentAngularVelocity() {
068        return shooterMotorRight.getVelocity().getValue();
069    }
070
071    @Override
072    protected void stopMotors() {
073        shooterMotorRight.stopMotor();
074        shooterMotorCenter.stopMotor();
075        shooterMotorLeft.stopMotor();
076        shooterMotorCenter.setControl(shooterFollowerController);
077        shooterMotorLeft.setControl(shooterFollowerController);
078    }
079
080    @Override
081    public void setVoltageOverride(Voltage voltage) {
082        this.voltageOverride = Optional.of(voltage);
083    }
084
085    private ControlRequest getShooterControl() {
086        if (voltageOverride.isPresent()) {
087            return new VoltageOut(voltageOverride.get());
088        }
089        return shooterController.withVelocity(getState().getTargetAngularVelocity().in(RotationsPerSecond));
090    }
091
092    @Override
093    public void periodic() {
094        if (Settings.EnabledSubsystems.SHOOTER.get()) {
095            shooterMotorRight.setControl(getShooterControl());
096            // leader first
097        } else {
098            stopMotors();
099        }
100        
101        shooterMotorRight.update(Settings.DT);
102        shooterMotorCenter.update(Settings.DT);
103        shooterMotorLeft.update(Settings.DT);
104        RobotVisualizer.getInstance().updateShooter(getCurrentAngularVelocity());
105        super.periodic();
106    }
107
108    public SysIdRoutine getShooterSysIdRoutine() {
109        return SysId.getRoutine(
110                Settings.Shooter.RAMP_RATE,
111                Settings.Shooter.STEP_VOLTAGE,
112                "Shooter",
113                voltage -> setVoltageOverride(voltage),
114                () -> shooterMotorLeft.getPosition().getValue(),
115                () -> shooterMotorLeft.getVelocity().getValue(),
116                () -> shooterMotorLeft.getMotorVoltage().getValue(),
117                getInstance());
118    }
119}