src/serving/eval.py

"""
Evaluation metrics for the RAG pipeline.

Three levels of evaluation:

1. **Retrieval quality** — does the retriever surface the right chunks?
   Precision@k, Recall@k, MRR, NDCG@k.

2. **Answer quality** — does RAG actually improve the model's answers?
   Exact Match (EM) and Token F1, the standard SQuAD metrics.

3. **RAGAS-style diagnostics** — where in the pipeline do failures occur?
   Four dimensions from the RAGAS framework (Retrieval-Augmented
   Generation Assessment), implemented with token-overlap F1 instead
   of LLM-based scoring for zero-cost, deterministic evaluation:

   - *Faithfulness*: Is the answer grounded in the retrieved context?
     Mean best-sentence F1 against context chunks.
   - *Answer Relevance*: Does the answer address the question asked?
     Token F1 between answer and question.
   - *Context Precision*: Are the retrieved chunks relevant to the query?
     Fraction of chunks that overlap meaningfully with the ground truth.
   - *Context Recall*: Does the context cover the ground truth answer?
     Token F1 between the combined context and the ground truth.

   Together these four metrics localise failures: low context precision
   → retriever returning noise; low context recall → missing relevant
   docs; low faithfulness → model hallucinating beyond context; low
   answer relevance → model not addressing the question.

EM checks if the predicted answer exactly matches the ground truth
(after normalization). Token F1 treats both answers as bags of words
and computes precision/recall/F1 — a softer metric that gives partial
credit for overlapping tokens.
"""

from __future__ import annotations

import json
import math
import re
import string
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np


@dataclass
class QueryLabel:
    """A query with its ground-truth relevant chunk identifiers."""
    query: str
    relevant: list[str]  # substrings that identify relevant chunks
    source: str = ""     # which document the answer comes from


@dataclass
class RetrievalResult:
    """Per-query evaluation result."""
    query: str
    precision_at_k: float
    recall_at_k: float
    reciprocal_rank: float
    ndcg_at_k: float
    retrieved_texts: list[str] = field(default_factory=list)
    relevance_flags: list[bool] = field(default_factory=list)


def _dcg(relevances: list[bool], k: int) -> float:
    """Discounted Cumulative Gain at k."""
    score = 0.0
    for i in range(min(k, len(relevances))):
        if relevances[i]:
            score += 1.0 / math.log2(i + 2)  # i+2 because rank is 1-indexed
    return score


def _ndcg(relevances: list[bool], k: int, num_relevant: int) -> float:
    """Normalized DCG: actual DCG / ideal DCG."""
    dcg = _dcg(relevances, k)
    # Ideal: all relevant docs at the top
    ideal_relevances = [True] * min(num_relevant, k) + [False] * max(0, k - num_relevant)
    idcg = _dcg(ideal_relevances, k)
    return dcg / idcg if idcg > 0 else 0.0


def evaluate_query(
    retrieved: list[tuple[str, float]],
    label: QueryLabel,
    k: int,
) -> RetrievalResult:
    """Evaluate a single query's retrieval results against ground truth."""
    texts = [text for text, _ in retrieved[:k]]

    # Mark each retrieved chunk as relevant if it contains any relevant substring
    relevance = []
    for text in texts:
        text_lower = text.lower()
        is_relevant = any(r.lower() in text_lower for r in label.relevant)
        relevance.append(is_relevant)

    # Precision@k
    num_relevant_retrieved = sum(relevance)
    precision = num_relevant_retrieved / k if k > 0 else 0.0

    # Recall@k
    total_relevant = len(label.relevant)
    # Count how many of the relevant substrings were found in any retrieved chunk
    found = set()
    for text in texts:
        text_lower = text.lower()
        for r in label.relevant:
            if r.lower() in text_lower:
                found.add(r.lower())
    recall = len(found) / total_relevant if total_relevant > 0 else 0.0

    # MRR: reciprocal rank of first relevant result
    rr = 0.0
    for i, rel in enumerate(relevance):
        if rel:
            rr = 1.0 / (i + 1)
            break

    # NDCG@k
    ndcg = _ndcg(relevance, k, total_relevant)

    return RetrievalResult(
        query=label.query,
        precision_at_k=precision,
        recall_at_k=recall,
        reciprocal_rank=rr,
        ndcg_at_k=ndcg,
        retrieved_texts=texts,
        relevance_flags=relevance,
    )


