"""Task 38 · sigmoid_gate_topk_renorm

MoE 路由门算子 (DeepSeek-V2/V3 / Qwen3-MoE 风格) - 3-kernel Triton 实现:

  K1 (_sigmoid_bias_kernel) :  sel = sigmoid(routed) + bias
  K2 (_topk_kernel)         :  indices = topk(sel, k)   (自写, 迭代 tl.argmax)
  K3 (_gate_finalize_kernel):  gather 原始 routed → cat shared
                               → sigmoid → renorm → scale → split 输出

国产 NPU 套路全套 (来自 D245 验证):
  - input 指针 cast 到 int16/int32 (防 NaN 转换)
  - stride / pid 用 int64 cast (防 overflow)
  - with torch.get_device_module(x.device).device(x.device): 切设备
  - enable_fp_fusion=False, num_warps=4
  - 不用 tl.sigmoid, 用 1.0 / (1.0 + tl.exp(-x))
"""
import torch
import triton
import triton.language as tl


# ============================================================
# Kernel 1: sel = sigmoid(routed) + bias
# ============================================================
@triton.jit
def _sigmoid_bias_kernel(
    routed_ptr,    # [T, N] fp16/bf16
    bias_ptr,      # [N] fp32
    sel_ptr,       # [T, N] fp32 (中间 buffer)
    N,
    stride_t,      # int (T 维度 stride)
    BLOCK_N: tl.constexpr,
):
    # input 指针 cast 防 NaN
    if routed_ptr.dtype.element_ty.primitive_bitwidth == 16:
        routed_ptr = routed_ptr.to(tl.pointer_type(tl.int16))

    pid_t = tl.program_id(0).to(tl.int64)
    pid_n = tl.program_id(1).to(tl.int64)

    offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)
    mask = offs < N

    x = tl.load(routed_ptr + pid_t * stride_t + offs, mask=mask, other=0).to(tl.float32)
    b = tl.load(bias_ptr + offs, mask=mask, other=0.0)

    # 不用 tl.sigmoid (国产 NPU 不支持), 用 1/(1+exp(-x))
    sig = 1.0 / (1.0 + tl.exp(-x))
    sel = sig + b

    tl.store(sel_ptr + pid_t * N + offs, sel, mask=mask)


