"""Losses: next-token cross-entropy (per position), controller regression/classification losses."""
from __future__ import annotations

import torch
import torch.nn.functional as F


def token_nll(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:
    """Per-position negative log-likelihood [B,T] (float32); ignored positions get 0."""
    B, T, V = logits.shape
    nll = F.cross_entropy(logits.float().reshape(B * T, V), targets.reshape(B * T), ignore_index=ignore_index, reduction="none")
    return nll.view(B, T)


def lm_loss(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:
    nll = token_nll(logits, targets, ignore_index)
    mask = (targets != ignore_index).float()
    return (nll * mask).sum() / mask.sum().clamp(min=1.0)


def controller_losses(preds: dict, g_obs: torch.Tensor, valid: torch.Tensor, huber_delta: float = 0.1):
    """preds: dict of [N] tensors for the logged (state, action) pairs; g_obs: [N] observed gains (detached);
    valid: [N] bool. Returns dict(gain=Huber, harm=BCE, nll=Gaussian NLL for sigma)."""
    g_obs = g_obs.detach()
    v = valid.float()
    n = v.sum().clamp(min=1.0)
    gain = (F.huber_loss(preds["g_hat"], g_obs, delta=huber_delta, reduction="none") * v).sum() / n
    harm_t = (g_obs < 0).float()
    harm = (F.binary_cross_entropy_with_logits(preds["harm_logit"], harm_t, reduction="none") * v).sum() / n
    log_sigma = preds["log_sigma"].clamp(-7, 3)
    nll = ((0.5 * ((g_obs - preds["g_hat"].detach()) ** 2) * torch.exp(-2 * log_sigma) + log_sigma) * v).sum() / n
    return {"gain": gain, "harm": harm, "nll": nll}