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        SmartDashboard.putData("Vision/Set Lab Pipeline", new SetPipeline(4));
038    }
039
040    public static LimelightVision getInstance() {
041        return instance;
042    }
043
044    private String[] names;
045
046    private BooleanSubscriber[] camerasEnabled;
047
048    private MegaTagMode megaTagMode;
049
050    public enum MegaTagMode {
051
052        MEGATAG1, MEGATAG2
053    }
054
055    public LimelightVision() {
056        names = new String[Cameras.LimelightCameras.length];
057        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
058            names[i] = Cameras.LimelightCameras[i].name();
059            Pose3d robotRelativePose = Cameras.LimelightCameras[i].location();
060            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));
061            LimelightHelpers.setRewindEnabled(names[i], true);
062        }
063        camerasEnabled = new BooleanSubscriber[Cameras.LimelightCameras.length];
064        for (int i = 0; i < camerasEnabled.length; i++) {
065            camerasEnabled[i] = DogLog.tunable("Vision/" + names[i] + " Is Enabled", true);
066            setIMUMode(1);
067            DogLog.log("Vision/" + names[i] + " Has Data", false);
068        }
069        setMegaTagMode(MegaTagMode.MEGATAG1);
070    }
071
072    public void setTagWhitelist(int[] ids) {
073        for (String name : names) {
074            LimelightHelpers.SetFiducialIDFiltersOverride(name, ids);
075        }
076    }
077
078    public void enable() {
079        EnabledSubsystems.VISION.getTopic().publish().set(true);
080    }
081
082    public void disable() {
083        EnabledSubsystems.VISION.getTopic().publish().set(false);
084    }
085
086    public void setCameraEnabled(String name, boolean enabled) {
087        for (int i = 0; i < names.length; i++) {
088            if (names[i].equals(name)) {
089                camerasEnabled[i].getTopic().publish().set(enabled);
090            }
091        }
092    }
093
094    public void setMegaTagMode(MegaTagMode mode) {
095        this.megaTagMode = mode;
096    }
097
098    public void setIMUMode(int mode) {
099        for (String name : names) {
100            LimelightHelpers.SetIMUMode(name, mode);
101        }
102    }
103
104    public void setPipeline(int pipeline) {
105        for (String name : names) {
106            LimelightHelpers.setPipelineIndex(name, pipeline);
107        }
108    }
109
110    /**
111     * Allows you to set the convergence speed of the internal LL IMU and robot gyro.
112     *
113     * @param assistValue, an double that sets the correction speed of the complementary filter for the IMU. IMU Mode 4
114     * uses the fusing of the internal IMU (1khz) with the external gyro reading as well. Higher values ranging towards 1
115     * indicate a faster convergence of internal IMU to the robot IMU mode. Defaults to 0.001.
116     */
117    public void setIMUAssistValue(double assistValue) {
118        for (String name : names) {
119            LimelightHelpers.SetIMUAssistAlpha(name, assistValue);
120        }
121    }
122
123    public IMUData[] getIMUData() {
124        IMUData[] data = new IMUData[Cameras.LimelightCameras.length];
125        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
126            data[i] = LimelightHelpers.getIMUData(Cameras.LimelightCameras[i].name());
127        }
128        return data;
129    }
130
131    public void captureRewind(double timeSecs) {
132        for (int i = 0; i < Cameras.LimelightCameras.length; i++) {
133            LimelightHelpers.triggerRewindCapture(names[i], timeSecs);
134        }
135    }
136
137    @Override
138    public void periodic() {
139        if (!EnabledSubsystems.VISION.get()) {
140            return;
141        }
142        for (int i = 0; i < names.length; i++) {
143            DogLog.log("Vision/" + names[i] + "/Heartbeat", LimelightHelpers.getHeartbeat(names[i]));
144            if (!camerasEnabled[i].get()) {
145                DogLog.log("Vision/" + names[i] + " Has Data", false);
146                continue;
147            }
148            String limelightName = names[i];
149            // Seed robot heading (used by MT2)
150            LimelightHelpers.SetRobotOrientation(limelightName, (CommandSwerveDrivetrain.getInstance().getPose().getRotation().getDegrees() + (Robot.isBlue() ? 0 : 180)) % 360, 0, 0, 0, 0, 0);
151            PoseEstimate poseEstimate;
152            // MegaTag switching
153            if (megaTagMode == MegaTagMode.MEGATAG1) {
154                poseEstimate = Robot.isBlue() ? LimelightHelpers.getBotPoseEstimate_wpiBlue(limelightName) : LimelightHelpers.getBotPoseEstimate_wpiRed(limelightName);
155            } else {
156                poseEstimate = Robot.isBlue() ? LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(limelightName) : LimelightHelpers.getBotPoseEstimate_wpiRed_MegaTag2(limelightName);
157            }
158            boolean notNull = false;
159            boolean withinAngularVelocityTolerance = false;
160            boolean poseAtOrigin = false;
161            // Adding to pose estimator
162            DogLog.log("Vision/Pose Not Estimate Null", poseEstimate != null);
163            if (poseEstimate != null && poseEstimate.tagCount > 0) {
164                notNull = true;
165                DogLog.log("Vision/Pose Estimate X", poseEstimate.pose.getX());
166                DogLog.log("Vision/Pose Estimate Y", poseEstimate.pose.getY());
167                DogLog.log("Vision/Pose Estimate Theta", poseEstimate.pose.getRotation().getDegrees());
168                DogLog.log("Vision/Tag Count", poseEstimate.tagCount);
169                if (poseEstimate.pose.equals(Settings.Vision.INVALID_POSITION)) {
170                    poseAtOrigin = true;
171                }
172                if (CommandSwerveDrivetrain.getInstance().getChassisSpeeds().omegaRadiansPerSecond < Settings.Vision.MAX_ANGULAR_VELOCITY_RAD_SEC) {
173                    withinAngularVelocityTolerance = true;
174                }
175                Boolean isValidPose = notNull && withinAngularVelocityTolerance && !poseAtOrigin;
176                DogLog.log("Vision/Pose at Origin?", poseAtOrigin);
177                DogLog.log("Vision/Within Angular Velocity", withinAngularVelocityTolerance);
178                DogLog.log("Vision/isValidPose", isValidPose);
179                DogLog.log("Vision/isWithinAngularVel", withinAngularVelocityTolerance);
180                DogLog.log("Vision/poseAtOrigin", poseAtOrigin);
181                Pose2d robotPose = poseEstimate.pose;
182                double timestamp = poseEstimate.timestampSeconds;
183                if (megaTagMode == MegaTagMode.MEGATAG1 && isValidPose) {
184                    CommandSwerveDrivetrain.getInstance().addVisionMeasurement(robotPose, timestamp, Settings.Vision.MT1_STDEVS);
185                } else if (megaTagMode == MegaTagMode.MEGATAG2 && isValidPose) {
186                    CommandSwerveDrivetrain.getInstance().addVisionMeasurement(robotPose, timestamp, Settings.Vision.MT2_STDEVS);
187                }
188                DogLog.log("Vision/Pose", robotPose);
189                DogLog.log("Vision/Pose X Component", robotPose.getX());
190                DogLog.log("Vision/Pose Y Component", robotPose.getY());
191                DogLog.log("Vision/Pose Theta (Degrees)", robotPose.getRotation().getDegrees());
192                DogLog.log("Vision/" + names[i] + " Has Data", true);
193                DogLog.log("Vision/MegaTag Mode", megaTagMode.toString());
194                // this yaw is seems to be the robot yaw passed into the LL
195                DogLog.forceNt.log("Vision/Pipeline", LimelightHelpers.getCurrentPipelineIndex(limelightName));
196                DogLog.log("Vision/Limelight Robot Yaw", LimelightHelpers.getIMUData(limelightName).robotYaw);
197                // this is just the yaw of the internal imu
198                DogLog.log("Vision/Limelight Yaw", LimelightHelpers.getIMUData(limelightName).Yaw);
199                DogLog.log("Vision/Has at least 2 tags", poseEstimate.tagCount >= 2);
200            }
201        }
202    }
203}