"""opsd_loss.py — Self-distillation loss, lifted from siyan-zhao/OPSD.

Original source: github.com/siyan-zhao/OPSD::OPSDTrainer.generalized_jsd_loss (MIT).
Verified self-contained via DeepWiki audit on 2026-05-25.

Mathematical reference:
- OPSD paper: Zhao et al., "Self-Distilled Reasoner: On-Policy Self-Distillation
  for LLMs", arXiv:2601.18734.
- SDPO paper: Hübotter et al., "Reinforcement Learning via Self-Distillation",
  arXiv:2601.20802 (formalizes the same loss as Composer 2.5's "Targeted RL with
  Textual Feedback").

The loss computes JSD/KL divergence between a teacher distribution (model
conditioned on privileged information / a hint) and a student distribution
(model on the original context). Both come from the SAME model — the teacher
is just "the model with hint inserted into context."

Composer 2.5 uses this with the privileged information being a "hint" inserted
at the error-turn site. We use the same loss; the data collator constructs
ctx_teacher = ctx_student + hint_at_error_turn for us.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def generalized_jsd_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    labels: torch.Tensor | None = None,
    beta: float = 0.5,
    temperature: float = 1.0,
    reduction: str = "batchmean",
    logits_are_probs: bool = False,
    top_k: int | None = None,
    token_clip: float | None = None,
) -> torch.Tensor:
    """Generalized Jensen-Shannon Divergence loss between student and teacher.

    Args:
        student_logits: (B, T, V) — student model logits at each token position.
        teacher_logits: (B, T, V) — teacher (= same model with hint context) logits.
        labels: (B, T) — token-level mask. Positions with label == -100 are ignored
            (standard HF padding/ignored convention). For Composer-style hint-distill,
            mask should be 1 at error-turn tokens AFTER the hint, 0 elsewhere.
        beta: in [0, 1]. 0 = forward KL (student → teacher); 1 = reverse KL
            (teacher → student); 0.5 = symmetric JSD (default, recommended).
        temperature: softens distributions; T > 1 encourages distribution-matching
            on broader tail probabilities. SDPO paper uses 1.0.
        reduction: "batchmean" (sum / batch_size, like torch.nn.KLDivLoss) or "sum".
        logits_are_probs: if True, inputs are already probabilities (skip softmax).
        top_k: restrict KL to top-k tokens of the teacher distribution.
            Saves compute on large vocabularies (Qwen3 vocab = 152K).
        token_clip: clip per-token JSD to this max. Stabilizes training.
            SDPO paper does NOT clip; OPSD code defaults to None (no clip).

    Returns:
        Scalar loss tensor.
    """
    # Temperature scaling
    if not logits_are_probs:
        student_logits = student_logits / temperature
        teacher_logits = teacher_logits / temperature

    # Top-k restriction (optional, for vocab-size compute savings)
    if top_k is not None:
        # Restrict to top-k tokens of teacher; renormalize both there.
        teacher_topk_vals, teacher_topk_idx = teacher_logits.topk(top_k, dim=-1)
        student_topk_vals = student_logits.gather(-1, teacher_topk_idx)
        student_log_probs = F.log_softmax(student_topk_vals, dim=-1)
        teacher_log_probs = F.log_softmax(teacher_topk_vals, dim=-1)
    else:
        student_log_probs = F.log_softmax(student_logits, dim=-1)
        teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)

    # KL / JSD computation
    if beta == 0.0:
        # Forward KL: KL(student || teacher)
        per_token_div = F.kl_div(
            student_log_probs, teacher_log_probs,
            reduction="none", log_target=True,
        ).sum(dim=-1)
    elif beta == 1.0:
        # Reverse KL: KL(teacher || student)
        per_token_div = F.kl_div(
            teacher_log_probs, student_log_probs,
            reduction="none", log_target=True,
        ).sum(dim=-1)
    else:
        # JSD (symmetric, beta = 0.5 default):
        #   M = 0.5 * (P + Q); JSD = 0.5 * (KL(P||M) + KL(Q||M))
        # Implementation via log-space mixture:
        #   log_m = logaddexp(log p, log q) - log 2
        log_mixture = torch.logaddexp(student_log_probs, teacher_log_probs) - torch.log(
            torch.tensor(2.0, device=student_logits.device)
        )
        kl_student_mixture = F.kl_div(
            log_mixture, student_log_probs, reduction="none", log_target=True
        ).sum(dim=-1)
        kl_teacher_mixture = F.kl_div(
            log_mixture, teacher_log_probs, reduction="none", log_target=True
        ).sum(dim=-1)
        per_token_div = beta * kl_student_mixture + (1.0 - beta) * kl_teacher_mixture

    # Optional per-token clip (stability)
    if token_clip is not None:
        per_token_div = per_token_div.clamp(max=token_clip)

    # Mask out ignored positions (labels == -100, the HF convention)
    if labels is not None:
        loss_mask = (labels != -100).float()
        per_token_div = per_token_div * loss_mask
        n_valid = loss_mask.sum().clamp(min=1.0)
    else:
        n_valid = torch.tensor(per_token_div.numel(), device=per_token_div.device, dtype=per_token_div.dtype)

    if reduction == "batchmean":
        # batchmean = sum over (B*T_valid) / B
        return per_token_div.sum() / per_token_div.shape[0]
    elif reduction == "sum":
        return per_token_div.sum()
    elif reduction == "mean":
        return per_token_div.sum() / n_valid
    elif reduction == "none":
        return per_token_div
    else:
        raise ValueError(f"Unknown reduction: {reduction}")


__all__ = ["generalized_jsd_loss"]