package frc.robot.subsystems;
import java.util.Optional;
import java.util.function.Supplier;

import com.ctre.phoenix6.Utils;
import com.ctre.phoenix6.swerve.SwerveDrivetrainConstants;
import com.ctre.phoenix6.swerve.SwerveModuleConstants;
import com.ctre.phoenix6.swerve.SwerveRequest;
import com.pathplanner.lib.auto.AutoBuilder;
import com.pathplanner.lib.config.PIDConstants;
import com.pathplanner.lib.config.RobotConfig;
import com.pathplanner.lib.controllers.PPHolonomicDriveController;

import edu.wpi.first.math.Matrix;
import edu.wpi.first.math.VecBuilder;
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.numbers.N1;
import edu.wpi.first.math.numbers.N3;
import edu.wpi.first.wpilibj.DriverStation;
import edu.wpi.first.wpilibj.DriverStation.Alliance;
import edu.wpi.first.wpilibj.smartdashboard.Field2d;
import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
import edu.wpi.first.wpilibj2.command.Command;
import edu.wpi.first.wpilibj2.command.Subsystem;
import frc.robot.Constants.DriveMode;
import frc.robot.generated.TunerConstants.TunerSwerveDrivetrain;
import frc.robot.Constants;

import edu.wpi.first.networktables.*;


/**
 * Class that extends the Phoenix 6 SwerveDrivetrain class and implements
 * Subsystem so it can easily be used in command-based projects.
 *
 * Generated by the 2026 Tuner X Swerve Project Generator
 * https://v6.docs.ctr-electronics.com/en/stable/docs/tuner/tuner-swerve/index.html
 */
public class CommandSwerveDrivetrain extends TunerSwerveDrivetrain implements Subsystem {

    /* Blue alliance sees forward as 0 degrees (toward red alliance wall) */
    private static final Rotation2d kBlueAlliancePerspectiveRotation = Rotation2d.kZero;
    /* Red alliance sees forward as 180 degrees (toward blue alliance wall) */
    private static final Rotation2d kRedAlliancePerspectiveRotation = Rotation2d.k180deg;
    /* Keep track if we've ever applied the operator perspective before or not */
    private boolean m_hasAppliedOperatorPerspective = false;

    private DriveMode m_driveMode = DriveMode.OTHER;
    
    private NetworkTable positioningNetworkTable;
    private NetworkTable ll_shooter_NT;
    private NetworkTable ll_intaker_NT;
    private NetworkTable driveNetworkTable;
    
    private final StructPublisher<Pose2d> posePub =
    NetworkTableInstance.getDefault()
        .getStructTopic("/Odometry/RobotPose", Pose2d.struct)
        .publish();

    private Field2d m_field;
    private Field2d shooter_ll_field;
    private Field2d intaker_ll_field;

    public CommandSwerveDrivetrain(
        SwerveDrivetrainConstants drivetrainConstants,
        SwerveModuleConstants<?, ?, ?>... modules
    ) {
        super(drivetrainConstants, modules);
        m_field = new Field2d();
        shooter_ll_field = new Field2d();
        intaker_ll_field = new Field2d();
        SmartDashboard.putData("Field", m_field);
        SmartDashboard.putData("PositioningShooterll", shooter_ll_field);
        SmartDashboard.putData("PositioningIntakerll", intaker_ll_field);
        LimelightHelpers.SetIMUMode(Constants.Limelight.LIMELIGHT_NAME_Shooter, 1);
        LimelightHelpers.SetIMUMode(Constants.Limelight.LIMELIGHT_NAME_Intaker, 1);
        positioningNetworkTable = NetworkTableInstance.getDefault().getTable("Positioning");
        driveNetworkTable = NetworkTableInstance.getDefault().getTable("Drive");
        ll_shooter_NT = NetworkTableInstance.getDefault().getTable("Positioning/limelight-shooter");
        ll_intaker_NT = NetworkTableInstance.getDefault().getTable("Positioning/limelight-intaker");

        configueAutoBuilder();

    }


