"""自定义 Metal 算子：indexed 量化线性 + fused MoE expert（实验性加速路径）。

这些 kernel 走 `mx.fast.metal_kernel`，按 `indices[pair]` 选择专家权重，
直接在 packed 量化权重上算 SwiGLU/down，避免反量化中间张量。仅在对应
CUSTOM_* 环境开关命中且 bit 宽匹配时启用，默认关闭（不改变模型数值）。
"""
from typing import Tuple  # noqa: F401  (保留以兼容历史 import)

import mlx.core as mx

from mlx_streaming import config
from mlx_streaming.config import parse_layers_env as _parse_layers_env

# kernel 编译缓存：按 (tile) / (hidden,moe_inter,...) 复用已编译 kernel。
_CUSTOM_QKERNEL_CACHE = {}
_CUSTOM_FUSED_MOE_CACHE = {}


def _custom_qproj_enabled(layer_idx: int, bits: int) -> bool:
    if not config.custom_qproj():
        return False
    if bits != config.custom_qproj_bits():
        return False
    layers = _parse_layers_env("CUSTOM_QPROJ_LAYERS")
    return layers is None or int(layer_idx) in layers


def _custom_fused_moe_enabled(layer_idx: int, proj_bits: dict) -> bool:
    if not config.custom_fused_moe():
        return False
    bits = config.custom_fused_moe_bits()
    if any(int(proj_bits[name]) != bits for name in ("gate_proj", "up_proj", "down_proj")):
        return False
    layers = _parse_layers_env("CUSTOM_FUSED_MOE_LAYERS")
    return layers is None or int(layer_idx) in layers


def _custom_qproj_targets() -> set[str]:
    spec = config.custom_qproj_targets()
    return {x.strip() for x in spec.split(",") if x.strip()}


def _custom_qlinear_indexed(x: mx.array, indices: mx.array, weight: mx.array,
                            scales: mx.array, biases: mx.array,
                            out_dim: int, in_dim: int, group_size: int,
                            bits: int, tile: int = 4) -> mx.array:
    """indexed custom qlinear: x[p,in] 使用 indices[p] 选择专家权重，输出 [p,out]。"""
    key = tile
    kernel = _CUSTOM_QKERNEL_CACHE.get(key)
    if kernel is None:
        source = r"""
            constexpr int lanes_per_row = 256 / rows_per_group;
            constexpr int block_size = rows_per_group * lanes_per_row;
            uint tid = thread_position_in_threadgroup.x;
            uint group_id = thread_position_in_grid.x / block_size;
            uint local_row = tid / lanes_per_row;
            uint lane = tid % lanes_per_row;
            uint global_row = group_id * rows_per_group + local_row;
            if (global_row >= pairs * out_dim) return;
            uint pair = global_row / out_dim;
            uint out_row = global_row % out_dim;
            uint expert = indices[pair];
            constexpr uint mask = (1u << bits) - 1u;
            constexpr int words_per_row = (in_dim * bits) / 32;
            constexpr int groups_per_row = in_dim / group_size;
            threadgroup float partial[block_size];
            float acc = 0.0f;
            for (int col = int(lane); col < in_dim; col += lanes_per_row) {
                int bit_offset = col * bits;
                int word_idx = bit_offset / 32;
                int shift = bit_offset % 32;
                uint base = (expert * out_dim + out_row) * words_per_row;
                uint word = weight[base + word_idx];
                uint q = (word >> shift);
                if (shift + bits > 32) {
                    uint next_word = weight[base + word_idx + 1];
                    q |= (next_word << (32 - shift));
                }
                q = q & mask;
                int g = col / group_size;
                uint sb = (expert * out_dim + out_row) * groups_per_row + g;
                float wv = float(q) * scales[sb] + biases[sb];
                acc += wv * x[pair * in_dim + col];
            }
            partial[tid] = acc;
            threadgroup_barrier(mem_flags::mem_threadgroup);
            for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
                if (lane < stride) {
                    partial[tid] += partial[tid + stride];
                }
                threadgroup_barrier(mem_flags::mem_threadgroup);
            }
            if (lane == 0) {
                y[global_row] = partial[0 + (local_row * lanes_per_row)];
            }
        """
        kernel = mx.fast.metal_kernel(
            name=f"custom_qlinear_indexed_tile{tile}",
            input_names=["x", "indices", "weight", "scales", "biases"],
            output_names=["y"],
            source=source,
        )
        _CUSTOM_QKERNEL_CACHE[key] = kernel
    pairs = int(indices.size)
    grid_groups = (pairs * out_dim + tile - 1) // tile
    (y,) = kernel(
        inputs=[x, indices.astype(mx.uint32), weight, scales, biases],
        output_shapes=[(pairs, out_dim)],
        output_dtypes=[mx.float32],
        grid=(grid_groups * 256, 1, 1),
        threadgroup=(256, 1, 1),
        template=[
            ("pairs", pairs),
            ("out_dim", out_dim),
            ("in_dim", in_dim),
            ("group_size", group_size),
            ("bits", bits),
            ("rows_per_group", tile),
        ],
    )
    return y


