#!/usr/bin/env python3
"""Arm 8 -- what a readout costs per decision when the server may batch.

    python3 run_arm8.py --levels 1,4,8 --out rows-addenda

Pre-registered in README "Addenda arms 4-11". The server is arm 2's with ONE
flag changed, `--max-num-seqs 8`; this driver changes one thing about the run,
the number of decisions in flight, and nothing about the decisions themselves.

WHY THREADS AND NOT AN ASYNC HTTP CLIENT. The point of the arm is the SERVER's
batching, not the client's IO model, and the one thing that must not vary
between arm 2 and arm 8 is the request itself. So each stream calls `run.py`'s
own `call_vllm` -- the same body, the same `chat_template_kwargs`, the same
`temperature 0`, the same `top_logprobs` -- through `asyncio.to_thread`, and the
concurrency is in how many of those are in flight at once. A hand-rolled async
client would have been a second request builder to keep in step with the first,
which is exactly the kind of drift this bench refuses elsewhere.

ACCURACY IS A FINDING HERE, NOT A METRIC. The items, prompts, seed and
temperature are arm 2's, and a greedy readout must not move because the server
batched it. The report therefore carries a per-level accuracy AND the ids of any
item whose choice differs from concurrency 1, because a difference is the whole
result and must not be averaged away.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import statistics
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import run as R              # noqa: E402
import run_addenda as A      # noqa: E402  -- for load_kit


def build_items() -> list[dict]:
    """Task (c)'s 108 items as (id, label, options, prompt) -- arm 2's bytes."""
    roster = {g["id"]: g["name"] for g in A.load_kit("manifest")["c_roster"]}
    out = []
    for item in A.load_kit("task_c"):
        options = [(gid, roster[gid]) for gid in item["options"]]
        state = ("A dinner table with six guests: %s."
                 % ", ".join(roster[g] for g in item["options"]))
        prompt = R.readout_prompt(state, R.INSTR_C % item["text"], options)
        out.append({"id": item["id"], "label": item["label"],
                    "options": options, "prompt": prompt,
                    "prompt_sha256": R.sha(prompt)})
    return out


def score(resp: dict, options: list) -> tuple[str, list, list]:
    """Arm 2's readout, verbatim in arithmetic: floor, /T, softmax, argmax."""
    raw, _detail = R.letter_scores(resp, len(options))
    z = [(R.FLOOR if v is None else v) / R.READOUT_T for v in raw]
    p = R.softmax(z)
    best = max(range(len(p)), key=p.__getitem__)
    return options[best][0], raw, p