    private void configueAutoBuilder(){
        
        RobotConfig config;
        try {
            config = RobotConfig.fromGUISettings();
        } catch (Exception e) {
            throw new RuntimeException("Failed to load PathPlanner RobotConfig", e);
        }

        AutoBuilder.configure(
            this::getPose,                  // 当前机器人 Pose
            this::resetPose,                // 重置里程计
            this::getRobotRelativeSpeeds,   // ROBOT RELATIVE speeds
            (speeds, feedforwards) -> driveRobotRelative(speeds),      // 用 speeds 驱动机器人
            new PPHolonomicDriveController(
                new PIDConstants(5.0, 0.0, 0.0),   // X/Y PID
                new PIDConstants(5.0, 0.0, 0.0)    // Theta PID
            ),
            config,
            () -> false,
            this
        );

    }
    public Pose2d getPose() {
        return getState().Pose;
    }

    public ChassisSpeeds getRobotRelativeSpeeds() {
        return getState().Speeds;
    }

    public void driveRobotRelative(ChassisSpeeds speeds) {
        ChassisSpeeds scaled = new ChassisSpeeds(
            speeds.vxMetersPerSecond,                 // X 不变
            speeds.vyMetersPerSecond,                 // Y 不变
            speeds.omegaRadiansPerSecond * 1.3        // 只放大旋转
        );

        setControl(
            new SwerveRequest.ApplyRobotSpeeds().withSpeeds(scaled)
        ); 
    }

<<<<<<< HEAD
=======
    



        
>>>>>>> a8acbce9dc302adc52e28e6f8c205677f50ea20e
    private static final double kMaxVisionJumpMeters = 1.0;
    private static final double kMaxOmegaDegPerSec = 360.0;

    /**
     * Returns a command that applies the specified control request to this swerve drivetrain.
     *
     * @param request Function returning the request to apply
     * @return Command to run
     */
    public Command applyRequest(Supplier<SwerveRequest> request) {
        return run(() -> this.setControl(request.get()));
    }

    public Command applyRequest(Supplier<SwerveRequest> requestSupplier, Supplier<DriveMode> driveModeSupplier) {
        return run(() -> {
            DriveMode mode = driveModeSupplier.get();
            m_driveMode = mode;
            this.setControl(requestSupplier.get());
        });
    }
    
    public void forceUsingLimelightMT2() {
        LimelightHelpers.PoseEstimate mt2 =
        LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(Constants.Limelight.LIMELIGHT_NAME_Shooter);

        if (mt2 == null || mt2.tagCount == 0) {
            return;
        }

        var stdDevs = VecBuilder.fill(0.7, 0.7, 9999999);
        addVisionMeasurement(mt2.pose, mt2.timestampSeconds, stdDevs);
    }
    
