"""Dense Metal attention and explicit BF16 rounding boundaries."""

from __future__ import annotations

import mlx.core as mx


def silu(x):
    # PyTorch evaluates SiLU in its accumulation dtype, then rounds once.
    y = x.astype(mx.float32)
    return (y * mx.sigmoid(y)).astype(x.dtype)


def rms_norm(x, weight=None, eps=1e-5, *, rounded=False):
    # Qwen explicitly rounds before the affine multiply. ATen nn.RMSNorm
    # (DiT/VAE) multiplies its affine in FP32 before rounding. MLX's low
    # precision fast RMSNorm implements Qwen's order, so upcast for ATen.
    if rounded:
        normalized = mx.fast.rms_norm(x, None, eps)
        return normalized if weight is None else normalized * weight
    return mx.fast.rms_norm(
        x.astype(mx.float32), None if weight is None else weight.astype(mx.float32), eps
    ).astype(x.dtype)


def rotate(x, cos, sin):
    n = cos.shape[-1]
    a = x[..., :n]
    half = n // 2
    rotated = mx.concatenate((-a[..., half:], a[..., :half]), axis=-1)
    out = a * cos.astype(x.dtype) + rotated * sin.astype(x.dtype)
    return mx.concatenate((out, x[..., n:]), axis=-1)


def attention(q, k, v, *, causal=False, query_chunk=0):
    """Inputs [batch, sequence, heads, dim]; GQA without materializing repeated KV."""
    q, k, v = (a.transpose(0, 2, 1, 3) for a in (q, k, v))
    if not query_chunk or q.shape[2] <= query_chunk:
        out = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=q.shape[-1] ** -0.5, mask="causal" if causal else None
        )
    else:
        rows = []
        for start in range(0, q.shape[2], query_chunk):
            qq = q[:, :, start : start + query_chunk]
            mask = None
            if causal:
                mask = (
                    mx.arange(start, start + qq.shape[2])[:, None]
                    >= mx.arange(k.shape[2])[None]
                )
            rows.append(
                mx.fast.scaled_dot_product_attention(
                    qq, k, v, scale=q.shape[-1] ** -0.5, mask=mask
                )
            )
            mx.eval(rows[-1])
        out = mx.concatenate(rows, axis=2)
    return out.transpose(0, 2, 1, 3)


def modulate(
    x, norm_weight, shift, scale, indices, eps=1e-5, *, kernel="native", threads=256
):
    if (
        kernel == "metal-exact"
        and x.ndim == 2
        and x.shape[1] == 5376
        and x.dtype == mx.bfloat16
    ):
        from .tuning import norm_modulate_exact, supported_device

        if supported_device() and 4096 <= x.shape[0] <= 9216:
            return norm_modulate_exact(x, norm_weight, shift, scale, indices, eps)
    if kernel in ("metal", "metal-qk"):
        from .kernels import norm_modulate

        return norm_modulate(x, norm_weight, shift, scale, indices, eps, threads)
    h = rms_norm(x, norm_weight, eps)
    return h * (mx.array(1, x.dtype) + scale[indices]) + shift[indices]


def gated_add(x, other, gate, indices):
    return x + gate[indices] * other


def swiglu(x):
    a, b = mx.split(x, 2, axis=-1)
    return silu(a) * b


def layer_norm(x, weight, bias, eps=1e-5):
    return mx.fast.layer_norm(
        x.astype(mx.float32), weight.astype(mx.float32), bias.astype(mx.float32), eps
    ).astype(x.dtype)