@dataclass
class EvalSummary:
    """Aggregated metrics across all queries."""
    num_queries: int
    mean_precision_at_k: float
    mean_recall_at_k: float
    mrr: float
    mean_ndcg_at_k: float
    hit_rate: float                # fraction of queries with >= 1 relevant chunk
    per_query: list[RetrievalResult]


def evaluate_retriever(
    pipeline,
    labels: list[QueryLabel],
    k: int = 3,
) -> EvalSummary:
    """Run all labeled queries through the pipeline and compute aggregate metrics."""
    results = []
    for label in labels:
        retrieved = pipeline.retrieve(label.query, top_k=k)
        result = evaluate_query(retrieved, label, k)
        results.append(result)

    n = len(results)
    hits = sum(1 for r in results if any(r.relevance_flags)) if n else 0
    return EvalSummary(
        num_queries=n,
        mean_precision_at_k=np.mean([r.precision_at_k for r in results]) if n else 0.0,
        mean_recall_at_k=np.mean([r.recall_at_k for r in results]) if n else 0.0,
        mrr=np.mean([r.reciprocal_rank for r in results]) if n else 0.0,
        mean_ndcg_at_k=np.mean([r.ndcg_at_k for r in results]) if n else 0.0,
        hit_rate=hits / n if n else 0.0,
        per_query=results,
    )


def load_labels(path: str | Path) -> list[QueryLabel]:
    """Load query labels from a JSON file."""
    with open(path) as f:
        data = json.load(f)
    return [QueryLabel(**item) for item in data]


# ---------------------------------------------------------------------------
# Answer Quality Metrics (SQuAD-style)
# ---------------------------------------------------------------------------

def normalize_answer(text: str) -> str:
    """Normalize answer text for comparison: lowercase, strip articles/punctuation/whitespace."""
    text = text.lower()
    # Remove articles
    text = re.sub(r"\b(a|an|the)\b", " ", text)
    # Remove punctuation
    text = text.translate(str.maketrans("", "", string.punctuation))
    # Collapse whitespace
    text = " ".join(text.split())
    return text


def exact_match(prediction: str, ground_truth: str) -> float:
    """1.0 if normalized prediction == normalized ground truth, else 0.0."""
    return 1.0 if normalize_answer(prediction) == normalize_answer(ground_truth) else 0.0


def token_f1(prediction: str, ground_truth: str) -> float:
    """Token-level F1 between predicted and ground-truth answers."""
    pred_tokens = normalize_answer(prediction).split()
    gold_tokens = normalize_answer(ground_truth).split()

    if not gold_tokens:
        return 1.0 if not pred_tokens else 0.0
    if not pred_tokens:
        return 0.0

    common = sum((Counter(pred_tokens) & Counter(gold_tokens)).values())
    if common == 0:
        return 0.0

    precision = common / len(pred_tokens)
    recall = common / len(gold_tokens)
    return 2 * precision * recall / (precision + recall)


# ---------------------------------------------------------------------------
# Structured Numeric Extraction & Matching
# ---------------------------------------------------------------------------

# Regex for numbers with optional sign, commas, decimals, and magnitude suffixes.
_NUMBER_RE = re.compile(
    r"[-+]?\$?\s*\d[\d,]*(?:\.\d+)?"          # base number (opt. sign, $, commas, decimal)
    r"(?:\s*(?:billion|million|thousand|B|M|K))?"  # optional magnitude suffix
    r"|[-+]?\d+(?:\.\d+)?%",                   # or a percentage like 12.3%
    re.IGNORECASE,
)

_MAGNITUDE_MAP = {
    "billion": 1e9, "b": 1e9,
    "million": 1e6, "m": 1e6,
    "thousand": 1e3, "k": 1e3,
}