# ============================================================
# Kernel 2: top-k 选 indices (自写, 迭代 argmax)
# ============================================================
@triton.jit
def _topk_kernel(
    sel_ptr,       # [T, N] fp32 (来自 K1)
    idx_ptr,       # [T, k] int32
    N,
    K: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid = tl.program_id(0).to(tl.int64)

    offs = tl.arange(0, BLOCK_N)
    mask = offs < N

    sel = tl.load(sel_ptr + pid * N + offs, mask=mask, other=-float('inf'))

    NEG_INF: tl.constexpr = float('-inf')
    for i in tl.static_range(K):
        # 找最大值的索引 (scalar)
        idx = tl.argmax(sel, axis=0)
        # 把 sel[idx] 设为 -inf (避免重复选)
        is_max = (offs == idx)
        sel = tl.where(is_max & mask, NEG_INF, sel)
        # 写 indices
        tl.store(idx_ptr + pid * K + i, idx.to(tl.int32))


# ============================================================
# Kernel 3: gather + cat + sigmoid + renorm + scale + split
# ============================================================
@triton.jit
def _gate_finalize_kernel(
    routed_ptr,         # [T, N] fp16/bf16
    shared_ptr,         # [T, S] fp16/bf16
    idx_ptr,            # [T, k] int32
    route_scale,        # float
    global_scale_ptr,   # [1] fp32
    routed_w_ptr,       # [T, k] output dtype (input dtype)
    shared_w_ptr,       # [T, S] output dtype (input dtype)
    N, S,
    stride_routed_t, stride_shared_t,
    stride_routed_w_t, stride_shared_w_t,
    K: tl.constexpr,
    BLOCK_S: tl.constexpr,  # next_pow2(S)
):
    # input 指针 cast 防 NaN, output 不 cast (output 是计算结果, 不需保 NaN bits)
    if routed_ptr.dtype.element_ty.primitive_bitwidth == 16:
        routed_ptr = routed_ptr.to(tl.pointer_type(tl.int16))
        shared_ptr = shared_ptr.to(tl.pointer_type(tl.int16))

    pid = tl.program_id(0).to(tl.int64)

    # ===== 加载 routed 的 K 个值 (按 indices gather) =====
    idx_offs = tl.arange(0, K)
    indices = tl.load(idx_ptr + pid * K + idx_offs).to(tl.int64)
    routed_vals = tl.load(routed_ptr + pid * stride_routed_t + indices).to(tl.float32)
    routed_sigmoid = 1.0 / (1.0 + tl.exp(-routed_vals))

    # ===== 加载 shared 的 S 个值 =====
    s_offs = tl.arange(0, BLOCK_S)
    s_mask = s_offs < S
    shared_vals = tl.load(
        shared_ptr + pid * stride_shared_t + s_offs, mask=s_mask, other=0
    ).to(tl.float32)
    shared_sigmoid = 1.0 / (1.0 + tl.exp(-shared_vals))

    # ===== 计算 sum (routed + shared 一起归一化) =====
    sum_routed = tl.sum(routed_sigmoid, axis=0)
    sum_shared = tl.sum(shared_sigmoid, axis=0)
    total_sum = sum_routed + sum_shared

    # ===== scale =====
    gs = tl.load(global_scale_ptr)
    inv = (route_scale * gs) / total_sum

    # ===== 写 routed_w =====
    routed_w = (routed_sigmoid * inv).to(routed_w_ptr.dtype.element_ty)
    tl.store(routed_w_ptr + pid * stride_routed_w_t + idx_offs, routed_w)

    # ===== 写 shared_w =====
    shared_w = (shared_sigmoid * inv).to(shared_w_ptr.dtype.element_ty)
    tl.store(shared_w_ptr + pid * stride_shared_w_t + s_offs, shared_w, mask=s_mask)


# ============================================================
# Python wrapper
# ============================================================
def _next_pow2(x):
    p = 1
    while p < x:
        p *= 2
    return p


def sigmoid_gate_topk_renorm(logits, k, n_shared_experts, route_scale, global_scale, bias):
    T, G = logits.shape
    N = G - n_shared_experts
    S = n_shared_experts

    if logits.ndim != 2:
        raise ValueError('logits must have shape [T, N+S]')
    if not isinstance(k, int) or k <= 0 or k > N:
        raise ValueError(f'k must be 0 < k <= N={N}, got {k}')
    if not isinstance(n_shared_experts, int) or n_shared_experts < 0 or n_shared_experts > S:
        raise ValueError(f'n_shared_experts must be 0 <= S={S}, got {n_shared_experts}')
    if bias.shape != (N,):
        raise ValueError(f'bias must have shape [{N}], got {tuple(bias.shape)}')
    if global_scale.numel() != 1:
        raise ValueError(f'global_scale must be a scalar tensor, got shape {tuple(global_scale.shape)}')
    if logits.device.type in ('cpu', 'meta', 'mps'):
        raise RuntimeError('a real Triton accelerator backend is required')

    routed = logits[:, :N]
    shared = logits[:, N:]

    # ===== 中间 buffer =====
    sel = torch.empty((T, N), dtype=torch.float32, device=logits.device)
    indices = torch.empty((T, k), dtype=torch.int32, device=logits.device)

    BLOCK_N = max(16, _next_pow2(N))
    BLOCK_S = max(16, _next_pow2(S))

    module = torch.get_device_module(logits.device)
    with module.device(logits.device):
        # ===== Kernel 1: sigmoid + bias =====
        grid1 = (T, triton.cdiv(N, BLOCK_N))
        _sigmoid_bias_kernel[grid1](
            routed, bias, sel, N, routed.stride(0),
            BLOCK_N=BLOCK_N,
            enable_fp_fusion=False, num_warps=4,
        )

        # ===== Kernel 2: topk =====
        grid2 = (T,)
        _topk_kernel[grid2](
            sel, indices, N, K=k, BLOCK_N=BLOCK_N,
            enable_fp_fusion=False, num_warps=4,
        )

        # ===== Kernel 3: gate finalize =====
        routed_w = torch.empty((T, k), dtype=logits.dtype, device=logits.device)
        shared_w = torch.empty((T, S), dtype=logits.dtype, device=logits.device)
        grid3 = (T,)
        _gate_finalize_kernel[grid3](
            routed, shared, indices, route_scale, global_scale,
            routed_w, shared_w, N, S,
            routed.stride(0), shared.stride(0),
            routed_w.stride(0), shared_w.stride(0),
            K=k, BLOCK_S=BLOCK_S,
            enable_fp_fusion=False, num_warps=4,
        )

    return routed_w, indices, shared_w


reference = sigmoid_gate_topk_renorm