"""Validate a continuous-vector KD cache and reject ambiguous KL target fields."""
from __future__ import annotations

import argparse
import hashlib
import json
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


CACHE_ROOT = ROOT / "data" / "processed_kd_clean_v1" / "cache"


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest().upper()


def softmax(values: np.ndarray) -> np.ndarray:
    shifted = values.astype(np.float64) - values.max(axis=1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=1, keepdims=True)


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


def rotate(values: np.ndarray, roll_deg: np.ndarray) -> np.ndarray:
    angle = np.deg2rad(roll_deg)
    result = values.copy()
    result[:, 0] = np.cos(angle) * values[:, 0] - np.sin(angle) * values[:, 1]
    result[:, 1] = np.sin(angle) * values[:, 0] + np.cos(angle) * values[:, 1]
    return result


def angular(first: np.ndarray, second: np.ndarray) -> np.ndarray:
    first = first / np.linalg.norm(first, axis=1, keepdims=True)
    second = second / np.linalg.norm(second, axis=1, keepdims=True)
    return np.rad2deg(np.arccos(np.clip(np.sum(first * second, axis=1), -1.0, 1.0)))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--participant", required=True)
    parser.add_argument("--tag", default="official448_pointkd")
    args = parser.parse_args()
    cache_path = CACHE_ROOT / f"{args.participant}.{args.tag}.h5"
    summary_path = CACHE_ROOT / f"{args.participant}.{args.tag}.summary.json"
    validation_path = CACHE_ROOT / f"{args.participant}.{args.tag}.validation.json"
    if validation_path.exists():
        raise FileExistsError(f"refusing to overwrite validation: {validation_path}")
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    errors: list[str] = []
    if file_sha256(cache_path) != summary["cache_sha256"]:
        errors.append("cache SHA-256 differs from summary")
    required = {
        "sample_id", "relative_frame_path", "source_index", "left_patches", "landmarks",
        "left_gaze", "left_roll_deg", "teacher_pitch_logits_raw", "teacher_yaw_logits_raw",
        "teacher_target_vector", "teacher_target_pitch_deg", "teacher_target_yaw_deg",
        "teacher_target_error_deg",
    }
    prohibited = {"teacher_pitch_logits", "teacher_yaw_logits", "teacher_aligned_vector"}
    metrics = {}
    with h5py.File(cache_path, "r") as handle:
        if required - set(handle):
            errors.append(f"missing required fields: {sorted(required - set(handle))}")
        if prohibited & set(handle):
            errors.append(f"ambiguous roll-rebinned KL fields present: {sorted(prohibited & set(handle))}")
        counts = {name: int(handle[name].shape[0]) for name in required if name in handle}
        if len(set(counts.values())) != 1:
            errors.append(f"row counts differ: {counts}")
        raw_pitch = handle["teacher_pitch_logits_raw"][:]
        raw_yaw = handle["teacher_yaw_logits_raw"][:]
        pitch = softmax(raw_pitch) @ np.arange(90) * 4.0 - 180.0
        yaw = softmax(raw_yaw) @ np.arange(90) * 4.0 - 180.0
        expected_vector = rotate(vectors(pitch, yaw), handle["left_roll_deg"][:])
        stored_vector = handle["teacher_target_vector"][:].astype(np.float64)
        metrics["max_teacher_target_vector_abs_diff"] = float(np.max(np.abs(expected_vector - stored_vector)))
        gaze_deg = np.rad2deg(handle["left_gaze"][:].astype(np.float64))
        recomputed_error = angular(stored_vector, vectors(gaze_deg[:, 0], gaze_deg[:, 1]))
        stored_error = handle["teacher_target_error_deg"][:]
        metrics["max_teacher_error_recompute_difference_deg"] = float(
            np.max(np.abs(recomputed_error - stored_error))
        )
        metrics["teacher_target_error_mean_deg"] = float(recomputed_error.mean())
        metrics["max_teacher_target_norm_error"] = float(
            np.max(np.abs(np.linalg.norm(stored_vector, axis=1) - 1.0))
        )
        numeric = [name for name in required if name in handle and handle[name].dtype.kind not in "OSU"]
        if any(not np.isfinite(handle[name][:]).all() for name in numeric):
            errors.append("non-finite numeric values found")
        if metrics["max_teacher_target_vector_abs_diff"] > 1e-6:
            errors.append("stored teacher target vectors do not reproduce")
        if metrics["max_teacher_error_recompute_difference_deg"] > 1e-5:
            errors.append("stored teacher errors do not reproduce")
        if handle.attrs.get("approved_distillation") != "continuous 3D vector loss only":
            errors.append("continuous-only distillation approval attribute is missing")
    report = {
        "schema": "mpiigaze-point-kd-cache-validation-v3",
        "participant": args.participant,
        "cache_path": str(cache_path.resolve()),
        "cache_sha256": file_sha256(cache_path),
        "rows": summary["accepted_rows"],
        "metrics": metrics,
        "errors": errors,
        "pass": not errors,
    }
    validation_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8", newline="\n")
    print(json.dumps(report, indent=2))
    if errors:
        raise SystemExit(1)


if __name__ == "__main__":
    main()