def extract_numbers(text: str) -> list[float]:
    """Extract all numeric values from *text*, normalising magnitudes.

    Handles currency symbols, commas, magnitude suffixes (billion/million/
    thousand/B/M/K), and percentages.  Returns a deduplicated list of
    floats in the order they appear.

    >>> extract_numbers("Revenue was $1.04 billion and costs $283 million")
    [1040000000.0, 283000000.0]
    >>> extract_numbers("grew 12.3% year-over-year")
    [12.3]
    """
    results: list[float] = []
    for match in _NUMBER_RE.finditer(text):
        raw = match.group(0)
        # Strip currency symbols and whitespace
        raw = re.sub(r"[$€£¥\s]", "", raw)
        # Handle percentages — strip % and return the number itself
        if raw.endswith("%"):
            try:
                results.append(float(raw[:-1]))
            except ValueError:
                pass
            continue
        # Detect magnitude suffix
        multiplier = 1.0
        lower = raw.lower()
        for suffix, mult in _MAGNITUDE_MAP.items():
            if lower.endswith(suffix):
                raw = raw[: -len(suffix)]
                multiplier = mult
                break
        # Strip commas
        raw = raw.replace(",", "")
        try:
            results.append(float(raw) * multiplier)
        except ValueError:
            pass
    # Deduplicate while preserving order
    seen: set[float] = set()
    deduped: list[float] = []
    for v in results:
        if v not in seen:
            seen.add(v)
            deduped.append(v)
    return deduped


def numeric_match(prediction: str, ground_truth: str, tolerance: float = 0.01) -> float:
    """Score whether the prediction contains the key number(s) from ground_truth.

    Extracts numbers from both strings.  For each ground-truth number,
    checks if any predicted number is within *tolerance* (relative).
    Returns the fraction of ground-truth numbers matched.

    A tolerance of 0.01 means 1% relative difference is accepted,
    handling rounding ($1.038B → $1.04B).

    Returns 1.0 if ground_truth contains no numbers (vacuously true).

    >>> numeric_match("Revenue was $1.04 billion", "$1,038,765 thousand")
    1.0
    >>> numeric_match("Revenue was $500 million", "$1,038,765 thousand")
    0.0
    """
    gt_nums = extract_numbers(ground_truth)
    if not gt_nums:
        return 1.0
    pred_nums = extract_numbers(prediction)
    if not pred_nums:
        return 0.0

    matched = 0
    for gt in gt_nums:
        for pred in pred_nums:
            if gt == 0:
                if pred == 0:
                    matched += 1
                    break
            elif abs(pred - gt) / abs(gt) <= tolerance:
                matched += 1
                break
    return matched / len(gt_nums)


@dataclass
class AnswerResult:
    """Per-question answer quality result."""
    question: str
    ground_truth: str
    prediction: str
    em: float
    f1: float
    used_rag: bool
    nm: float = 0.0  # numeric match score


@dataclass
class AnswerQualitySummary:
    """Aggregated answer quality metrics."""
    num_questions: int
    mean_em: float
    mean_f1: float
    mean_nm: float
    per_question: list[AnswerResult]


def evaluate_answer_quality(results: list[AnswerResult]) -> AnswerQualitySummary:
    """Compute aggregate EM, F1, and numeric match from per-question results."""
    n = len(results)
    return AnswerQualitySummary(
        num_questions=n,
        mean_em=np.mean([r.em for r in results]) if n else 0.0,
        mean_f1=np.mean([r.f1 for r in results]) if n else 0.0,
        mean_nm=np.mean([r.nm for r in results]) if n else 0.0,
        per_question=results,
    )


# ---------------------------------------------------------------------------
# RAGAS-style Diagnostic Metrics
# ---------------------------------------------------------------------------
#
# These implement the four RAGAS dimensions using token-overlap F1 instead
# of LLM-based scoring.  This gives deterministic, zero-cost evaluation
# that can run on every query without a model call.


def _split_sentences(text: str) -> list[str]:
    """Split on sentence-ending punctuation followed by space or EOL."""
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [s.strip() for s in parts if s.strip()]


def faithfulness(answer: str, context_chunks: list[str]) -> float:
    """Score how faithfully the answer reflects the retrieved context.

    Splits the answer into sentences and scores each against all chunks
    using token F1.  The faithfulness score is the mean of the best F1
    each sentence achieves.  Sentences under 4 words are skipped (headers,
    transitions).

    High faithfulness = answer sticks to what the context says.
    Low faithfulness = answer contains claims not in the context
    (potential hallucination).

    Args:
        answer: The model's generated answer.
        context_chunks: Retrieved chunk texts.

    Returns:
        Mean best-sentence grounding score in [0.0, 1.0].
    """
    if not answer or not context_chunks:
        return 0.0

    sentences = _split_sentences(answer)
    scores: list[float] = []

    for sent in sentences:
        if len(sent.split()) < 4:
            continue
        best_f1 = max(
            (token_f1(sent, chunk) for chunk in context_chunks),
            default=0.0,
        )
        scores.append(best_f1)

    return float(np.mean(scores)) if scores else 0.0


