src/eval/metrics.py

"""Numeric metrics for predicted vs gold cascade chains.

The judge route in ``src/eval/evaluator.py`` produces qualitative grades
(evidence_level / timing_alignment) by reading the news; this module
takes the orthogonal route — compare the predicted chain directly
against a reference (gold) chain and produce precision / recall / F1.

The gold chain is the test-event cascade chain extracted by step 03 of
the data pipeline (``data/processed/cascade_chains/{event_id}.json``).
That chain was extracted with full news as input; the predictor only
sees event metadata + RAG-retrieved historical chains, so the
information gap is large enough that the gold side serves as a
defensible reference.

Bumping ``ALGORITHM_VERSION`` invalidates every cached gold evaluation
via ``GoldEvaluator``'s fingerprint. Do this whenever the matching
algorithm or any default threshold changes in a way that would alter
prior numeric results.
"""

from __future__ import annotations

from typing import Protocol

import numpy as np

from src.models.schemas import CascadeChain, CascadeNode, NodeMatch, NodeMatchCandidate

ALGORITHM_VERSION = "v10"  # v0.6 issue A — BFS predictor re-tune; v9 caches were produced under BFS knobs calibrated for v0.4 retrieval distribution and must regenerate
DEFAULT_COSINE_THRESHOLD = 0.35  # was 0.5 (v2); recalibrated for new gold's named-entity density


class SupportsEmbedTexts(Protocol):
    def embed_texts(self, texts: list[str]) -> list[list[float]]: ...


def match_nodes(
    predicted: list[CascadeNode],
    gold: list[CascadeNode],
    embedder: SupportsEmbedTexts,
    threshold: float = DEFAULT_COSINE_THRESHOLD,
) -> list[NodeMatch]:
    """Greedy node alignment with a hard same-domain filter.

    Returns the list of accepted matches only. Thin wrapper around
    ``match_nodes_with_diagnostics`` that drops the candidate list.
    """
    matches, _ = match_nodes_with_diagnostics(
        predicted, gold, embedder, threshold=threshold
    )
    return matches


def match_nodes_with_diagnostics(
    predicted: list[CascadeNode],
    gold: list[CascadeNode],
    embedder: SupportsEmbedTexts,
    threshold: float = DEFAULT_COSINE_THRESHOLD,
) -> tuple[list[NodeMatch], list[NodeMatchCandidate]]:
    """Same matching as ``match_nodes`` but also surfaces every same-domain pair.

    For every (predicted, gold) pair where the domains match, compute the
    cosine similarity between description embeddings. Sort candidates by
    similarity, then greedily accept pairs whose similarity ≥ threshold
    and where neither side has been claimed yet.

    Returns a tuple ``(matches, candidates)`` where ``matches`` is the
    accepted alignment (same as ``match_nodes``), and ``candidates`` is
    every same-domain pair (including those below threshold or crowded
    out by greedy selection), sorted by descending cosine. Used by
    --dump-match-debug to support offline threshold sweeps.

    Both inputs typically have ≤40 nodes (extraction cap), so the
    O(P·G) pairing is fine; Hungarian assignment is in the backlog.
    """
    if not predicted or not gold:
        return [], []

    p_texts = [n.description for n in predicted]
    g_texts = [n.description for n in gold]
    # Single batched call — cheaper than per-pair embed.
    p_vecs = np.array(embedder.embed_texts(p_texts), dtype=float)
    g_vecs = np.array(embedder.embed_texts(g_texts), dtype=float)

    # sentence-transformers returns L2-normalized vectors when
    # normalize_embeddings=True (Embedder default), so dot product == cosine.
    # Renormalize defensively to keep this module independent of caller.
    p_vecs = _l2_normalize(p_vecs)
    g_vecs = _l2_normalize(g_vecs)
    sim = p_vecs @ g_vecs.T  # shape (P, G)

    # Collect ALL same-domain pairs first (no threshold filter) so the
    # diagnostics output can show the full distribution, including the
    # near-misses that --threshold sweeps are interested in.
    all_pairs: list[tuple[float, int, int]] = []
    for i, p_node in enumerate(predicted):
        for j, g_node in enumerate(gold):
            if p_node.domain != g_node.domain:
                continue  # hard same-domain filter — applies to diag output too
            score = float(sim[i, j])
            all_pairs.append((score, i, j))

    all_pairs.sort(key=lambda x: x[0], reverse=True)

    matched_p: set[int] = set()
    matched_g: set[int] = set()
    matches: list[NodeMatch] = []
    accepted_pairs: set[tuple[int, int]] = set()

    for score, i, j in all_pairs:
        if score < threshold:
            break  # candidates are sorted desc; nothing further can pass
        if i in matched_p or j in matched_g:
            continue
        matched_p.add(i)
        matched_g.add(j)
        accepted_pairs.add((i, j))
        matches.append(
            NodeMatch(
                p_id=predicted[i].id,
                g_id=gold[j].id,
                cosine=round(score, 4),
                severity_match=(predicted[i].severity == gold[j].severity),
            )
        )

    candidates = [
        NodeMatchCandidate(
            p_id=predicted[i].id,
            g_id=gold[j].id,
            cosine=round(score, 4),
            severity_match=(predicted[i].severity == gold[j].severity),
            accepted=((i, j) in accepted_pairs),
        )
        for score, i, j in all_pairs
    ]
    return matches, candidates