    private void fuseLimelightMT2() {
        Pose2d current = getState().Pose;

        LimelightHelpers.PoseEstimate mt2_shooter =
        getLimelightPoseEstimate(Constants.Limelight.LIMELIGHT_NAME_Shooter);
        
        LimelightHelpers.PoseEstimate mt2_intaker =
        getLimelightPoseEstimate(Constants.Limelight.LIMELIGHT_NAME_Intaker);

        boolean shooterAvailable = ifLimelightPositioningAvailable(mt2_shooter, current);
        boolean intakerAvailable = ifLimelightPositioningAvailable(mt2_intaker, current);
        ll_shooter_NT.getEntry("available").setBoolean(shooterAvailable);
        ll_intaker_NT.getEntry("available").setBoolean(intakerAvailable);
        ll_shooter_NT.getEntry("TagCount").setNumber(mt2_shooter != null ? mt2_shooter.tagCount : 0);
        ll_intaker_NT.getEntry("TagCount").setNumber(mt2_intaker != null ? mt2_intaker.tagCount : 0);
        
        LimelightHelpers.PoseEstimate mt2_final = null;

        if(!shooterAvailable && !intakerAvailable) {
            positioningNetworkTable.getEntry("MetaTag2Available").setBoolean(false);
            positioningNetworkTable.getEntry("Usingll").setString("None");
            //positioningNetworkTable.getEntry("MetaTag2DistanceDifference").setDouble(Double.NaN);
            
            return;
        }else if (shooterAvailable && !intakerAvailable) {
            double shooterDistance = current.getTranslation().getDistance(mt2_shooter.pose.getTranslation());
            ll_shooter_NT.getEntry("Distance").setDouble(shooterDistance);
            ll_shooter_NT.getEntry("X").setDouble(mt2_shooter.pose.getX());
            ll_shooter_NT.getEntry("Y").setDouble(mt2_shooter.pose.getY());
            positioningNetworkTable.getEntry("MetaTag2Available").setBoolean(true);
            positioningNetworkTable.getEntry("Usingll").setString("Shooter");
            //positioningNetworkTable.getEntry("MetaTag2DistanceDifference").setDouble(Double.NaN);
            mt2_final = mt2_shooter;
        }else if (!shooterAvailable && intakerAvailable) {
            double intakerDistance = current.getTranslation().getDistance(mt2_intaker.pose.getTranslation());
            ll_intaker_NT.getEntry("Distance").setDouble(intakerDistance);
            ll_intaker_NT.getEntry("X").setDouble(mt2_intaker.pose.getX());
            ll_intaker_NT.getEntry("Y").setDouble(mt2_intaker.pose.getY());
            positioningNetworkTable.getEntry("MetaTag2Available").setBoolean(true);
            positioningNetworkTable.getEntry("Usingll").setString("Intaker");
            //positioningNetworkTable.getEntry("MetaTag2DistanceDifference").setDouble(Double.NaN);
            mt2_final = mt2_intaker;
        }else {
            // 两个都可用时的策略：这里简单选更近的，你也
            double shooterDistance = current.getTranslation().getDistance(mt2_shooter.pose.getTranslation());
            double intakerDistance = current.getTranslation().getDistance(mt2_intaker.pose.getTranslation());
            double difference = mt2_shooter.pose.getTranslation().getDistance(mt2_intaker.pose.getTranslation());

            ll_shooter_NT.getEntry("Distance").setDouble(shooterDistance);
            ll_intaker_NT.getEntry("Distance").setDouble(intakerDistance);
            ll_shooter_NT.getEntry("X").setDouble(mt2_shooter.pose.getX());
            ll_shooter_NT.getEntry("Y").setDouble(mt2_shooter.pose.getY());
            ll_intaker_NT.getEntry("X").setDouble(mt2_intaker.pose.getX());
            ll_intaker_NT.getEntry("Y").setDouble(mt2_intaker.pose.getY());
            positioningNetworkTable.getEntry("MetaTag2DistanceDifference").setDouble(difference);
            positioningNetworkTable.getEntry("MetaTag2XDifference").setDouble(Math.abs(mt2_shooter.pose.getX() - mt2_intaker.pose.getX()));
            positioningNetworkTable.getEntry("MetaTag2YDifference").setDouble(Math.abs(mt2_shooter.pose.getY() - mt2_intaker.pose.getY()));
            positioningNetworkTable.getEntry("Usingll").setString(shooterDistance <= intakerDistance ? "ChosenShooter" : "ChosenIntaker");

            mt2_final = shooterDistance <= intakerDistance ? mt2_shooter : mt2_intaker;
        }
        
        Pose2d mt2Pose = mt2_final.pose;

        double distance = current.getTranslation().getDistance(mt2Pose.getTranslation());
        Rotation2d deltaRotation = current.getRotation().minus(mt2Pose.getRotation());

        positioningNetworkTable.getEntry("CurrentPoseX").setDouble(current.getX());
        positioningNetworkTable.getEntry("CurrentPoseY").setDouble(current.getY());
        positioningNetworkTable.getEntry("MetaTag2PoseX").setDouble(mt2Pose.getX());
        positioningNetworkTable.getEntry("MetaTag2PoseY").setDouble(mt2Pose.getY());
        positioningNetworkTable.getEntry("PoseDistance").setDouble(distance);
        positioningNetworkTable.getEntry("RotationDifference").setDouble(Math.abs(deltaRotation.getDegrees()));

        // 4) 设置测量噪声（示例：theta 基本不信任就给极大）
        // 你也可以根据 mt2.avgTagDist / mt2.tagSpan 动态调
        var stdDevs = VecBuilder.fill(0.7, 0.7, 9999999);

        // 5) 融合：注意 timestampSeconds 的处理
        addVisionMeasurement(mt2_final.pose, mt2_final.timestampSeconds, stdDevs);
        positioningNetworkTable.getEntry("MetaTag2Available").setBoolean(true);
    }

