#!/usr/bin/env python3
"""Audit the Qwen235B real-machine surfaces before comparing Frontier.

This script intentionally does not consume simulator output.  It establishes
whether each real response surface is complete and discriminative enough to
support a later claim about simulator config selection.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from pathlib import Path
from typing import Any


SCHEMA = "frontier-multicase-ground-truth-v0"
EXPECTED_PROBES = 6


def load_json(path: Path) -> dict[str, Any]:
    with path.open(encoding="utf-8") as source:
        value = json.load(source)
    if not isinstance(value, dict):
        raise ValueError(f"expected JSON object: {path}")
    return value


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def config_from_result(result: dict[str, Any]) -> dict[str, int]:
    flags = result["config_patch"]["flag_patch"]
    tp = int(flags["tensor-parallel-size"])
    dp = int(flags.get("data-parallel-size", 1))
    return {
        "tp": tp,
        "dp": dp,
        "ep": int(flags.get("expert-parallel-size", 1)),
        "mns": int(flags["max-num-seqs"]),
        "mbt": int(flags["max-num-batched-tokens"]),
        "gpu_count": tp * dp,
    }


def cell_id(config: dict[str, int]) -> str:
    topology = f"tp{config['tp']}"
    if config["dp"] != 1 or config["ep"] != 1:
        topology += f"_dp{config['dp']}_ep{config['ep']}"
    return f"{topology}_mns{config['mns']}_mbt{config['mbt']}"


def trial_record(case: str, path: Path) -> dict[str, Any]:
    result = load_json(path)
    config = config_from_result(result)
    score = float(result["best_request_rate"]) / config["gpu_count"]
    probes = result.get("probes", [])
    infeasible_above = [
        float(probe.get("payload", probe)["request_rate"]) / config["gpu_count"]
        for probe in probes
        if not probe["feasible"]
        and float(probe.get("payload", probe)["request_rate"])
        / config["gpu_count"]
        > score
    ]
    upper_bound = min(infeasible_above) if infeasible_above else None
    probe_count = len(result.get("probes", []))
    primary_result = result.get("best_source") == "primary_search"
    no_probe_failure = not bool(result.get("completed_with_probe_failure", False))
    fully_valid = (
        result.get("status") == "completed"
        and probe_count == EXPECTED_PROBES
        and primary_result
        and no_probe_failure
    )
    return {
        "case": case,
        "cell_id": cell_id(config),
        **config,
        "score_req_s_per_gpu": score,
        "capacity_lower_bound_req_s_per_gpu": score,
        "capacity_upper_bound_req_s_per_gpu": upper_bound,
        "capacity_bracket_width_req_s_per_gpu": (
            upper_bound - score if upper_bound is not None else None
        ),
        "best_request_rate_req_s": float(result["best_request_rate"]),
        "best_sampling_u": float(result["best_sampling_u"]),
        "best_pass_rate": float(result["best_pass_rate"]),
        "probe_count": probe_count,
        "best_source": result.get("best_source"),
        "completed_with_probe_failure": bool(
            result.get("completed_with_probe_failure", False)
        ),
        "fully_valid": fully_valid,
        "result_path": str(path),
        "result_sha256": sha256_file(path),
    }


def summarize_case(case: str, rows: list[dict[str, Any]]) -> dict[str, Any]:
    if not rows:
        raise ValueError(f"no rows for case: {case}")
    ids = [row["cell_id"] for row in rows]
    if len(ids) != len(set(ids)):
        duplicates = sorted(cell for cell in set(ids) if ids.count(cell) > 1)
        raise ValueError(f"duplicate cells for {case}: {duplicates}")

    best = max(row["score_req_s_per_gpu"] for row in rows)
    tolerance = max(1e-12, best * 1e-9)
    top = [
        row["cell_id"]
        for row in rows
        if math.isclose(row["score_req_s_per_gpu"], best, abs_tol=tolerance)
    ]
    distinct_scores = []
    for score in sorted({row["score_req_s_per_gpu"] for row in rows}, reverse=True):
        if not any(math.isclose(score, seen, abs_tol=tolerance) for seen in distinct_scores):
            distinct_scores.append(score)

    max_lower_bound = max(row["capacity_lower_bound_req_s_per_gpu"] for row in rows)
    possibly_optimal = [
        row["cell_id"]
        for row in rows
        if row["capacity_upper_bound_req_s_per_gpu"] is None
        or row["capacity_upper_bound_req_s_per_gpu"] + tolerance >= max_lower_bound
    ]

    total_pairs = len(rows) * (len(rows) - 1) // 2
    tied_pairs = sum(
        1
        for left_index, left in enumerate(rows)
        for right in rows[left_index + 1 :]
        if math.isclose(
            left["score_req_s_per_gpu"],
            right["score_req_s_per_gpu"],
            abs_tol=tolerance,
        )
    )
    return {
        "case": case,
        "cell_count": len(rows),
        "fully_valid_cell_count": sum(bool(row["fully_valid"]) for row in rows),
        "invalid_cells": [row["cell_id"] for row in rows if not row["fully_valid"]],
        "best_score_req_s_per_gpu": best,
        "top_set": sorted(top),
        "top_set_size": len(top),
        "random_top_set_hit_rate": len(top) / len(rows),
        "distinct_score_count": len(distinct_scores),
        "distinct_scores_req_s_per_gpu": distinct_scores,
        "possibly_optimal_set_from_search_brackets": sorted(possibly_optimal),
        "possibly_optimal_set_size": len(possibly_optimal),
        "pair_count": total_pairs,
        "tied_pair_count": tied_pairs,
        "informative_pair_count": total_pairs - tied_pairs,
        "informative_pair_fraction": (
            (total_pairs - tied_pairs) / total_pairs if total_pairs else 0.0
        ),
    }


def markdown_report(metrics: dict[str, Any], rows: list[dict[str, Any]]) -> str:
    lines = [
        "# Qwen235B ground-truth audit",
        "",
        "Objective: maximum SLO-feasible offered request throughput per GPU.",
        "This report contains real-machine data only; it makes no Frontier match claim.",
        "",
        "| case | valid cells | score levels | top-set size | random top-set hit | informative pairs |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for case in metrics["cases"]:
        lines.append(
            f"| {case['case']} | {case['fully_valid_cell_count']}/{case['cell_count']} "
            f"| {case['distinct_score_count']} | {case['top_set_size']}/{case['cell_count']} "
            f"| {case['random_top_set_hit_rate']:.1%} "
            f"| {case['informative_pair_count']}/{case['pair_count']} "
            f"({case['informative_pair_fraction']:.1%}) |"
        )

    lines.extend(["", "## Cells", ""])
    for case in metrics["cases"]:
        lines.extend(
            [
                f"### {case['case']}",
                "",
                "| cell | capacity bracket (req/s/GPU) | valid | probes | source |",
                "|---|---:|---:|---:|---|",
            ]
        )
        for row in sorted(
            (row for row in rows if row["case"] == case["case"]),
            key=lambda row: row["cell_id"],
        ):
            upper = row["capacity_upper_bound_req_s_per_gpu"]
            bracket = (
                f"[{row['capacity_lower_bound_req_s_per_gpu']:.9f}, "
                f"{upper:.9f})"
                if upper is not None
                else f"[{row['capacity_lower_bound_req_s_per_gpu']:.9f}, +inf)"
            )
            lines.append(
                f"| {row['cell_id']} | {bracket} "
                f"| {'yes' if row['fully_valid'] else 'no'} | {row['probe_count']} "
                f"| {row['best_source']} |"
            )
        lines.extend(
            [
                "",
                f"Top set: `{', '.join(case['top_set'])}`.",
                f"Possibly optimal under binary-search brackets: "
                f"`{', '.join(case['possibly_optimal_set_from_search_brackets'])}`.",
                "",
            ]
        )

    lines.extend(
        [
            "## Interpretation guardrail",
            "",
            "A Frontier top-set hit is insufficient by itself because the surfaces contain "
            "large ties. The later comparison must report selected-config regret and "
            "tie-aware pairwise ranking, and must keep invalid real cells visible.",
            "",
        ]
    )
    return "\n".join(lines)


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as target:
        writer = csv.DictWriter(target, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--prefill-root", action="append", type=Path, required=True)
    parser.add_argument("--decode-root", action="append", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cases = {"qwen235b_prefill_only": args.prefill_root, "qwen235b_decode_only": args.decode_root}
    rows = []
    for case, roots in cases.items():
        for root in roots:
            paths = sorted(root.glob("store/*/trials/trial-*/result.json"))
            if not paths:
                raise ValueError(f"no result files below {root}")
            rows.extend(trial_record(case, path) for path in paths)

    summaries = [
        summarize_case(case, [row for row in rows if row["case"] == case])
        for case in cases
    ]
    metrics = {"schema": SCHEMA, "cases": summaries}
    args.output_dir.mkdir(parents=True, exist_ok=True)
    write_csv(args.output_dir / "cells.csv", rows)
    (args.output_dir / "metrics.json").write_text(
        json.dumps(metrics, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    (args.output_dir / "report.md").write_text(
        markdown_report(metrics, rows), encoding="utf-8"
    )


if __name__ == "__main__":
    main()