"""Inference-only fused SwiGLU selection for dense MLP layers."""

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

from astrai.extension.backend.linear import linear
from astrai.extension.dispatch import env_mode
from astrai.extension.loader import is_available
from astrai.extension.ops.swiglu import bf16_swiglu


def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
    # Keep the existing linear backend in the fallback chain. This preserves
    # any independently qualified GEMV batches instead of making the fusion
    # decision suppress linear-level optimizations.
    return linear(x, up_weight) * F.silu(linear(x, gate_weight))


def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
    return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())


def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
    return not (
        torch.is_grad_enabled()
        or not x.is_cuda
        or x.dtype != torch.bfloat16
        or up_weight.dtype != torch.bfloat16
        or gate_weight.dtype != torch.bfloat16
        or x.ndim not in (1, 2)
        or up_weight.ndim != 2
        or gate_weight.ndim != 2
        or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
        or up_weight.shape != gate_weight.shape
        or x.shape[-1] != up_weight.shape[1]
        or x.shape[-1] % 8 != 0
        or x.device != up_weight.device
        or x.device != gate_weight.device
        or not x.is_contiguous()
        or not up_weight.is_contiguous()
        or not gate_weight.is_contiguous()
        # The fused kernel reads all streams as uint4; contiguous-but-offset
        # views are routed to the unfused chain instead of failing.
        or (x.data_ptr() & 15) != 0
        or (up_weight.data_ptr() & 15) != 0
        or (gate_weight.data_ptr() & 15) != 0
        or not is_available("bf16_swiglu")
    )


def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
    """Apply the dense-MLP SwiGLU projection with a safe torch fallback.

    ``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain; ``1`` forces
    the fused primitive for supported inputs; ``auto`` (the default) uses
    the fused primitive for decode batches with M in ``{1, ..., 8}``. The
    fused kernel reads x once and covers both projections plus the SiLU
    gate-multiply in a single launch, measured 9-15% faster than the
    unfused chain per MLP call on L20 with L2-thrashing weight rotation.
    """
    if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
        return _fused_swiglu(x, up_weight, gate_weight)
    return _unfused_swiglu(x, up_weight, gate_weight)


__all__ = ["swiglu"]