"""Evaluation metrics for the own classifier (pure Python, standard library only, no repo imports).

Spec: docs/market-psychology-own-classifier-training.md, sections 5 and 7. Runs on the pod (Python 3.11+)
and locally (``classifier_eval.py``). Labels are option strings; a distribution is a dict option ->
probability in taxonomy option order; ``tax`` is anything with ``names`` and ``fields`` (``taxonomy.Taxonomy``).

Definitions (fixed here, where the spec leaves room):
- **Null gold labels** are skipped by every per-field metric (they are masked in training). d3-style
  end-to-end scoring (``end_to_end``) instead counts a null as "not the class", as
  ``d3_calibration.outcomes`` does; its callers build the rows.
- **F1 of a class** = 2TP / (2TP + FP + FN). Precision = TP / (TP + FP) is None without predictions;
  recall = TP / (TP + FN) is None without gold support.
- **Macro-F1** is the unweighted mean F1 over every class of the field's taxonomy options, except a class
  with no gold support and no prediction (2TP + FP + FN = 0), whose F1 is undefined and which is left out.
  A class with gold support or predictions always counts, so a missed or only falsely predicted class
  contributes 0. A label outside the given classes raises ValueError.
- **Top label** of a distribution: the highest probability, ties to the earlier option (``taxonomy.top``).
- **Top-label ECE**: confidence = the top label's probability, correct = top label == gold; 15 equal-width
  bins over [0, 1], bin = min(floor(confidence x 15), 14), so bin k holds [k/15, (k+1)/15) and the last
  bin also holds 1.0; ECE = sum over non-empty bins of (bin size / n) x |accuracy - mean confidence|.
- **Class-wise ECE** (one against the rest): the same bins over p(class), outcome gold == class.
- **Log-loss**: mean -ln(max(p(gold), 1e-12)) over non-null gold (the floor is ``calibrate.FLOOR``).
- **Shares**: predicted share = mean p(class), gold share = share of gold == class, both over the rows
  whose gold is not null.
- **Exact match**: rows whose gold is complete (no null in the scored fields) with every field equal;
  rows with a null gold field are excluded and counted.
- **Bootstrap**: author-group percentile bootstrap as in ``d3_calibration.bootstrap``: sorted groups,
  ``random.Random(seed)``, each replicate draws as many groups as there are with replacement and
  concatenates their ids; ``fn(ids)`` returning None is skipped; the interval is
  [values[floor(a/2 (m-1))], values[floor((1-a/2)(m-1))]] over the m sorted kept values. A paired
  statistic is any ``fn`` that computes both sides on the same drawn ids (for example F1(student) -
  F1(d3)).
"""
from __future__ import annotations

from collections import Counter, defaultdict
import math
import random

VERSION = 'classifier-metrics-v1'
FLOOR = 1e-12
BINS = 15
REPS = 2000
SEED = 'classifier-eval-20261017'


# ------------------------------------------------------------------ helpers
def top(dist: dict) -> str:
    """Highest-probability option; ties go to the earlier option (as ``taxonomy.top``)."""
    best = None
    for k, v in dist.items():
        if best is None or v > dist[best]:
            best = k
    return best


def softmax(logits, t: float = 1.0) -> list:
    scaled = [x / t for x in logits]
    m = max(scaled)
    exp = [math.exp(x - m) for x in scaled]
    total = sum(exp)
    return [e / total for e in exp]


def _pairs(gold, other):
    gold, other = list(gold), list(other)
    if len(gold) != len(other):
        raise ValueError('Gold and predictions differ in length: %d, %d' % (len(gold), len(other)))
    return [(g, o) for g, o in zip(gold, other) if g is not None]


def _ratio(a, b):
    return a / b if b else None


def _bin(p: float, bins: int) -> int:
    return min(max(int(math.floor(p * bins)), 0), bins - 1)


def _ece(points, bins: int):
    """points: (confidence, outcome 0/1)."""
    if not points:
        return None
    total = defaultdict(lambda: [0, 0.0, 0.0])
    for conf, hit in points:
        b = total[_bin(conf, bins)]
        b[0] += 1
        b[1] += conf
        b[2] += hit
    n = len(points)
    return sum(c / n * abs(acc / c - s / c) for c, s, acc in total.values())


