2025RidgeScape/src/main/java/frc4388/robot/subsystems/SwerveDrive.java

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

package frc4388.robot.subsystems;

import java.util.Optional;
import java.util.function.DoubleSupplier;

import com.ctre.phoenix6.Utils;
import com.ctre.phoenix6.hardware.CANcoder;
import com.ctre.phoenix6.hardware.TalonFX;
import com.ctre.phoenix6.swerve.SwerveDrivetrain;
import com.ctre.phoenix6.swerve.SwerveRequest;

import edu.wpi.first.math.geometry.Pose2d;
import edu.wpi.first.math.geometry.Rotation2d;
import edu.wpi.first.math.geometry.Translation2d;
import edu.wpi.first.math.kinematics.ChassisSpeeds;
import edu.wpi.first.math.kinematics.SwerveModuleState;
import edu.wpi.first.networktables.GenericEntry;
import edu.wpi.first.wpilibj.shuffleboard.BuiltInLayouts;
import edu.wpi.first.wpilibj.shuffleboard.BuiltInWidgets;
import edu.wpi.first.wpilibj.shuffleboard.Shuffleboard;
import edu.wpi.first.wpilibj.shuffleboard.ShuffleboardLayout;
import edu.wpi.first.wpilibj.smartdashboard.Field2d;
import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
// import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
import frc4388.robot.Constants.SwerveDriveConstants;
import frc4388.robot.Constants.VisionConstants;
import frc4388.utility.Status;
import frc4388.utility.Subsystem;
import frc4388.utility.Status.ReportLevel;
import edu.wpi.first.wpilibj.DriverStation;

import com.pathplanner.lib.controllers.PPHolonomicDriveController;
import com.pathplanner.lib.config.PIDConstants;
import com.pathplanner.lib.auto.AutoBuilder;
import com.pathplanner.lib.commands.PathPlannerAuto;
import com.pathplanner.lib.config.RobotConfig;

public class SwerveDrive extends Subsystem {
    private SwerveDrivetrain<TalonFX, TalonFX, CANcoder> swerveDriveTrain;

    private Vision vision;
  
    private int gear_index = SwerveDriveConstants.STARTING_GEAR;
    private boolean stopped = false;

    public double speedAdjust = SwerveDriveConstants.MAX_SPEED_MEETERS_PER_SEC * SwerveDriveConstants.GEARS[gear_index];
    public double rotSpeedAdjust = SwerveDriveConstants.MAX_ROT_SPEED;
    public double autoSpeedAdjust = SwerveDriveConstants.MAX_SPEED_MEETERS_PER_SEC * 0.25; // cap auto performance to 25%
  
    public double rotTarget = 0.0;
    public Rotation2d orientRotTarget = new Rotation2d();
    public ChassisSpeeds chassisSpeeds = new ChassisSpeeds();

    /** Creates a new SwerveDrive. */
    public SwerveDrive(SwerveDrivetrain<TalonFX, TalonFX, CANcoder> swerveDriveTrain, Vision vision) {
	// public SwerveDrive(SwerveDrivetrain<TalonFX, TalonFX, CANcoder> swerveDriveTrain) {
	super();

	this.swerveDriveTrain = swerveDriveTrain;
	this.vision = vision;

	RobotConfig config;
	try{
	    config = RobotConfig.fromGUISettings();
	} catch (Exception e) {
	    // Handle exception as needed
	    config = null;
	}
        // DoubleSupplier a = () -> 1.d;
        AutoBuilder.configure(
			      () -> {
                    var pose = swerveDriveTrain.samplePoseAt(Utils.getCurrentTimeSeconds()).orElse(new Pose2d());
                    // pose = new Pose2d(pose.getX(), pose.getY(), pose.getRotation().times(-1));
                    return pose;//getRotation().times(-1)
                  }, // Robot pose supplier
			      swerveDriveTrain::resetPose, // Method to reset odometry (will be called if your auto has a starting pose)
			      () -> swerveDriveTrain.getState().Speeds, // ChassisSpeeds supplier. MUST BE ROBOT RELATIVE
			      (speeds, feedforwards) -> swerveDriveTrain.setControl(new SwerveRequest.ApplyRobotSpeeds()
										    .withSpeeds(speeds)
										    ), // Method that will drive the robot given ROBOT RELATIVE ChassisSpeeds. Also optionally outputs individual module feedforwards
			      new PPHolonomicDriveController( // PPHolonomicController is the built in path following controller for holonomic drive trains
							      new PIDConstants(5.0, 0.0, 0.0), // Translation PID constants
							      new PIDConstants(5.0, 0.0, 0.0) // Rotation PID constants
							      ),
			      config, // The robot configuration
			      () -> {
				  // Boolean supplier that controls when the path will be mirrored for the red alliance
				  // This will flip the path being followed to the red side of the field.
				  // THE ORIGIN WILL REMAIN ON THE BLUE SIDE

				  var alliance = DriverStation.getAlliance();
				  if (alliance.isPresent()) {
				      return alliance.get() == DriverStation.Alliance.Red;
				  }
				  return false;
			      },
			      this // Reference to this subsystem to set requirements
			      );

	    }

