"""Audit cached teacher logits under plausible pitch/yaw coordinate conventions."""
from pathlib import Path
import sys

ROOT = Path(r"E:\Gaze_estimation")
sys.path.insert(0, str(ROOT / ".codex_deps"))
import h5py
import numpy as np
import pandas as pd

OUT = ROOT / "artifacts" / "kd-teacher-trap-diagnostic"


def expectation(logits):
    x = logits - logits.max(axis=1, keepdims=True)
    p = np.exp(x); p /= p.sum(axis=1, keepdims=True)
    return (p * np.arange(90)).sum(axis=1) * 2.0 - 90.0


def vectors(pitch_deg, yaw_deg):
    pitch = np.deg2rad(pitch_deg); yaw = np.deg2rad(yaw_deg)
    return np.column_stack([-np.cos(pitch) * np.sin(yaw), -np.sin(pitch), -np.cos(pitch) * np.cos(yaw)])


rows = []
for subject in ("p01", "p08", "p11"):
    path = ROOT / "data" / "processed" / f"{subject}_v16.h5"
    with h5py.File(path, "r") as f:
        a = expectation(f["teacher_pitch_logits"][:])
        b = expectation(f["teacher_yaw_logits"][:])
        gt = np.rad2deg(f["left_gaze"][:])
    gt_vec = vectors(gt[:, 0], gt[:, 1])
    for swap in (False, True):
        base_p, base_y = (b, a) if swap else (a, b)
        for pitch_sign in (-1, 1):
            for yaw_sign in (-1, 1):
                p = pitch_sign * base_p; y = yaw_sign * base_y
                axis = np.abs(np.column_stack([p, y]) - gt).mean(axis=1)
                pred_vec = vectors(p, y)
                angular = np.rad2deg(np.arccos(np.clip((pred_vec * gt_vec).sum(axis=1), -1, 1)))
                rows.append({
                    "subject": subject, "swap_pitch_yaw": swap,
                    "pitch_sign": pitch_sign, "yaw_sign": yaw_sign,
                    "samples": len(gt), "axis_mae_deg": axis.mean(),
                    "angular_3d_deg": angular.mean(),
                })
df = pd.DataFrame(rows)
summary = df.groupby(["swap_pitch_yaw", "pitch_sign", "yaw_sign"], as_index=False)[["axis_mae_deg", "angular_3d_deg"]].mean()
summary["rank_axis"] = summary.axis_mae_deg.rank(method="min").astype(int)
summary["rank_3d"] = summary.angular_3d_deg.rank(method="min").astype(int)
df.to_csv(OUT / "teacher_coordinate_by_subject.csv", index=False)
summary.sort_values("axis_mae_deg").to_csv(OUT / "teacher_coordinate_summary.csv", index=False)
print(summary.sort_values("axis_mae_deg").to_string(index=False))