"""Rebalance the Model-V dataset by downsampling all-correct (-1) examples.

ProcessBench F1 = harmonic_mean(acc_error, acc_correct), so a training set skewed
toward -1 biases V to over-accept, crushing acc_error and thus F1. This makes a
~balanced default (correct:erroneous = --ratio) and preserves the original.

    python scripts/rebalance_verifier.py            # train.jsonl + val.jsonl, ratio 1.0
    python scripts/rebalance_verifier.py --ratio 1.5

Original files are moved to *_full.jsonl; the balanced set takes the canonical name.
"""
import argparse
import json
import random
from pathlib import Path


def rebalance(path: Path, ratio: float, seed: int) -> tuple[int, int, int]:
    rows = [json.loads(l) for l in path.read_text().splitlines() if l.strip()]
    err = [r for r in rows if r.get("first_error_index", -1) != -1]
    cor = [r for r in rows if r.get("first_error_index", -1) == -1]
    rng = random.Random(seed)
    rng.shuffle(cor)
    keep_cor = cor[: int(round(len(err) * ratio))]
    balanced = err + keep_cor
    rng.shuffle(balanced)

    full = path.with_name(path.stem + "_full.jsonl")
    if not full.exists():
        path.rename(full)  # preserve original once
    with open(path, "w") as f:
        for r in balanced:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    return len(err), len(keep_cor), len(balanced)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default="data/verifier")
    ap.add_argument("--ratio", type=float, default=1.0, help="correct:erroneous ratio to keep")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    for name in ("train.jsonl", "val.jsonl"):
        p = Path(args.dir) / name
        if not p.exists():
            continue
        e, c, t = rebalance(p, args.ratio, args.seed)
        print(f"{name}: {t} rows (erroneous {e}, all-correct {c}); original -> {p.stem}_full.jsonl")
    return 0


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