def compute_metrics(
    predicted: CascadeChain,
    gold: CascadeChain,
    matches: list[NodeMatch],
) -> dict:
    """Aggregate node-level matches into the per-event metric dict.

    Keys:
      - precision, recall, f1: node-level (matched / total)
      - domain_jaccard: |P_dom ∩ G_dom| / |P_dom ∪ G_dom|
      - severity_match_rate: fraction of matches with identical severity
        (None when no matches)
    """
    p_total = len(predicted.cascade_events)
    g_total = len(gold.cascade_events)
    m = len(matches)

    precision = m / p_total if p_total else 0.0
    recall = m / g_total if g_total else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0

    p_domains = {n.domain for n in predicted.cascade_events}
    g_domains = {n.domain for n in gold.cascade_events}
    union = p_domains | g_domains
    domain_jaccard = (
        len(p_domains & g_domains) / len(union) if union else 0.0
    )

    severity_match_rate = (
        sum(1 for x in matches if x.severity_match) / m if m else None
    )

    return {
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "f1": round(f1, 4),
        "domain_jaccard": round(domain_jaccard, 4),
        "severity_match_rate": (
            round(severity_match_rate, 4) if severity_match_rate is not None else None
        ),
    }


def _l2_normalize(arr: np.ndarray) -> np.ndarray:
    """Row-wise L2 normalization, safe against zero rows."""
    norms = np.linalg.norm(arr, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return arr / norms


# v0.5 issue #65 — domain-bag (category-level) matching.
# Drops the description-cosine threshold from match_nodes_with_diagnostics;
# pairs pred and gold within each domain by multiset min count.
# Defended in technical_report/v0.5/metric_refactor_report.md §2.

# Sentinel for the cosine field on domain-bag NodeMatch. Cosine is in
# [0, 1] for L2-normalized embeddings, so -1.0 is unambiguously "no
# embedding similarity was computed" and JSON-serializes cleanly
# (unlike NaN, which round-trips to null and fails pydantic float
# validation on reload).
_DOMAIN_BAG_COSINE_SENTINEL: float = -1.0


def match_nodes_domain_bag(
    predicted: list[CascadeNode],
    gold: list[CascadeNode],
) -> list[NodeMatch]:
    """Domain-bag matching — pair pred and gold nodes within each domain
    by multiset min, without requiring description cosine similarity.

    For each domain present in both predicted and gold, this emits
    ``min(pred_count_d, gold_count_d)`` NodeMatch objects. The 1-to-1
    constraint is preserved (each pred / gold node appears in at most
    one match), but the cosine threshold is dropped.

    Severity_match flag is computed by greedy multiset alignment per
    severity bucket: pairs that share severity within a domain are
    paired first; remaining nodes are then paired across severities
    (severity_match=False for these). The cosine field is set to
    ``_DOMAIN_BAG_COSINE_SENTINEL`` (-1.0) since no embedding
    similarity is computed here.
    """
    from collections import defaultdict

    pred_by_dom: dict[str, list[CascadeNode]] = defaultdict(list)
    gold_by_dom: dict[str, list[CascadeNode]] = defaultdict(list)
    for n in predicted:
        pred_by_dom[n.domain].append(n)
    for n in gold:
        gold_by_dom[n.domain].append(n)

    matches: list[NodeMatch] = []
    domains = set(pred_by_dom) & set(gold_by_dom)
    for d in sorted(domains):
        ps = pred_by_dom[d]
        gs = gold_by_dom[d]

        sev_pred: dict[str, list[CascadeNode]] = defaultdict(list)
        sev_gold: dict[str, list[CascadeNode]] = defaultdict(list)
        for n in ps:
            sev_pred[n.severity].append(n)
        for n in gs:
            sev_gold[n.severity].append(n)

        used_p_ids: set[str] = set()
        used_g_ids: set[str] = set()
        # Pass 1: pair within same severity (severity_match=True)
        for sev in sorted(set(sev_pred) | set(sev_gold)):
            k = min(len(sev_pred[sev]), len(sev_gold[sev]))
            for i in range(k):
                pn = sev_pred[sev][i]
                gn = sev_gold[sev][i]
                used_p_ids.add(pn.id)
                used_g_ids.add(gn.id)
                matches.append(
                    NodeMatch(
                        p_id=pn.id, g_id=gn.id,
                        cosine=_DOMAIN_BAG_COSINE_SENTINEL,
                        severity_match=True,
                    )
                )
        # Pass 2: cross-severity leftover (up to multiset min)
        leftover_p = [n for n in ps if n.id not in used_p_ids]
        leftover_g = [n for n in gs if n.id not in used_g_ids]
        for pn, gn in zip(leftover_p, leftover_g):
            matches.append(
                NodeMatch(
                    p_id=pn.id, g_id=gn.id,
                    cosine=_DOMAIN_BAG_COSINE_SENTINEL,
                    severity_match=(pn.severity == gn.severity),
                )
            )
    return matches


def compute_metrics_domain_bag(
    matches: list[NodeMatch],
    predicted: list[CascadeNode],
    gold: list[CascadeNode],
) -> tuple[float, float, float, float, float | None]:
    """Compute (precision, recall, f1, domain_jaccard, severity_match_rate)
    for domain-bag matches. severity_match_rate is None when no matches.
    """
    m = len(matches)
    p_count = len(predicted)
    g_count = len(gold)

    precision = m / p_count if p_count else 0.0
    recall = m / g_count if g_count else 0.0
    f1 = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) else 0.0
    )

    pred_doms = {n.domain for n in predicted}
    gold_doms = {n.domain for n in gold}
    union = pred_doms | gold_doms
    domain_jaccard = (
        len(pred_doms & gold_doms) / len(union) if union else 0.0
    )

    sev_rate = (
        sum(1 for x in matches if x.severity_match) / m if m else None
    )
    return (precision, recall, f1, domain_jaccard, sev_rate)


