#!/usr/bin/env python3
"""Combines the measured and researched per-core/per-card/per-machine rates recorded in
`docs/rig-capacity.md` with the rig's actual card, machine and core counts to print
absolute capacity: sampling-tier candidates per hour, capacity-tier escalations per hour,
verification-tier cores available for compilation and the inner loop, total power draw
against the budget, and which stage binds first.

Counts live in a gitignored local file, never in this script or in the repository.
`rig-inventory.example.json` shows the shape with placeholder values; copy it to
`rig-inventory.local.json` and fill in the actual counts.

Usage: bench/capacity.py [path-to-inventory.json]   # defaults to rig-inventory.local.json
"""
import json
import sys
from pathlib import Path

BENCH_DIR = Path(__file__).resolve().parent

# Rates below are sourced from docs/rig-capacity.md; re-run the benchmarks there
# (oracle-scaling.sh, run-inner-loop.sh, compile-batch.sh) and the GPU-tier research and
# update both places together if the numbers drift.

# Sampling tier: candidates per second per card, at fleet scale. Data-parallel replicas
# scale close to linearly, so this is a fleet-wide per-card rate, not a single-card peak.
SAMPLING_CANDIDATES_PER_S_PER_CARD = 0.47

# Capacity tier: escalation calls per second per machine (the larger model invoked on
# cases that survive N failed samples), at an 8k-token prompt / 1.5k-token output.
CAPACITY_ESCALATIONS_PER_S_PER_MACHINE = 0.02

# Verification tier: measured on the 16-core dev box, per core, under contention from a
# concurrent WebGPU CTS sweep (docs/rig-capacity.md reports the load alongside). Likely a
# lower bound for a quiet, dedicated verification machine.
#
# cases/s per core, harness::compare::compare() against cached dumps, at 16 threads (full
# box saturation, the closest stand-in for a fully loaded verification machine).
INNER_LOOP_CASES_PER_S_PER_CORE = 17058
#
# candidates/s per machine, batch layout, `cargo build --release` uncached at N=1000 (the
# realistic case: freshly generated candidates, not yet in any cache). Treated as
# per-machine because a single batch compile is bound by one machine, not additive across
# machines beyond running one batch per machine at a time.
COMPILE_CANDIDATES_PER_S_PER_MACHINE = 1000 / 6.95


def bottleneck(stages):
    """`stages` maps a stage name to a throughput in candidates/hour (or None if
    unmeasured). Returns the name and value of the smallest measured stage."""
    measured = {k: v for k, v in stages.items() if v is not None}
    if not measured:
        return None, None
    name = min(measured, key=measured.get)
    return name, measured[name]


def main():
    path = Path(sys.argv[1]) if len(sys.argv) > 1 else BENCH_DIR / "rig-inventory.local.json"
    if not path.exists():
        sys.exit(
            f"{path} not found; copy rig-inventory.example.json to rig-inventory.local.json "
            "and fill in the rig's actual counts"
        )
    inv = json.loads(path.read_text())

    v = inv["verification_tier"]
    verification_cores = v["total_cores"]
    verification_machines = v["machines"]
    verification_watts = verification_machines * v["machine_watts"]

    s = inv["sampling_tier"]
    sampling_watts = s["cards"] * s["card_watts"]
    candidates_per_hour = s["cards"] * SAMPLING_CANDIDATES_PER_S_PER_CARD * 3600

    c = inv["capacity_tier"]
    capacity_watts = c["machines"] * c["machine_watts"]
    escalations_per_hour = c["machines"] * CAPACITY_ESCALATIONS_PER_S_PER_MACHINE * 3600

    extra_watts = 0.0
    extra = inv.get("extra_tier")
    if extra:
        extra_watts = extra.get("cards", 0) * extra.get("card_watts", 0.0)

    total_watts = verification_watts + sampling_watts + capacity_watts + extra_watts
    budget_watts = inv["power_budget_w"]

    inner_loop_per_hour = None
    if INNER_LOOP_CASES_PER_S_PER_CORE is not None:
        inner_loop_per_hour = verification_cores * INNER_LOOP_CASES_PER_S_PER_CORE * 3600
    compile_per_hour = None
    if COMPILE_CANDIDATES_PER_S_PER_MACHINE is not None:
        compile_per_hour = (
            verification_machines * COMPILE_CANDIDATES_PER_S_PER_MACHINE * 3600
        )

    stage, stage_per_hour = bottleneck(
        {
            "sampling (candidate generation)": candidates_per_hour,
            "verification (batch compilation)": compile_per_hour,
            "verification (inner-loop evaluation)": inner_loop_per_hour,
        }
    )

    result = {
        "verification_tier_cores": verification_cores,
        "verification_tier_machines": verification_machines,
        "sampling_tier_cards": s["cards"],
        "candidates_per_hour": candidates_per_hour,
        "capacity_tier_machines": c["machines"],
        "escalations_per_hour": escalations_per_hour,
        "compile_candidates_per_hour": compile_per_hour,
        "inner_loop_cases_per_hour": inner_loop_per_hour,
        "bottleneck_stage": stage,
        "bottleneck_candidates_per_hour": stage_per_hour,
        "total_power_w": total_watts,
        "power_budget_w": budget_watts,
        "power_headroom_w": budget_watts - total_watts,
        "power_binds_before_hardware_count": total_watts > budget_watts,
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()