    private LimelightHelpers.PoseEstimate getLimelightPoseEstimate(String limelightName) {
        if (Constants.Limelight.UsingMetaTag2) {
           return LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(limelightName);
        }
        return LimelightHelpers.getBotPoseEstimate_wpiBlue(limelightName);
    }
    private boolean ifLimelightPositioningAvailable(LimelightHelpers.PoseEstimate mt2, Pose2d currentPose) {
        // 1) 没看到 tag：拒绝
        if (mt2 == null || mt2.tagCount == 0) {
            return false;
        }
        // 2) 转太快：拒绝（角速度你可以从 Pigeon2 或 drivetrain state 拿）
        double omegaDegPerSec = getPigeon2().getAngularVelocityZWorld().getValueAsDouble(); // 若你有这个信号
        if (Math.abs(omegaDegPerSec) > kMaxOmegaDegPerSec) {
            System.out.println("NoUsingMetaTag2ForTooFastRotationSpeed");
            return false;
        }      

        // 3) 跳变太大：拒绝（Phoenix6 推荐的稳定策略）
        //Pose2d current = getState().Pose;
        if (currentPose.getTranslation().getDistance(mt2.pose.getTranslation()) > kMaxVisionJumpMeters) {
            System.out.println("NoUsingMetaTag2ForTooChangingPose");
            return false;
        }

        return true;
    }
<<<<<<< HEAD
=======

>>>>>>> a8acbce9dc302adc52e28e6f8c205677f50ea20e
    public String driveToAprilTag(){
        //Rotation2d targetHeading = getPose().getRotation();
        if(!LimelightSupplier.isTargetVisible() ){
            return "No Target";
        }
        boolean ifRightTag = false;
        int trueId = LimelightSupplier.getAprilTagID();
        for(int id : new int[]{16,32}){
            if(id == trueId){
                ifRightTag = true;
<<<<<<< HEAD
                System.out.println("Target get.");
=======
>>>>>>> a8acbce9dc302adc52e28e6f8c205677f50ea20e
                break;
            }
        }
        if(!ifRightTag){
<<<<<<< HEAD
            System.out.println("No Target.");
=======
>>>>>>> a8acbce9dc302adc52e28e6f8c205677f50ea20e
            return "No Target";
        }
        double tx = LimelightSupplier.getTX() ;
        double targetTz = -LimelightSupplier.getTargetTZ();
        double ry = LimelightSupplier.getTargetRotationY();

        double multiply = Constants.Limelight.AutoClimbAlignSpeed;
        if(Math.abs(tx)<Constants.Limelight.AutoClimbAlignTolerance){
            tx = 0;
        }
        else{
            tx *= multiply ;
        }
        
        if(Math.abs(targetTz)<Constants.Limelight.AutoClimbAlignTolerance){
            targetTz = 0;
        }
        else{
            targetTz*= multiply ;
            targetTz*= 5;
        }

        if(Math.abs(ry)<Constants.Limelight.AutoClimbAlignTolerance){
            ry = 0;
        }
        else{
            ry/=10;
            ry = Math.min(ry, Constants.Limelight.AutoClimbmaxAngularVelocity);
            ry = Math.max(ry, -Constants.Limelight.AutoClimbmaxAngularVelocity);
        }
        
        
        positioningNetworkTable.getEntry("interestPointsClimb tx").setDouble(tx);
        positioningNetworkTable.getEntry("interestPointsClimb targetTz").setDouble(targetTz);
        positioningNetworkTable.getEntry("interestPointsClimb ry").setDouble(ry);
        

        ChassisSpeeds scaled = new ChassisSpeeds(
            tx,                 // X 不变
            targetTz,                 // Y 不变
            ry
        );
        driveRobotRelative(scaled);
        // drive(
        //     new Translation2d(tx, targetTz),
        //     ry,
        //     false,
        //     false
        // );

        return "Driving to AprilTag";
    }
    @Override
    public void periodic() {

        Pose2d pose = getState().Pose;
        m_field.setRobotPose(pose);
        posePub.set(pose);
        shooter_ll_field.setRobotPose(LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(Constants.Limelight.LIMELIGHT_NAME_Shooter).pose);
        intaker_ll_field.setRobotPose(LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(Constants.Limelight.LIMELIGHT_NAME_Intaker).pose);
        driveNetworkTable.getEntry("YawFromState").setDouble(getState().Pose.getRotation().getDegrees());
        driveNetworkTable.getEntry("DriveMode").setString(m_driveMode.name());
        driveNetworkTable.getEntry("IsFieldCentric").setBoolean(m_driveMode == DriveMode.FIELD_CENTRIC);

        // 1) 每帧喂给 LL：机器人当前 yaw（度），其余先全 0
        LimelightHelpers.SetRobotOrientation(
            Constants.Limelight.LIMELIGHT_NAME_Shooter,
            pose.getRotation().getDegrees(), // 关键：要符合 wpiBlue 约定
            0, 0, 0, 0, 0
        );
        LimelightHelpers.SetRobotOrientation(
            Constants.Limelight.LIMELIGHT_NAME_Intaker,
            pose.getRotation().getDegrees(), // 关键：要符合 wpiBlue 约定
            0, 0, 0, 0, 0
        );

        fuseLimelightMT2();

        if (!m_hasAppliedOperatorPerspective || DriverStation.isDisabled()) {
            DriverStation.getAlliance().ifPresent(allianceColor -> {
                setOperatorPerspectiveForward(
                    allianceColor == Alliance.Red
                        ? kRedAlliancePerspectiveRotation
                        : kBlueAlliancePerspectiveRotation
                );
                m_hasAppliedOperatorPerspective = true;
            });
        }
    }        
        
        
            /**
     * Adds a vision measurement to the Kalman Filter. This will correct the odometry pose estimate
     * while still accounting for measurement noise.
     *
     * @param visionRobotPoseMeters The pose of the robot as measured by the vision camera.
     * @param timestampSeconds The timestamp of the vision measurement in seconds.
     */
    @Override
    public void addVisionMeasurement(Pose2d visionRobotPoseMeters, double timestampSeconds) {
        super.addVisionMeasurement(visionRobotPoseMeters, Utils.fpgaToCurrentTime(timestampSeconds));
    }

