"""Tune a selector instead of retraining what it selects from.

Mis-selection is fixed at the selector: a router, a retriever or a verifier.
`tune` takes a decision log whose records carry `scores` ({candidate: score})
and a `target`, and fits three cheap repairs:

1. Per-candidate standardisation. Experts that each report their own
   confidence are not on one scale; a habitually confident candidate wins
   picks it should lose. Each candidate's scores are z-scored with its own
   mean and spread.
2. Deferral. When the top two candidates are close, the pick is a coin toss.
   A margin threshold τ is chosen so that the accepted picks meet a target ε,
   and the rest are deferred (to a verifier, a larger model or a person).
3. A shortlist. The smallest k whose top-k contains the target often enough,
   so a verifier only has to judge k candidates instead of the whole pool.

Every choice is fitted on a calibration half and judged on a held-out half,
never on the decisions it was fitted to: a filter must be estimated rather
than trained on the data it judges. Records alternate between the halves in
log order, so the split is reproducible.
"""
from dataclasses import dataclass, field
from math import sqrt
from typing import Dict, List, Optional, Sequence, Tuple

from .estimate import Rate
from .log import Decision
from .stats import diff_ci


@dataclass
class Selector:
    """A fitted selection rule: z-scores (if `standardise`) plus a deferral margin."""
    standardise: bool
    stats: Dict[str, Tuple[float, float]] = field(default_factory=dict)  # candidate -> (mean, sd)
    fallback: Tuple[float, float] = (0.0, 1.0)
    margin: Optional[float] = None  # defer when top1 − top2 < margin

    def adjusted(self, scores: Dict[str, float]) -> Dict[str, float]:
        if not self.standardise:
            return dict(scores)
        out = {}
        for c, s in scores.items():
            mu, sd = self.stats.get(c, self.fallback)
            out[c] = (s - mu) / sd
        return out

    def ranked(self, scores: Dict[str, float]) -> List[str]:
        adj = self.adjusted(scores)
        return sorted(adj, key=lambda c: (-adj[c], c))

    def select(self, scores: Dict[str, float]) -> Tuple[Optional[str], float]:
        """Return (choice, margin); choice is None when the pick should be deferred."""
        adj = self.adjusted(scores)
        order = sorted(adj, key=lambda c: (-adj[c], c))
        gap = adj[order[0]] - adj[order[1]] if len(order) > 1 else float("inf")
        if self.margin is not None and gap < self.margin:
            return None, gap
        return order[0], gap

    def to_json(self):
        return {"standardise": self.standardise, "stats": {c: list(v) for c, v in self.stats.items()},
                "fallback": list(self.fallback), "margin": self.margin}


def _usable(ds: Sequence[Decision]) -> List[Decision]:
    return [d for d in ds if d.scores and d.target is not None and d.target in d.scores and len(d.scores) > 1]


def split(ds: Sequence[Decision]) -> Tuple[List[Decision], List[Decision]]:
    u = _usable(ds)
    return u[0::2], u[1::2]


def _mean_sd(xs: List[float]) -> Tuple[float, float]:
    m = sum(xs) / len(xs)
    sd = sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1)) if len(xs) > 1 else 0.0
    return m, sd


def fit_standardiser(ds: Sequence[Decision]) -> Selector:
    per: Dict[str, List[float]] = {}
    every: List[float] = []
    for d in ds:
        for c, s in d.scores.items():
            per.setdefault(c, []).append(s)
            every.append(s)
    gm, gsd = _mean_sd(every)
    fallback = (gm, gsd if gsd > 0 else 1.0)
    stats = {}
    for c, xs in per.items():
        m, sd = _mean_sd(xs)
        stats[c] = (m, sd) if len(xs) > 1 and sd > 0 else fallback
    return Selector(True, stats, fallback)


def eps_of(sel: Selector, ds: Sequence[Decision]) -> Rate:
    k = sum(1 for d in ds if sel.ranked(d.scores)[0] != d.target)
    return Rate.of(k, len(ds))


def topk_miss(sel: Selector, ds: Sequence[Decision], k: int) -> Rate:
    miss = sum(1 for d in ds if d.target not in sel.ranked(d.scores)[:k])
    return Rate.of(miss, len(ds))


def _gap(sel: Selector, d: Decision) -> float:
    adj = sel.adjusted(d.scores)
    a, b = sorted(adj.values(), reverse=True)[:2]
    return a - b


def fit_margin(sel: Selector, ds: Sequence[Decision], target_eps: float) -> Optional[float]:
    """Smallest margin whose accepted picks have a Wilson upper bound on ε ≤ target_eps."""
    pts = sorted(((_gap(sel, d), sel.ranked(d.scores)[0] != d.target) for d in ds), key=lambda t: -t[0])
    best = None
    wrong = 0
    for i, (g, w) in enumerate(pts):
        wrong += w
        # Only cut between distinct gaps, so every decision at this gap is accepted.
        if i + 1 < len(pts) and pts[i + 1][0] == g:
            continue
        if Rate.of(wrong, i + 1).hi <= target_eps:
            best = g
    return best


