"""The verifier: leave-one-out fit of a candidate solver on attested pairs.

One object, reused everywhere: grammar selection, CEGIS failure feedback, and
(offline) RL reward. score = sqrt(EM * chrF) on held-out attested pairs,
minus an MDL penalty so the simplest adequate grammar wins ties.

A "candidate" is anything with `predict(src: str) -> str` for the relevant
direction; grammars, analogy baselines, and raw LLM outputs all fit.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Callable, List, Optional, Sequence, Tuple

from .metrics import chrf, exact_match

Predictor = Callable[[str], Optional[str]]


@dataclass
class Verdict:
    em: float
    chrf: float
    mdl: float
    failures: List[Tuple[str, str, str]] = field(default_factory=list)  # (src, gold, pred)

    @property
    def fit(self) -> float:
        return math.sqrt(max(self.em, 0.0) * max(self.chrf, 0.0))

    @property
    def score(self) -> float:
        return self.fit - self.mdl

    def __repr__(self) -> str:
        return f"Verdict(em={self.em:.3f}, chrf={self.chrf:.3f}, mdl={self.mdl:.4f}, n_fail={len(self.failures)})"


def evaluate(
    predict: Predictor,
    pairs: Sequence[Tuple[str, str]],
    mdl_cost: float = 0.0,
    mdl_weight: float = 0.002,
) -> Verdict:
    """Score `predict` on attested (src, gold) pairs. A None/empty prediction
    scores 0 on both metrics for that pair. mdl_cost is the grammar's
    description length (see dsl.grammar.Grammar.mdl); weighted lightly so it
    only breaks ties."""
    if not pairs:
        return Verdict(0.0, 0.0, mdl_cost * mdl_weight)
    em_sum = chrf_sum = 0.0
    failures = []
    for src, gold in pairs:
        pred = predict(src) or ""
        e, c = exact_match(pred, gold), chrf(pred, gold)
        em_sum += e
        chrf_sum += c
        if e < 1.0:
            failures.append((src, gold, pred))
    n = len(pairs)
    return Verdict(em_sum / n, chrf_sum / n, mdl_cost * mdl_weight, failures)


LOO_MAX_FOLDS = 12


def leave_one_out(
    fit_predict: Callable[[Sequence[Tuple[str, str]]], Predictor],
    pairs: Sequence[Tuple[str, str]],
    mdl_cost: float = 0.0,
    max_folds: int = LOO_MAX_FOLDS,
) -> Verdict:
    """True LOO for candidates that are *fit* from pairs (analogy, alignment
    baselines): refit without pair i, predict pair i. For a fixed grammar
    (already synthesized), use `evaluate` directly — the LLM saw the pairs,
    but the grammar either reproduces them or it doesn't.

    Refitting is O(pairs^2)+ per fold (alignment rebuild), so folds are capped
    at `max_folds` evenly-spaced held-out pairs — an unbiased estimate is all
    the selection needs, and the 30-minute budget cannot afford exact LOO on
    40-pair puzzles."""
    if not pairs:
        return Verdict(0.0, 0.0, 0.0)
    n = len(pairs)
    if n <= max_folds:
        fold_idx = range(n)
    else:
        step = n / max_folds
        fold_idx = sorted({int(k * step) for k in range(max_folds)})
    em_sum = chrf_sum = 0.0
    failures = []
    for i in fold_idx:
        src, gold = pairs[i]
        held_in = [p for j, p in enumerate(pairs) if j != i]
        pred = fit_predict(held_in)(src) or ""
        e, c = exact_match(pred, gold), chrf(pred, gold)
        em_sum += e
        chrf_sum += c
        if e < 1.0:
            failures.append((src, gold, pred))
    k = len(list(fold_idx))
    return Verdict(em_sum / k, chrf_sum / k, mdl_cost * 0.002, failures)


def select_best(
    candidates: Sequence[Tuple[str, Predictor, float]],
    pairs: Sequence[Tuple[str, str]],
) -> Tuple[Optional[str], Optional[Predictor], Verdict]:
    """Pick the best (name, predictor, mdl_cost) by verifier score.
    Ties broken by lower MDL (already in score), then earlier order
    (candidates should be ordered by prior preference: symbolic first)."""
    best: Tuple[Optional[str], Optional[Predictor], Verdict] = (None, None, Verdict(0, 0, 0))
    best_score = -1e9
    for name, pred, mdl in candidates:
        v = evaluate(pred, pairs, mdl)
        if v.score > best_score + 1e-9:
            best_score = v.score
            best = (name, pred, v)
    return best