    /**
     * Adds a vision measurement to the Kalman Filter. This will correct the odometry pose estimate
     * while still accounting for measurement noise.
     * <p>
     * Note that the vision measurement standard deviations passed into this method
     * will continue to apply to future measurements until a subsequent call to
     * {@link #setVisionMeasurementStdDevs(Matrix)} or this method.
     *
     * @param visionRobotPoseMeters The pose of the robot as measured by the vision camera.
     * @param timestampSeconds The timestamp of the vision measurement in seconds.
     * @param visionMeasurementStdDevs Standard deviations of the vision pose measurement
     *     in the form [x, y, theta]ᵀ, with units in meters and radians.
     */
    @Override
    public void addVisionMeasurement(
        Pose2d visionRobotPoseMeters,
        double timestampSeconds,
        Matrix<N3, N1> visionMeasurementStdDevs
    ) {
        super.addVisionMeasurement(visionRobotPoseMeters, Utils.fpgaToCurrentTime(timestampSeconds), visionMeasurementStdDevs);
    }

    /**
     * Return the pose at a given timestamp, if the buffer is not empty.
     *
     * @param timestampSeconds The timestamp of the pose in seconds.
     * @return The pose at the given timestamp (or Optional.empty() if the buffer is empty).
     */
    @Override
    public Optional<Pose2d> samplePoseAt(double timestampSeconds) {
        return super.samplePoseAt(Utils.fpgaToCurrentTime(timestampSeconds));
    }
}