def answer_relevance(answer: str, question: str) -> float:
    """Score how relevant the answer is to the question.

    Uses token F1 between the answer and question.  A relevant answer
    reuses question terms and addresses the topic; an irrelevant answer
    discusses something unrelated.

    This is a lightweight proxy for the LLM-based RAGAS answer relevance
    (which generates synthetic questions from the answer and measures
    similarity).  Token overlap captures whether the answer stays on-topic.

    Args:
        answer: The model's generated answer.
        question: The user's original question.

    Returns:
        Token F1 in [0.0, 1.0].
    """
    if not answer or not question:
        return 0.0
    return token_f1(answer, question)


def context_precision(
    context_chunks: list[str],
    ground_truth: str,
    threshold: float = 0.10,
) -> float:
    """Fraction of retrieved chunks that are relevant to the ground truth.

    A chunk is considered relevant if its token F1 against the ground
    truth exceeds the threshold.  High context precision means the
    retriever isn't returning noise; low means many retrieved chunks
    are irrelevant.

    Args:
        context_chunks: Retrieved chunk texts.
        ground_truth: The expected answer.
        threshold: Minimum token F1 to consider a chunk relevant.

    Returns:
        Precision score in [0.0, 1.0].
    """
    if not context_chunks:
        return 0.0

    relevant = sum(
        1 for chunk in context_chunks
        if token_f1(chunk, ground_truth) >= threshold
    )
    return relevant / len(context_chunks)


def context_recall(context_chunks: list[str], ground_truth: str) -> float:
    """Score how well the combined context covers the ground truth.

    Concatenates all retrieved chunks and computes token F1 against the
    ground truth.  High recall means the context contains the information
    needed to answer; low recall means relevant information is missing.

    Args:
        context_chunks: Retrieved chunk texts.
        ground_truth: The expected answer.

    Returns:
        Token F1 in [0.0, 1.0].
    """
    if not context_chunks or not ground_truth:
        return 0.0

    combined = " ".join(context_chunks)
    return token_f1(combined, ground_truth)


@dataclass
class RAGASResult:
    """Per-query RAGAS diagnostic scores."""
    question: str
    faithfulness: float
    answer_relevance: float
    context_precision: float
    context_recall: float


@dataclass
class RAGASSummary:
    """Aggregated RAGAS scores across all queries."""
    num_queries: int
    mean_faithfulness: float
    mean_answer_relevance: float
    mean_context_precision: float
    mean_context_recall: float
    per_query: list[RAGASResult]


def evaluate_ragas(results: list[RAGASResult]) -> RAGASSummary:
    """Compute aggregate RAGAS scores from per-query results."""
    n = len(results)
    return RAGASSummary(
        num_queries=n,
        mean_faithfulness=float(np.mean([r.faithfulness for r in results])) if n else 0.0,
        mean_answer_relevance=float(np.mean([r.answer_relevance for r in results])) if n else 0.0,
        mean_context_precision=float(np.mean([r.context_precision for r in results])) if n else 0.0,
        mean_context_recall=float(np.mean([r.context_recall for r in results])) if n else 0.0,
        per_query=results,
    )


# ---------------------------------------------------------------------------
# Error Taxonomy
# ---------------------------------------------------------------------------
#
# Classifies *why* a question failed, not just *that* it failed.
# Each failed question gets exactly one category — the most upstream
# failure, since fixing that would likely fix downstream symptoms.

class ErrorType:
    """Constants for error taxonomy categories."""
    EMPTY_REFUSAL = "empty_refusal"
    RETRIEVAL_MISS = "retrieval_miss"
    FORMAT_MISMATCH = "format_mismatch"
    WRONG_EXTRACTION = "wrong_extraction"
    HALLUCINATION = "hallucination"
    CORRECT = "correct"


# Patterns that indicate the model refused to answer or produced nothing.
_REFUSAL_PATTERNS = re.compile(
    r"not\s+found|not\s+available|i\s+don.?t\s+(?:know|have)|"
    r"cannot\s+(?:find|determine)|no\s+(?:information|data|context)",
    re.IGNORECASE,
)


