"""Losses for stage ① alignment training (T3.1).

- :func:`lm_loss`: masked, shifted cross-entropy — computes LM loss only on the
  answer span (``labels != -100``), with the standard right-shift.
- :func:`infonce_contrastive_loss`: symmetric CLIP-style InfoNCE between two
  batches of representations. Used to align two *views* of the same time series
  (original vs lightly-augmented), forcing the TS encoder to be discriminative
  rather than collapse to a constant (which would let the LLM ignore the TS
  tokens).

Total stage ① loss (composed by ``train/stage1.py``):

    loss = lm_loss + lambda * contrastive_loss    (lambda default 0.1)

Design ref: §4.1 / §4.2 — "扰动一致性/InfoNCE". We implement InfoNCE in TS-
embedding space (encoder/projector output) rather than answer-embedding space,
because (a) it is cheap (no LLM forward needed for the augmented view) and
(b) the alignment goal is precisely that distinct series produce distinct TS
representations; the LM term already binds TS→answer. T3.3 (>50% answer change
under *strong* perturbation) is then evaluated separately, using perturbations
larger than the light augmentations used as InfoNCE positives.
"""
from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import Tensor


def lm_loss(logits: Tensor, labels: Tensor, ignore_index: int = -100) -> Tensor:
    """Masked, shifted cross-entropy.

    ``logits``  : [B, S, V]
    ``labels``  : [B, S]  (``ignore_index`` masks non-answer positions)
    Returns a scalar.
    """
    if labels.numel() == 0:
        return logits.new_zeros(())
    # standard causal shift: predict token t+1 from position t
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    if (shift_labels != ignore_index).sum() == 0:
        return logits.new_zeros(())
    loss = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=ignore_index,
        reduction="mean",
    )
    return loss


def infonce_contrastive_loss(
    z1: Tensor, z2: Tensor, temperature: float = 0.07
) -> Tensor:
    """Symmetric InfoNCE (CLIP-style) between two sets of representations.

    Positive pair = (z1[i], z2[i]); negatives = all j != i in the batch.
    Inputs are L2-normalized internally, so scale is ignored.

    ``z1``, ``z2`` : [B, D]
    Returns a scalar.
    """
    z1n = F.normalize(z1, dim=-1)
    z2n = F.normalize(z2, dim=-1)
    logits = z1n @ z2n.t() / temperature          # [B, B]
    targets = torch.arange(z1.size(0), device=z1.device)
    loss_12 = F.cross_entropy(logits, targets)
    loss_21 = F.cross_entropy(logits.t(), targets)
    return 0.5 * (loss_12 + loss_21)