Training

"""Inspectable objectives used to train retrieval and RAG control policies.

This module does not train a neural network.  It computes the scalar losses and
sampling diagnostics that a framework such as PyTorch would differentiate.
Keeping the mathematics in plain Python makes temperature, negative selection,
distillation, preference optimization, and policy-gradient assumptions visible
in the notebooks.
"""

import math
from dataclasses import dataclass
from typing import Iterable, List, Mapping, Optional, Sequence, Set, Tuple


@dataclass(frozen=True)
class NegativeExample:
    """A candidate negative plus provenance needed to detect false negatives."""

    identifier: str
    score: float
    source_id: str = ""
    answer_ids: Tuple[str, ...] = ()
    teacher_relevance: Optional[float] = None


@dataclass(frozen=True)
class NegativeMiningReport:
    """Selected hard negatives and candidates excluded as possible positives."""

    selected: Tuple[NegativeExample, ...]
    excluded_false_negatives: Tuple[NegativeExample, ...]
    excluded_easy: Tuple[NegativeExample, ...]


@dataclass(frozen=True)
class PolicyLoss:
    """REINFORCE loss with the returns and advantages used to obtain it."""

    loss: float
    returns: Tuple[float, ...]
    advantages: Tuple[float, ...]


def logsumexp(values: Sequence[float]) -> float:
    """Numerically stable ``log(sum(exp(values)))``."""

    if not values:
        raise ValueError("logsumexp requires at least one value")
    maximum = max(values)
    return maximum + math.log(sum(math.exp(value - maximum) for value in values))


def softmax(logits: Sequence[float], temperature: float = 1.0) -> Tuple[float, ...]:
    """Temperature-scaled categorical probabilities."""

    if temperature <= 0:
        raise ValueError("temperature must be positive")
    if not logits:
        return ()
    scaled = [value / temperature for value in logits]
    normalizer = logsumexp(scaled)
    return tuple(math.exp(value - normalizer) for value in scaled)


def contrastive_loss(
    positive_score: float,
    negative_scores: Sequence[float],
    temperature: float = 1.0,
) -> float:
    """InfoNCE loss for one query, one positive, and explicit negatives."""

    probabilities = softmax((positive_score, *negative_scores), temperature)
    return -math.log(max(probabilities[0], 1e-300))


def in_batch_contrastive_loss(
    similarities: Sequence[Sequence[float]],
    positive_indices: Optional[Sequence[int]] = None,
    temperature: float = 1.0,
    valid_mask: Optional[Sequence[Sequence[bool]]] = None,
) -> float:
    """Mean row-wise contrastive loss with optional false-negative masking.

    ``similarities[i][j]`` is the query-i/document-j score.  The usual paired
    batch has ``positive_indices[i] == i``.  ``valid_mask`` can suppress known
    alternate positives from each denominator while retaining the designated
    positive.
    """

    if not similarities:
        raise ValueError("similarity matrix must be non-empty")
    width = len(similarities[0])
    if width == 0 or any(len(row) != width for row in similarities):
        raise ValueError("similarity matrix must be rectangular and non-empty")
    positives = tuple(range(len(similarities))) if positive_indices is None else tuple(positive_indices)
    if len(positives) != len(similarities):
        raise ValueError("one positive index is required per query")
    if valid_mask is not None and (
        len(valid_mask) != len(similarities) or any(len(row) != width for row in valid_mask)
    ):
        raise ValueError("valid mask must match the similarity matrix")

    losses: List[float] = []
    for row_index, (row, positive) in enumerate(zip(similarities, positives)):
        if positive < 0 or positive >= width:
            raise ValueError("positive index is outside the matrix")
        retained = []
        positive_position = -1
        for column, score in enumerate(row):
            valid = valid_mask is None or valid_mask[row_index][column] or column == positive
            if valid:
                if column == positive:
                    positive_position = len(retained)
                retained.append(score)
        probabilities = softmax(retained, temperature)
        losses.append(-math.log(max(probabilities[positive_position], 1e-300)))
    return sum(losses) / len(losses)


def pairwise_hinge_loss(positive_score: float, negative_score: float, margin: float = 1.0) -> float:
    """Margin ranking loss ``max(0, margin - s+ + s-)``."""

    if margin < 0:
        raise ValueError("margin must be non-negative")
    return max(0.0, margin - positive_score + negative_score)


def pairwise_logistic_loss(positive_score: float, negative_score: float) -> float:
    """Smooth pairwise ranking loss ``softplus(s- - s+)``."""

    difference = negative_score - positive_score
    if difference > 0:
        return difference + math.log1p(math.exp(-difference))
    return math.log1p(math.exp(difference))


def listwise_cross_entropy(
    student_logits: Sequence[float],
    relevance: Sequence[float],
    temperature: float = 1.0,
) -> float:
    """ListNet-style cross entropy between relevance and model distributions."""

    if len(student_logits) != len(relevance) or not student_logits:
        raise ValueError("student and relevance lists must be equally sized and non-empty")
    target = softmax(relevance, temperature)
    predicted = softmax(student_logits, temperature)
    return -sum(expected * math.log(max(actual, 1e-300)) for expected, actual in zip(target, predicted))


