#!/usr/bin/env python3
"""Audit the complete failure-retaining 27-cell Transformer T2 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.transformer_crossover_t2 import (
    DEFAULT_SELECTOR,
    DEPTHS,
    METHODS,
    registry_sha256,
    selector_report,
    t2_jobs,
)
from experiments.crossover_hardware import assert_hardware_report


TRAIN_HASH = (
    "6ec305602a99ac2802745a134e1f5e33e2231b4855525b00b9aebb730ac2626f"
)
VALIDATION_HASH = (
    "d37d30cc0c8327c270d493299c3dca54135f6d5f1c9ef60cda78076e311204b1"
)
EXPECTED_TRAIN_TOKENS = 5000 * 32 * 64
EXPECTED_VALIDATION_TOKENS = 1742 * 64
PARAMETERS = {4: 813_568, 8: 1_602_048, 12: 2_390_528}


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 expected_work(record, job):
    work = record["work"]
    method = job["method"]
    depth = job["depth"]
    train = EXPECTED_TRAIN_TOKENS
    validation = EXPECTED_VALIDATION_TOKENS
    presentations = 2 * train if method in ("pepita", "ff") else train
    candidate = 65 * validation if method == "ff" else 0
    relaxation = 0
    local_vjp = 0
    if method == "dualprop":
        relaxation = 16 * train
        local_vjp = train * (depth + 1) * 16
    elif method == "ep":
        relaxation = 24 * train + 20 * validation
        local_vjp = train * (depth + 1) * 24
    assert work["ordinary_training_tokens"] == train
    assert work["ordinary_validation_tokens"] == validation
    assert work["training_token_presentations"] == presentations
    assert work["candidate_token_presentations"] == candidate
    assert work["relaxation_token_passes"] == relaxation
    assert work["local_vjp_token_evaluations"] == local_vjp
    assert work["logical_task_loss_queries"] == 0
    assert work["completed_optimizer_steps"] == 5000
    assert work["forward_parameter_count"] == PARAMETERS[depth]
    assert work["full_forward_token_passes"] > 0
    assert work["enumerated_full_forward_macs"] > 0
    if method in ("dfa", "pepita", "ff", "dualprop", "ep"):
        assert (
            work["local_block_token_evaluations"] > 0
            or method == "ff"
        )
    return work


def audit_completed(job, manifest, source, hardware_policy):
    assert manifest["output_exists"] is True
    assert manifest["output_sha256"] == sha256(job["output"])
    record = read_json(job["output"])
    provenance = record["provenance"]
    assert provenance["git_commit"] == source["git_commit"]
    assert provenance["git_tracked_dirty"] is False
    args = record["args"]
    expected_args = {
        "method": job["method"],
        "depth": job["depth"],
        "width": 128,
        "heads": 4,
        "mlp_ratio": 4,
        "context_length": 64,
        "batch_size": 32,
        "eval_batch_size": 32,
        "train_steps": 5000,
        "lr": job["rate"],
        "schedule": "cosine",
        "min_lr": job["rate"] * 0.1,
        "warmup_steps": 100,
        "weight_decay": 0.1,
        "run_seed": 0,
        "model_seed": 2027,
        "loader_seed": 0,
        "negative_seed": 5001,
        "ep_sign_seed": 5002,
        "traffic_ratio": 4.0,
        "ff_threshold": 2.0,
        "ff_score_from_layer": 1,
        "ep_beta": 0.5,
        "ep_dt": 0.5,
        "ep_free_steps": 20,
        "ep_nudge_steps": 4,
        "dp_alpha": 0.0,
        "dp_beta": 0.1,
        "dp_inference_passes": 16,
        "eval_every": 0,
        "max_val_batches": 0,
    }
    for key, expected in expected_args.items():
        assert args[key] == expected, (
            f"{job['experiment_name']}: argument drift in {key}"
        )
    assert record["protocol_family"] == (
        "transformer_local_learning_crossover"
    )
    dataset = record["dataset"]
    assert dataset["train_sha256"] == TRAIN_HASH
    assert dataset["validation_sha256"] == VALIDATION_HASH
    architecture = record["architecture"]
    assert architecture["depth"] == job["depth"]
    assert architecture["width"] == 128
    assert architecture["heads"] == 4
    assert architecture["context_length"] == 64
    assert architecture["forward_parameter_count"] == PARAMETERS[job["depth"]]
    evaluation = record["evaluation_protocol"]
    assert evaluation["split"] == "validation"
    assert evaluation["test_evaluations"] == 0
    assert evaluation["test_used_for_selection"] is False
    assert len(record["validation"]) == 1
    final = record["final"]
    assert final["tokens"] == EXPECTED_VALIDATION_TOKENS
    assert final["step"] <= 5000
    first_nonfinite = record["first_nonfinite_step"]
    nll = float(final["nll"])
    perplexity = float(final["perplexity"])
    accuracy = float(final["accuracy"])
    is_finite = (
        first_nonfinite is None
        and final["step"] == 5000
        and math.isfinite(nll)
        and math.isfinite(perplexity)
        and math.isfinite(accuracy)
    )
    work = None
    if is_finite:
        work = expected_work(record, job)
    else:
        assert record["work"]["ordinary_training_tokens"] <= (
            EXPECTED_TRAIN_TOKENS
        )
        assert record["work"]["logical_task_loss_queries"] == 0
    hardware = record["hardware"]
    locked_hardware = manifest["hardware_lock"]
    assert_hardware_report(locked_hardware, hardware_policy)
    assert (
        hardware["cuda_visible_devices"]
        == locked_hardware["cuda_visible_devices"]
    )
    assert hardware["device_name"] == locked_hardware["physical_gpu_name"]
    assert hardware["peak_memory_allocated_bytes"] is not None
    assert hardware["peak_memory_reserved_bytes"] is not None
    return {
        "finite": is_finite,
        "first_nonfinite_step": first_nonfinite,
        "completed_optimizer_steps":
            int(record["work"]["completed_optimizer_steps"]),
        "final_validation_nll": nll if math.isfinite(nll) else None,
        "final_validation_perplexity":
            perplexity if math.isfinite(perplexity) else None,
        "final_validation_accuracy":
            accuracy if math.isfinite(accuracy) else None,
        "forward_parameter_count":
            int(architecture["forward_parameter_count"]),
        "feedback_parameter_count":
            int(architecture["feedback_parameter_count"]),
        "peak_memory_allocated_bytes":
            int(hardware["peak_memory_allocated_bytes"]),
        "total_wall_seconds": float(record["total_wall_seconds"]),
        "work": work if work is not None else record["work"],
    }


def audit_job(job, source, selector, hardware_policy):
    manifest_path = job["output"] + ".manifest.json"
    if not os.path.isfile(manifest_path):
        raise AssertionError(f"missing T2 manifest: {job['experiment_name']}")
    manifest = read_json(manifest_path)
    for key in (
        "stage",
        "method",
        "architecture",
        "depth",
        "rate",
        "experiment_name",
        "output",
        "timeout_seconds",
        "command",
    ):
        assert manifest[key] == job[key], (
            f"{job['experiment_name']}: manifest drift in {key}"
        )
    assert manifest["source"] == source
    assert manifest["selector"] == selector
    hardware = manifest["hardware_lock"]
    assert_hardware_report(hardware, hardware_policy)
    common = {
        "cell_id": f"transformer{job['depth']}::{job['method']}",
        "method": job["method"],
        "depth": job["depth"],
        "rate": job["rate"],
        "status": manifest["status"],
        "manifest": os.path.relpath(manifest_path, ROOT),
        "driver_wall_seconds": float(manifest["driver_wall_seconds"]),
        "physical_gpu_index": hardware["physical_gpu_index"],
        "physical_gpu_uuid": hardware["physical_gpu_uuid"],
        "output_sha256": manifest["output_sha256"],
    }
    if manifest["status"] == "completed":
        return {
            **common,
            **audit_completed(job, manifest, source, hardware_policy),
        }
    assert manifest["status"] in {
        "timeout", "nonzero_exit", "missing_output"
    }
    if manifest["output_exists"]:
        assert manifest["output_sha256"] == sha256(job["output"])
    else:
        assert manifest["output_sha256"] is None
    return {
        **common,
        "finite": False,
        "first_nonfinite_step": None,
        "completed_optimizer_steps": None,
        "final_validation_nll": None,
        "final_validation_perplexity": None,
        "final_validation_accuracy": None,
        "forward_parameter_count": None,
        "feedback_parameter_count": None,
        "peak_memory_allocated_bytes": None,
        "total_wall_seconds": None,
        "work": None,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--selector", default=DEFAULT_SELECTOR)
    parser.add_argument(
        "--out", default="results/transformer_crossover/t2_audit.json"
    )
    args = parser.parse_args()
    selector = selector_report(args.selector)
    jobs = t2_jobs(selector["selected_rates"])
    launch_path = os.path.join(
        ROOT, "results", "transformer_crossover", "t2_launch.json"
    )
    assert os.path.isfile(launch_path), "missing Transformer T2 launch lock"
    launch = read_json(launch_path)
    assert launch["stage"] == "t2"
    assert launch["selector"] == selector
    assert launch["registry_sha256"] == registry_sha256(jobs)
    assert launch["num_jobs"] == 27
    hardware_policy = launch["hardware_policy"]
    assert (
        launch["allowed_physical_gpus"]
        == hardware_policy["allowed_physical_gpu_indices"]
    )
    source = launch["source"]
    records = [
        audit_job(job, source, selector, hardware_policy) for job in jobs
    ]
    assert len(records) == 27
    assert len({record["cell_id"] for record in records}) == 27
    assert {
        (record["method"], record["depth"]) for record in records
    } == {(method, depth) for method in METHODS for depth in DEPTHS}
    failures = [
        record["cell_id"] for record in records if not record["finite"]
    ]
    report = {
        "audit_status": "passed",
        "stage": "transformer_crossover_t2",
        "complete_grid": True,
        "failure_retaining": True,
        "num_expected_cells": 27,
        "num_audited_cells": len(records),
        "num_finite_cells": len(records) - len(failures),
        "failed_or_incomplete_cells": failures,
        "test_policy": "none",
        "source": source,
        "selector": selector,
        "hardware_policy": hardware_policy,
        "launch_lock": {
            "path": os.path.relpath(launch_path, ROOT),
            "sha256": sha256(launch_path),
            "registry_sha256": launch["registry_sha256"],
        },
        "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(
            {
                "audit_status": report["audit_status"],
                "num_audited_cells": len(records),
                "num_finite_cells": report["num_finite_cells"],
                "failed_or_incomplete_cells": failures,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()