"""Research kernels. No default dispatch until full-output validation."""

from functools import lru_cache

import mlx.core as mx

from experiments.v2_mpp import gemm


class TunedFC2:
    def __init__(self, layer):
        self.original = layer
        self.bias = layer.bias
        self.weight = layer.weight
        self.zero = mx.zeros(self.weight.shape[0], self.weight.dtype)

    @property
    def nbytes(self):
        return self.original.nbytes

    def __call__(self, x):
        if x.ndim != 2 or x.dtype != mx.bfloat16 or self.weight.shape != (5376, 14336):
            return self.original(x)
        return gemm(
            x,
            self.weight,
            self.zero if self.bias is None else self.bias,
            64,
            128,
            512,
            4,
            8,
        )


@lru_cache(None)
def _qk():
    return mx.fast.metal_kernel(
        name="h3_r3_qk_native_reduce",
        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;
        #pragma unroll
        for (uint j = 0; j < 4; ++j) {
            uint c = lane * 4 + j;
            float q = float(qkv[qb+c]), k = float(qkv[kb+c]);
            qsum += q*q; ksum += k*k;
        }
        float qi = metal::precise::rsqrt(simd_sum(qsum)/float(dim)+epsilon[0]);
        float ki = metal::precise::rsqrt(simd_sum(ksum)/float(dim)+epsilon[0]);
        #pragma unroll
        for (uint j = 0; j < 4; ++j) {
            uint c = lane * 4 + j;
            T qn = T(float(qw[c])*(float(qkv[qb+c])*qi));
            T kn = T(float(kw[c])*(float(qkv[kb+c])*ki));
            if (c < rot) {
                uint partner = (c+rot/2)%rot;
                T qother = T(float(qw[partner])*(float(qkv[qb+partner])*qi));
                T kother = T(float(kw[partner])*(float(qkv[kb+partner])*ki));
                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);
            }
            uint offset = HEAD_MAJOR ? (head*qkv_shape[0]+token)*dim+c : row*dim+c;
            qout[offset] = qn; kout[offset] = kn;
        }
        """,
    )


def qk_native_reduce(qkv, qw, kw, cos, sin, eps=1e-5, *, head_major=False):
    n, _, heads, dim = qkv.shape
    if dim != 128:
        raise ValueError("Research kernel only supports D128")
    shape = (1, heads, n, dim) if head_major else (1, n, heads, dim)
    return _qk()(
        inputs=[
            qkv,
            qw,
            kw,
            cos.reshape(n, -1),
            sin.reshape(n, -1),
            mx.array([eps], mx.float32),
        ],
        template=[("T", qkv.dtype), ("HEAD_MAJOR", head_major)],
        grid=(32, n * heads, 1),
        threadgroup=(32, 1, 1),
        output_shapes=[shape, shape],
        output_dtypes=[qkv.dtype, qkv.dtype],
    )


@lru_cache(None)
def _norm():
    return mx.fast.metal_kernel(
        name="h3_r3_norm_mod_native_reduce",
        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 r = 0; r < width; r += TG*4) {
            #pragma unroll
            for (uint j = 0; j < 4; ++j) {
                uint c = r + tid*4+j;
                if (c < width) { 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 (warp == 0) partial[lane] = 0;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (lane == 0) partial[warp] = sum;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        threadgroup float inverse;
        if (warp == 0) {
            float s = simd_sum(partial[lane]);
            if (lane == 0) inverse = metal::precise::rsqrt(s/float(width)+epsilon[0]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        uint index = indices[row]*width;
        for (uint r = 0; r < width; r += TG*4) {
            #pragma unroll
            for (uint j = 0; j < 4; ++j) {
                uint c = r+tid*4+j;
                if (c < width) {
                    T affine = T(float(weight[c])*(float(x[row*width+c])*inverse));
                    T factor = T(T(1)+scale[index+c]);
                    T product = T(affine*factor);
                    out[row*width+c] = T(product+shift[index+c]);
                }
            }
        }
        """,
    )


def norm_native_reduce(x, weight, shift, scale, indices, eps=1e-5, threads=1024):
    return _norm()(
        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]