"""W4A16 weight-only quantized GEMM (AWQ/GPTQ-style, group_size=128, bf16).

Fused unpack + dequant + GEMM. V1: Triton kernel for all shapes.

Layout trick: instead of interleaving nibbles along K in registers, load the
packed (K//2, N) uint8 tile once and do TWO dots per group:
    acc += x_even @ dequant(lo_nibbles) + x_odd @ dequant(hi_nibbles)
Each tile spans exactly one 128-row group (64 packed rows), so scales/zeros
are scalar-per-column within the tile and broadcast cleanly.
"""
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.jit
def _w4a16_kernel(
    X, WQ, SC, ZO, OUT,
    M, N, K,
    BM: tl.constexpr, BN: tl.constexpr,
    GROUP_M: tl.constexpr, NSTAGE: tl.constexpr,
):
    # Each k-tile covers KP packed rows = 2*KP k values = one 128 group (KP=64).
    KP: tl.constexpr = 64
    pid = tl.program_id(0)
    grid_m = tl.cdiv(M, BM)
    grid_n = tl.cdiv(N, BN)
    width = GROUP_M * grid_n
    group_id = pid // width
    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
    pid_m = group_id * GROUP_M + (pid % group_size)
    pid_n = (pid % width) // group_size

    rm = pid_m * BM + tl.arange(0, BM)
    rn = pid_n * BN + tl.arange(0, BN)
    rk = tl.arange(0, KP)
    rk2 = tl.arange(0, 2 * KP)
    m_mask = rm < M

    KH = K // 2
    x_ptr = X + rm[:, None] * K + rk2[None, :]
    w_ptr = WQ + rk[:, None] * N + rn[None, :]
    g_ptr = SC + rn  # + group*N as we advance
    z_ptr = ZO + rn

    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for g in tl.range(0, KH, KP, num_stages=NSTAGE):
        # x tile: (BM, 2*KP) contiguous k, split even/odd k -> two dots
        a = tl.load(x_ptr, mask=m_mask[:, None], other=0.0)
        xe, xo = tl.split(tl.reshape(a, (BM, KP, 2)))
        w = tl.load(w_ptr)
        s = tl.load(g_ptr)
        z = tl.load(z_ptr)
        lo = (w & 0xF).to(tl.bfloat16)
        hi = ((w >> 4) & 0xF).to(tl.bfloat16)
        b_e = (lo - z) * s
        b_o = (hi - z) * s
        acc = tl.dot(xe, b_e, acc)
        acc = tl.dot(xo, b_o, acc)
        x_ptr += 2 * KP
        w_ptr += KP * N
        g_ptr += N
        z_ptr += N

    o_ptr = OUT + rm[:, None] * N + rn[None, :]
    tl.store(o_ptr, acc.to(tl.bfloat16), mask=m_mask[:, None])


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
        Kh = K // 2
        self.register_buffer("w_q", torch.zeros(Kh, 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))

        # config pick (tuned later)
        if M <= 16:
            self.BM, self.BN, self.nw, self.ns = 16, 128, 4, 4
        elif M <= 32:
            self.BM, self.BN, self.nw, self.ns = 32, 128, 4, 4
        else:
            self.BM, self.BN, self.nw, self.ns = 64, 128, 4, 3

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M = x.shape[0]
        out = torch.empty((M, self.N), dtype=torch.bfloat16, device=x.device)
        grid = (triton.cdiv(M, self.BM) * triton.cdiv(self.N, self.BN),)
        _w4a16_kernel[grid](
            x, self.w_q, self.scales, self.zeros, out,
            M, self.N, self.K,
            BM=self.BM, BN=self.BN, GROUP_M=8, NSTAGE=self.ns,
            num_warps=self.nw,
        )
        return out


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]