"""W4A16 weight-only int4 quantized GEMM for B200 (SM100 Blackwell).

AWQ/GPTQ-style asymmetric int4 with explicit zero-points and per-group bf16
scales. Fuses unpack + dequant + GEMM in a single Triton kernel pass so the
int4 weight stream is read from HBM exactly once.

Packing: w_q is (K//2, N) uint8. Low nibble of byte kh = weight at even K row
(2*kh); high nibble = odd K row (2*kh+1). Group size 128 along K, so 64 packed
rows per group. scales/zeros are (K//128, N) bf16.

Dequant per group: w_bf[k,n] = (nibble[k,n] - zeros[k//128,n]) * scales[k//128,n]
"""
from __future__ import annotations

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

GROUP_SIZE = 128


def _cfg(bm, bn, w, s):
    return triton.Config({"BLOCK_M": bm, "BLOCK_N": bn}, num_warps=w, num_stages=s)


_CONFIGS = [
    _cfg(16, 128, 4, 3),
    _cfg(16, 256, 4, 3),
    _cfg(16, 256, 8, 4),
    _cfg(32, 128, 4, 3),
    _cfg(32, 256, 8, 4),
    _cfg(64, 128, 4, 4),
    _cfg(64, 256, 8, 4),
    _cfg(128, 128, 8, 4),
    _cfg(128, 256, 8, 4),
]


@triton.autotune(configs=_CONFIGS, key=["M", "N", "K"])
@triton.jit
def _w4a16_kernel(
    x_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wq_kh, stride_wq_n,
    stride_s_g, stride_s_n,
    stride_z_g, stride_z_n,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, GK: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    m_mask = offs_m < M
    n_mask = offs_n < N

    KH_PER_G: tl.constexpr = GK // 2  # 64 packed rows per group
    n_groups = K // GK
    offs_kh = tl.arange(0, KH_PER_G)  # 0..63

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for g in range(n_groups):
        kh = g * KH_PER_G + offs_kh  # (64,) packed rows for this group
        # Packed int4 weights: (64, BLOCK_N)
        wq = tl.load(
            wq_ptr + kh[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n,
            mask=n_mask[None, :], other=0,
        ).to(tl.int32)
        lo = (wq & 0xF).to(tl.float32)         # even K row (2*kh)
        hi = ((wq >> 4) & 0xF).to(tl.float32)  # odd  K row (2*kh+1)

        s = tl.load(s_ptr + g * stride_s_g + offs_n * stride_s_n,
                    mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_ptr + g * stride_z_g + offs_n * stride_z_n,
                    mask=n_mask, other=0.0).to(tl.float32)

        lo_deq = ((lo - z[None, :]) * s[None, :]).to(tl.bfloat16)
        hi_deq = ((hi - z[None, :]) * s[None, :]).to(tl.bfloat16)

        # Activations for even/odd K rows.
        k_even = 2 * kh
        k_odd = 2 * kh + 1
        x_even = tl.load(
            x_ptr + offs_m[:, None] * stride_xm + k_even[None, :] * stride_xk,
            mask=m_mask[:, None], other=0.0,
        ).to(tl.bfloat16)
        x_odd = tl.load(
            x_ptr + offs_m[:, None] * stride_xm + k_odd[None, :] * stride_xk,
            mask=m_mask[:, None], other=0.0,
        ).to(tl.bfloat16)

        acc += tl.dot(x_even, lo_deq, out_dtype=tl.float32)
        acc += tl.dot(x_odd, hi_deq, out_dtype=tl.float32)

    out = acc.to(tl.bfloat16)
    tl.store(
        out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on,
        out, mask=m_mask[:, None] & n_mask[None, :],
    )


def _w4a16_gemm(x, w_q, scales, zeros, N, K):
    M = x.shape[0]
    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"]))
    _w4a16_kernel[grid](
        x, w_q, scales, zeros, out,
        M, N, K,
        x.stride(0), x.stride(1),
        w_q.stride(0), w_q.stride(1),
        scales.stride(0), scales.stride(1),
        zeros.stride(0), zeros.stride(1),
        out.stride(0), out.stride(1),
        GK=GROUP_SIZE,
    )
    return out


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 and K % 2 == 0
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size
        self.register_buffer("w_q", torch.zeros(K // 2, N, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(n_groups, N, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n_groups, N, dtype=torch.bfloat16))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.contiguous()
        return _w4a16_gemm(x, self.w_q, self.scales, self.zeros, self.N, self.K)


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]