"""Verifier exploitability audit from a Best-of-N log.

A Best-of-N log has one JSON object per line, one line per scored candidate:

    problem   which problem the candidate answers (required)
    score     the verifier's / reward model's score (required)
    correct   whether the candidate is actually right (required)
    pool      optional label for the pool the verifier faced (e.g. the prefilter)
    seen      optional: true if the verifier was trained or tuned on this problem

Everything is measured on the data, not assumed:

- problem ε: the fraction of mixed problems (at least one right and one wrong
  candidate) on which some wrong candidate scores at least as high as every
  right one. Ties count against the verifier.
- candidate ε: the fraction of wrong candidates, in mixed problems, that score
  at least as high as their problem's best right candidate. This is the
  per-candidate exploiter rate in the verifier-bottleneck model.
- the Best-of-N curve: expected accuracy of argmax over a random N-subset of
  each problem's candidates, computed exactly rather than by sampling. Its
  peak N* is the safe budget, and the drop from N* to the largest N measures
  how far more sampling hurts.
- ε* = p/(1−p): below it, consensus among high scorers neutralises the
  collapse; above it, buy a less exploitable verifier.

Exploitability belongs to a verifier and its pool, so ε is reported per
`pool`. A verifier must be judged on problems it was not fitted to, so when
`seen` is present the audit compares seen against unseen problems.
"""
import json
from dataclasses import dataclass, field
from math import comb, sqrt
from typing import Dict, Iterable, List, Optional

from .estimate import Rate
from .stats import Z95, diff_ci


@dataclass
class Candidate:
    problem: str
    score: float
    correct: bool
    pool: Optional[str] = None
    seen: Optional[bool] = None


def parse_candidates(lines: Iterable[str]) -> List[Candidate]:
    out = []
    for i, line in enumerate(lines, 1):
        line = line.strip()
        if not line:
            continue
        try:
            o = json.loads(line)
            c = Candidate(str(o["problem"]), float(o["score"]), bool(o["correct"]),
                          None if o.get("pool") is None else str(o["pool"]),
                          None if o.get("seen") is None else bool(o["seen"]))
        except (ValueError, KeyError, TypeError):
            raise ValueError(f"line {i} is not a candidate record (needs problem, score, correct)")
        out.append(c)
    if not out:
        raise ValueError("no candidates found")
    return out


def read_candidates(path) -> List[Candidate]:
    with open(path) as f:
        return parse_candidates(f)


def group(cands: List[Candidate]) -> Dict[str, List[Candidate]]:
    g: Dict[str, List[Candidate]] = {}
    for c in cands:
        g.setdefault(c.problem, []).append(c)
    return g


def _ranked(cs: List[Candidate]) -> List[bool]:
    """Correctness in descending score order, wrong before right on ties."""
    return [c.correct for c in sorted(cs, key=lambda c: (-c.score, c.correct))]


def bon_problem(order: List[bool], n: int) -> float:
    """Exact P(argmax of a random n-subset is right) for one problem."""
    m = len(order)
    total = comb(m, n)
    return sum(comb(m - 1 - r, n - 1) for r, ok in enumerate(order) if ok) / total


def _exploited(cs: List[Candidate]) -> Optional[bool]:
    right = [c.score for c in cs if c.correct]
    wrong = [c.score for c in cs if not c.correct]
    if not right or not wrong:
        return None
    return max(wrong) >= max(right)


def problem_eps(problems: Dict[str, List[Candidate]]) -> Rate:
    flags = [f for f in (_exploited(cs) for cs in problems.values()) if f is not None]
    return Rate.of(sum(flags), len(flags))


def candidate_eps(problems: Dict[str, List[Candidate]]) -> Rate:
    k = n = 0
    for cs in problems.values():
        right = [c.score for c in cs if c.correct]
        wrong = [c.score for c in cs if not c.correct]
        if not right or not wrong:
            continue
        best = max(right)
        n += len(wrong)
        k += sum(1 for w in wrong if w >= best)
    return Rate.of(k, n)


@dataclass
class CurvePoint:
    n: int
    accuracy: float
    problems: int