    // public void oneModuleTest(SwerveModule module, Translation2d leftStick, Translation2d rightStick){
    //   // double ang = Math.atan2(rightStick.getY(), rightStick.getX());
    //   // rightStick.getAngle()
    //   double speed = Math.sqrt(Math.pow(leftStick.getX(), 2) + Math.pow(leftStick.getY(), 2));
    //   // System.out.println(ang);
    //   // module.go(ang);
    //   // Rotation2d rot = Rotation2d.fromRadians(ang);
    //   Rotation2d rot = new Rotation2d(rightStick.getX(), rightStick.getY());
    //   SwerveModuleState state = new SwerveModuleState(speed, rot);
    //   module.setDesiredState(state);
    // }

    public void driveWithInput(Translation2d leftStick, Translation2d rightStick, boolean fieldRelative) {
	if (rightStick.getNorm() < 0.05 && leftStick.getNorm() < 0.05 && stopped == false) // if no imput and the swerve drive is still going:
	    stopModules(); // stop the swerve
    
	if (rightStick.getNorm() < 0.05 && leftStick.getNorm() < 0.05) //if no imput
	    return; // don't bother doing swerve drive math and return early.

	leftStick = leftStick.rotateBy(Rotation2d.fromDegrees(SwerveDriveConstants.FORWARD_OFFSET));
    
    if (fieldRelative) {
      swerveDriveTrain.setControl(new SwerveRequest.FieldCentric()
        .withVelocityX(leftStick.getX()*speedAdjust)
        .withVelocityY(-leftStick.getY()*speedAdjust)
        .withRotationalRate(rightStick.getX()*rotSpeedAdjust)
      );
      // double rot = 0;
      
	    // ! drift correction
	    // dependant on if the new odomitry system acounts for rotational drift, this may not be needed.
	    //   if (rightStick.getNorm() > 0.05) {
	    //     rotTarget = swerveDriveTrain.getRotation3d().toRotation2d().getDegrees();
	    //     swerveDriveTrain.setControl(new SwerveRequest.FieldCentric()
	    //       .withVelocityX(leftStick.getX()*speedAdjust)
	    //       .withVelocityY(leftStick.getY()*speedAdjust)
	    //       .withRotationalRate(rightStick.getY()*rotSpeedAdjust)
	    //     );
      
	    //   //  SmartDashboard.putBoolean("drift correction", false);
	    //     stopped = false;
	    //   } else if(leftStick.getNorm() > 0.05) {
	    //     if (!stopped) {
	    //       stopModules();
	    //       stopped = true;
	    //     }

	    //  //   SmartDashboard.putBoolean("drift correction", true);
        
	    //     // rot = ((rotTarget - gyro.getAngle()) / 360) * SwerveDriveConstants.ROT_CORRECTION_SPEED;

	    //   }

	    //   // Use the left joystick to set speed. Apply a cubic curve and the set max speed.
	    //   Translation2d speed = leftStick.times(leftStick.getNorm() * speedAdjust);
	    //   // Translation2d cubedSpeed = new Translation2d(Math.pow(speed.getX(), 3.00), Math.pow(speed.getY(), 3.00));

	    //   // Convert field-relative speeds to robot-relative speeds.
	    //   // chassisSpeeds = chassisSpeeds.
	    //   chassisSpeeds = ChassisSpeeds.fromFieldRelativeSpeeds(-1 * speed.getX(), -1 * speed.getY(), (-1 * rightStick.getX() * rotSpeedAdjust) - rot_correction, gyro.getRotation2d().times(-1));
	} else {      // Create robot-relative speeds.
	    swerveDriveTrain.setControl(new SwerveRequest.RobotCentric()
					.withVelocityX(leftStick.getX()*speedAdjust)
					.withVelocityY(-leftStick.getY()*speedAdjust)
					.withRotationalRate(rightStick.getX()*rotSpeedAdjust)
					);
	}
    }

    public void driveWithInputOrientation(Translation2d leftStick, Translation2d rightStick) { // there is no practical reason to have a robot relitive version of this, and no pre provided version 
	if (rightStick.getNorm() < 0.05 && leftStick.getNorm() < 0.05 && stopped == false) // if no imput and the swerve drive is still going:
	    stopModules(); // stop the swerve
  
	if (rightStick.getNorm() < 0.05 && leftStick.getNorm() < 0.05) //if no imput
	    return; // don't bother doing swerve drive math and return early.

	leftStick.rotateBy(Rotation2d.fromDegrees(SwerveDriveConstants.FORWARD_OFFSET));
    
	swerveDriveTrain.setControl(new SwerveRequest.FieldCentricFacingAngle()
				    .withVelocityX(leftStick.getX()*speedAdjust)
				    .withVelocityY(leftStick.getY()*speedAdjust)
				    .withTargetDirection(rightStick.getAngle())
				    ); 
    }

