#!/usr/bin/env python3
"""Diagnose ego-cam ArUco detection on the saved raw frame.

Loads /tmp/cal_raw_ego_camera.png and tries detectMarkers with:
  1. DEFAULT DetectorParameters (what our auto-cal uses)
  2. PERMISSIVE params (looser thresholds, corner refinement on)
  3. Preprocessing: adaptive contrast (CLAHE) + detect
  4. Downscale-then-detect (some markers detect better at smaller scale)
Reports marker count + IDs for each. Saves annotated PNG.
"""
import cv2
import sys
from pathlib import Path
import numpy as np

# Load the saved raw ego frame
IN_PATH = Path("/tmp/cal_raw_ego_camera.png")
if not IN_PATH.exists():
    print(f"missing {IN_PATH}", file=sys.stderr); sys.exit(1)

bgr = cv2.imread(str(IN_PATH))
if bgr is None:
    print("couldn't decode PNG", file=sys.stderr); sys.exit(1)
print(f"[dx] frame shape: {bgr.shape}")
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
d = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)

# Look for arm1 IDs (217-225) and arm2 IDs (226-234)
_ARM1_IDS = set(range(217, 226))
_ARM2_IDS = set(range(226, 235))

def _classify(ids):
    if ids is None:
        return 0, 0, 0, []
    ids_flat = ids.flatten()
    a1 = sum(1 for i in ids_flat if int(i) in _ARM1_IDS)
    a2 = sum(1 for i in ids_flat if int(i) in _ARM2_IDS)
    return len(ids_flat), a1, a2, sorted(int(i) for i in ids_flat)


def _run(name, params, prep=None):
    g = gray if prep is None else prep(gray)
    corners, ids, rejected = cv2.aruco.detectMarkers(g, d, parameters=params)
    n, a1, a2, id_list = _classify(ids)
    print(f"[{name:>28}] total={n:2d}  arm1={a1}  arm2={a2}  "
          f"rejected={len(rejected)}  ids={id_list}")
    return corners, ids


# 1. Default params
p_default = cv2.aruco.DetectorParameters()
c1, i1 = _run("default", p_default)

# 2. Permissive params
p_perm = cv2.aruco.DetectorParameters()
p_perm.cornerRefinementMethod = cv2.aruco.CORNER_REFINE_SUBPIX
p_perm.minMarkerPerimeterRate = 0.01     # default 0.03 — allow smaller markers
p_perm.maxMarkerPerimeterRate = 4.0
p_perm.polygonalApproxAccuracyRate = 0.05   # default 0.03 — more tolerant
p_perm.maxErroneousBitsInBorderRate = 0.5   # default 0.35 — more border noise ok
p_perm.errorCorrectionRate = 0.8            # default 0.6 — recover more bits
p_perm.adaptiveThreshWinSizeMin = 3
p_perm.adaptiveThreshWinSizeMax = 45         # default 23
p_perm.adaptiveThreshWinSizeStep = 4         # default 10
c2, i2 = _run("permissive", p_perm)

# 3. CLAHE contrast preprocess + permissive
def _clahe(g):
    return cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(g)
c3, i3 = _run("CLAHE + permissive", p_perm, prep=_clahe)

# 4. Detect on 2x downscaled (sometimes helps)
def _down2x(g):
    return cv2.resize(g, (g.shape[1] // 2, g.shape[0] // 2),
                       interpolation=cv2.INTER_AREA)
c4, i4 = _run("2x downscale + default", p_default, prep=_down2x)
c5, i5 = _run("2x downscale + permissive", p_perm, prep=_down2x)

# 5. Detect on 0.5x upscaled
def _up2x(g):
    return cv2.resize(g, (g.shape[1] * 2, g.shape[0] * 2),
                       interpolation=cv2.INTER_CUBIC)
c6, i6 = _run("2x upscale + default", p_default, prep=_up2x)
c7, i7 = _run("2x upscale + permissive", p_perm, prep=_up2x)

# Pick whichever variant found the most for the annotated dump.
best_corners, best_ids, best_name = None, None, None
best_n = -1
for name, corners, ids in [
        ("default", c1, i1), ("permissive", c2, i2),
        ("clahe_permissive", c3, i3),
        ("2x_down_default", c4, i4), ("2x_down_permissive", c5, i5),
        ("2x_up_default", c6, i6), ("2x_up_permissive", c7, i7)]:
    n = 0 if ids is None else len(ids.flatten())
    if n > best_n:
        best_n = n
        best_corners, best_ids, best_name = corners, ids, name

if best_ids is not None and best_n > 0:
    # If the best was on a scaled image, annotate original with rescaled corners.
    annotated = bgr.copy()
    if best_name in ("2x_down_default", "2x_down_permissive"):
        scale = 2.0
        corners_up = [c * scale for c in best_corners]
    elif best_name in ("2x_up_default", "2x_up_permissive"):
        scale = 0.5
        corners_up = [c * scale for c in best_corners]
    else:
        corners_up = best_corners
    cv2.aruco.drawDetectedMarkers(annotated, corners_up, best_ids)
    out = Path("/tmp/ego_diagnostic_annotated.png")
    cv2.imwrite(str(out), annotated)
    print(f"\n[dx] best variant '{best_name}' found {best_n} markers — "
          f"annotated dump at {out}")
else:
    print("\n[dx] NO variant found any markers on this frame")