async def one_level(arm, items: list[dict], level: int, path: str) -> dict:
    sem = asyncio.Semaphore(level)
    rows: list[dict] = []

    async def do(item):
        async with sem:
            t0 = time.monotonic()
            try:
                resp = await asyncio.to_thread(
                    R.call_vllm, arm, item["prompt"], predict=1)
            except R.SeatError as exc:
                return {"id": item["id"], "label": item["label"],
                        "refused": "seat-error", "error": str(exc)[:300],
                        "choice": None, "seconds": time.monotonic() - t0}
            choice, raw, p = score(resp, item["options"])
            return {"id": item["id"], "label": item["label"], "choice": choice,
                    "correct": choice == item["label"],
                    "letter_logprobs": raw,
                    "floored": [i for i, v in enumerate(raw) if v is None],
                    "probs": {k: p[i] for i, (k, _) in enumerate(item["options"])},
                    "prompt_sha256": item["prompt_sha256"],
                    "seconds": resp["_seconds"],
                    "prompt_tokens": resp.get("prompt_eval_count"),
                    "eval_tokens": resp.get("eval_count"),
                    "concurrency": level, "stamp": R.utc()}

    started = time.monotonic()
    rows = await asyncio.gather(*(do(i) for i in items))
    wall = time.monotonic() - started
    with open(path, "w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    lat = sorted(r["seconds"] for r in rows if r.get("seconds"))
    reached = [r for r in rows if not r.get("refused")]
    return {"concurrency": level, "n": len(rows), "reached": len(reached),
            "refused": len(rows) - len(reached),
            "correct": sum(1 for r in reached if r.get("correct")),
            "accuracy": (sum(1 for r in reached if r.get("correct")) / len(reached)
                         if reached else None),
            "wall_seconds": wall,
            "decisions_per_second": len(reached) / wall if wall else None,
            "median_seconds": statistics.median(lat) if lat else None,
            "p95_seconds": (lat[min(len(lat) - 1, int(0.95 * len(lat)))]
                            if lat else None),
            "rows_path": os.path.basename(path)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--levels", default="1,4,8")
    ap.add_argument("--out", default=os.path.join(HERE, "rows-addenda"))
    a = ap.parse_args(argv)
    os.makedirs(a.out, exist_ok=True)
    arm = R.ARMS["openjev-fp8-readout"]
    items = build_items()
    levels = [int(x) for x in a.levels.split(",")]
    reports = []
    baseline: dict[str, str] = {}
    for lv in levels:
        base = os.path.join(a.out, "arm8.c.conc%d" % lv)
        stamp = R.utc()
        # `box` is an arm-3 field. benchbox deliberately stays on the ARM-2 copy of
        # run.py for every addenda arm, so that all of arms 4-11 are built by ONE
        # instrument rather than half by each; its decision path is byte-identical
        # to the canonical copy's (verified by AST hash over readout_prompt,
        # decide, letter_scores, call_vllm, summarise and every calibration
        # constant). So this reads the field if it is there and names the box it
        # is actually running on if it is not -- it never types a box name it did
        # not get from somewhere.
        box = getattr(arm, "box", None) or os.uname().nodename
        print("ANNOUNCE  box %s · cards %s · arm 8 (--max-num-seqs 8) · task c "
              "readout · concurrency %d · %d items · start %s · expect ~%d min"
              % (box, ", ".join(c[:16] for c in arm.cards), lv, len(items),
                 stamp, max(1, 6 // lv)), flush=True)
        idle = R.idle_reading(arm.cards, base + ".idle.csv")
        with R.Boards(arm.cards, base + ".watts.csv"):
            rep = asyncio.run(one_level(arm, items, lv, base + ".jsonl"))
        energy = R.Boards.integrate(base + ".watts.csv")
        rep.update(started=stamp, ended=R.utc(), idle=idle, energy=energy, box=box,
                   cards=list(arm.cards), model=arm.model,
                   server="arm 2's serve.sh with --max-num-seqs 8, all else identical")
        if rep["reached"]:
            net = energy["joules"] - (idle.get("mean_watts") or 0) * energy["seconds"]
            rep["energy"]["joules_per_decision"] = energy["joules"] / rep["reached"]
            rep["energy"]["net_joules_per_decision"] = net / rep["reached"]
        # THE FINDING CHECK, AND WHY IT IS NOT JOINED ON `id`.
        # task (c)'s kit has 108 distinct lines but only 50 distinct id STRINGS:
        # `build_kit` composes the id from (guest, edition, course, line_index)
        # and two wall editions can produce the same tuple. The texts differ,
        # every row is a real decision, and no measured number is wrong -- but
        # the id is NOT a key, and a dict keyed on it silently compares 50 of
        # 108 rows. So the check joins by POSITION (asyncio.gather preserves
        # input order, and these rows are written in that order) and ASSERTS
        # that the two runs' prompt sha256s line up before comparing a choice.
        # A misalignment raises rather than quietly reporting agreement.
        cur = [json.loads(l) for l in open(base + ".jsonl", encoding="utf-8") if l.strip()]
        if not baseline:
            baseline = cur
            rep["differs_from_conc1"] = []
        else:
            if len(cur) != len(baseline):
                raise SystemExit("arm 8: row counts differ, %d vs %d"
                                 % (len(cur), len(baseline)))
            mism = [i for i, (a, b) in enumerate(zip(baseline, cur))
                    if a["prompt_sha256"] != b["prompt_sha256"]]
            if mism:
                raise SystemExit("arm 8: prompts not aligned at positions %s" % mism[:5])
            rep["differs_from_conc1"] = [
                {"position": i, "id": b["id"], "conc1": a.get("choice"),
                 "this": b.get("choice")}
                for i, (a, b) in enumerate(zip(baseline, cur))
                if a.get("choice") != b.get("choice")]
        rep["compared_rows"] = len(cur)
        rep["join"] = "by position, with prompt_sha256 asserted equal (the kit's id is not unique)"
        rep["n_differs_from_conc1"] = len(rep["differs_from_conc1"])
        with open(base + ".report.json", "w", encoding="utf-8") as fh:
            json.dump(rep, fh, indent=1, sort_keys=True, ensure_ascii=False)
        reports.append(rep)
        print(json.dumps({k: rep[k] for k in
                          ("concurrency", "accuracy", "median_seconds", "p95_seconds",
                           "decisions_per_second", "n_differs_from_conc1")},
                         default=str), flush=True)
    with open(os.path.join(a.out, "arm8.summary.json"), "w") as fh:
        json.dump(reports, fh, indent=1, sort_keys=True, default=str)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())