@dataclass
class Audit:
    candidates: int
    problems: int
    mixed: int
    p: float  # share of candidates that are right
    problem_eps: Rate
    candidate_eps: Rate
    eps_star: Optional[float]
    curve: List[CurvePoint]
    n_star: int
    peak: float
    at_max: float
    decline: float  # peak − accuracy at the largest N, same problems
    decline_lo: float
    decline_hi: float
    by_pool: Dict[str, Rate] = field(default_factory=dict)
    seen: Optional[Rate] = None
    unseen: Optional[Rate] = None
    seen_gap: Optional[tuple] = None  # (d, lo, hi) unseen − seen problem ε
    verdict: str = ""
    notes: List[str] = field(default_factory=list)


def _grid(nmax: int) -> List[int]:
    g, n = [], 1
    while n < nmax:
        g.append(n)
        n *= 2
    return g + [nmax]


EXPLOITED = "more samples hurt"
AT_RISK = "exploits above ε*"
SAFE = "no collapse at this budget"
THIN = "not enough mixed problems"


def audit(cands: List[Candidate]) -> Audit:
    problems = group(cands)
    orders = {k: _ranked(v) for k, v in problems.items()}
    nmax = min(len(o) for o in orders.values())
    curve = []
    for n in _grid(nmax):
        accs = [bon_problem(o, n) for o in orders.values()]
        curve.append(CurvePoint(n, sum(accs) / len(accs), len(accs)))
    best = max(curve, key=lambda c: (round(c.accuracy, 12), -c.n))
    last = curve[-1]
    # Paired over problems: accuracy at N* minus accuracy at the largest N.
    ds = [bon_problem(o, best.n) - bon_problem(o, last.n) for o in orders.values()]
    mean = sum(ds) / len(ds)
    sd = sqrt(sum((d - mean) ** 2 for d in ds) / (len(ds) - 1)) if len(ds) > 1 else 0.0
    half = Z95 * sd / sqrt(len(ds))

    p = sum(c.correct for c in cands) / len(cands)
    pe, ce = problem_eps(problems), candidate_eps(problems)
    eps_star = p / (1 - p) if p < 1 else None

    a = Audit(len(cands), len(problems), pe.n, p, pe, ce, eps_star, curve, best.n, best.accuracy,
              last.accuracy, mean, mean - half, mean + half)

    pools = sorted({c.pool for c in cands if c.pool is not None})
    for name in pools:
        a.by_pool[name] = problem_eps({k: v for k, v in problems.items() if v[0].pool == name})

    if any(c.seen is not None for c in cands):
        s = problem_eps({k: v for k, v in problems.items() if v[0].seen})
        u = problem_eps({k: v for k, v in problems.items() if v[0].seen is False})
        a.seen, a.unseen = s, u
        if s.n and u.n:
            a.seen_gap = diff_ci(s.k, s.n, u.k, u.n)

    if pe.n < 30:
        a.verdict = THIN
        a.notes.append(f"Only {pe.n} problems have both a right and a wrong candidate; ε intervals are wide.")
    elif a.decline_lo > 0:
        a.verdict = EXPLOITED
        a.notes.append(f"Accuracy peaks at N={best.n} and falls by {mean:.3f} by N={last.n}. "
                       "Cap the budget at N* and spend the surplus on a less exploitable verifier.")
    elif eps_star is not None and ce.lo > eps_star:
        a.verdict = AT_RISK
        a.notes.append(f"Candidate ε {ce.p:.3f} is above ε* = p/(1−p) = {eps_star:.3f}: exploits can win a "
                       "plurality, so consensus will not rescue larger budgets. Audit again at larger N.")
    else:
        a.verdict = SAFE
        if eps_star is not None and ce.k:
            a.notes.append(f"Candidate ε {ce.p:.3f} is below ε* = p/(1−p) = {eps_star:.3f}: if the curve turns "
                           "down at larger N, selecting by consensus among high scorers should hold it.")
    if a.seen_gap and a.seen_gap[1] > 0:
        a.notes.append(f"Unseen problems are exploited more often than seen ones (+{a.seen_gap[0]:.3f}). "
                       "An audit on the verifier's own training data understates ε; report the unseen figure.")
    if len(a.by_pool) > 1:
        a.notes.append("ε differs by pool: it belongs to the verifier and the pool it faces. "
                       "Calibrate on the post-prefilter pool used in production.")
    return a