"""Same-position denoising objectives. No autoregressive shift is used."""
from __future__ import annotations
import torch
import torch.nn.functional as F


def _separate_answer_padding_loss(
    selected_ce, example_ids, counts, sampled_t, token_weights,
    selected_answer_mask, selected_padding_mask, answer_lengths, padding_lengths,
    answer_padding_weights, compute_unweighted_metric,
):
    """Normalize each region per example before applying fixed mixing weights.

    Region lengths count all eligible positions, not just this draw's masks.
    Empty regions/draws contribute zero; coefficients are never redistributed.
    Include eligible examples with zero masks in the batch average so masking
    fewer positions does not condition the objective on a nonempty draw.
    """
    if sampled_t is None:
        raise ValueError("Separate answer/padding loss requires sampled_t")
    if any(value is None for value in (
        selected_answer_mask, selected_padding_mask, answer_lengths, padding_lengths,
    )):
        raise ValueError("Separate answer/padding loss requires region masks and lengths")
    weighted_ce = selected_ce if token_weights is None else selected_ce * token_weights
    inverse_t_ce = weighted_ce / sampled_t[example_ids].to(weighted_ce.dtype).clamp_min(1e-8)
    component_losses = []
    valid = (answer_lengths + padding_lengths) > 0
    valid_count = valid.sum()
    for selected_mask, lengths in (
        (selected_answer_mask, answer_lengths), (selected_padding_mask, padding_lengths),
    ):
        sums = torch.zeros(counts.shape[0], device=weighted_ce.device, dtype=weighted_ce.dtype).scatter_add(
            0, example_ids, inverse_t_ce * selected_mask,
        )
        per_example = sums / lengths.clamp_min(1)
        component_losses.append((per_example * valid).sum() / valid_count.clamp_min(1))
    answer_loss, padding_loss = component_losses
    answer_weight, padding_weight = answer_padding_weights
    loss = answer_weight * answer_loss + padding_weight * padding_loss
    metrics = {
        "weighted_loss": loss.detach(),
        "answer_loss": answer_loss.detach(),
        "padding_loss": padding_loss.detach(),
        "valid_examples": valid_count.detach(),
        "supervised_tokens": counts.sum().detach(),
    }
    if compute_unweighted_metric:
        metrics["unweighted_masked_token_ce"] = (selected_ce.sum() / counts.sum().clamp_min(1)).detach()
    return loss, metrics


def _finish_selected_loss(
    selected_ce_sums,
    selected_ce_total,
    counts,
    sampled_t,
    normalization_mask,
    compute_unweighted_metric,
):
    """Reduce already-selected token losses with the configured weighting."""
    valid = counts > 0
    per_example = selected_ce_sums / counts.clamp_min(1)
    if sampled_t is None:
        weighted = per_example
    else:
        if normalization_mask is None:
            raise ValueError("normalization_mask is required when sampled_t is provided")
        response_lengths = normalization_mask.sum(dim=1).clamp_min(1)
        weighted = selected_ce_sums / sampled_t.to(selected_ce_sums.device).clamp_min(1e-8) / response_lengths

    valid_count = valid.sum()
    loss = (weighted * valid).sum() / valid_count.clamp_min(1)
    metrics = {
        "weighted_loss": loss.detach(),
        "valid_examples": valid_count.detach(),
        "supervised_tokens": counts.sum().detach(),
    }
    if compute_unweighted_metric:
        metrics["unweighted_masked_token_ce"] = (
            selected_ce_total / counts.sum().clamp_min(1)
        ).detach()
    return loss, metrics


