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

Fuses int4 unpack + per-group dequant + GEMM in a single Triton kernel, modeled
on the group-quant mixed-input matmul pattern (see torchao hqq _mixed_mm_kernel):
  * packed uint8 weights (low nibble = even-K, high nibble = odd-K)
  * dequant w = (q - z) * s with per-group (group=128) bf16 scales/zeros
  * a single tl.dot per K-iteration with fp32 accumulation
  * group-M L2 swizzle, autotuned over IO-bound (small M) and compute-bound tiles
"""
from __future__ import annotations

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

GROUP_SIZE = 128


def _io_configs():
    cfgs = []
    for num_stages in (4, 5):
        for bn in (64, 128, 256):
            for bk in (64, 128):
                warps = 4 if bn <= 128 else 8
                cfgs.append(Config({"BLOCK_M": 16, "BLOCK_N": bn, "BLOCK_K": bk,
                                    "GROUP_M": 8}, num_stages=num_stages, num_warps=warps))
    for num_stages in (4, 5):
        cfgs.append(Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128,
                            "GROUP_M": 8}, num_stages=num_stages, num_warps=4))
        cfgs.append(Config({"BLOCK_M": 32, "BLOCK_N": 256, "BLOCK_K": 128,
                            "GROUP_M": 8}, num_stages=num_stages, num_warps=8))
    return cfgs


def _compute_configs():
    cfgs = []
    for bm, bn, bk, st, w in (
        (64, 128, 128, 4, 4),
        (64, 256, 128, 4, 4),
        (128, 128, 128, 4, 4),
        (128, 256, 128, 3, 8),
        (128, 128, 64, 4, 4),
        (256, 128, 128, 3, 8),
        (64, 128, 64, 5, 4),
        (128, 256, 64, 3, 8),
    ):
        cfgs.append(Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk,
                            "GROUP_M": 8}, num_stages=st, num_warps=w))
    return cfgs


@triton.jit
def _w4a16_kernel(
    A, B, scales_ptr, zeros_ptr, C,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    stride_sg, stride_sn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
    QGROUP: tl.constexpr,
):
    pid = tl.program_id(0)
    grid_m = tl.cdiv(M, BLOCK_M)
    grid_n = tl.cdiv(N, BLOCK_N)
    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

    ram = tl.max_contiguous(tl.multiple_of(pid_m * BLOCK_M + tl.arange(0, BLOCK_M), BLOCK_M), BLOCK_M)
    rbn = tl.max_contiguous(tl.multiple_of(pid_n * BLOCK_N + tl.arange(0, BLOCK_N), BLOCK_N), BLOCK_N)

    rk_full = tl.arange(0, BLOCK_K)
    rk_half = tl.arange(0, BLOCK_K // 2)

    A += ram[:, None] * stride_am + rk_full[None, :] * stride_ak
    B += rk_half[:, None] * stride_bk + rbn[None, :] * stride_bn

    offs_scale_n = pid_n * BLOCK_N * stride_sn + tl.arange(0, BLOCK_N) * stride_sn

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

    for k in range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(A)
        qb = tl.load(B)

        g = (k * BLOCK_K) // QGROUP
        scales = tl.load(scales_ptr + offs_scale_n + g * stride_sg)
        zeros = tl.load(zeros_ptr + offs_scale_n + g * stride_sg)

        # unsigned int4 unpack (values 0..15)
        qb_lo = qb & 0xF
        qb_hi = qb >> 4
        dq_b = tl.join(
            qb_lo.to(tl.float16).to(tl.bfloat16),
            qb_hi.to(tl.float16).to(tl.bfloat16),
        ).permute(0, 2, 1).reshape(BLOCK_K, BLOCK_N)

        dq_b = (dq_b - zeros[None, :]) * scales[None, :]
        acc += tl.dot(a, dq_b, out_dtype=tl.float32, input_precision="ieee")

        A += BLOCK_K * stride_ak
        B += (BLOCK_K // 2) * stride_bk

    acc = acc.to(tl.bfloat16)
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    C = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(C, acc, mask=mask)


_w4a16_io = triton.autotune(_io_configs(), key=["M", "N", "K"])(_w4a16_kernel)
_w4a16_compute = triton.autotune(_compute_configs(), key=["M", "N", "K"])(_w4a16_kernel)


def _launch(kernel, x, wq, scales, zeros, out, M, N, K):
    grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
    kernel[grid](
        x, wq, scales, zeros, out,
        M, N, K,
        x.stride(0), x.stride(1),
        wq.stride(0), wq.stride(1),
        out.stride(0), out.stride(1),
        scales.stride(0), scales.stride(1),
        QGROUP=GROUP_SIZE,
    )


def w4a16_gemm(x, wq, scales, zeros, M, N, K):
    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
    # Dispatch: tiny-M is IO-bound; large-M is compute-bound.
    if M <= 64:
        _launch(_w4a16_io, x, wq, scales, zeros, out, M, N, K)
    else:
        _launch(_w4a16_compute, x, wq, scales, zeros, out, M, N, K)
    return out


class Model(nn.Module):
    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        self.register_buffer("w_q", torch.zeros((K // 2, N), dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros((K // 128, N), dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros((K // 128, 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.M, self.N, self.K)


M = 1
N = 12288
K = 4096


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


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