from __future__ import annotations

import torch

from ..models.student_t import masked_student_t_nll, gaussian_nll


def pinball_loss(y: torch.Tensor, mask: torch.Tensor, q: torch.Tensor, levels: list[float]) -> torch.Tensor:
    """y, mask: [B, H]; q: [B, H, Q]."""
    lv = torch.tensor(levels, device=y.device, dtype=torch.float32).view(1, 1, -1)
    diff = y.float().unsqueeze(-1) - q
    loss = torch.maximum(lv * diff, (lv - 1) * diff).mean(-1)
    m = mask.float()
    return (loss * m).sum() / m.sum().clamp_min(1.0)


def forecast_loss(out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], head: str,
                  levels: list[float] | None = None, aux_vol_weight: float = 0.0) -> tuple[torch.Tensor, dict]:
    y, m = batch["y"].float(), batch["y_mask"]
    if head == "student_t":
        loss = masked_student_t_nll(y, m, out)
    elif head == "gaussian":
        loss = gaussian_nll(y, m, out)
    else:
        loss = pinball_loss(y, m, out["quantiles"], levels)
    logs = {"nll": float(loss.detach())}
    if aux_vol_weight > 0 and head != "quantile":
        # auxiliary: predicted variance of the H-sum vs realised squared sum (QLIKE-like, stabilised)
        var = (out["scale"] ** 2 * (out["nu"] / (out["nu"] - 2)).clamp(max=50)).sum(-1) if head == "student_t" \
            else (out["scale"] ** 2).sum(-1)
        rv = ((y * m.float()) ** 2).sum(-1)
        aux = (torch.log(var + 1e-12) + rv / (var + 1e-12)).mean()
        loss = loss + aux_vol_weight * aux
        logs["aux_vol"] = float(aux.detach())
    return loss, logs