"""FlagOS S2 赛道一 · hc_head (s2t1op055) · 国际A(NVIDIA)/国际B(AMD)通用版 · v2
DeepSeek-V4 "hc_head" LM-head 混合器: RMSNorm + 线性混合 + sigmoid 门控 + 加权求和
v2 修复（从 extend_attention.py v2 经验）:
  1) num_stages=1 显式 (头号嫌疑 — 国产 NPU + 部分国际卡多级流水支持不完整)
  2) num_warps 4 → 8 (大 BLOCK_D=1024 + HC_MULT=4 循环, 8 warps 更稳)
  3) pid int64 → int32 (D ≤ 7168, T 不会超 2.1B, 长度算术走 int32)
  4) 指针 cast 防 NaN (D245 验证)
"""
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)

    # 指针 cast 防 NaN（D245 验证）
    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))

    # ===== Pass 1: 算 squared sum + 算 mixes (linear projection) =====
    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

    # ===== sigmoid gate =====
    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  # [HC_MULT]

    # ===== Pass 2: 加权求和 (collapse hc_mult) =====
    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):
    """hc_head: DeepSeek-V4 HC head reduction for LM-head mixer.

    Computes gates from the RMS-normalized flattened HC residual
    and returns out = sum_i gate_i * residual_i, collapsing hc_mult streams.

    Args:
        x: [T, hc_mult, hidden_size] bfloat16
        hc_fn: [hc_mult, hc_mult * hidden_size] float32
        hc_scale: [1] float32
        hc_base: [hc_mult] float32
        norm_eps, hc_eps: float

    Returns:
        [T, hidden_size] bfloat16
    """
    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

    BLOCK_D = min(1024, 1 << max(5, (HC_DIM - 1).bit_length()))

    grid = (T,)
    _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=8, num_stages=1, enable_fp_fusion=False,
    )
    return out


reference = hc_head
