#!/usr/bin/env python3
"""Join the three audited families into the exact 81-cell X1 panel."""
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_registry import METHODS, primary_cells


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


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 finite_number(value):
    return (
        isinstance(value, (int, float))
        and not isinstance(value, bool)
        and math.isfinite(float(value))
    )


def assert_plain_header(report):
    assert report["gate"] == "pass"
    assert report["stage"] == "plain_cnn_p2"
    assert report["num_expected_records"] == 27
    assert report["num_audited_records"] == 27
    assert report["missing_experiments"] == []
    assert report["test_policy"] == "none"
    assert len(report["records"]) == 27


def assert_native_header(report, expected_stage):
    assert report["audit_status"] == "passed"
    assert report["stage"] == expected_stage
    assert report["complete_grid"] is True
    assert report["failure_retaining"] is True
    assert report["num_expected_cells"] == 27
    assert report["num_audited_cells"] == 27
    assert report["test_policy"] == "none"
    assert len(report["records"]) == 27


def plain_record(record):
    raw_accuracy = record.get("best_validation_accuracy")
    valid_accuracy = finite_number(raw_accuracy)
    if valid_accuracy:
        assert 0.0 <= float(raw_accuracy) <= 100.0
    is_finite = (
        record.get("status") == "completed"
        and record.get("outcome") == "finite"
        and valid_accuracy
    )
    hardware = record.get("physical_gpu") or {}
    return {
        "cell_id": f"{record['architecture']}::{record['method']}",
        "status": record.get("status"),
        "finite": is_finite,
        "validation_metric_name": "top1_accuracy",
        "validation_metric": (
            float(raw_accuracy) / 100.0 if valid_accuracy else None
        ),
        "best_validation_accuracy": (
            float(raw_accuracy) / 100.0 if valid_accuracy else None
        ),
        "final_validation_accuracy": (
            float(record["final_validation_accuracy"]) / 100.0
            if finite_number(record.get("final_validation_accuracy"))
            else None
        ),
        "validation_loss": (
            float(record["final_validation_loss"])
            if finite_number(record.get("final_validation_loss"))
            else None
        ),
        "wall_seconds": (
            float(record["driver_wall_seconds"])
            if finite_number(record.get("driver_wall_seconds"))
            else None
        ),
        "peak_memory_mib": (
            float(hardware["peak_process_gpu_memory_mib"])
            if finite_number(hardware.get("peak_process_gpu_memory_mib"))
            else None
        ),
        "work": record.get("work"),
    }


def resnet_record(record):
    raw_accuracy = record.get("best_validation_accuracy")
    valid_accuracy = finite_number(raw_accuracy)
    is_finite = record.get("finite") is True and valid_accuracy
    memory = record.get("peak_memory_allocated_bytes")
    return {
        "cell_id": record["cell_id"],
        "status": record.get("status"),
        "finite": is_finite,
        "validation_metric_name": "top1_accuracy",
        "validation_metric": (
            float(raw_accuracy) if valid_accuracy else None
        ),
        "best_validation_accuracy": (
            float(raw_accuracy) if valid_accuracy else None
        ),
        "final_validation_accuracy": (
            float(record["final_validation_accuracy"])
            if finite_number(record.get("final_validation_accuracy"))
            else None
        ),
        "validation_loss": (
            float(record["final_validation_loss"])
            if finite_number(record.get("final_validation_loss"))
            else None
        ),
        "wall_seconds": (
            float(record["driver_wall_seconds"])
            if finite_number(record.get("driver_wall_seconds"))
            else None
        ),
        "peak_memory_mib": (
            float(memory) / (1024.0 ** 2)
            if finite_number(memory) else None
        ),
        "work": record.get("work"),
    }


