"""Prove and export the recovered p01 legacy-H5-row to MPIIGaze-frame mapping."""
from __future__ import annotations

import csv
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


LEGACY = ROOT / "data" / "processed" / "p01_v16.h5"
CLEAN = ROOT / "data" / "processed_kd_clean_v1" / "cache" / "p01.full.h5"
OUT = ROOT / "artifacts" / "kd-teacher-trap-diagnostic"


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 decode(values: np.ndarray) -> list[str]:
    return [value.decode("utf-8") if isinstance(value, bytes) else str(value) for value in values]


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 expectation(values: np.ndarray) -> np.ndarray:
    return softmax(values) @ np.arange(90, dtype=np.float64) * 2.0 - 90.0


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 angular_error(prediction: np.ndarray, target: np.ndarray) -> np.ndarray:
    prediction /= np.linalg.norm(prediction, axis=1, keepdims=True)
    target /= np.linalg.norm(target, axis=1, keepdims=True)
    return np.rad2deg(np.arccos(np.clip(np.sum(prediction * target, axis=1), -1.0, 1.0)))


def main() -> None:
    with h5py.File(LEGACY, "r") as legacy, h5py.File(CLEAN, "r") as clean:
        legacy_count = int(legacy["left_gaze"].shape[0])
        clean_count = int(clean["left_gaze"].shape[0])
        left_equal = bool(np.array_equal(legacy["left_gaze"][:], clean["left_gaze"][:]))
        right_equal = bool(np.array_equal(legacy["right_gaze"][:], clean["right_gaze"][:]))
        landmark_delta = np.abs(
            legacy["landmarks"][:].astype(np.float64) - clean["landmarks"][:].astype(np.float64)
        )
        row_landmark_max = landmark_delta.reshape(clean_count, -1).max(axis=1)
        paths = decode(clean["relative_frame_path"][:])
        ids = decode(clean["sample_id"][:])
        image_hashes = decode(clean["raw_image_sha256"][:])
        source_indices = clean["source_index"][:].astype(np.int64)

        target_deg = np.rad2deg(clean["left_gaze"][:].astype(np.float64))
        target_vector = vectors(target_deg[:, 0], target_deg[:, 1])
        legacy_pitch = expectation(legacy["teacher_pitch_logits"][:])
        legacy_yaw = expectation(legacy["teacher_yaw_logits"][:])
        corrected_raw_pitch = expectation(clean["teacher_pitch_logits_raw"][:])
        corrected_raw_yaw = expectation(clean["teacher_yaw_logits_raw"][:])
        corrected_aligned_vector = clean["teacher_aligned_vector"][:].astype(np.float64)
        legacy_error = angular_error(vectors(legacy_pitch, legacy_yaw), target_vector.copy())
        corrected_raw_error = angular_error(
            vectors(corrected_raw_pitch, corrected_raw_yaw), target_vector.copy()
        )
        corrected_aligned_error = angular_error(corrected_aligned_vector, target_vector.copy())

        csv_path = OUT / "p01_legacy_row_to_frame_mapping.csv"
        with csv_path.open("w", encoding="utf-8", newline="") as stream:
            writer = csv.DictWriter(
                stream,
                fieldnames=(
                    "legacy_h5_row", "source_index", "sample_id", "relative_frame_path",
                    "raw_image_sha256", "left_gaze_exact", "right_gaze_exact",
                    "landmark_row_max_abs_diff",
                ),
            )
            writer.writeheader()
            for index in range(clean_count):
                writer.writerow(
                    {
                        "legacy_h5_row": index,
                        "source_index": int(source_indices[index]),
                        "sample_id": ids[index],
                        "relative_frame_path": paths[index],
                        "raw_image_sha256": image_hashes[index],
                        "left_gaze_exact": bool(
                            np.array_equal(legacy["left_gaze"][index], clean["left_gaze"][index])
                        ),
                        "right_gaze_exact": bool(
                            np.array_equal(legacy["right_gaze"][index], clean["right_gaze"][index])
                        ),
                        "landmark_row_max_abs_diff": float(row_landmark_max[index]),
                    }
                )

        mapping_pass = bool(
            legacy_count == clean_count
            and left_equal
            and right_equal
            and np.max(row_landmark_max) <= 1e-6
            and np.all(np.diff(source_indices) > 0)
            and len(ids) == len(set(ids))
        )
        report = {
            "legacy_h5": str(LEGACY),
            "legacy_h5_sha256": file_sha256(LEGACY),
            "clean_h5": str(CLEAN),
            "clean_h5_sha256": file_sha256(CLEAN),
            "legacy_rows": legacy_count,
            "clean_rows": clean_count,
            "left_gaze_all_rows_exact": left_equal,
            "right_gaze_all_rows_exact": right_equal,
            "landmarks_global_max_abs_diff": float(np.max(landmark_delta)),
            "landmarks_global_mean_abs_diff": float(np.mean(landmark_delta)),
            "landmark_rows_with_max_abs_diff_le_1e_6": int(np.sum(row_landmark_max <= 1e-6)),
            "source_indices_strictly_increasing": bool(np.all(np.diff(source_indices) > 0)),
            "unique_sample_ids": len(set(ids)),
            "mapping_csv": str(csv_path),
            "mapping_csv_sha256": file_sha256(csv_path),
            "mapping_status": "RECOVERED_FOR_P01" if mapping_pass else "NOT_PROVEN",
            "first_five": [
                {
                    "legacy_h5_row": index,
                    "source_index": int(source_indices[index]),
                    "sample_id": ids[index],
                    "relative_frame_path": paths[index],
                    "raw_image_sha256": image_hashes[index],
                }
                for index in range(min(5, clean_count))
            ],
            "teacher_comparison_against_roll_corrected_left_gaze": {
                "legacy_cached_teacher_3d_error_mean_deg": float(legacy_error.mean()),
                "corrected_loader_raw_teacher_3d_error_mean_deg": float(corrected_raw_error.mean()),
                "corrected_loader_roll_aligned_teacher_3d_error_mean_deg": float(
                    corrected_aligned_error.mean()
                ),
                "corrected_aligned_minus_legacy_error_deg": float(
                    corrected_aligned_error.mean() - legacy_error.mean()
                ),
                "legacy_vs_corrected_pitch_prediction_pearson": float(
                    np.corrcoef(legacy_pitch, corrected_raw_pitch)[0, 1]
                ),
                "legacy_vs_corrected_yaw_prediction_pearson": float(
                    np.corrcoef(legacy_yaw, corrected_raw_yaw)[0, 1]
                ),
            },
            "interpretation": (
                "The identical accepted-row count, exact left/right gaze arrays, and <=1e-6 "
                "landmark agreement establish the p01 legacy row order. Patch differences do "
                "not invalidate identity because patch-processing variants were already known."
            ),
        }
    output = OUT / "p01_legacy_row_mapping_audit.json"
    output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8", newline="\n")
    print(json.dumps(report, indent=2))
    if not mapping_pass:
        raise SystemExit(1)


if __name__ == "__main__":
    main()