import edu.wpi.first.math.Matrix;
import edu.wpi.first.math.VecBuilder;
import edu.wpi.first.math.numbers.*;

// ── Base constants ─────────────────────────────────────────────────────────
// Start with these values and iterate based on observed performance.
// BASE_XY_STDDEV: position stddev at 1m with 1 tag
// Decreasing this makes vision more aggressive at correcting odometry.
private static final double BASE_XY_STDDEV = 0.03;  // meters at 1m/1tag

// Theta stddev — effectively ignore vision heading.
// Only change this if you have multi-tag with very close range AND
// you've verified vision heading is more accurate than your gyro.
private static final double THETA_STDDEV = 9999.0;   // radians (effectively infinite)

/**
 * Computes vision measurement standard deviations based on detection quality.
 *
 * @param distanceMeters  Distance to the nearest detected tag (meters).
 *                        From: target.getBestCameraToTarget().getTranslation().getNorm()
 * @param numTags         Number of tags used in this estimate.
 *                        From: estimate.targetsUsed.size()
 * @return Matrix<N3,N1> of [x stddev, y stddev, theta stddev]
 */
private Matrix<N3, N1> computeVisionStdDevs(double distanceMeters, int numTags) {
    // XY error scales with distance² and improves proportionally with tag count
    double xyStdDev = BASE_XY_STDDEV * (distanceMeters * distanceMeters) / numTags;

    // Cap to prevent extreme values at long range from blowing up the filter.
    // Measurements beyond 7m or so should be rejected in Lesson 8 anyway.
    xyStdDev = Math.min(xyStdDev, 5.0);

    return VecBuilder.fill(xyStdDev, xyStdDev, THETA_STDDEV);
}

// Usage in periodic() — extract distance from the detected targets:
estimateOpt.ifPresent(estimate -> {
    int numTags = estimate.targetsUsed.size();

    // Average distance across all detected tags (better than using only the closest)
    double totalDist = estimate.targetsUsed.stream()
        .mapToDouble(t -> t.getBestCameraToTarget().getTranslation().getNorm())
        .sum();
    double avgDist = totalDist / numTags;

    Matrix<N3, N1> stdDevs = computeVisionStdDevs(avgDist, numTags);

    // Pass to drivetrain with both pose and per-measurement stddevs
    m_drive.addVisionMeasurement(
        estimate.estimatedPose.toPose2d(),
        estimate.timestampSeconds,
        stdDevs
    );

    // Log for tuning iteration (Lesson 7)
    SmartDashboard.putNumber("Vision/StdDevXY", stdDevs.get(0, 0));
    SmartDashboard.putNumber("Vision/AvgTagDist", avgDist);
    SmartDashboard.putNumber("Vision/NumTags", numTags);
});

public class VisionConstants {

    // ── Camera geometry ───────────────────────────────────────────────────────
    // Re-measure after any change to camera mounting.
    public static final Transform3d ROBOT_TO_FRONT_CAMERA =
        new Transform3d(new Translation3d(0.30, -0.10, 0.50),
                       new Rotation3d(0, Units.degreesToRadians(-15), 0));

    // ── Field ─────────────────────────────────────────────────────────────────
    public static final double FIELD_LENGTH_METERS = 16.54;
    public static final double FIELD_WIDTH_METERS  = 8.21;

    // ── Kalman filter — odometry trust ────────────────────────────────────────
    // Measured from actual robot drift tests (5m path, 5 runs).
    // Increase if your robot slips more (carpet-to-carpet variation).
    public static final Matrix<N3, N1> ODOMETRY_STD_DEVS =
        VecBuilder.fill(0.1, 0.1, 0.1);

    // ── Kalman filter — vision trust formula ──────────────────────────────────
    // xyStdDev = BASE_XY_STDDEV * distance² / numTags
    //
    // START HERE for tuning. Lower this → trust vision more aggressively.
    // Typical range: 0.01 (very aggressive) to 0.1 (very conservative).
    // Raise if vision is adding noise instead of correcting drift.
    public static final double BASE_XY_STDDEV = 0.03;

    // Heading from vision is rarely better than the gyro.
    // Set to a large number to effectively ignore vision heading.
    // Only reduce if: multi-tag at close range AND gyro has known drift issues.
    public static final double VISION_THETA_STDDEV = 9999.0;

    // Maximum distance at which to accept a vision measurement.
    // Beyond this, stddev is so large that measurements add noise.
    // Also used as an explicit filter in Lesson 8.
    public static final double MAX_VISION_DISTANCE_METERS = 5.0;
}