def _custom_fused_moe_indexed(x: mx.array, indices: mx.array,
                              gate_w: mx.array, gate_s: mx.array, gate_b: mx.array,
                              up_w: mx.array, up_s: mx.array, up_b: mx.array,
                              down_w: mx.array, down_s: mx.array, down_b: mx.array,
                              hidden: int, moe_inter: int, group_size: int,
                              bits: int) -> mx.array:
    """fused MoE expert: x[p] 用 indices[p] 选择专家，输出 down 后的 hidden。"""
    lanes_per_row = config.custom_fused_moe_lanes()
    block_size = config.custom_fused_moe_block()
    key = (hidden, moe_inter, group_size, bits, lanes_per_row, block_size)
    kernel = _CUSTOM_FUSED_MOE_CACHE.get(key)
    if kernel is None:
        source = r"""
            uint tid = thread_position_in_threadgroup.x;
            uint pair = thread_position_in_grid.x / block_size;
            if (pair >= pairs) return;
            uint expert = indices[pair];
            constexpr uint rows_per_step = block_size / lanes_per_row;
            uint local_row = tid / lanes_per_row;
            uint row_lane = tid % lanes_per_row;
            constexpr uint mask = (1u << bits) - 1u;
            constexpr int gu_words_per_row = (hidden * bits) / 32;
            constexpr int gu_groups_per_row = hidden / group_size;
            constexpr int down_words_per_row = (moe_inter * bits) / 32;
            constexpr int down_groups_per_row = moe_inter / group_size;
            threadgroup float act[1024];
            threadgroup float gate_part[block_size];
            threadgroup float up_part[block_size];

            // 多个 lane 协作计算同一 row，避免每个 dot 完全串行。
            for (uint row_base = 0; row_base < moe_inter; row_base += rows_per_step) {
                uint row = row_base + local_row;
                float gate_acc = 0.0f;
                float up_acc = 0.0f;
                if (row < moe_inter) {
                    for (uint col = row_lane; col < hidden; col += lanes_per_row) {
                        int bit_offset = int(col * bits);
                        int word_idx = bit_offset / 32;
                        int shift = bit_offset % 32;
                        uint gu_base = (expert * moe_inter + row) * gu_words_per_row;
                        uint qg = gate_w[gu_base + word_idx] >> shift;
                        uint qu = up_w[gu_base + word_idx] >> shift;
                        if (shift + bits > 32) {
                            qg |= gate_w[gu_base + word_idx + 1] << (32 - shift);
                            qu |= up_w[gu_base + word_idx + 1] << (32 - shift);
                        }
                        qg &= mask;
                        qu &= mask;
                        uint g = col / group_size;
                        uint sb = (expert * moe_inter + row) * gu_groups_per_row + g;
                        float xv = x[pair * hidden + col];
                        gate_acc += (float(qg) * gate_s[sb] + gate_b[sb]) * xv;
                        up_acc += (float(qu) * up_s[sb] + up_b[sb]) * xv;
                    }
                }
                gate_part[tid] = gate_acc;
                up_part[tid] = up_acc;
                threadgroup_barrier(mem_flags::mem_threadgroup);
                for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
                    if (row_lane < stride) {
                        gate_part[tid] += gate_part[tid + stride];
                        up_part[tid] += up_part[tid + stride];
                    }
                    threadgroup_barrier(mem_flags::mem_threadgroup);
                }
                if (row_lane == 0 && row < moe_inter) {
                    float gate_v = gate_part[tid];
                    float up_v = up_part[tid];
                    float sig = 1.0f / (1.0f + exp(-gate_v));
                    act[row] = gate_v * sig * up_v;
                }
                threadgroup_barrier(mem_flags::mem_threadgroup);
            }

            // down projection 输出最终 hidden。
            for (uint row_base = 0; row_base < hidden; row_base += rows_per_step) {
                uint row = row_base + local_row;
                float acc = 0.0f;
                if (row < hidden) {
                    for (uint col = row_lane; col < moe_inter; col += lanes_per_row) {
                        int bit_offset = int(col * bits);
                        int word_idx = bit_offset / 32;
                        int shift = bit_offset % 32;
                        uint base = (expert * hidden + row) * down_words_per_row;
                        uint q = down_w[base + word_idx] >> shift;
                        if (shift + bits > 32) {
                            q |= down_w[base + word_idx + 1] << (32 - shift);
                        }
                        q &= mask;
                        uint g = col / group_size;
                        uint sb = (expert * hidden + row) * down_groups_per_row + g;
                        acc += (float(q) * down_s[sb] + down_b[sb]) * act[col];
                    }
                }
                gate_part[tid] = acc;
                threadgroup_barrier(mem_flags::mem_threadgroup);
                for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
                    if (row_lane < stride) {
                        gate_part[tid] += gate_part[tid + stride];
                    }
                    threadgroup_barrier(mem_flags::mem_threadgroup);
                }
                if (row_lane == 0 && row < hidden) {
                    y[pair * hidden + row] = gate_part[tid];
                }
                threadgroup_barrier(mem_flags::mem_threadgroup);
            }
        """
        kernel = mx.fast.metal_kernel(
            name=f"custom_fused_moe_h{hidden}_i{moe_inter}_b{bits}_l{lanes_per_row}",
            input_names=[
                "x", "indices",
                "gate_w", "gate_s", "gate_b",
                "up_w", "up_s", "up_b",
                "down_w", "down_s", "down_b",
            ],
            output_names=["y"],
            source=source,
        )
        _CUSTOM_FUSED_MOE_CACHE[key] = kernel
    pairs = int(indices.size)
    (y,) = kernel(
        inputs=[
            x, indices.astype(mx.uint32),
            gate_w, gate_s, gate_b,
            up_w, up_s, up_b,
            down_w, down_s, down_b,
        ],
        output_shapes=[(pairs, hidden)],
        output_dtypes=[mx.float32],
        grid=(pairs * block_size, 1, 1),
        threadgroup=(block_size, 1, 1),
        template=[
            ("pairs", pairs),
            ("hidden", hidden),
            ("moe_inter", moe_inter),
            ("group_size", group_size),
            ("bits", bits),
            ("lanes_per_row", lanes_per_row),
            ("block_size", block_size),
        ],
    )
    return y