"""W4A16 weight-only quantized GEMM for H100 (SM90).

Fused unpack + dequant + GEMM.  Triton kernel with per-group (128) bf16
scales/zeros applied inside the K-loop, tensor-core dot in bf16 with fp32
accumulation.

  x:      (M, K)          bf16
  w_q:    (K // 2, N)     uint8   low nibble = even k row, high nibble = odd k row
  scales: (K // 128, N)   bf16
  zeros:  (K // 128, N)   bf16
  out:    (M, N)          bf16
"""
from __future__ import annotations

import torch
import torch.nn as nn
import triton
import triton.language as tl

OP_TYPE = "gemm_w4a16"
SUPPORTED_PRECISIONS = ["int4_bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

GROUP_SIZE = 128


# ---------------------------------------------------------------------------
# Triton fused kernel
# ---------------------------------------------------------------------------
@triton.jit
def _w4a16_gemm_kernel(
    x_ptr, w_ptr, s_ptr, z_ptr, y_ptr,
    M, N, K,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BM)
    num_pid_n = tl.cdiv(N, BN)

    # L2-friendly rasterization (grouped along M).
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_kh = tl.arange(0, BK // 2)  # packed-byte rows (two k per byte)
    mask_m = offs_m < M

    # x tiles for even / odd k lanes (two half-K dots per group tile)
    x_even_ptrs = x_ptr + offs_m[:, None] * K + (2 * offs_kh)[None, :]
    x_odd_ptrs = x_even_ptrs + 1
    # packed tile: (BK // 2, BN) uint8, contiguous along N
    w_ptrs = w_ptr + offs_kh[:, None] * N + offs_n[None, :]

    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for k0 in range(0, K, BK):
        b = tl.load(w_ptrs)                                           # (BK//2, BN) u8
        lo = (b & 0xF).to(tl.float32)
        hi = ((b >> 4) & 0xF).to(tl.float32)
        g = k0 // 128  # one quant group per BK tile (BK == GROUP_SIZE)
        s = tl.load(s_ptr + g * N + offs_n).to(tl.float32)            # (BN,)
        z = tl.load(z_ptr + g * N + offs_n).to(tl.float32)            # (BN,)
        w_even = ((lo - z[None, :]) * s[None, :]).to(tl.bfloat16)
        w_odd = ((hi - z[None, :]) * s[None, :]).to(tl.bfloat16)
        x_even = tl.load(x_even_ptrs, mask=mask_m[:, None], other=0.0)
        x_odd = tl.load(x_odd_ptrs, mask=mask_m[:, None], other=0.0)
        acc = tl.dot(x_even, w_even, acc)
        acc = tl.dot(x_odd, w_odd, acc)
        x_even_ptrs += BK
        x_odd_ptrs += BK
        w_ptrs += (BK // 2) * N

    y = acc.to(tl.bfloat16)
    y_ptrs = y_ptr + offs_m[:, None] * N + offs_n[None, :]
    tl.store(y_ptrs, y, mask=mask_m[:, None])


def _pick_config(M: int, N: int):
    # (BM, BN, num_warps, num_stages)
    if M <= 16:
        return 16, 128, 4, 4
    if M <= 32:
        return 32, 128, 4, 4
    return 64, 256, 8, 4


def w4a16_gemm(x, w_q, scales, zeros, group_size: int = GROUP_SIZE):
    M, K = x.shape
    Khalf, N = w_q.shape
    assert Khalf * 2 == K
    BM, BN, nw, ns = _pick_config(M, N)
    y = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    grid = (triton.cdiv(M, BM) * triton.cdiv(N, BN),)
    _w4a16_gemm_kernel[grid](
        x, w_q, scales, zeros, y,
        M, N, K,
        BM=BM, BN=BN, BK=group_size, GROUP_M=8,
        num_warps=nw, num_stages=ns,
    )
    return y


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K % group_size == 0
        assert K % 2 == 0
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size
        # Buffers match the reference exactly so load_state_dict(strict=True) works.
        self.register_buffer("w_q", torch.empty((K // 2, N), dtype=torch.uint8))
        self.register_buffer("scales", torch.empty((n_groups, N), dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.empty((n_groups, N), dtype=torch.bfloat16))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return w4a16_gemm(x, self.w_q, self.scales, self.zeros, self.group_size)


M = 1
N = 12288
K = 4096


def get_inputs():
    x = torch.randn(M, K, dtype=torch.bfloat16)
    return [x]


def get_init_inputs():
    return [M, N, K]