"""Metrics matching the official eval notebook (iolai-2026-workshop).

Lives in solver/ (not eval/) because the RUNTIME needs it: the verifier and
the fallback ladder score candidate answers with chrF/EM at inference time.
The dev-only eval/ package re-exports from here.

Semantics:

- EM: case-insensitive exact string equality after .strip() ONLY — punctuation
  and internal whitespace are significant. Gold items may carry alternatives
  (list of accepted strings); a hit on any alternative counts.
- chrF: sacrebleu CHRF() defaults (char n-grams 1..6, beta=2, whitespace not
  in n-grams, epsilon-smoothed per order), 0..1 here (notebook prints 0..100).
  Max over alternatives.
- Aggregate: per-item average of each metric; geometric mean reported as the
  headline (the competition combines EM and chrF; the notebook prints both).

chrF is implemented in pure python replicating sacrebleu's algorithm so the
eval sandbox needs no dependency; tests/test_scorer.py checks parity against
sacrebleu when it is installed.
"""

from __future__ import annotations

import math
import re
import unicodedata
from collections import Counter
from typing import Dict, List, Optional, Sequence, Union

CHRF_NGRAM_ORDER = 6
CHRF_BETA = 2.0
_EPS = 1e-16

Gold = Union[str, Sequence[str]]  # a gold item: one string or alternatives


def _alts(gold: Gold) -> List[str]:
    if isinstance(gold, str):
        return [gold]
    return [str(a) for a in gold] or [""]


def normalize_answer(s: str) -> str:
    """Official EM normalization: strip + lowercase. Nothing else — final
    punctuation and internal spacing are significant."""
    return str(s).strip().lower()


def exact_match(pred: str, gold: Gold) -> float:
    p = normalize_answer(pred)
    return 1.0 if any(p == normalize_answer(a) for a in _alts(gold)) else 0.0


def _char_ngrams(s: str, n: int) -> Counter:
    return Counter(s[i : i + n] for i in range(len(s) - n + 1))


def _chrf_single(pred: str, gold: str, n_order: int = CHRF_NGRAM_ORDER,
                 beta: float = CHRF_BETA) -> float:
    """Exact replication of sacrebleu CHRF defaults (whitespace stripped,
    effective-order smoothing): precision/recall averaged over orders where
    BOTH sides have n-grams; hypothesis counts are zeroed for orders the
    reference lacks. Returns [0, 1] (sacrebleu reports x100)."""
    pred_s = "".join(str(pred).split())
    gold_s = "".join(str(gold).split())
    avg_prec = avg_rec = 0.0
    effective = 0
    for n in range(1, n_order + 1):
        gn = _char_ngrams(gold_s, n)
        pn = _char_ngrams(pred_s, n)
        n_ref = sum(gn.values())
        n_hyp = sum(pn.values()) if gn else 0  # sacrebleu: no ref => no hyp hits
        if n_hyp > 0 and n_ref > 0:
            overlap = sum((pn & gn).values())
            avg_prec += overlap / n_hyp
            avg_rec += overlap / n_ref
            effective += 1
    if effective == 0:
        return 0.0
    avg_prec /= effective
    avg_rec /= effective
    if avg_prec + avg_rec == 0:
        return 0.0
    b2 = beta * beta
    return (1 + b2) * avg_prec * avg_rec / (b2 * avg_prec + avg_rec)


def chrf(pred: str, gold: Gold) -> float:
    return max(_chrf_single(pred, a) for a in _alts(gold))


def item_scores(pred: str, gold: Gold) -> Dict[str, float]:
    return {"em": exact_match(pred, gold), "chrf": chrf(pred, gold)}


def score_submission(
    preds: Sequence[Sequence[str]],
    golds: Sequence[Sequence[Gold]],
    weights: Optional[Sequence[Sequence[float]]] = None,
) -> Dict[str, float]:
    """Score a full submission.

    preds/golds: per row, a list of items; each gold item is a string or a
    list of accepted alternatives. weights: optional per-item point values;
    uniform if None. Length mismatches within a row are penalized: missing
    preds score 0, extra preds are ignored (notebook lines preds up by
    position exactly the same way).
    """
    total_w = 0.0
    em_w = 0.0
    chrf_w = 0.0
    n_items = 0
    for ri, (prow, grow) in enumerate(zip(preds, golds)):
        wrow = list(weights[ri]) if weights is not None else [1.0] * len(grow)
        for ii, gold in enumerate(grow):
            w = wrow[ii] if ii < len(wrow) else 1.0
            pred = prow[ii] if ii < len(prow) else ""
            total_w += w
            em_w += w * exact_match(pred, gold)
            chrf_w += w * chrf(pred, gold)
            n_items += 1
    if total_w == 0:
        return {"em": 0.0, "chrf": 0.0, "score": 0.0, "n_items": 0}
    em_avg = em_w / total_w
    chrf_avg = chrf_w / total_w
    return {
        "em": em_avg,
        "chrf": chrf_avg,
        "score": math.sqrt(em_avg * chrf_avg),
        "n_items": n_items,
    }