#!/usr/bin/env python3
"""Select predeclared simulator-lattice anchors for blind real confirmation."""

from __future__ import annotations

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


RATES = (0.10, 0.20, 0.40, 0.80, 1.20, 1.60, 2.40, 3.20)
CONFIG_NAMES = {
    f"tp{tp}_mns{mns}_mbt{mbt}"
    for tp in (4, 8)
    for mns in (64, 128)
    for mbt in (8192, 16384)
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--frontier-freeze", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--selection-slo", default="tpot_150ms")
    return parser.parse_args()


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def select_rates(loads: list[dict[str, Any]], slo: str) -> list[float]:
    ordered = sorted(loads, key=lambda item: float(item["offered_request_rate"]))
    rates = [float(item["offered_request_rate"]) for item in ordered]
    feasible = [bool(item["slos"][slo]["feasible"]) for item in ordered]
    selected = {rates[0]}
    for index in range(len(rates) - 1):
        if feasible[index] != feasible[index + 1]:
            selected.update((rates[index], rates[index + 1]))
    if len(selected) == 1:
        if all(feasible):
            selected.update(rates[-2:])
        elif not any(feasible):
            selected.update(rates[:2])
    return sorted(selected)


def warmup_requests(rate: float) -> int:
    return min(32, max(4, math.ceil(rate * 20.0)))


def main() -> None:
    args = parse_args()
    freeze_path = args.frontier_freeze.resolve()
    freeze = json.loads(freeze_path.read_text())
    if freeze.get("schema") != "frontier-qwen235b-t0-surface-v1":
        raise ValueError("unexpected Frontier freeze schema")
    if freeze.get("status") != "frozen_before_real_surface":
        raise ValueError("Frontier surface is not frozen")
    results = freeze.get("config_results") or []
    if len(results) != 8 or any(len(item.get("loads") or []) != 8 for item in results):
        raise ValueError("Frontier surface is incomplete")
    names = {item.get("config", {}).get("name") for item in results}
    if names != CONFIG_NAMES:
        raise ValueError(f"Frontier config set mismatch: {names}")
    for item in results:
        rates = tuple(sorted(float(load["offered_request_rate"]) for load in item["loads"]))
        if rates != RATES:
            raise ValueError(f"Frontier rate lattice mismatch for {item['config']['name']}: {rates}")

    cells = []
    total_expected_seconds = 0.0
    for item in results:
        config = item["config"]
        rates = select_rates(item["loads"], args.selection_slo)
        # Every anchor gets an independent server in both rounds.  The estimate
        # includes server startup, target-rate warmup and conservative drain
        # allowances for both the discarded and measured request streams.
        expected_seconds = 2 * sum(
            120.0
            + (warmup_requests(rate) - 1) / rate
            + 60.0
            + 63.0 / rate
            + 60.0
            for rate in rates
        )
        total_expected_seconds += expected_seconds * int(config["tp"])
        cells.append(
            {
                "config": config,
                "rates": rates,
                "rounds": 2,
                "requests_per_anchor": 64,
                "anchor_isolation": "fresh_server_per_rate_per_round",
                "target_rate_warmup_requests": {
                    f"{rate:.2f}": warmup_requests(rate) for rate in rates
                },
                "expected_wall_seconds": expected_seconds,
                "expected_h20_gpu_hours": expected_seconds * int(config["tp"]) / 3600.0,
                "expansion_rule": "if real labels do not bracket a transition, expand to the next frozen lattice anchor and repeat both directions",
            }
        )
    payload = {
        "schema": "qwen235b-t0-real-plan-v1",
        "frontier_freeze": {"path": str(freeze_path), "sha256": sha256(freeze_path)},
        "selection_slo": args.selection_slo,
        "selection_timing": "rate anchors frozen after complete simulator surface and before any accepted real surface cell",
        "execution_protocol_amendment": {
            "timing": "after excluded multi-rate-per-server diagnostic and before any accepted real surface cell",
            "reason": "observed cross-anchor GPU/kernel/batch warm-state leakage",
            "contract": "fresh server and target-rate warmup for every config-rate-round",
        },
        "strict_preregistered_slo": "tpot_40ms",
        "post_pilot_sensitivities": ["tpot_120ms", "tpot_150ms", "tpot_180ms"],
        "cells": cells,
        "expected_total_h20_gpu_hours": total_expected_seconds / 3600.0,
        "hard_timeout_hours_per_cell": 2.0,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"cells": len(cells), "expected_total_h20_gpu_hours": payload["expected_total_h20_gpu_hours"]}, sort_keys=True))


if __name__ == "__main__":
    main()