from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.attention import SDPBackend

try:
    from fla.modules.layernorm import rms_norm as fused_rms_norm
except ImportError:
    fused_rms_norm = None


class RMSNorm(nn.Module):
    """Zero-centered RMSNorm as used by Qwen3-Next.

    On CUDA this runs FLA's single-pass Triton kernel with the ``1 + weight``
    scale kept in FP32, which matches the eager FP32 math to rounding.
    """

    def __init__(self, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        scale = 1.0 + self.weight.float()
        if fused_rms_norm is not None and hidden_states.is_cuda:
            return fused_rms_norm(hidden_states, scale, None, eps=self.eps)
        input_dtype = hidden_states.dtype
        normalized = hidden_states.float()
        normalized = normalized * torch.rsqrt(
            normalized.square().mean(dim=-1, keepdim=True) + self.eps
        )
        return (normalized * scale).to(dtype=input_dtype)


def linear_cross_entropy_reference(
    states: torch.Tensor,
    weight: torch.Tensor,
    targets: torch.Tensor,
    *,
    z_loss_coefficient: float = 0.0,
    ignore_index: int = -100,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Eager ``(objective, cross_entropy, z_loss)`` over materialized logits."""

    logits = F.linear(states, weight)
    cross_entropy = F.cross_entropy(logits, targets, ignore_index=ignore_index)
    reduced_dtype = torch.promote_types(logits.dtype, torch.float32)
    if z_loss_coefficient > 0.0:
        # Keep logits in their training dtype for logsumexp. Casting the full
        # [tokens, vocab] tensor to fp32 materializes it again in the autograd
        # graph; upcast only the reduced [tokens] partition.
        log_partition = torch.logsumexp(logits, dim=-1).to(reduced_dtype)
        z_loss = log_partition.square()[targets.ne(ignore_index)].mean()
    else:
        z_loss = cross_entropy.new_zeros((), dtype=reduced_dtype)
    objective = cross_entropy + z_loss_coefficient * z_loss
    return objective, cross_entropy.detach(), z_loss.detach()


def linear_token_logprobs_reference(
    states: torch.Tensor,
    weight: torch.Tensor,
    targets: torch.Tensor,
    *,
    ignore_index: int = -100,
) -> torch.Tensor:
    """Eager fp32 ``log_softmax(states @ weight.T)[row, target]``; ignored rows are 0."""

    valid = targets.ne(ignore_index)
    logprobs = F.linear(states, weight).float().log_softmax(dim=-1)
    picked = logprobs.gather(1, targets.masked_fill(~valid, 0)[:, None]).squeeze(1)
    return picked.masked_fill(~valid, 0.0)


def softcap(values: torch.Tensor, cap: float) -> torch.Tensor:
    return cap * torch.tanh(values / cap)


def situ_glu_reference(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
    """``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``."""

    gate, up = gate_up.chunk(2, dim=-1)
    return softcap(gate, gate_cap) * torch.sigmoid(gate) * softcap(up, up_cap)


def attention_sdpa_backends(device: torch.device) -> list[SDPBackend]:
    if device.type != "cuda":
        return [SDPBackend.MATH]
    return [
        SDPBackend.CUDNN_ATTENTION,
        SDPBackend.FLASH_ATTENTION,
        SDPBackend.EFFICIENT_ATTENTION,
    ]