# ------------------------------------------------------------------ per field
def accuracy(gold, pred):
    pairs = _pairs(gold, pred)
    return _ratio(sum(1 for g, p in pairs if g == p), len(pairs))


def per_class(gold, pred, classes) -> dict:
    """{class: {support, predicted, tp, precision, recall, f1}} for every given class."""
    classes = list(classes)
    known = set(classes)
    hits, support, predicted = Counter(), Counter(), Counter()
    for g, p in _pairs(gold, pred):
        if g not in known or p not in known:
            raise ValueError('Label outside the classes: %r / %r' % (g, p))
        support[g] += 1
        predicted[p] += 1
        if g == p:
            hits[g] += 1
    out = {}
    for c in classes:
        tp, s, n = hits[c], support[c], predicted[c]
        out[c] = {'support': s, 'predicted': n, 'tp': tp, 'precision': _ratio(tp, n), 'recall': _ratio(tp, s),
                  'f1': _ratio(2 * tp, n + s)}  # 2TP + FP + FN = predicted + support
    return out


def macro_f1(gold, pred, classes):
    """Mean F1 over the classes with gold support or predictions (see the module docstring)."""
    scores = [v['f1'] for v in per_class(gold, pred, classes).values() if v['f1'] is not None]
    return sum(scores) / len(scores) if scores else None


def log_loss(gold, dists):
    pairs = _pairs(gold, dists)
    if not pairs:
        return None
    return sum(-math.log(max(d[g], FLOOR)) for g, d in pairs) / len(pairs)


def ece(gold, dists, bins: int = BINS):
    """Top-label expected calibration error."""
    points = []
    for g, d in _pairs(gold, dists):
        best = top(d)
        points.append((d[best], 1.0 if best == g else 0.0))
    return _ece(points, bins)


def class_ece(gold, dists, cls: str, bins: int = BINS):
    """One-against-the-rest calibration error of p(cls)."""
    return _ece([(d[cls], 1.0 if g == cls else 0.0) for g, d in _pairs(gold, dists)], bins)


def shares(gold, dists, cls: str) -> dict:
    pairs = _pairs(gold, dists)
    n = len(pairs)
    predicted = _ratio(sum(d[cls] for _, d in pairs), n)
    observed = _ratio(sum(1 for g, _ in pairs if g == cls), n)
    return {'n': n, 'predicted_share': predicted, 'gold_share': observed,
            'difference': None if predicted is None else predicted - observed}


def key_class(gold, dists, cls: str) -> dict:
    """Precision, recall, F1 (top label), class-wise ECE and predicted against gold share for one class."""
    gold, dists = list(gold), list(dists)
    pred = [top(d) for d in dists]
    known = set(gold) | set(pred) | {cls}
    known.discard(None)
    pc = per_class(gold, pred, sorted(known))[cls]
    out = dict(pc, class_ece=class_ece(gold, dists, cls))
    out.update(shares(gold, dists, cls))
    return out


def field_summary(gold, dists, classes) -> dict:
    """Accuracy, macro-F1, log-loss, top-label ECE and per-class figures of one field."""
    gold, dists = list(gold), list(dists)
    pred = [top(d) for d in dists]
    pc = per_class(gold, pred, classes)
    scores = [v['f1'] for v in pc.values() if v['f1'] is not None]
    return {'n': sum(1 for g in gold if g is not None), 'nulls': sum(1 for g in gold if g is None),
            'accuracy': accuracy(gold, pred), 'macro_f1': sum(scores) / len(scores) if scores else None,
            'macro_f1_classes': len(scores), 'log_loss': log_loss(gold, dists), 'ece': ece(gold, dists),
            'per_class': pc}


def exact_match(gold_rows, pred_rows, fields) -> dict:
    """gold_rows, pred_rows: aligned lists of {field: option|None} / {field: option}."""
    gold_rows, pred_rows = list(gold_rows), list(pred_rows)
    if len(gold_rows) != len(pred_rows):
        raise ValueError('Gold and predictions differ in length')
    n = matches = excluded = 0
    for g, p in zip(gold_rows, pred_rows):
        if any(g.get(f) is None for f in fields):
            excluded += 1
            continue
        n += 1
        matches += all(g[f] == p[f] for f in fields)
    return {'n': n, 'matches': matches, 'rate': _ratio(matches, n), 'excluded_incomplete_gold': excluded}