# v0.7 issue A — bootstrap CI telemetry.
# Spec: docs/superpowers/specs/2026-06-07-v07-bootstrap-ci-design.md
# CI is reported as additional AGGREGATE telemetry; it does NOT enter the
# KEEP/ROLLBACK acceptance panel (amendment §3.1 ±0.005 tolerance stays
# authoritative). See spec §1.2 for the framing decision.


def bootstrap_macro_ci(
    per_event_values: list[float | None],
    n_resamples: int = 1000,
    confidence: float = 0.95,
    seed: int = 42,
) -> dict[str, float | int]:
    """Event-level bootstrap CI for a single macro-averaged metric.

    Resamples ``per_event_values`` with replacement ``n_resamples`` times,
    takes the mean of each resample, and returns the requested-confidence
    percentile interval over those means. ``None`` entries are dropped
    before resampling; ``n`` in the returned dict reflects survivors.

    Returns a dict with keys: ``mean``, ``ci_low``, ``ci_high``, ``n``,
    ``n_resamples``, ``confidence``.

    Raises ``ValueError`` when no non-None values remain — silent NaN
    masks a real programmer error (e.g. an empty evaluable-event set).
    Degenerate inputs (n=1 or zero variance) collapse the CI onto the
    point estimate; this is well-defined and exercised by the test suite.
    """
    filtered = [float(v) for v in per_event_values if v is not None]
    if not filtered:
        raise ValueError(
            "per_event_values must contain at least one non-None entry"
        )
    arr = np.asarray(filtered, dtype=float)
    if arr.size == 1:
        v = float(arr[0])
        return {
            "mean": v,
            "ci_low": v,
            "ci_high": v,
            "n": 1,
            "n_resamples": int(n_resamples),
            "confidence": float(confidence),
        }
    rng = np.random.default_rng(seed)
    samples = rng.choice(arr, size=(n_resamples, arr.size), replace=True)
    means = samples.mean(axis=1)
    alpha = (1.0 - confidence) / 2.0
    lo, hi = np.percentile(means, [100.0 * alpha, 100.0 * (1.0 - alpha)])
    return {
        "mean": float(arr.mean()),
        "ci_low": float(lo),
        "ci_high": float(hi),
        "n": int(arr.size),
        "n_resamples": int(n_resamples),
        "confidence": float(confidence),
    }


