"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 燧原(Enflame)特化版 · v2
算法层同通用版 · 燧原小芯片套路（参考 chunk_scaled_dot_kkt_enflame.py 模式）:
  1) torch.get_device_module(x.device).device(x.device) wrapper (国产 NPU 必需)
  2) enable_fp_fusion=False (D245 验证)
  3) BLOCK_D: 1024 → 512 (小芯片 + 小 SRAM, 跟 chunk_enflame BK=32 一致思路)
  4) num_warps: 8 → 4 (小芯片 warp 资源少)
  5) num_stages=1 (国产 NPU 多级流水支持不完整)
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _hc_head_fwd_kernel(
    x_ptr, hc_fn_ptr, hc_scale_ptr, hc_base_ptr, out_ptr,
    T,
    HC_DIM: tl.constexpr, D: tl.constexpr, HC_MULT: tl.constexpr,
    BLOCK_D: tl.constexpr,
    norm_eps, hc_eps,
):
    pid = tl.program_id(0).to(tl.int32)

    if x_ptr.dtype.element_ty.primitive_bitwidth == 16:
        x_ptr = x_ptr.to(tl.pointer_type(tl.int16))
    if hc_fn_ptr.dtype.element_ty.primitive_bitwidth == 32:
        hc_fn_ptr = hc_fn_ptr.to(tl.pointer_type(tl.int32))

    sqr_sum = tl.zeros((), dtype=tl.float32)
    mixes = tl.zeros((HC_MULT,), dtype=tl.float32)

    for d_off in range(0, HC_DIM, BLOCK_D):
        d_idx = d_off + tl.arange(0, BLOCK_D)
        mask = d_idx < HC_DIM
        x_val = tl.load(x_ptr + pid * HC_DIM + d_idx, mask=mask, other=0.0).to(tl.float32)
        sqr_sum += tl.sum(x_val * x_val, axis=0)
        for m in tl.static_range(HC_MULT):
            fn_val = tl.load(hc_fn_ptr + m * HC_DIM + d_idx, mask=mask, other=0.0)
            mixes[m] += tl.sum(x_val * fn_val, axis=0)

    rsqrt = 1.0 / tl.sqrt(sqr_sum / HC_DIM + norm_eps)
    mixes = mixes * rsqrt

    scale = tl.load(hc_scale_ptr)
    bases = tl.load(hc_base_ptr + tl.arange(0, HC_MULT))
    pre = 1.0 / (1.0 + tl.exp(-(mixes * scale + bases))) + hc_eps

    for d_off in range(0, D, BLOCK_D):
        d_idx = d_off + tl.arange(0, BLOCK_D)
        mask = d_idx < D
        accum = tl.zeros((BLOCK_D,), dtype=tl.float32)
        for m in tl.static_range(HC_MULT):
            x_val = tl.load(x_ptr + pid * HC_DIM + m * D + d_idx, mask=mask, other=0.0).to(tl.float32)
            accum += pre[m] * x_val
        tl.store(out_ptr + pid * D + d_idx, accum.to(out_ptr.dtype.element_ty), mask=mask)


def hc_head(x, hc_fn, hc_scale, hc_base, norm_eps, hc_eps):
    if x.ndim != 3:
        raise ValueError('x must have shape [T, hc_mult, hidden_size]')
    if x.device.type in ('cpu', 'meta', 'mps'):
        raise RuntimeError('a real Triton accelerator backend is required')

    T, hc_mult, D = x.shape
    HC_DIM = hc_mult * D

    if hc_fn.shape != (hc_mult, HC_DIM):
        raise ValueError(f'hc_fn must have shape [{hc_mult}, {HC_DIM}], got {tuple(hc_fn.shape)}')
    if hc_scale.numel() != 1:
        raise ValueError(f'hc_scale must be scalar, got shape {tuple(hc_scale.shape)}')
    if hc_base.shape != (hc_mult,):
        raise ValueError(f'hc_base must have shape [{hc_mult}], got {tuple(hc_base.shape)}')

    x_flat = x.view(T, HC_DIM)
    out = torch.empty((T, D), dtype=x.dtype, device=x.device)
    if T == 0:
        return out

    # 燧原小 SRAM → 小 BLOCK
    BLOCK_D = min(512, 1 << max(5, (HC_DIM - 1).bit_length()))

    grid = (T,)
    with torch.get_device_module(x.device).device(x.device):
        _hc_head_fwd_kernel[grid](
            x_flat, hc_fn, hc_scale, hc_base, out,
            T,
            HC_DIM=HC_DIM, D=D, HC_MULT=hc_mult,
            BLOCK_D=BLOCK_D,
            norm_eps=norm_eps, hc_eps=hc_eps,
            num_warps=4, num_stages=1, enable_fp_fusion=False,
        )
    return out


reference = hc_head