def transformer_record(record):
    raw_nll = record.get("final_validation_nll")
    valid_nll = finite_number(raw_nll)
    is_finite = record.get("finite") is True and valid_nll
    memory = record.get("peak_memory_allocated_bytes")
    return {
        "cell_id": record["cell_id"],
        "status": record.get("status"),
        "finite": is_finite,
        "validation_metric_name": "negative_log_likelihood",
        "validation_metric": float(raw_nll) if valid_nll else None,
        "final_validation_nll": (
            float(raw_nll) if valid_nll else None
        ),
        "final_validation_perplexity": (
            float(record["final_validation_perplexity"])
            if finite_number(record.get("final_validation_perplexity"))
            else None
        ),
        "final_validation_accuracy": (
            float(record["final_validation_accuracy"])
            if finite_number(record.get("final_validation_accuracy"))
            else None
        ),
        "wall_seconds": (
            float(record["driver_wall_seconds"])
            if finite_number(record.get("driver_wall_seconds"))
            else None
        ),
        "peak_memory_mib": (
            float(memory) / (1024.0 ** 2)
            if finite_number(memory) else None
        ),
        "work": record.get("work"),
    }


def aggregate(plain, resnet, transformer):
    assert_plain_header(plain)
    assert_native_header(resnet, "resnet_crossover_r2")
    assert_native_header(transformer, "transformer_crossover_t2")
    normalized = (
        [plain_record(record) for record in plain["records"]]
        + [resnet_record(record) for record in resnet["records"]]
        + [transformer_record(record) for record in transformer["records"]]
    )
    identifiers = [record["cell_id"] for record in normalized]
    assert len(identifiers) == 81
    assert len(set(identifiers)) == 81, "duplicate crossover cell"
    registry = {cell["cell_id"]: cell for cell in primary_cells()}
    actual = set(identifiers)
    expected = set(registry)
    assert actual == expected, (
        f"crossover registry mismatch: missing={sorted(expected - actual)}, "
        f"unexpected={sorted(actual - expected)}"
    )
    records = []
    for record in normalized:
        cell = registry[record["cell_id"]]
        records.append({
            **record,
            "architecture": cell["architecture"],
            "family": cell["family"],
            "scale_index": cell["scale_index"],
            "size": cell["size"],
            "method": cell["method"],
        })
    records.sort(key=lambda row: (
        row["family"], row["scale_index"], METHODS.index(row["method"])
    ))
    for architecture in {row["architecture"] for row in records}:
        methods = {
            row["method"] for row in records
            if row["architecture"] == architecture
        }
        assert methods == set(METHODS), (
            f"{architecture} does not contain all nine methods"
        )
    failures = [
        record["cell_id"] for record in records if not record["finite"]
    ]
    return {
        "audit_status": "passed",
        "stage": "crossover_x1_81",
        "complete_grid": True,
        "failure_retaining": True,
        "num_architecture_scale_points": 9,
        "num_methods": 9,
        "num_expected_cells": 81,
        "num_audited_cells": len(records),
        "num_finite_cells": len(records) - len(failures),
        "failed_or_incomplete_cells": failures,
        "test_policy": "none",
        "records": records,
    }


def input_lock(path):
    absolute = os.path.abspath(path)
    return {"path": absolute, "sha256": sha256(absolute)}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--plain", required=True,
        help="complete audit JSON emitted by the external plain-CNN P2",
    )
    parser.add_argument(
        "--resnet", default="results/resnet_crossover/r2_audit.json"
    )
    parser.add_argument(
        "--transformer",
        default="results/transformer_crossover/t2_audit.json",
    )
    parser.add_argument(
        "--out", default="results/crossover_81_audit.json"
    )
    args = parser.parse_args()
    report = aggregate(
        read_json(args.plain),
        read_json(args.resnet),
        read_json(args.transformer),
    )
    report["inputs"] = {
        "plain_cnn": input_lock(args.plain),
        "resnet": input_lock(args.resnet),
        "transformer": input_lock(args.transformer),
    }
    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({
        "audit_status": report["audit_status"],
        "num_audited_cells": report["num_audited_cells"],
        "num_finite_cells": report["num_finite_cells"],
        "failed_or_incomplete_cells":
            report["failed_or_incomplete_cells"],
    }, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()