001/**
002 * ********************** PROJECT RON ************************
003 */
004/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */
005/* Use of this source code is governed by an MIT-style license */
006/* that can be found in the repository LICENSE file.           */
007/**
008 * ***********************************************************
009 */
010package com.stuypulse.robot.subsystems.vision;
011
012import static edu.wpi.first.units.Units.*;
013import com.stuypulse.robot.Robot;
014import com.stuypulse.robot.commands.vision.SetPipeline;
015import com.stuypulse.robot.constants.Cameras;
016import com.stuypulse.robot.constants.Settings;
017import com.stuypulse.robot.constants.Settings.EnabledSubsystems;
018import com.stuypulse.robot.subsystems.swerve.CommandSwerveDrivetrain;
019import com.stuypulse.robot.util.vision.LimelightHelpers;
020import com.stuypulse.robot.util.vision.LimelightHelpers.IMUData;
021import com.stuypulse.robot.util.vision.LimelightHelpers.PoseEstimate;
022import edu.wpi.first.math.geometry.Pose2d;
023import edu.wpi.first.math.geometry.Pose3d;
024import edu.wpi.first.networktables.BooleanSubscriber;
025import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
026import edu.wpi.first.wpilibj2.command.SubsystemBase;
027import dev.doglog.DogLog;
028
029public class LimelightVision extends SubsystemBase {
030
031    private static final LimelightVision instance;
032
033    static {
034        instance = new LimelightVision();
035        SmartDashboard.putData("Vision/Set Cloudy Pipeline", new SetPipeline(0));
036        SmartDashboard.putData("Vision/Set Sunny Pipeline", new SetPipeline(1));
037    }
038
039    public static LimelightVision getInstance() {
040        return instance;
041    }
042
043    private String[] names;
044
045    private BooleanSubscriber[] camerasEnabled;
046
047    private MegaTagMode megaTagMode;
048
049    public enum MegaTagMode {
050
051        MEGATAG1, MEGATAG2
052    }
053
054    public LimelightVision() {
055        names = new String[Cameras.LimelightCameras.length];
056        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
057            names[i] = Cameras.LimelightCameras[i].name();
058            Pose3d robotRelativePose = Cameras.LimelightCameras[i].location();
059            LimelightHelpers.setCameraPose_RobotSpace(names[i], robotRelativePose.getX(), robotRelativePose.getY(), robotRelativePose.getZ(), robotRelativePose.getRotation().getMeasureX().in(Degrees), robotRelativePose.getRotation().getMeasureY().in(Degrees), robotRelativePose.getRotation().getMeasureZ().in(Degrees));
060            LimelightHelpers.setRewindEnabled(names[i], true);
061        }
062        camerasEnabled = new BooleanSubscriber[Cameras.LimelightCameras.length];
063        for (int i = 0; i < camerasEnabled.length; i++) {
064            camerasEnabled[i] = DogLog.tunable("Vision/" + names[i] + " Is Enabled", true);
065            setIMUMode(1);
066            DogLog.log("Vision/" + names[i] + " Has Data", false);
067        }
068        setMegaTagMode(MegaTagMode.MEGATAG1);
069    }
070
071    public void setTagWhitelist(int[] ids) {
072        for (String name : names) {
073            LimelightHelpers.SetFiducialIDFiltersOverride(name, ids);
074        }
075    }
076
077    public void enable() {
078        EnabledSubsystems.VISION.getTopic().publish().set(true);
079    }
080
081    public void disable() {
082        EnabledSubsystems.VISION.getTopic().publish().set(false);
083    }
084
085    public void setCameraEnabled(String name, boolean enabled) {
086        for (int i = 0; i < names.length; i++) {
087            if (names[i].equals(name)) {
088                camerasEnabled[i].getTopic().publish().set(enabled);
089            }
090        }
091    }
092
093    public void setMegaTagMode(MegaTagMode mode) {
094        this.megaTagMode = mode;
095    }
096
097    public void setIMUMode(int mode) {
098        for (String name : names) {
099            LimelightHelpers.SetIMUMode(name, mode);
100        }
101    }
102
103    public void setPipeline(int pipeline) {
104        for (String name : names) {
105            LimelightHelpers.setPipelineIndex(name, pipeline);
106        }
107    }
108
109    /**
110     * Allows you to set the convergence speed of the internal LL IMU and robot gyro.
111     *
112     * @param assistValue, an double that sets the correction speed of the complementary filter for the IMU. IMU Mode 4
113     * uses the fusing of the internal IMU (1khz) with the external gyro reading as well. Higher values ranging towards 1
114     * indicate a faster convergence of internal IMU to the robot IMU mode. Defaults to 0.001.
115     */
116    public void setIMUAssistValue(double assistValue) {
117        for (String name : names) {
118            LimelightHelpers.SetIMUAssistAlpha(name, assistValue);
119        }
120    }
121
122    public IMUData[] getIMUData() {
123        IMUData[] data = new IMUData[Cameras.LimelightCameras.length];
124        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
125            data[i] = LimelightHelpers.getIMUData(Cameras.LimelightCameras[i].name());
126        }
127        return data;
128    }
129
130    public void captureRewind(double timeSecs) {
131        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
132            LimelightHelpers.triggerRewindCapture(names[i], timeSecs);
133        }
134    }
135
136    @Override
137    public void periodic() {
138        if (!EnabledSubsystems.VISION.get()) {
139            return;
140        }
141        for (int i = 0; i < names.length; i++) {
142            DogLog.log("Vision/" + names[i] + "/Heartbeat", LimelightHelpers.getHeartbeat(names[i]));
143            if (!camerasEnabled[i].get()) {
144                DogLog.log("Vision/" + names[i] + " Has Data", false);
145                continue;
146            }
147            String limelightName = names[i];
148            // Seed robot heading (used by MT2)
149            LimelightHelpers.SetRobotOrientation(limelightName, (CommandSwerveDrivetrain.getInstance().getPose().getRotation().getDegrees() + (Robot.isBlue() ? 0 : 180)) % 360, 0, 0, 0, 0, 0);
150            PoseEstimate poseEstimate;
151            // MegaTag switching
152            if (megaTagMode == MegaTagMode.MEGATAG1) {
153                poseEstimate = Robot.isBlue() ? LimelightHelpers.getBotPoseEstimate_wpiBlue(limelightName) : LimelightHelpers.getBotPoseEstimate_wpiRed(limelightName);
154            } else {
155                poseEstimate = Robot.isBlue() ? LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(limelightName) : LimelightHelpers.getBotPoseEstimate_wpiRed_MegaTag2(limelightName);
156            }
157            boolean notNull = false;
158            boolean withinAngularVelocityTolerance = false;
159            boolean poseAtOrigin = false;
160            // Adding to pose estimator
161            SmartDashboard.putBoolean("Vision/Pose Not Estimate Null", poseEstimate != null);
162            if (poseEstimate != null && poseEstimate.tagCount > 0) {
163                notNull = true;
164                SmartDashboard.putNumber("Vision/Pose Estimate X", poseEstimate.pose.getX());
165                SmartDashboard.putNumber("Vision/Pose Estimate Y", poseEstimate.pose.getY());
166                SmartDashboard.putNumber("Vision/Pose Estimate Theta", poseEstimate.pose.getRotation().getDegrees());
167                SmartDashboard.putNumber("Vision/Tag Count", poseEstimate.tagCount);
168                if (poseEstimate.pose.equals(Settings.Vision.INVALID_POSITION)) {
169                    poseAtOrigin = true;
170                }
171                if (CommandSwerveDrivetrain.getInstance().getChassisSpeeds().omegaRadiansPerSecond < Settings.Vision.MAX_ANGULAR_VELOCITY_RAD_SEC) {
172                    withinAngularVelocityTolerance = true;
173                }
174                Boolean isValidPose = notNull && withinAngularVelocityTolerance && !poseAtOrigin;
175                SmartDashboard.putBoolean("Vision/Pose at Origin?", poseAtOrigin);
176                SmartDashboard.putBoolean("Vision/Within Angular Velocity", withinAngularVelocityTolerance);
177                DogLog.log("Vision/isValidPose", isValidPose);
178                DogLog.log("Vision/isWithinAngularVel", withinAngularVelocityTolerance);
179                DogLog.log("Vision/poseAtOrigin", poseAtOrigin);
180                Pose2d robotPose = poseEstimate.pose;
181                double timestamp = poseEstimate.timestampSeconds;
182                if (megaTagMode == MegaTagMode.MEGATAG1 && isValidPose) {
183                    CommandSwerveDrivetrain.getInstance().addVisionMeasurement(robotPose, timestamp, Settings.Vision.MT1_STDEVS);
184                } else if (megaTagMode == MegaTagMode.MEGATAG2 && isValidPose) {
185                    CommandSwerveDrivetrain.getInstance().addVisionMeasurement(robotPose, timestamp, Settings.Vision.MT2_STDEVS);
186                }
187                DogLog.log("Vision/Pose", robotPose);
188                DogLog.log("Vision/Pose X Component", robotPose.getX());
189                DogLog.log("Vision/Pose Y Component", robotPose.getY());
190                DogLog.log("Vision/Pose Theta (Degrees)", robotPose.getRotation().getDegrees());
191                DogLog.log("Vision/" + names[i] + " Has Data", true);
192                DogLog.log("Vision/MegaTag Mode", megaTagMode.toString());
193                // this yaw is seems to be the robot yaw passed into the LL
194                DogLog.forceNt.log("Vision/Pipeline", LimelightHelpers.getCurrentPipelineIndex(limelightName));
195                DogLog.log("Vision/Limelight Robot Yaw", LimelightHelpers.getIMUData(limelightName).robotYaw);
196                // this is just the yaw of the internal imu
197                DogLog.log("Vision/Limelight Yaw", LimelightHelpers.getIMUData(limelightName).Yaw);
198                DogLog.log("Vision/Has at least 2 tags", poseEstimate.tagCount >= 2);
199            }
200        }
201    }
202}