001/************************ PROJECT ZIXI ************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved.*/
003/* This work is licensed under the terms of the MIT license.  */
004/**************************************************************/
005package com.stuypulse.robot.subsystems.vision;
006
007import static org.wpilib.units.Units.Degrees;
008
009import org.wpilib.math.geometry.Pose3d;
010import org.wpilib.math.geometry.Rotation2d;
011import org.wpilib.math.geometry.Rotation3d;
012import org.wpilib.math.geometry.Transform3d;
013import org.wpilib.math.util.Units;
014import org.wpilib.networktables.DoubleArrayPublisher;
015import org.wpilib.networktables.DoubleArraySubscriber;
016import org.wpilib.networktables.DoubleSubscriber;
017import org.wpilib.networktables.IntegerPublisher;
018import org.wpilib.networktables.IntegerSubscriber;
019import org.wpilib.networktables.NetworkTable;
020import org.wpilib.networktables.NetworkTableInstance;
021import org.wpilib.networktables.TimestampedDoubleArray;
022import org.wpilib.system.RobotController;
023
024import java.util.Arrays;
025import java.util.HashSet;
026import java.util.LinkedList;
027import java.util.List;
028import java.util.Set;
029import java.util.function.Supplier;
030
031/** IO implementation for real Limelight hardware. */
032public class VisionIOLimelight implements VisionIO {
033    private final Supplier<Rotation2d> rotationSupplier;
034    private final DoubleArrayPublisher orientationPublisher;
035    private final IntegerPublisher pipelinePublisher;
036    private final IntegerPublisher imuModePublisher;
037    private final DoubleArrayPublisher aprilTagWhitelistPublisher;
038
039    private final DoubleSubscriber latencySubscriber;
040    private final DoubleSubscriber txSubscriber;
041    private final DoubleSubscriber tySubscriber;
042    private final DoubleArraySubscriber megatag1Subscriber;
043    private final DoubleArraySubscriber megatag2Subscriber;
044    private final IntegerSubscriber pipelineSubscriber;
045
046    private MegaTagMode megaTagMode;
047
048    /**
049     * Creates a new VisionIOLimelight.
050     *
051     * @param name The configured name of the Limelight.
052     * @param rotationSupplier Supplier for the current estimated rotation, used for MegaTag 2.
053     */
054    public VisionIOLimelight(
055            String name, Transform3d robotToCamera, Supplier<Rotation2d> rotationSupplier) {
056        NetworkTable table = NetworkTableInstance.getDefault().getTable(name);
057        this.rotationSupplier = rotationSupplier;
058        orientationPublisher = table.getDoubleArrayTopic("robot_orientation_set").publish();
059        pipelinePublisher = table.getIntegerTopic("pipeline").publish();
060        imuModePublisher = table.getIntegerTopic("imumode_set").publish();
061        aprilTagWhitelistPublisher = table.getDoubleArrayTopic("fiducial_id_filters_set").publish();
062        latencySubscriber = table.getDoubleTopic("tl").subscribe(0.0);
063        txSubscriber = table.getDoubleTopic("tx").subscribe(0.0);
064        tySubscriber = table.getDoubleTopic("ty").subscribe(0.0);
065        megatag1Subscriber =
066                table.getDoubleArrayTopic("botpose_wpiblue").subscribe(new double[] {});
067        megatag2Subscriber =
068                table.getDoubleArrayTopic("botpose_orb_wpiblue").subscribe(new double[] {});
069        pipelineSubscriber = table.getIntegerTopic("pipeline").subscribe(0);
070
071        megaTagMode = MegaTagMode.MEGATAG_1;
072
073        table.getDoubleArrayTopic("camerapose_robotspace_set")
074                .publish()
075                .accept(
076                        new double[] {
077                            robotToCamera.getX(),
078                            robotToCamera.getY(),
079                            robotToCamera.getZ(),
080                            robotToCamera.getRotation().getMeasureX().in(Degrees),
081                            robotToCamera.getRotation().getMeasureY().in(Degrees),
082                            robotToCamera.getRotation().getMeasureZ().in(Degrees)
083                        });
084    }
085
086    @Override
087    public void updateInputs(VisionIOInputs inputs) {
088        // Update connection status based on whether an update has been seen in the last
089        // 250ms
090        inputs.connected =
091                ((RobotController.getMonotonicTime() - latencySubscriber.getLastChange()) / 1000)
092                        < 250;
093
094        inputs.currentPipeline = (int) pipelineSubscriber.get();
095
096        // Update target observation
097        inputs.latestTargetObservation =
098                new TargetObservation(
099                        Rotation2d.fromDegrees(txSubscriber.get()),
100                        Rotation2d.fromDegrees(tySubscriber.get()));
101
102        // Update orientation for MegaTag 2
103        orientationPublisher.accept(
104                new double[] {rotationSupplier.get().getDegrees(), 0.0, 0.0, 0.0, 0.0, 0.0});
105        NetworkTableInstance.getDefault()
106                .flush(); // Increases network traffic but recommended by Limelight
107
108        // Read new pose observations from NetworkTables
109        Set<Integer> tagIds = new HashSet<>();
110        List<PoseObservation> poseObservations = new LinkedList<>();
111
112        if (megaTagMode == MegaTagMode.MEGATAG_1) {
113            for (TimestampedDoubleArray rawSample : megatag1Subscriber.readQueue()) {
114                if (rawSample.value.length == 0) continue;
115                for (int i = 11; i < rawSample.value.length; i += 7) {
116                    tagIds.add((int) rawSample.value[i]);
117                }
118                poseObservations.add(
119                        new PoseObservation(
120                                // Timestamp, based on server timestamp of publish and latency
121                                rawSample.timestamp * 1.0e-6 - rawSample.value[6] * 1.0e-3,
122
123                                // 3D pose estimate
124                                parsePose(rawSample.value),
125
126                                // Ambiguity, using only the first tag because ambiguity isn't
127                                // applicable for
128                                // multitag
129                                rawSample.value.length >= 18 ? rawSample.value[17] : 0.0,
130
131                                // Tag count
132                                (int) rawSample.value[7],
133
134                                // Average tag distance
135                                rawSample.value[9],
136
137                                // Observation type
138                                PoseObservationType.MEGATAG_1));
139            }
140        }
141
142        if (megaTagMode == MegaTagMode.MEGATAG_2) {
143            for (TimestampedDoubleArray rawSample : megatag2Subscriber.readQueue()) {
144                if (rawSample.value.length == 0) continue;
145                for (int i = 11; i < rawSample.value.length; i += 7) {
146                    tagIds.add((int) rawSample.value[i]);
147                }
148                poseObservations.add(
149                        new PoseObservation(
150                                // Timestamp, based on server timestamp of publish and latency
151                                rawSample.timestamp * 1.0e-6 - rawSample.value[6] * 1.0e-3,
152
153                                // 3D pose estimate
154                                parsePose(rawSample.value),
155
156                                // Ambiguity, zeroed because the pose is already disambiguated
157                                0.0,
158
159                                // Tag count
160                                (int) rawSample.value[7],
161
162                                // Average tag distance
163                                rawSample.value[9],
164
165                                // Observation type
166                                PoseObservationType.MEGATAG_2));
167            }
168        }
169
170        // Save pose observations to inputs object
171        inputs.poseObservations = new PoseObservation[poseObservations.size()];
172        for (int i = 0; i < poseObservations.size(); i++) {
173            inputs.poseObservations[i] = poseObservations.get(i);
174        }
175
176        // Save tag IDs to inputs objects
177        inputs.tagIds = new int[tagIds.size()];
178        int i = 0;
179        for (int id : tagIds) {
180            inputs.tagIds[i++] = id;
181        }
182    }
183
184    /** Parses the 3D pose from a Limelight botpose array. */
185    private static Pose3d parsePose(double[] rawLLArray) {
186        return new Pose3d(
187                rawLLArray[0],
188                rawLLArray[1],
189                rawLLArray[2],
190                new Rotation3d(
191                        Units.degreesToRadians(rawLLArray[3]),
192                        Units.degreesToRadians(rawLLArray[4]),
193                        Units.degreesToRadians(rawLLArray[5])));
194    }
195
196    @Override
197    public void applyOutputs(VisionIOOutputs outputs) {
198        pipelinePublisher.accept(outputs.pipeline);
199
200        this.megaTagMode = outputs.megaTagMode;
201
202        aprilTagWhitelistPublisher.accept(
203                Arrays.stream(outputs.aprilTagWhitelist).asDoubleStream().toArray());
204        imuModePublisher.accept(outputs.imuMode);
205    }
206}