#!/usr/bin/env python3
"""Analyze real, calibrated, old-profile, and new-profile P1 probe outcomes."""

from __future__ import annotations

import argparse
import csv
import json
import statistics
from pathlib import Path
from typing import Any


MODES = ("historical-calibrated", "historical-profile-only", "vllm020-profile-only")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--controller-state", type=Path, required=True)
    parser.add_argument("--calibrated", type=Path, required=True)
    parser.add_argument("--old-profile-only", type=Path, required=True)
    parser.add_argument("--new-profile-only", type=Path, required=True)
    parser.add_argument("--output-json", type=Path, required=True)
    parser.add_argument("--output-csv", type=Path, required=True)
    return parser.parse_args()


def load_results(path: Path) -> dict[tuple[str, str], dict[str, Any]]:
    payload = json.loads(path.read_text())
    if payload["status"] != "PASS" or len(payload["results"]) != 12:
        raise ValueError(f"expected 12 passing simulator probes in {path}")
    return {(row["cell"], row["role"]): row for row in payload["results"]}


def main() -> None:
    args = parse_args()
    controller = json.loads(args.controller_state.read_text())
    real: dict[tuple[str, str], dict[str, Any]] = {}
    for cell, value in controller["cells"].items():
        tp = int(value["tp"])
        for run in value["runs"]:
            if run["role"] in ("low1", "high1"):
                real[(cell, run["role"])] = {
                    **run,
                    "tp": tp,
                    "offered_req_s_per_gpu": int(run["selected_count"]) / 60 / tp,
                }
    if len(real) != 12:
        raise ValueError(f"expected 12 real P1 probes, found {len(real)}")

    modes = {
        "historical-calibrated": load_results(args.calibrated),
        "historical-profile-only": load_results(args.old_profile_only),
        "vllm020-profile-only": load_results(args.new_profile_only),
    }
    rows: list[dict[str, Any]] = []
    summaries: dict[str, Any] = {}
    for mode in MODES:
        predicted = modes[mode]
        false_feasible = 0
        false_infeasible = 0
        pass_errors: list[float] = []
        capacity_lower_bounds: dict[str, float] = {}
        for key in sorted(real):
            real_row = real[key]
            sim_row = predicted[key]
            scorer = sim_row["scorer"]
            sim_feasible = bool(scorer["slo"]["feasible"])
            real_feasible = bool(real_row["feasible"])
            false_feasible += int(sim_feasible and not real_feasible)
            false_infeasible += int(real_feasible and not sim_feasible)
            pass_error = abs(float(scorer["slo"]["pass_rate"]) - float(real_row["pass_rate"]))
            pass_errors.append(pass_error)
            rows.append(
                {
                    "mode": mode,
                    "cell": key[0],
                    "role": key[1],
                    "real_feasible": real_feasible,
                    "sim_feasible": sim_feasible,
                    "real_pass_rate": float(real_row["pass_rate"]),
                    "sim_pass_rate": float(scorer["slo"]["pass_rate"]),
                    "pass_rate_absolute_error": pass_error,
                    "offered_req_s_per_gpu": float(real_row["offered_req_s_per_gpu"]),
                    "sim_throughput_req_s_per_gpu": float(
                        scorer["throughput_requests_per_second_per_gpu"]
                    ),
                }
            )
            if sim_feasible:
                capacity_lower_bounds[key[0]] = max(
                    capacity_lower_bounds.get(key[0], 0.0),
                    float(real_row["offered_req_s_per_gpu"]),
                )
        agreement = 12 - false_feasible - false_infeasible
        summaries[mode] = {
            "probe_classification": {
                "agreement": agreement,
                "accuracy": agreement / 12,
                "false_feasible": false_feasible,
                "false_infeasible": false_infeasible,
            },
            "pass_rate_mae": statistics.mean(pass_errors),
            "feasible_probe_count": sum(
                bool(row["scorer"]["slo"]["feasible"]) for row in predicted.values()
            ),
            "p1_capacity_lower_bounds_req_s_per_gpu": {
                cell: capacity_lower_bounds.get(cell, 0.0)
                for cell in sorted(controller["cells"])
            },
            "rank_identifiable": bool(capacity_lower_bounds),
        }

    output = {
        "schema": "frontier-qwen30-p1-profile-ablation.v1",
        "scope": {
            "cells": 6,
            "probes_per_cell": 2,
            "roles": ["low1", "high1"],
            "reading": "held-out boundary classification, not a complete capacity sweep",
        },
        "sources": {
            "controller_state": str(args.controller_state.resolve()),
            "historical_calibrated": str(args.calibrated.resolve()),
            "historical_profile_only": str(args.old_profile_only.resolve()),
            "vllm020_profile_only": str(args.new_profile_only.resolve()),
        },
        "summaries": summaries,
        "rows": rows,
    }
    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    args.output_json.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
    args.output_csv.parent.mkdir(parents=True, exist_ok=True)
    with args.output_csv.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    print(args.output_json)


if __name__ == "__main__":
    main()