def selected_denoising_loss(
    selected_logits,
    selected_labels,
    example_ids,
    counts,
    sampled_t=None,
    normalization_mask=None,
    *,
    compute_unweighted_metric=True,
    token_weights=None,
    answer_padding_weights=None,
    selected_answer_mask=None,
    selected_padding_mask=None,
    answer_lengths=None,
    padding_lengths=None,
):
    """Compute the objective when the LM head emitted supervised positions only."""
    selected_ce = F.cross_entropy(selected_logits, selected_labels, reduction="none")
    if answer_padding_weights is not None:
        return _separate_answer_padding_loss(
            selected_ce, example_ids, counts, sampled_t, token_weights,
            selected_answer_mask, selected_padding_mask, answer_lengths, padding_lengths,
            answer_padding_weights, compute_unweighted_metric,
        )
    weighted_ce = selected_ce if token_weights is None else selected_ce * token_weights
    selected_ce_sums = torch.zeros(
        counts.shape[0], device=selected_logits.device, dtype=weighted_ce.dtype
    ).scatter_add(0, example_ids, weighted_ce)
    return _finish_selected_loss(
        selected_ce_sums,
        selected_ce.sum(),
        counts,
        sampled_t,
        normalization_mask,
        compute_unweighted_metric,
    )


def masked_denoising_loss(
    logits,
    labels,
    loss_mask,
    sampled_t=None,
    normalization_mask=None,
    *,
    compute_unweighted_metric=True,
    sparse_positions=True,
    token_weights=None,
    answer_padding_weights=None,
    answer_mask=None,
    padding_mask=None,
):
    """Compute masked CE with optional LLaDA-style inverse-t weighting.

    Returns a differentiable scalar and aggregate metrics. Examples without selected
    tokens are excluded rather than changing another example's denominator. When
    sampled_t is provided, normalization_mask must represent the complete eligible
    response, not only the positions selected for corruption. Optional token_weights
    correct frontier positions by t / p(position) before inverse-t reduction;
    unweighted_masked_token_ce always reports the original CE.
    """
    counts = loss_mask.sum(dim=1)
    if answer_padding_weights is not None:
        if answer_mask is None or padding_mask is None:
            raise ValueError("Separate answer/padding loss requires answer_mask and padding_mask")
        answer_mask = answer_mask & ~padding_mask
        example_ids, token_ids = loss_mask.nonzero(as_tuple=True)
        return selected_denoising_loss(
            logits[example_ids, token_ids], labels[example_ids, token_ids], example_ids,
            counts, sampled_t, normalization_mask,
            compute_unweighted_metric=compute_unweighted_metric,
            token_weights=None if token_weights is None else token_weights[example_ids, token_ids],
            answer_padding_weights=answer_padding_weights,
            selected_answer_mask=answer_mask[example_ids, token_ids],
            selected_padding_mask=padding_mask[example_ids, token_ids],
            answer_lengths=answer_mask.sum(dim=1), padding_lengths=padding_mask.sum(dim=1),
        )
    if sparse_positions:
        # Computing CE over [batch, sequence, vocabulary] wastes a large softmax
        # on positions excluded from the objective. Reuse one set of selected
        # indices for the logits, labels, and per-example reduction.
        example_ids, token_ids = loss_mask.nonzero(as_tuple=True)
        selected_ce = F.cross_entropy(
            logits[example_ids, token_ids], labels[example_ids, token_ids], reduction="none"
        )
        weighted_ce = selected_ce if token_weights is None else selected_ce * token_weights[example_ids, token_ids]
        return _finish_selected_loss(
            torch.zeros(
                logits.shape[0], device=logits.device, dtype=weighted_ce.dtype
            ).scatter_add(0, example_ids, weighted_ce),
            selected_ce.sum(),
            counts,
            sampled_t,
            normalization_mask,
            compute_unweighted_metric,
        )
    else:
        # Structured all-token training has nothing to compact; avoid copying
        # the entire logits tensor through advanced indexing in that mode.
        token_ce = F.cross_entropy(logits.transpose(1, 2), labels, reduction="none")
        selected_ce_sums = (token_ce * loss_mask).sum(dim=1)
        weighted_ce_sums = selected_ce_sums if token_weights is None else (token_ce * loss_mask * token_weights).sum(dim=1)
        return _finish_selected_loss(
            weighted_ce_sums,
            selected_ce_sums.sum(),
            counts,
            sampled_t,
            normalization_mask,
            compute_unweighted_metric,
        )