#!/usr/bin/env python3
"""Small structural tests for the cross-config transfer diagnostic."""

from __future__ import annotations

import importlib.util
from pathlib import Path


HERE = Path(__file__).resolve().parent
SPEC = importlib.util.spec_from_file_location(
    "analyze_residual_transfer", HERE / "analyze_residual_transfer.py"
)
assert SPEC is not None and SPEC.loader is not None
MODULE = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)


def example(cell: str, level: str, value: float) -> dict:
    tp_text, mns_text = cell.split("_")
    return {
        "cell": cell,
        "level": level,
        "tp": int(tp_text[2:]),
        "mns": int(mns_text[3:]),
        "pass_rate_residual": value - 0.5,
        "real_pass_rate_rep1": value,
        "sim_pass_rate": 0.5,
        "offered_req_s_per_gpu": 1.0 if level == "low" else 2.0,
        "real_feasible": value >= 0.95,
        "sim_feasible": False,
        "state_residual": {"values": {"batch.mean": value}},
        "engine": {
            "common": {"batch": {"mean": value, "max": value, "cv": 0.0}},
            "engine_only": {"kv": value},
        },
    }


def main() -> None:
    cells = ("tp1_mns8", "tp1_mns16", "tp2_mns8", "tp2_mns16")
    examples = [
        example(cell, level, 0.1 + index / 10.0)
        for level in ("low", "high")
        for index, cell in enumerate(cells)
    ]
    rows = MODULE.transitions(examples)
    assert len(rows) == 48
    assert any(row["source_level"] != row["target_level"] for row in rows)
    assert all(row["source_cell"] != row["target_cell"] for row in rows)
    assert len(rows[0]["hybrid_telemetry"]) > len(rows[0]["hybrid_base"])
    assert len(rows[0]["direct_telemetry"]) > len(rows[0]["direct_base"])
    for feature, target in (
        ("hybrid_base", "target_residual_delta"),
        ("hybrid_telemetry", "target_residual_delta"),
        ("direct_base", "target_real_pass_rate_delta"),
        ("direct_telemetry", "target_real_pass_rate_delta"),
    ):
        prediction = MODULE.grouped_predictions(
            rows, feature_name=feature, target_name=target, regularization=1.0
        )
        assert len(prediction) == len(rows)
    print("telemetry residual transfer: PASS")


if __name__ == "__main__":
    main()