"""Paired permutation test between two eval JSONL sidecars.

Compares per-query scores from two `eval.run_eval` sidecar files so
"P@1 went 0.88 → 0.91 on 25 queries" is not treated as a win.

    python -m eval.pvalue \\
        --a eval/results/baseline.jsonl \\
        --b eval/results/new.jsonl \\
        --metric rr

Exit 0 even when the difference is not significant — this is a report,
not a gate. No third-party deps; `random.Random(seed)` is enough.
"""
from __future__ import annotations

import argparse
import json
import random
from pathlib import Path


def load_sidecar(path: Path) -> list[dict]:
    rows: list[dict] = []
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def paired_permutation(
    a: list[float],
    b: list[float],
    n_resamples: int = 10000,
    seed: int = 0,
) -> dict:
    """Two-sided paired permutation test on per-query scores.

    Null: each pair is exchangeable (randomly flipping the sign of
    A_i - B_i). p_value is the fraction of permutations whose
    |mean diff| is at least as large as the observed |mean(A-B)|.
    """
    if len(a) != len(b):
        raise ValueError(f"paired lengths differ: {len(a)} vs {len(b)}")
    if not a:
        raise ValueError("no paired queries to compare")
    diffs = [x - y for x, y in zip(a, b)]
    n = len(diffs)
    observed = sum(diffs) / n
    abs_obs = abs(observed)
    rng = random.Random(seed)
    extreme = 0
    for _ in range(n_resamples):
        total = 0.0
        for d in diffs:
            total += d if rng.random() < 0.5 else -d
        if abs(total / n) >= abs_obs - 1e-15:
            extreme += 1
    p_value = extreme / n_resamples
    return {
        "A_mean": sum(a) / n,
        "B_mean": sum(b) / n,
        "Diff(A-B)": observed,
        "p_value": p_value,
        "significant": p_value < 0.05,
        "n": n,
        "n_resamples": n_resamples,
    }


def _index(rows: list[dict], metric: str) -> dict[tuple[str, str], float]:
    """Map (retriever, query) -> score."""
    out: dict[tuple[str, str], float] = {}
    for row in rows:
        retriever = str(row.get("retriever") or "")
        query = str(row.get("query") or "")
        if metric == "p_at_1":
            score = float(row.get("p_at_1") or 0)
        else:
            score = float(row.get("rr") or 0)
        out[(retriever, query)] = score
    return out


def compare(
    rows_a: list[dict],
    rows_b: list[dict],
    metric: str = "rr",
    retriever: str | None = None,
    n_resamples: int = 10000,
    seed: int = 0,
) -> list[dict]:
    """Join on (retriever, query). One result dict per shared retriever."""
    ia, ib = _index(rows_a, metric), _index(rows_b, metric)
    retrievers = sorted({r for r, _ in ia} & {r for r, _ in ib})
    if retriever:
        retrievers = [r for r in retrievers if r == retriever]
        if not retrievers:
            raise ValueError(f"retriever {retriever!r} not in both sidecars")
    reports = []
    for name in retrievers:
        queries = sorted({q for r, q in ia if r == name} & {q for r, q in ib if r == name})
        if not queries:
            continue
        a_scores = [ia[(name, q)] for q in queries]
        b_scores = [ib[(name, q)] for q in queries]
        report = paired_permutation(a_scores, b_scores, n_resamples=n_resamples, seed=seed)
        report["retriever"] = name
        report["metric"] = metric
        reports.append(report)
    if not reports:
        raise ValueError("no overlapping (retriever, query) pairs")
    return reports


def render(reports: list[dict]) -> str:
    lines = ["# Permutation test", ""]
    for r in reports:
        sig = "yes" if r["significant"] else "no"
        lines += [
            f"## `{r['retriever']}` ({r['metric']}, n={r['n']})",
            "",
            f"- A_mean: `{r['A_mean']:.4f}`",
            f"- B_mean: `{r['B_mean']:.4f}`",
            f"- Diff(A-B): `{r['Diff(A-B)']:.4f}`",
            f"- p_value: `{r['p_value']:.4f}` ({r['n_resamples']} resamples)",
            f"- significant (p < 0.05): **{sig}**",
            "",
        ]
    return "\n".join(lines)


def main() -> int:
    p = argparse.ArgumentParser(description="Paired permutation test on two eval JSONL sidecars.")
    p.add_argument("--a", type=Path, required=True, help="sidecar JSONL (system A)")
    p.add_argument("--b", type=Path, required=True, help="sidecar JSONL (system B)")
    p.add_argument("--metric", choices=("rr", "p_at_1"), default="rr")
    p.add_argument("--retriever", default=None, help="restrict to one retriever name")
    p.add_argument("--n-resamples", type=int, default=10000)
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()
    reports = compare(
        load_sidecar(args.a),
        load_sidecar(args.b),
        metric=args.metric,
        retriever=args.retriever,
        n_resamples=args.n_resamples,
        seed=args.seed,
    )
    print(render(reports), end="")
    return 0


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