#!/usr/bin/env python3
from __future__ import annotations

from analyze_pilot_e2e import expanded_top_k, replay


def candidate(
    cell: str,
    sim_score: float,
    real_feasible: bool,
    probability: float,
) -> dict[str, object]:
    return {
        "cell": cell,
        "level": "high",
        "sim_throughput_req_s_per_gpu": sim_score,
        "real_goodput_req_s_per_gpu": sim_score,
        "real_feasible": real_feasible,
        "setup_h20_hours": 0.1,
        "full_trial_h20_hours": 0.05,
        "prefix_h20_hours": 0.01,
        "instrument_probability": probability,
    }


def main() -> None:
    candidates = [
        candidate("a", 3.0, True, 0.5),
        candidate("b", 2.0, False, 0.01),
        candidate("c", 2.0, True, 0.99),
    ]
    shortlist = expanded_top_k(candidates, 2)
    assert [item["cell"] for item in shortlist] == ["a", "b", "c"]
    result = replay(
        shortlist,
        probability_key="instrument_probability",
        oracle_goodput=3.0,
        common_failure_h20_hours=0.02,
    )
    assert result["selected_cell"] == "a"
    assert result["false_accept"] == 0
    assert result["false_reject"] == 0
    assert result["early_accept"] == 1
    assert result["early_reject"] == 1
    assert result["online_h20_hours"] > 0
    print("fidelity pilot e2e: PASS")


if __name__ == "__main__":
    main()