"""Stage-b reward-path ceiling test: a PERFECT policy must score ~100%.

House rule from a paid-for incident (mistral-prm-ppo, 2026-08-14): two PPO
runs measured a broken answer-extractor -- a perfect policy scored 0/100.
No GRPO rollout launches until this driver certifies the EXACT reward path
(extract_answer -> OutcomeVerifier -> CompositeReward + format instrument)
at >= --floor on every gold source, with every miss listed in the receipt
for eyeballing.

CPU-only, no model: each (problem, gold) becomes
`build_perfect_completion(gold, eos)` -- byte-identical to what a policy
flawlessly imitating its SFT targets emits, EOS included (rollout decodes
with skip_special_tokens=False).

Sources: gsm8k (1,319), math500 (500), deepmath (--deepmath-n sampled
evenly across shards; stage-b's actual gold distribution -- the reason this
must re-run here and not only on the stage-a sets).

    python scripts/ceiling_test.py --deepmath-n 5000
"""

import argparse
import json
import os
import sys
from pathlib import Path

# Run from a checkout without an install: put the repo root on sys.path first.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.data.schema import extract_answer, register_answer_suffixes
from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test, read_math500
from align.data.stageb_sources import StageBSources, iter_bigmath, iter_deepmath
from align.eval.roundtrip import build_perfect_completion
from align.eval.verifier import OutcomeVerifier, VerifierConfig
from align.rl.format_instrument import format_scores
from align.rl.rewards import CompositeReward

EOS = "</s>"


#: The corpora live in two hubs: the stage-a pins under sources.DEFAULT_HUB
#: (the voxtral cache this project inherited), the stage-b downloads under
#: the rlhf hf_cache. Resolved per-source-set, not via one env var.
STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub"


def _gold_sources(args) -> dict[str, list[str]]:
    out: dict[str, list[str]] = {}
    sources = DataSources.from_local_hub()  # DEFAULT_HUB: gsm8k/math500 pins
    if args.gsm8k:
        out["gsm8k"] = [gsm8k_gold_answer(r) for r in read_gsm8k_test(sources)]
    if args.math500:
        out["math500"] = [r["answer"] for r in read_math500(sources)]
    if args.deepmath_n or args.bigmath_n:
        sb = StageBSources.from_local_hub(
            hub=args.stageb_hub,
            need_deepmath=bool(args.deepmath_n),
            need_openmath=False,
            need_bigmath=bool(args.bigmath_n),
        )
        for key, it, total, n, field in (
            ("deepmath", iter_deepmath, 103_110, args.deepmath_n, "final_answer"),
            ("bigmath", iter_bigmath, 251_122, args.bigmath_n, "answer"),
        ):
            if not n:
                continue
            golds: list[str] = []
            for i, row in enumerate(it(sb)):
                # even sample across the stream without materializing the corpus
                if i % max(1, total // n) == 0:
                    golds.append(str(row[field]))
                if len(golds) >= n:
                    break
            out[key] = golds
    return out


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--floor", type=float, default=0.98)
    parser.add_argument("--gsm8k", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--math500", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--deepmath-n", type=int, default=5000)
    parser.add_argument("--bigmath-n", type=int, default=5000)
    parser.add_argument(
        "--workers", type=int, default=4, help="verifier procs; keep <=4 on the shared box"
    )
    parser.add_argument("--out", default="outputs/eval/ceiling_test.json")
    parser.add_argument("--stageb-hub", default=STAGEB_HUB)
    args = parser.parse_args()

    register_answer_suffixes(EOS)  # BEFORE any extraction (golden.py's incident)
    verifier = OutcomeVerifier(VerifierConfig(workers=args.workers))
    reward = CompositeReward(format_weight=0.1)  # overlong off: length is not under test

    report: dict = {"floor": args.floor, "eos": EOS, "sources": {}}
    overall_pass = True
    for name, golds in _gold_sources(args).items():
        completions = [build_perfect_completion(g, EOS) for g in golds]
        preds = [extract_answer(c) for c in completions]
        extract_ok = sum(1 for p in preds if p is not None)
        outcomes = verifier.verify_batch(golds, preds)
        correct = [o.name == "CORRECT" for o in outcomes]
        fscores = [format_scores(c)["format/score"] for c in completions]
        rewards = [
            reward(accuracy=1.0 if ok else 0.0, format_score=f, response_len=0)["reward/total"]
            for ok, f in zip(correct, fscores)
        ]
        ceiling = sum(correct) / len(golds) if golds else float("nan")
        misses = [
            {"gold": g, "pred": p, "outcome": o.name}
            for g, p, o, ok in zip(golds, preds, outcomes, correct)
            if not ok
        ]
        passed = ceiling >= args.floor
        overall_pass &= passed
        report["sources"][name] = {
            "n": len(golds),
            "extract_ok_rate": extract_ok / len(golds),
            "verify_correct_rate": ceiling,
            "mean_format_score": sum(fscores) / len(fscores),
            "mean_reward_total": sum(rewards) / len(rewards),
            "passed": passed,
            "misses": misses[:50],
            "misses_total": len(misses),
        }
        print(
            f"{name}: n={len(golds)} extract_ok={extract_ok / len(golds):.4f} "
            f"ceiling={ceiling:.4f} misses={len(misses)} -> {'PASS' if passed else 'FAIL'}",
            flush=True,
        )

    report["overall_pass"] = bool(overall_pass)
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    tmp = args.out + ".partial"
    with open(tmp, "w") as f:
        json.dump(report, f, indent=1)
    os.replace(tmp, args.out)
    print(f"receipt: {args.out} | overall_pass={overall_pass}")
    if not overall_pass:
        print(
            "CEILING FAIL: eyeball every miss in the receipt before touching the "
            "verifier -- a representational-equivalence gap and a broken extractor "
            "look identical in the rate and completely different in the misses.",
            file=sys.stderr,
        )
    return 0 if overall_pass else 1


if __name__ == "__main__":
    sys.exit(main())