    public boolean rotateToTarget(double angle) {
	swerveDriveTrain.setControl(new SwerveRequest.FieldCentricFacingAngle()
				    .withVelocityX(0)
				    .withVelocityY(0)
				    .withTargetDirection(Rotation2d.fromDegrees(angle))
				    );

	if (Math.abs(angle - getGyroAngle()) < 5.0) {
	    return true;
	}

	return false;
    }

    public void driveWithInputRotation(Translation2d leftStick, Rotation2d rot) {
	// if (leftStick.getNorm() < 0.05 && stopped == false) // if no imput and the swerve drive is still going:
	//   stopModules(); // stop the swerve
    
	// if (leftStick.getNorm() < 0.05) //if no imput
	//   return; // don't bother doing swerve drive math and return early.

	leftStick = leftStick.rotateBy(Rotation2d.fromDegrees(SwerveDriveConstants.FORWARD_OFFSET));
    
	swerveDriveTrain.setControl(new SwerveRequest.FieldCentricFacingAngle()
				    .withVelocityX(leftStick.getX()*-speedAdjust)
				    .withVelocityY(leftStick.getY()*speedAdjust)
				    .withTargetDirection(rot)
				    );
	// double
    }

    public double getGyroAngle() {
	return swerveDriveTrain.getRotation3d().getAngle();
    }

    public void resetGyro() {
	swerveDriveTrain.tareEverything();
    }

  public void stopModules() {
    swerveDriveTrain.setControl(new SwerveRequest.SwerveDriveBrake());
  }

    @Override
    public void periodic() {
	// This method will be called once per scheduler run\
	SmartDashboard.putNumber("Gyro", getGyroAngle());
	SmartDashboard.putNumber("RotTartget", rotTarget);

	double time = Vision.getTime();

	vision.setLastOdomPose(swerveDriveTrain.samplePoseAt(time));

	if(vision.isTag()){
	    swerveDriveTrain.addVisionMeasurement(vision.getPose2d(), time);
	}

	// if(e.isPresent())
    }

    private void reset_index() {
	gear_index = SwerveDriveConstants.STARTING_GEAR; // however we wish to initialize the gear (What gear does the robot start in?)
    }

    public void shiftDown() {
	if (gear_index == -1 || gear_index >= SwerveDriveConstants.GEARS.length) reset_index(); // If outof bounds: reset index
	int i = gear_index - 1;
	if (i == -1) i = 0;
	setPercentOutput(SwerveDriveConstants.GEARS[i]);
	gear_index = i;
    }

    public void shiftUp() {
	if (gear_index == -1 || gear_index >= SwerveDriveConstants.GEARS.length) reset_index(); // If outof bounds: reset index
	int i = gear_index + 1;
	if (i == SwerveDriveConstants.GEARS.length) i = SwerveDriveConstants.GEARS.length - 1;
	setPercentOutput(SwerveDriveConstants.GEARS[i]);
	gear_index = i;
    }

    public void setPercentOutput(double speed) {
	speedAdjust = SwerveDriveConstants.MAX_SPEED_MEETERS_PER_SEC * speed;
	gear_index = -1;
    }

    public void setToSlow() {
	setPercentOutput(SwerveDriveConstants.SLOW_SPEED);
	gear_index = 0;
    }

    public void setToFast() {
	setPercentOutput(SwerveDriveConstants.FAST_SPEED);
	gear_index = 1;
    }

    public void setToTurbo() {
	setPercentOutput(SwerveDriveConstants.TURBO_SPEED);
	gear_index = 2;
    }

    public void shiftUpRot() {
	rotSpeedAdjust = SwerveDriveConstants.ROTATION_SPEED;
    }

    public void shiftDownRot() {
	rotSpeedAdjust = SwerveDriveConstants.MIN_ROT_SPEED;
    }

    @Override
    public String getSubsystemName() {
	return "Swerve Drive Controller";
    }

    ShuffleboardLayout subsystemLayout = Shuffleboard.getTab("Subsystems")
	.getLayout(getSubsystemName(), BuiltInLayouts.kList)
	.withSize(2, 2);

    GenericEntry sbGyro = subsystemLayout
	.add("Gyro angle", 0)
	.withWidget(BuiltInWidgets.kGyro)
	.getEntry();

    GenericEntry sbShiftState = subsystemLayout
	.add("Shift State", 0)
	.withWidget(BuiltInWidgets.kNumberBar)
	.getEntry();

    @Override
    public void queryStatus() {
	sbGyro.setDouble(getGyroAngle());
	sbShiftState.setDouble(this.speedAdjust);
  
	//TODO: Add more status things
    }

    @Override
    public Status diagnosticStatus() {
	Status status = new Status();

	status.addReport(ReportLevel.ERROR, "Don't know how to diganose new CTRE swerve systems. please check under the CAN(t) section for more detailed information about the swerves there.");
    
	return status;
    }
}