"""Experimental Metal RMSNorm + indexed AdaLN, with BF16 operation boundaries.

One threadgroup owns one row. The tuning harness checks output error as well as
latency for 128/256/512 threads; this is opt-in because reduction order differs.
"""

from functools import lru_cache

import mlx.core as mx


@lru_cache(maxsize=8)
def _kernel(threads):
    if threads not in (128, 256, 512):
        raise ValueError("norm threads must be 128, 256 or 512")
    return mx.fast.metal_kernel(
        name=f"h3_norm_mod_{threads}",
        input_names=["x", "weight", "shift", "scale", "indices", "epsilon"],
        output_names=["out"],
        source=r"""
        uint tid = thread_position_in_threadgroup.x;
        uint row = threadgroup_position_in_grid.y;
        uint width = x_shape[1];
        float sum = 0;
        for (uint c = tid; c < width; c += TG) {
            float a = float(x[row * width + c]);
            sum += a * a;
        }
        sum = simd_sum(sum);
        threadgroup float partial[32];
        uint lane = tid % 32, warp = tid / 32;
        if (lane == 0) partial[warp] = sum;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (warp == 0) {
            float v = lane < TG / 32 ? partial[lane] : 0.0f;
            v = simd_sum(v);
            if (lane == 0) partial[0] = metal::rsqrt(v / float(width) + epsilon[0]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        float inv = partial[0];
        uint index = indices[row] * width;
        for (uint c = tid; c < width; c += TG) {
            T affine = T(float(x[row * width + c]) * inv * float(weight[c]));
            T factor = T(T(1) + scale[index + c]);
            T product = T(affine * factor);
            out[row * width + c] = T(product + shift[index + c]);
        }
        """,
    )


def norm_modulate(x, weight, shift, scale, indices, eps=1e-5, threads=256):
    if x.ndim != 2 or x.shape[-1] != weight.shape[0]:
        raise ValueError("Expected [rows,width] and a matching norm weight")
    return _kernel(threads)(
        inputs=[
            x,
            weight,
            shift,
            scale,
            indices.astype(mx.uint32),
            mx.array([eps], mx.float32),
        ],
        template=[("T", x.dtype), ("TG", threads)],
        grid=(threads, x.shape[0], 1),
        threadgroup=(threads, 1, 1),
        output_shapes=[x.shape],
        output_dtypes=[x.dtype],
    )[0]


@lru_cache(maxsize=1)
def _qk_kernel():
    return mx.fast.metal_kernel(
        name="h3_qk_norm_rope",
        input_names=["qkv", "qw", "kw", "cosine", "sine", "epsilon"],
        output_names=["qout", "kout"],
        source=r"""
        uint lane = thread_position_in_threadgroup.x;
        uint row = threadgroup_position_in_grid.y;
        uint heads = qkv_shape[2], dim = qkv_shape[3];
        uint token = row / heads, head = row % heads;
        uint qb = token * 3 * heads * dim + head * dim;
        uint kb = qb + heads * dim;
        uint rot = cosine_shape[1];
        float qsum = 0, ksum = 0;
        for (uint c = lane; c < dim; c += 32) {
            float q = float(qkv[qb+c]), k = float(qkv[kb+c]);
            qsum += q*q; ksum += k*k;
        }
        float qi = metal::rsqrt(simd_sum(qsum)/float(dim)+epsilon[0]);
        float ki = metal::rsqrt(simd_sum(ksum)/float(dim)+epsilon[0]);
        for (uint c = lane; c < dim; c += 32) {
            T qn = T(float(qkv[qb+c])*qi*float(qw[c]));
            T kn = T(float(qkv[kb+c])*ki*float(kw[c]));
            if (c < rot) {
                uint partner = (c+rot/2)%rot;
                T qother = T(float(qkv[qb+partner])*qi*float(qw[partner]));
                T kother = T(float(qkv[kb+partner])*ki*float(kw[partner]));
                if (c < rot/2) { qother = -qother; kother = -kother; }
                T qcos = T(qn*cosine[token*rot+c]);
                T kcos = T(kn*cosine[token*rot+c]);
                T qsin = T(qother*sine[token*rot+c]);
                T ksin = T(kother*sine[token*rot+c]);
                qn = T(qcos+qsin); kn = T(kcos+ksin);
            }
            qout[row*dim+c] = qn;
            kout[row*dim+c] = kn;
        }
        """,
    )


def qk_norm_rope(qkv, qw, kw, cos, sin, eps=1e-5):
    """Fused head RMSNorm/partial RoPE; qkv [S,3,H,D], outputs [1,S,H,D]."""
    if qkv.ndim != 4 or qkv.shape[1] != 3:
        raise ValueError("QKV must have shape [sequence,3,heads,dim]")
    n, _, heads, dim = qkv.shape
    shape = (1, n, heads, dim)
    rot = cos.shape[-1]
    return _qk_kernel()(
        inputs=[
            qkv,
            qw,
            kw,
            cos.reshape(n, rot),
            sin.reshape(n, rot),
            mx.array([eps], mx.float32),
        ],
        template=[("T", qkv.dtype)],
        grid=(32, n * heads, 1),
        threadgroup=(32, 1, 1),
        output_shapes=[shape, shape],
        output_dtypes=[qkv.dtype, qkv.dtype],
    )