def aggregate_with_ci(
    evaluations: "list[GoldEvaluation]",
    metrics_to_ci: tuple[str, ...] = (
        "category_recall",
        "category_severity_match_rate",
        "cosine_f1",
        "category_f1",
    ),
    bootstrap_seed: int = 42,
    outlier_event_ids: set[str] | None = None,
) -> dict[str, dict]:
    """Compute mean + bootstrap CI for each macro metric across an
    evaluation set, grouping by event first and averaging across seeds.

    Pipeline:
      1. Drop evaluations whose event_id is in ``outlier_event_ids``.
      2. Group surviving evaluations by event_id; for each metric build
         the list of seed-level values.
      3. Per event, average across seeds → one float per event.
      4. Feed the per-event means into ``bootstrap_macro_ci``.

    Returns a dict ``{metric_name: ci_dict}`` where each ``ci_dict`` is
    the output of ``bootstrap_macro_ci`` plus two extra keys:
    ``per_event_values`` (event_id → seed-mean) and
    ``per_event_seed_values`` (event_id → list of seed values).

    Cosine F1 is stored on the GoldEvaluation schema as ``f1`` (legacy
    name); the alias map below resolves that without making callers care.

    Caches missing a metric (e.g. category_severity_match_rate is None
    for pre-v0.5 caches) contribute no value for that metric. A metric
    with zero contributing events is omitted from the output and logged.
    """
    import logging
    from collections import defaultdict

    log = logging.getLogger(__name__)
    outlier_set = outlier_event_ids or set()

    # GoldEvaluation stores cosine F1 under the legacy `f1` attribute name.
    # All other metric names match their attribute names 1:1.
    metric_attr_map = {
        "cosine_f1": "f1",
        "category_f1": "category_f1",
        "category_recall": "category_recall",
        "category_severity_match_rate": "category_severity_match_rate",
        "category_precision": "category_precision",
        "precision": "precision",
        "recall": "recall",
        "domain_jaccard": "domain_jaccard",
    }

    grouped: dict[str, dict[str, list[float]]] = defaultdict(
        lambda: defaultdict(list)
    )
    for ev in evaluations:
        if ev.event_id in outlier_set:
            continue
        for m in metrics_to_ci:
            attr = metric_attr_map.get(m, m)
            v = getattr(ev, attr, None)
            if v is None:
                continue
            grouped[ev.event_id][m].append(float(v))

    out: dict[str, dict] = {}
    for m in metrics_to_ci:
        per_event_seed_values = {
            eid: per_metric[m]
            for eid, per_metric in grouped.items()
            if per_metric.get(m)
        }
        if not per_event_seed_values:
            log.warning(
                "aggregate_with_ci: no evaluable events for metric '%s'", m
            )
            continue
        per_event_means = {
            eid: sum(vs) / len(vs)
            for eid, vs in per_event_seed_values.items()
        }
        ci = bootstrap_macro_ci(
            list(per_event_means.values()),
            seed=bootstrap_seed,
        )
        ci["per_event_values"] = per_event_means
        ci["per_event_seed_values"] = per_event_seed_values
        out[m] = ci
    return out