# ------------------------------------------------------------------ end to end (d3_calibration terms)
def binary(truth, flags) -> dict:
    """Hit rate (recall), false-alarm rate, precision and F1 of 0/1 decisions against 0/1 truth."""
    truth, flags = [bool(x) for x in truth], [bool(x) for x in flags]
    if len(truth) != len(flags):
        raise ValueError('Truth and decisions differ in length')
    tp = sum(1 for t, f in zip(truth, flags) if t and f)
    fp = sum(1 for t, f in zip(truth, flags) if f and not t)
    pos = sum(truth)
    neg = len(truth) - pos
    fn = pos - tp
    return {'n': len(truth), 'positives': pos, 'predicted': tp + fp, 'tp': tp, 'fp': fp, 'fn': fn,
            'hit_rate': _ratio(tp, pos), 'false_alarm_rate': _ratio(fp, neg), 'precision': _ratio(tp, tp + fp),
            'f1': _ratio(2 * tp, 2 * tp + fp + fn)}


def end_to_end(rows: dict, ids, field: str, cls: str, side: str, reference: str = 'astra') -> dict:
    """Rates of one side's end-to-end labels against the reference, in ``d3_calibration.outcomes`` rows:
    ``rows[i][side][field]`` is the label inside that side's primary population, else None (not the class)."""
    ids = list(ids)
    out = binary([rows[i][reference][field] == cls for i in ids], [rows[i][side][field] == cls for i in ids])
    out['reference_share'] = _ratio(out['positives'], out['n'])
    out['side_share'] = _ratio(out['predicted'], out['n'])
    return out


def f1_of(truth, flags):
    return binary(truth, flags)['f1']


# ------------------------------------------------------------------ bootstrap
def bootstrap(ids, group_of, fn, reps: int = REPS, seed=SEED, level: float = 0.95) -> dict:
    """Author-group percentile bootstrap of fn(ids). ``group_of``: dict or callable id -> group."""
    get = group_of.__getitem__ if isinstance(group_of, dict) else group_of
    groups = defaultdict(list)
    for i in ids:
        groups[get(i)].append(i)
    keys, rng, values = sorted(groups), random.Random(seed), []
    for _ in range(reps):
        draw = [i for k in (rng.choice(keys) for _ in keys) for i in groups[k]]
        v = fn(draw)
        if v is not None:
            values.append(v)
    values.sort()
    low = round((1.0 - level) / 2, 12)  # exactly 0.025 and 0.975 at level 0.95, as d3_calibration.bootstrap
    high = round(1.0 - low, 12)
    interval = None
    if values:
        m = len(values) - 1
        interval = [values[int(low * m)], values[int(high * m)]]
    return {'interval': interval, 'level': level, 'reps': reps, 'reps_used': len(values), 'groups': len(keys),
            'seed': seed}


# ------------------------------------------------------------------ checkpoint selection
def selection_loss(pred_rows, gold_rows, tax, temps=None):
    """Mean over fields of the mean log-loss (non-null gold only; fields without any are skipped).

    ``pred_rows``: prediction rows with ``logits``; ``gold_rows``: export rows with ``id`` and ``targets``
    (or a dict id -> targets). Every gold row needs a prediction. Uncalibrated unless ``temps`` is given.
    """
    gold = gold_rows if isinstance(gold_rows, dict) else {r['id']: r['targets'] for r in gold_rows}
    preds = {}
    for p in pred_rows:
        if p['id'] in preds:
            raise ValueError('Duplicate prediction: ' + str(p['id']))
        preds[p['id']] = p
    missing = [i for i in gold if i not in preds]
    if missing:
        raise ValueError('%d gold rows have no prediction' % len(missing))
    temps = temps or {}
    losses = []
    for f in tax.names:
        options = tax.fields[f]
        total, n = 0.0, 0
        for i, targets in gold.items():
            g = targets.get(f)
            if g is None:
                continue
            p = softmax(preds[i]['logits'][f], temps.get(f, 1.0))[options.index(g)]
            total -= math.log(max(p, FLOOR))
            n += 1
        if n:
            losses.append(total / n)
    if not losses:
        raise ValueError('No field has a non-null gold label')
    return sum(losses) / len(losses)