def _normalize_numbers(text: str) -> str:
    """Normalise numeric formats for fairer comparison.

    Strips currency symbols, commas, and common suffixes so that
    "$1,038,765 thousand" and "$1.04 billion" both become cleaner
    numeric tokens.
    """
    text = text.lower()
    # Remove currency symbols
    text = re.sub(r"[$€£¥]", "", text)
    # Remove commas in numbers: 1,038,765 -> 1038765
    text = re.sub(r"(\d),(\d)", r"\1\2", text)
    # Expand common magnitude suffixes
    text = re.sub(r"(\d+(?:\.\d+)?)\s*billion", lambda m: str(int(float(m.group(1)) * 1_000_000)), text)
    text = re.sub(r"(\d+(?:\.\d+)?)\s*million", lambda m: str(int(float(m.group(1)) * 1_000)), text)
    text = re.sub(r"(\d+(?:\.\d+)?)\s*thousand", lambda m: str(int(float(m.group(1)))), text)
    # Remove "approximately", "about", etc.
    text = re.sub(r"\b(?:approximately|about|roughly|around|~)\b", "", text)
    return " ".join(text.split())


def classify_error(
    prediction: str,
    ground_truth: str,
    context_chunks: list[str],
    f1_score: float,
    faithfulness_score: float = 0.0,
    success_threshold: float = 0.20,
    context_relevance_threshold: float = 0.10,
) -> str:
    """Classify why a RAG answer failed.

    Returns one of the ``ErrorType`` constants.  The classification
    follows a decision tree that identifies the most upstream failure:

    1. If F1 >= threshold → CORRECT (not a failure)
    2. If prediction is empty or a refusal → EMPTY_REFUSAL
    3. If no retrieved chunk is relevant to the ground truth → RETRIEVAL_MISS
    4. If number-normalised F1 is much higher → FORMAT_MISMATCH
    5. If faithfulness is low (answer not grounded) → HALLUCINATION
    6. Otherwise → WRONG_EXTRACTION (context had it, model got wrong fact)

    Args:
        prediction: The model's answer.
        ground_truth: The expected answer.
        context_chunks: Retrieved chunk texts (empty list if no RAG).
        f1_score: Pre-computed token F1 for this question.
        faithfulness_score: Pre-computed faithfulness (from RAGAS).
        success_threshold: F1 above this is considered correct.
        context_relevance_threshold: Min chunk-vs-ground-truth F1 to
            consider a chunk relevant.

    Returns:
        One of the ``ErrorType`` string constants.
    """
    # 1. Already correct
    if f1_score >= success_threshold:
        return ErrorType.CORRECT

    # 2. Empty or refusal
    stripped = prediction.strip()
    if not stripped or len(stripped) < 5 or _REFUSAL_PATTERNS.search(stripped):
        return ErrorType.EMPTY_REFUSAL

    # 3. Retrieval miss — no chunk relevant to the ground truth
    if context_chunks:
        best_chunk_f1 = max(
            (token_f1(chunk, ground_truth) for chunk in context_chunks),
            default=0.0,
        )
        if best_chunk_f1 < context_relevance_threshold:
            return ErrorType.RETRIEVAL_MISS
    else:
        return ErrorType.RETRIEVAL_MISS

    # 4. Format mismatch — normalise numbers and re-check
    norm_f1 = token_f1(
        _normalize_numbers(prediction),
        _normalize_numbers(ground_truth),
    )
    if norm_f1 >= success_threshold and norm_f1 - f1_score > 0.05:
        return ErrorType.FORMAT_MISMATCH

    # 5. Hallucination — answer not grounded in retrieved context
    if faithfulness_score < 0.10:
        return ErrorType.HALLUCINATION

    # 6. Wrong extraction — context had it, model pulled wrong fact
    return ErrorType.WRONG_EXTRACTION


def compute_error_distribution(
    errors: list[str],
) -> dict[str, int]:
    """Count occurrences of each error type.

    Args:
        errors: List of ``ErrorType`` constants, one per query.

    Returns:
        Dict mapping error type to count.
    """
    dist: dict[str, int] = {}
    for e in errors:
        dist[e] = dist.get(e, 0) + 1
    return dist