#!/usr/bin/env python3
"""Compare serialized training diagnostics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def _compare(left: np.ndarray, right: np.ndarray, atol: float, rtol: float):
    delta = np.abs(left.astype(np.float64) - right.astype(np.float64))
    passed = np.isclose(left, right, atol=atol, rtol=rtol)
    return {
        "status": "PASS" if bool(np.all(passed)) else "FAIL",
        "shape": list(left.shape),
        "max_abs": float(delta.max()) if delta.size else 0.0,
        "mean_abs": float(delta.mean()) if delta.size else 0.0,
        "outside_tolerance": int((~passed).sum()),
        "elements": int(delta.size),
        "atol": atol,
        "rtol": rtol,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("reference", type=Path)
    parser.add_argument("candidate", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    with np.load(args.reference) as left, np.load(args.candidate) as right:
        stages = {
            "logits": _compare(left["logits"], right["logits"], 1e-5, 1e-5),
            "loss": _compare(left["loss"], right["loss"], 1e-5, 1e-5),
            "gradients": _compare(left["gradients"], right["gradients"], 1e-4, 1e-5),
            "parameters_after_step": _compare(
                left["parameters"], right["parameters"], 1e-5, 1e-5
            ),
        }
        if "history" in left.files and "history" in right.files:
            stages["history"] = _compare(left["history"], right["history"], 1e-5, 1e-5)
        if "exp_avg" in left.files and "exp_avg" in right.files:
            stages["optimizer_exp_avg"] = _compare(
                left["exp_avg"], right["exp_avg"], 1e-5, 1e-5
            )
            stages["optimizer_exp_avg_sq"] = _compare(
                left["exp_avg_sq"], right["exp_avg_sq"], 1e-5, 1e-5
            )
            stages["optimizer_steps"] = _compare(
                left["optimizer_steps"], right["optimizer_steps"], 0.0, 0.0
            )
    report = {"schema_version": 1, "stages": stages}
    report["overall_status"] = (
        "PASS"
        if all(stage["status"] == "PASS" for stage in stages.values())
        else "FAIL"
    )
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "training_step.json").write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n"
    )
    lines = [
        "# One-step training parity",
        "",
        f"Overall status: **{report['overall_status']}**",
        "",
        "| Stage | Status | Max abs | Outside tolerance |",
        "|---|---|---:|---:|",
    ]
    for name, stage in stages.items():
        lines.append(
            f"| {name} | {stage['status']} | {stage['max_abs']} | "
            f"{stage['outside_tolerance']} |"
        )
    (args.output / "training_step.md").write_text("\n".join(lines) + "\n")
    print(json.dumps({"overall_status": report["overall_status"]}, indent=2))
    return 0 if report["overall_status"] == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())