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 com.stuypulse.robot.constants.FieldConstants;
008
009import org.wpilib.math.geometry.Pose3d;
010import org.wpilib.math.geometry.Rotation2d;
011import org.wpilib.math.geometry.Transform3d;
012
013import java.util.HashSet;
014import java.util.LinkedList;
015import java.util.List;
016import java.util.Optional;
017import java.util.Set;
018import org.photonvision.PhotonCamera;
019import org.photonvision.targeting.MultiTargetPNPResult;
020import org.photonvision.targeting.PhotonPipelineResult;
021import org.photonvision.targeting.PhotonTrackedTarget;
022
023/** IO implementation for real PhotonVision hardware. */
024public class VisionIOPhoton implements VisionIO {
025    protected final PhotonCamera camera;
026    protected final Transform3d robotToCamera;
027
028    /**
029     * Creates a new VisionIOPhotonVision.
030     *
031     * @param name The configured name of the camera.
032     * @param robotToCamera The 3D position of the camera relative to the robot.
033     */
034    public VisionIOPhoton(String name, Transform3d robotToCamera) {
035        camera = new PhotonCamera(name);
036        this.robotToCamera = robotToCamera;
037    }
038
039    @Override
040    public void updateInputs(VisionIOInputs inputs) {
041        inputs.connected = camera.isConnected();
042
043        // Read new camera observations
044        Set<Short> tagIds = new HashSet<>();
045        List<PoseObservation> poseObservations = new LinkedList<>();
046        for (PhotonPipelineResult result : camera.getAllUnreadResults()) {
047            // Update latest target observation
048            if (result.hasTargets()) {
049                inputs.latestTargetObservation =
050                        new TargetObservation(
051                                Rotation2d.fromDegrees(result.getBestTarget().getYaw()),
052                                Rotation2d.fromDegrees(result.getBestTarget().getPitch()));
053            } else {
054                inputs.latestTargetObservation =
055                        new TargetObservation(new Rotation2d(), new Rotation2d());
056            }
057
058            // Add pose observation
059            if (result.multitagResult.isPresent()) { // Multitag result
060                MultiTargetPNPResult multitagResult = result.multitagResult.get();
061
062                // Calculate robot pose
063                Transform3d fieldToCamera = multitagResult.estimatedPose.best;
064                Transform3d fieldToRobot = fieldToCamera.plus(robotToCamera.inverse());
065                Pose3d robotPose =
066                        new Pose3d(fieldToRobot.getTranslation(), fieldToRobot.getRotation());
067
068                // Calculate average tag distance
069                double totalTagDistance = 0.0;
070                for (PhotonTrackedTarget target : result.targets) {
071                    totalTagDistance += target.bestCameraToTarget.getTranslation().getNorm();
072                }
073
074                // Add tag IDs
075                tagIds.addAll(multitagResult.fiducialIDsUsed);
076
077                // Add observation
078                poseObservations.add(
079                        new PoseObservation(
080                                result.getTimestampSeconds(), // Timestamp
081                                robotPose, // 3D pose estimate
082                                multitagResult.estimatedPose.ambiguity, // Ambiguity
083                                multitagResult.fiducialIDsUsed.size(), // Tag count
084                                totalTagDistance / result.targets.size(), // Average tag distance
085                                PoseObservationType.PHOTONVISION)); // Observation type
086
087            } else if (!result.targets.isEmpty()) { // Single tag result
088                PhotonTrackedTarget target = result.targets.get(0);
089
090                // Calculate robot pose
091                Optional<Pose3d> tagPose =
092                        FieldConstants.FIELD_LAYOUT.getTagPose(target.fiducialId);
093                if (tagPose.isPresent()) {
094                    Transform3d fieldToTarget =
095                            new Transform3d(
096                                    tagPose.get().getTranslation(), tagPose.get().getRotation());
097                    Transform3d cameraToTarget = target.bestCameraToTarget;
098                    Transform3d fieldToCamera = fieldToTarget.plus(cameraToTarget.inverse());
099                    Transform3d fieldToRobot = fieldToCamera.plus(robotToCamera.inverse());
100                    Pose3d robotPose =
101                            new Pose3d(fieldToRobot.getTranslation(), fieldToRobot.getRotation());
102
103                    // Add tag ID
104                    tagIds.add((short) target.fiducialId);
105
106                    // Add observation
107                    poseObservations.add(
108                            new PoseObservation(
109                                    result.getTimestampSeconds(), // Timestamp
110                                    robotPose, // 3D pose estimate
111                                    target.poseAmbiguity, // Ambiguity
112                                    1, // Tag count
113                                    cameraToTarget
114                                            .getTranslation()
115                                            .getNorm(), // Average tag distance
116                                    PoseObservationType.PHOTONVISION)); // Observation type
117                }
118            }
119        }
120
121        // Save pose observations to inputs object
122        inputs.poseObservations = new PoseObservation[poseObservations.size()];
123        for (int i = 0; i < poseObservations.size(); i++) {
124            inputs.poseObservations[i] = poseObservations.get(i);
125        }
126
127        // Save tag IDs to inputs objects
128        inputs.tagIds = new int[tagIds.size()];
129        int i = 0;
130        for (int id : tagIds) {
131            inputs.tagIds[i++] = id;
132        }
133    }
134
135    @Override
136    public void applyOutputs(VisionIOOutputs outputs) {
137        camera.setPipelineIndex(outputs.pipeline);
138    }
139}