"""The pre-registered primary contrast: paired McNemar between two eval reports.

Frozen pre-registration (docs/RESULTS.md): DPO Arm B result vs new-SFT,
GSM8K-full n=1,319, paired exact McNemar, alpha=0.05. This driver is pure
arithmetic over two reports' per_problem rows -- it computes what was
pre-registered and nothing else. The statistic itself is
`align.eval.metrics.mcnemar`, the SAME implementation the pre-registration
froze (exact binomial in the small-discordant regime); this script must
never reimplement it.

Pairing is enforced, not assumed: rows join on problem_id within --source,
and a partial join REFUSES by default -- a silently dropped problem is a
silently different experiment. pass@1 is correct[0], matching how the
banked SFT report's headline numbers were computed.

    python scripts/primary_contrast.py \
        --arm outputs/eval/arm-b-final.json \
        --baseline outputs/eval/sft-final.json \
        --source gsm8k --expect-n 1319 \
        --label "dpo_arm_b_response vs sft-final"
"""

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.eval.metrics import mcnemar, mcnemar_exact, wilson_interval


def _pass1_by_id(report: dict, source: str) -> dict[str, bool]:
    out: dict[str, bool] = {}
    for row in report["per_problem"]:
        if row.get("source") != source:
            continue
        pid = row["problem_id"]
        if pid in out:
            raise ValueError(f"duplicate problem_id {pid!r} in one report; join is ill-defined")
        out[pid] = bool(row["correct"][0])
    return out


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--arm", required=True, help="eval report JSON for the arm under test")
    parser.add_argument("--baseline", required=True, help="eval report JSON for the baseline side")
    parser.add_argument("--source", default="gsm8k")
    parser.add_argument("--expect-n", type=int, default=1319, help="refuse if the paired n differs")
    parser.add_argument("--alpha", type=float, default=0.05)
    parser.add_argument("--label", required=True, help="human label recorded in the receipt")
    parser.add_argument("--out", default=None, help="default: outputs/eval/primary_contrast.json")
    parser.add_argument(
        "--allow-partial",
        action="store_true",
        help="permit a partial join (still recorded loudly); NOT the pre-registered path",
    )
    args = parser.parse_args()

    arm_report = json.load(open(args.arm))
    base_report = json.load(open(args.baseline))
    arm = _pass1_by_id(arm_report, args.source)
    base = _pass1_by_id(base_report, args.source)

    shared = sorted(set(arm) & set(base))
    missing = len(set(arm) ^ set(base))
    if (missing or len(shared) != args.expect_n) and not args.allow_partial:
        print(
            f"REFUSED: paired n={len(shared)} (expected {args.expect_n}), "
            f"{missing} unmatched problem_ids. The pre-registered contrast is "
            "full-set paired; a partial join is a different experiment. "
            "Pass --allow-partial only with a dated RESULTS.md entry explaining why.",
            file=sys.stderr,
        )
        return 2

    # Discordants: b = arm right where baseline wrong; c = the reverse.
    b = sum(1 for p in shared if arm[p] and not base[p])
    c = sum(1 for p in shared if base[p] and not arm[p])
    both_right = sum(1 for p in shared if arm[p] and base[p])
    both_wrong = len(shared) - b - c - both_right
    # PRIMARY p-value: the exact binomial at any discordant count, because
    # the frozen pre-registration's WORDING is "paired exact McNemar" while
    # the frozen mcnemar() implementation switches to chi-square at 25+
    # discordants. Resolved by the 2026-08-22 dated RESULTS.md entry, decided
    # BEFORE any arm data existed; mcnemar()'s output is recorded alongside.
    stat_exact, p_exact = mcnemar_exact(b, c)
    stat_frozen, p_frozen = mcnemar(b, c)
    frozen_branch = "exact_binomial" if (b + c) < 25 else "chi_square_continuity_corrected"

    n = len(shared)
    arm_k = b + both_right
    base_k = c + both_right
    receipt = {
        "label": args.label,
        "source": args.source,
        "n_paired": n,
        "unmatched": missing,
        "arm": {"report": args.arm, "pass@1": arm_k / n, "ci95": wilson_interval(arm_k, n)},
        "baseline": {
            "report": args.baseline,
            "pass@1": base_k / n,
            "ci95": wilson_interval(base_k, n),
        },
        "table": {
            "both_right": both_right,
            "both_wrong": both_wrong,
            "arm_only": b,
            "baseline_only": c,
        },
        "mcnemar": {
            "method": "exact_binomial",
            "statistic": stat_exact,
            "p_two_sided": p_exact,
            "alpha": args.alpha,
            "frozen_impl": {
                "statistic": stat_frozen,
                "p_two_sided": p_frozen,
                "branch": frozen_branch,
            },
        },
        "significant": bool(p_exact < args.alpha),
        "direction": "arm" if b > c else ("baseline" if c > b else "tie"),
    }

    out_path = args.out or os.path.join("outputs", "eval", "primary_contrast.json")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    tmp = out_path + ".partial"
    with open(tmp, "w") as f:
        json.dump(receipt, f, indent=1)
    os.replace(tmp, out_path)

    print(f"{args.label} on {args.source}, n={n} paired")
    print(f"  arm      pass@1 = {receipt['arm']['pass@1']:.4f}  {args.arm}")
    print(f"  baseline pass@1 = {receipt['baseline']['pass@1']:.4f}  {args.baseline}")
    print(f"  discordant: arm-only {b}, baseline-only {c} (concordant {both_right}+{both_wrong})")
    print(
        f"  exact McNemar p = {p_exact:.5f} (alpha {args.alpha}) -> significant = "
        f"{receipt['significant']}  [frozen impl {frozen_branch}: p = {p_frozen:.5f}]"
    )
    print(f"receipt written: {out_path}")
    return 0


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