#!/usr/bin/env python3
"""Audit the three-depth raw-KP mixed-traffic control."""
import argparse
import hashlib
import json
import math
import os
import sys


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from experiments.crossover_hardware import assert_hardware_report  # noqa: E402
from experiments.kp_raw_traffic_scaling import (  # noqa: E402
    DEPTHS,
    RESULT_ROOT,
    jobs,
    registry_sha256,
)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def read_json(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


def finite(value):
    return (
        isinstance(value, (int, float))
        and not isinstance(value, bool)
        and math.isfinite(float(value))
    )


def audit_job(job, launch):
    manifest_path = job["output"] + ".manifest.json"
    if not os.path.isfile(manifest_path):
        return {
            "depth": job["depth"],
            "architecture": job["architecture"],
            "status": "missing_manifest",
            "finite": False,
            "final_validation_accuracy": None,
        }
    manifest = read_json(manifest_path)
    assert manifest["source"] == launch["source"]
    assert manifest["command"] == job["command"]
    assert manifest["hardware_lock"]["physical_gpu_name"] == "NVIDIA RTX A6000"
    assert_hardware_report(manifest["hardware_lock"], launch["hardware_policy"])
    if manifest["status"] != "completed" or not manifest["output_exists"]:
        return {
            "depth": job["depth"],
            "architecture": job["architecture"],
            "status": manifest["status"],
            "finite": False,
            "final_validation_accuracy": None,
            "driver_wall_seconds": manifest["driver_wall_seconds"],
        }
    assert sha256(job["output"]) == manifest["output_sha256"]
    record = read_json(job["output"])
    args = record["args"]
    expected = {
        "mode": "kp_traffic",
        "traffic_rule": "raw",
        "traffic_ratio": 4.0,
        "traffic_seed": 5000,
        "traffic_calibration_examples": 64,
        "predictor_mode": "closed_form",
        "predictor_warmup_steps": 1,
        "predictor_every": 0,
        "neutral_projection": 1,
        "depth": job["depth"],
        "width": 16,
        "seed": 0,
        "loader_seed": 0,
        "split_seed": 2027,
        "epochs": 200,
        "eval_split": "validation",
        "eval_every": 1,
        "lr": 0.1,
        "output_lr": 0.1,
    }
    for key, wanted in expected.items():
        assert args[key] == wanted, (key, args[key], wanted)
    assert record["provenance"] == {
        "git_commit": launch["source"]["git_commit"],
        "git_tracked_dirty": False,
    }
    assert record["split"]["train_examples"] == 45_000
    assert record["split"]["validation_examples"] == 5_000
    assert record["split"]["test_examples"] == 10_000
    assert record["evaluation_protocol"]["test_evaluations"] == 0
    assert record["counters"]["ordinary_examples"] == 9_000_000
    assert record["counters"]["traffic_calibration_examples"] == 64
    assert record["counters"]["predictor_update_examples"] == 64
    assert record["counters"]["neutral_projection_examples"] == 9_000_000
    assert record["counters"]["logical_batch_loss_queries"] == 0
    final = record["final"]
    complete = len(record["epochs"]) == 200
    is_finite = bool(
        complete
        and final["finite"] is True
        and finite(final["accuracy"])
        and finite(final["loss"])
    )
    return {
        "depth": job["depth"],
        "architecture": job["architecture"],
        "status": "completed" if complete else "incomplete_trajectory",
        "finite": is_finite,
        "final_validation_accuracy": (
            float(final["accuracy"]) if finite(final["accuracy"]) else None
        ),
        "final_validation_loss": (
            float(final["loss"]) if finite(final["loss"]) else None
        ),
        "first_nonfinite_step": record.get("first_nonfinite_step"),
        "total_wall_seconds": record["timing"]["total_timed_wall_s"],
        "peak_memory_allocated_bytes": record["hardware"][
            "peak_memory_allocated_bytes"
        ],
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--out", default=os.path.join(RESULT_ROOT, "r1_audit.json")
    )
    args = parser.parse_args()
    launch_path = os.path.join(RESULT_ROOT, "r1_launch.json")
    launch = read_json(launch_path)
    registered = jobs()
    assert launch["registry_sha256"] == registry_sha256(registered)
    records = [audit_job(job, launch) for job in registered]
    report = {
        "audit_status": "passed",
        "stage": "raw_kp_traffic_scaling_r1",
        "complete_registry": True,
        "num_expected_cells": len(DEPTHS),
        "num_audited_cells": len(records),
        "num_finite_cells": sum(record["finite"] for record in records),
        "failure_retaining": True,
        "test_policy": "none",
        "records": records,
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()