def fit_margin_cost(sel: Selector, ds: Sequence[Decision], defer_cost: float, error_cost: float) -> Optional[float]:
    """Margin that minimises defer_cost · deferred + error_cost · accepted mistakes.

    None means deferring nothing is cheapest."""
    pts = sorted(((_gap(sel, d), sel.ranked(d.scores)[0] != d.target) for d in ds), key=lambda t: -t[0])
    total_wrong = sum(w for _, w in pts)
    best_cost, best = error_cost * total_wrong, None  # accept everything
    wrong = 0
    for i, (g, w) in enumerate(pts):
        wrong += w
        if i + 1 < len(pts) and pts[i + 1][0] == g:
            continue
        if i + 1 == len(pts):
            break  # accepting everything is the baseline above
        cost = error_cost * wrong + defer_cost * (len(pts) - i - 1)
        if cost < best_cost:
            # Defer everything below the next gap down: margin sits just above it.
            best_cost, best = cost, (g + pts[i + 1][0]) / 2
    return best


def cost_per_decision(sel: Selector, ds: Sequence[Decision], defer_cost: float, error_cost: float) -> float:
    total = 0.0
    for d in ds:
        choice, _ = sel.select(d.scores)
        total += defer_cost if choice is None else error_cost * (choice != d.target)
    return total / len(ds)


@dataclass
class Tuning:
    calibration: int
    held_out: int
    raw_cal: Rate
    std_cal: Rate
    raw_held: Rate
    std_held: Rate
    standardise: bool
    selector: Selector
    target_eps: float
    margin: Optional[float]
    coverage: Optional[Rate]  # held-out share of picks accepted
    accepted_eps: Optional[Rate]  # held-out ε among accepted picks
    shortlist: Optional[int]
    shortlist_miss: Optional[Rate]
    notes: List[str] = field(default_factory=list)
    cost_held: Optional[float] = None  # held-out cost per decision with deferral (cost mode)
    cost_held_base: Optional[float] = None  # same, without deferral


def tune(ds: Sequence[Decision], target_eps: float = 0.02, max_k: int = 10,
         defer_cost: Optional[float] = None, error_cost: Optional[float] = None) -> Tuning:
    """Fit on calibration, report on held-out. With `defer_cost` and `error_cost`
    the margin minimises expected cost instead of meeting `target_eps`."""
    cost_mode = defer_cost is not None and error_cost is not None
    cal, held = split(ds)
    if len(cal) < 2 or not held:
        raise ValueError("need at least three decisions with `scores` that include the `target`")
    raw = Selector(False)
    std = fit_standardiser(cal)
    rc, sc = eps_of(raw, cal), eps_of(std, cal)
    rh, sh = eps_of(raw, held), eps_of(std, held)
    use_std = sc.k < rc.k  # chosen on calibration only
    base = std if use_std else raw

    margin = (fit_margin_cost(base, cal, defer_cost, error_cost) if cost_mode
              else fit_margin(base, cal, target_eps))
    cov = acc_eps = None
    if margin is not None:
        sel = Selector(base.standardise, base.stats, base.fallback, margin)
        picks = [(sel.select(d.scores)[0], d.target) for d in held]
        kept = [(c, t) for c, t in picks if c is not None]
        cov = Rate.of(len(kept), len(held))
        acc_eps = Rate.of(sum(1 for c, t in kept if c != t), len(kept))
    else:
        sel = base

    shortlist = miss = None
    pool = min(len(d.scores) for d in cal)
    for k in range(1, min(max_k, pool) + 1):
        if topk_miss(base, cal, k).hi <= target_eps:
            shortlist, miss = k, topk_miss(base, held, k)
            break

    t = Tuning(len(cal), len(held), rc, sc, rh, sh, use_std, sel, target_eps, margin, cov, acc_eps, shortlist, miss)
    if cost_mode:
        t.cost_held = cost_per_decision(sel, held, defer_cost, error_cost)
        t.cost_held_base = cost_per_decision(base, held, defer_cost, error_cost)
    head = sh if use_std else rh
    if use_std:
        d, lo, hi = diff_ci(rh.k, rh.n, sh.k, sh.n)
        if hi < 0:
            t.notes.append(f"Standardising each candidate's scores cuts held-out ε from {rh.p:.3f} to {sh.p:.3f}: "
                           "the candidates were not reporting on one scale.")
        else:
            t.notes.append(f"Standardised scores won on calibration, but the held-out change ({d:+.3f}, "
                           f"95% CI [{lo:.3f}, {hi:.3f}]) is not significant.")
    if cost_mode:
        if margin is None:
            t.notes.append(f"At a deferral cost of {defer_cost:g} and an error cost of {error_cost:g}, "
                           "deferring nothing is cheapest on calibration.")
        else:
            t.notes.append(f"Defer below a gap of {margin:.3g}: held-out cost per decision {t.cost_held:.4g} "
                           f"vs {t.cost_held_base:.4g} without deferral (coverage {cov.p:.0%}).")
    elif margin is None:
        t.notes.append(f"No margin reaches ε ≤ {target_eps:g} on calibration; collect more labelled decisions "
                       "or fix the scores themselves.")
    elif cov is not None and acc_eps is not None and not cost_mode:
        t.notes.append(f"Defer when the top-two gap is below {margin:.3g}: on held-out decisions that keeps "
                       f"{cov.p:.0%} of picks at ε {acc_eps.p:.3f} (vs {head.p:.3f} with no deferral).")
        if acc_eps.n and acc_eps.lo > target_eps:
            t.notes.append("Held-out ε among accepted picks is above target: the calibration half was not "
                           "representative. Refit on more data before deploying.")
    if shortlist:
        t.notes.append(f"A verifier choosing among the top {shortlist} sees the target {1 - miss.p:.1%} of the time "
                       "on held-out decisions.")
    if len(cal) < 200:
        t.notes.append(f"Only {len(cal)} calibration decisions: intervals are wide.")
    return t