def kl_distillation_loss(
    student_logits: Sequence[float],
    teacher_logits: Sequence[float],
    temperature: float = 1.0,
) -> float:
    """Teacher-to-student KL divergence with the conventional ``T²`` scale."""

    if len(student_logits) != len(teacher_logits) or not student_logits:
        raise ValueError("student and teacher lists must be equally sized and non-empty")
    student = softmax(student_logits, temperature)
    teacher = softmax(teacher_logits, temperature)
    divergence = sum(
        target * (math.log(max(target, 1e-300)) - math.log(max(actual, 1e-300)))
        for target, actual in zip(teacher, student)
    )
    return divergence * temperature * temperature


def binary_cross_entropy(logit: float, label: float) -> float:
    """Stable logistic loss for a relevance or retrieve/stop label."""

    if not 0.0 <= label <= 1.0:
        raise ValueError("label must be between zero and one")
    return max(logit, 0.0) - logit * label + math.log1p(math.exp(-abs(logit)))


def dpo_loss(
    policy_chosen_logp: float,
    policy_rejected_logp: float,
    reference_chosen_logp: float,
    reference_rejected_logp: float,
    beta: float = 0.1,
) -> float:
    """Direct Preference Optimization loss for one chosen/rejected pair."""

    if beta <= 0:
        raise ValueError("beta must be positive")
    policy_margin = policy_chosen_logp - policy_rejected_logp
    reference_margin = reference_chosen_logp - reference_rejected_logp
    logit = beta * (policy_margin - reference_margin)
    return binary_cross_entropy(logit, 1.0)


def discounted_returns(rewards: Sequence[float], discount: float = 1.0) -> Tuple[float, ...]:
    """Right-to-left discounted returns for a retrieval trajectory."""

    if not 0.0 <= discount <= 1.0:
        raise ValueError("discount must be between zero and one")
    running = 0.0
    output: List[float] = []
    for reward in reversed(rewards):
        running = reward + discount * running
        output.append(running)
    return tuple(reversed(output))


def reinforce_loss(
    action_log_probabilities: Sequence[float],
    rewards: Sequence[float],
    baseline: Optional[Sequence[float]] = None,
    discount: float = 1.0,
) -> PolicyLoss:
    """Monte-Carlo policy-gradient loss for search/query/stop actions."""

    if len(action_log_probabilities) != len(rewards) or not rewards:
        raise ValueError("one action log-probability is required per non-empty reward")
    returns = discounted_returns(rewards, discount)
    baselines = tuple(0.0 for _ in rewards) if baseline is None else tuple(baseline)
    if len(baselines) != len(rewards):
        raise ValueError("baseline must have one value per action")
    advantages = tuple(value - base for value, base in zip(returns, baselines))
    loss = -sum(logp * advantage for logp, advantage in zip(action_log_probabilities, advantages))
    return PolicyLoss(loss=loss / len(rewards), returns=returns, advantages=advantages)


def false_negative_mask(
    query_positive_sets: Sequence[Iterable[str]],
    candidate_ids: Sequence[str],
) -> Tuple[Tuple[bool, ...], ...]:
    """Build a denominator mask that removes known alternate positives.

    ``False`` means the candidate should not act as a negative for that query.
    A designated positive can still be retained by
    :func:`in_batch_contrastive_loss` through its explicit positive index.
    """

    masks = []
    for positives in query_positive_sets:
        positive_set = set(positives)
        masks.append(tuple(candidate not in positive_set for candidate in candidate_ids))
    return tuple(masks)


def mine_hard_negatives(
    candidates: Sequence[NegativeExample],
    positive_source_ids: Iterable[str] = (),
    positive_answer_ids: Iterable[str] = (),
    k: int = 5,
    minimum_score: float = -math.inf,
    teacher_positive_threshold: float = 0.5,
) -> NegativeMiningReport:
    """Choose high-scoring negatives while quarantining likely false negatives.

    Source identity, answer aliases, and an optional teacher label are three
    independent reasons a retrieved item may be an unlabeled positive.  The
    report preserves excluded items so mining quality can be audited rather
    than silently trusting the training labels.
    """

    if k < 0:
        raise ValueError("k must be non-negative")
    positive_sources: Set[str] = set(positive_source_ids)
    positive_answers: Set[str] = set(positive_answer_ids)
    false_negatives: List[NegativeExample] = []
    eligible: List[NegativeExample] = []
    easy: List[NegativeExample] = []
    for candidate in candidates:
        overlaps_answer = bool(set(candidate.answer_ids) & positive_answers)
        teacher_positive = (
            candidate.teacher_relevance is not None
            and candidate.teacher_relevance >= teacher_positive_threshold
        )
        same_source = bool(candidate.source_id and candidate.source_id in positive_sources)
        if overlaps_answer or teacher_positive or same_source:
            false_negatives.append(candidate)
        elif candidate.score < minimum_score:
            easy.append(candidate)
        else:
            eligible.append(candidate)
    eligible.sort(key=lambda item: (-item.score, item.identifier))
    selected = eligible[:k]
    easy.extend(eligible[k:])
    easy.sort(key=lambda item: (-item.score, item.identifier))
    false_negatives.sort(key=lambda item: (-item.score, item.identifier))
    return NegativeMiningReport(tuple(selected), tuple(false_negatives), tuple(easy))