"""Autogenerated Triton kernel for row-wise layer normalization."""

from __future__ import annotations

import torch
import warnings

try:
    import triton
    import triton.language as tl
except ImportError:  # pragma: no cover - optional dependency
    triton = None
    tl = None


BLOCK_SIZE = 128
EPSILON = 1e-05


if triton is not None:

    @triton.jit
    def layer_norm_kernel(x_ptr, gamma_ptr, beta_ptr, output_ptr, stride, hidden_size, *, BLOCK_SIZE: tl.constexpr):
        row = tl.program_id(axis=0)
        offsets = row * stride + tl.arange(0, BLOCK_SIZE)
        mask = offsets < (row + 1) * stride

        x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
        mean = tl.sum(x, axis=0) / hidden_size
        centered = x - mean
        variance = tl.sum(centered * centered, axis=0) / hidden_size
        inv_std = tl.rsqrt(variance + EPSILON)
        gamma = tl.load(gamma_ptr + offsets, mask=mask, other=1.0)
        beta = tl.load(beta_ptr + offsets, mask=mask, other=0.0)
        normalized = centered * inv_std * gamma + beta
        tl.store(output_ptr + offsets, normalized, mask=mask)


def run_kernel(x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
    """Apply layer normalization with learned scale and bias."""

    if x.dim() != 2:
        raise ValueError("Input must be a 2D tensor [batch, hidden]")
    batch, hidden = x.shape
    if gamma.shape != (hidden,) or beta.shape != (hidden,):
        raise ValueError("Gamma/Beta shapes must match the hidden dimension")

    fallback = torch.nn.functional.layer_norm(
        x, (hidden,), weight=gamma, bias=beta, eps=EPSILON
    )

    if triton is None:
        return fallback

    output = torch.empty_like(x)
    grid = (batch,)
    try:
        layer_norm_kernel[grid](
            x,
            gamma,
            beta,
            output,
            x.stride(0),
            hidden,
            BLOCK_SIZE=BLOCK_SIZE,
        )
    except Exception as exc:  # pragma: no cover - triton runtime issues
        warnings.warn(
            f"Triton layer_norm kernel failed; falling back to torch: {exc}",
            RuntimeWarning,
        )
        return fallback
    return output