#!/usr/bin/env python3
"""Deterministic completeness and failure-retention tests for X1 joiner."""
import copy

from analyze_crossover_81 import aggregate
from crossover_registry import METHODS


PLAIN_ARCHITECTURES = ("minicnn", "vgglike", "vgg16")
RESNET_DEPTHS = (20, 32, 56)
TRANSFORMER_DEPTHS = (4, 8, 12)


def plain_report():
    records = []
    for architecture in PLAIN_ARCHITECTURES:
        for method in METHODS:
            records.append({
                "architecture": architecture,
                "method": method,
                "status": "completed",
                "outcome": "finite",
                "best_validation_accuracy": 75.0,
                "final_validation_accuracy": 74.0,
                "final_validation_loss": 1.0,
                "driver_wall_seconds": 10.0,
                "physical_gpu": {
                    "peak_process_gpu_memory_mib": 1000.0,
                },
                "work": {"synthetic": True},
            })
    return {
        "gate": "pass",
        "stage": "plain_cnn_p2",
        "num_expected_records": 27,
        "num_audited_records": 27,
        "missing_experiments": [],
        "test_policy": "none",
        "records": records,
    }


def native_report(family):
    if family == "resnet":
        stage = "resnet_crossover_r2"
        cell_ids = [
            f"resnet{depth}::{method}"
            for depth in RESNET_DEPTHS for method in METHODS
        ]
    else:
        stage = "transformer_crossover_t2"
        cell_ids = [
            f"transformer{depth}::{method}"
            for depth in TRANSFORMER_DEPTHS for method in METHODS
        ]
    records = []
    for cell_id in cell_ids:
        common = {
            "cell_id": cell_id,
            "status": "completed",
            "finite": True,
            "driver_wall_seconds": 12.0,
            "peak_memory_allocated_bytes": 1024 ** 3,
            "work": {"synthetic": True},
        }
        if family == "resnet":
            common.update({
                "best_validation_accuracy": 0.8,
                "final_validation_accuracy": 0.79,
                "final_validation_loss": 0.9,
            })
        else:
            common.update({
                "final_validation_nll": 2.0,
                "final_validation_perplexity": 7.389,
                "final_validation_accuracy": 0.3,
            })
        records.append(common)
    return {
        "audit_status": "passed",
        "stage": stage,
        "complete_grid": True,
        "failure_retaining": True,
        "num_expected_cells": 27,
        "num_audited_cells": 27,
        "test_policy": "none",
        "records": records,
    }


def must_fail(plain, resnet, transformer):
    try:
        aggregate(plain, resnet, transformer)
    except (AssertionError, KeyError):
        return
    raise AssertionError("invalid 81-cell panel was accepted")


def main():
    plain = plain_report()
    resnet = native_report("resnet")
    transformer = native_report("transformer")
    report = aggregate(plain, resnet, transformer)
    assert report["num_audited_cells"] == 81
    assert report["num_finite_cells"] == 81
    assert report["failed_or_incomplete_cells"] == []

    missing = copy.deepcopy(plain)
    missing["records"].pop()
    must_fail(missing, resnet, transformer)

    duplicate = copy.deepcopy(transformer)
    duplicate["records"][-1] = copy.deepcopy(duplicate["records"][0])
    must_fail(plain, resnet, duplicate)

    failed = copy.deepcopy(resnet)
    cell = failed["records"][0]
    cell.update({
        "status": "timeout",
        "finite": False,
        "best_validation_accuracy": None,
        "final_validation_accuracy": None,
        "final_validation_loss": None,
    })
    report = aggregate(plain, failed, transformer)
    assert report["num_audited_cells"] == 81
    assert report["num_finite_cells"] == 80
    assert report["failed_or_incomplete_cells"] == ["resnet20::bp"]
    print("crossover 81-cell smoke: all checks passed")


if __name__ == "__main__":
    main()