"""W4A16 weight-only int4 GEMM (AWQ/GPTQ-style asymmetric int4) for SM120.

Fused unpack + dequant + GEMM in a single Triton pass.

Layout / scheme (must match reference.py exactly for load_state_dict strict=True):
  x:      (M, K)        bf16
  w_q:    (K//2, N)     uint8   -- two int4 packed per byte:
                                   low nibble  = even-K row  (k = 2*kh)
                                   high nibble = odd-K  row  (k = 2*kh+1)
  scales: (K//128, N)   bf16
  zeros:  (K//128, N)   bf16    -- stored as float zero-point
  out:    (M, N)        bf16

Dequant (per group of 128 along K):
  w_bf[k,n] = (unpack(w_q)[k,n] - zeros[k//128,n]) * scales[k//128,n]

Trick: the K reduction is order-independent, so split it into even/odd k:
  y = sum_kh x[2kh]*w_lo[kh] + sum_kh x[2kh+1]*w_hi[kh]
with w_lo = low nibble, w_hi = high nibble. Skips int4 interleaving.
BLOCK_K = 128 so one K-block == one quant group == one scale/zero vector.
"""
from __future__ import annotations

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

GROUP_SIZE = 128


def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    K, N = w_q.shape
    assert K % 2 == 0
    lo = w_q[[REDACTED: IP]].to(torch.uint8) & 0xF
    hi = w_q[[REDACTED: IP]].to(torch.uint8) & 0xF
    return (lo | (hi << 4)).contiguous()


def _direct_configs():
    cfgs = []
    for bm in (16, 32, 64, 128, 256):
        for bn in (64, 128, 256):
            for w, s in ((4, 3), (4, 4), (8, 3), (8, 4)):
                cfgs.append(
                    triton.Config(
                        {"BLOCK_M": bm, "BLOCK_N": bn, "GROUP_M": 8},
                        num_warps=w, num_stages=s,
                    )
                )
    return cfgs


@triton.autotune(configs=_direct_configs(), key=["M", "N", "K"])
@triton.jit
def _w4a16_direct(
    x_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wkh, stride_wn,
    stride_sg, stride_sn,
    stride_zg, stride_zn,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, GROUP_M: tl.constexpr,
):
    BLOCK_K: tl.constexpr = 128
    HK: tl.constexpr = BLOCK_K // 2  # 64 packed bytes per K-block

    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    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 * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_kh = tl.arange(0, HK)

    m_mask = offs_m < M
    n_mask = offs_n < N

    num_k_blocks = tl.cdiv(K, BLOCK_K)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for kb in range(0, num_k_blocks):
        k0 = kb * BLOCK_K
        kh0 = k0 // 2
        offs_ke = k0 + 2 * offs_kh
        offs_ko = offs_ke + 1

        x_e = tl.load(
            x_ptr + offs_m[:, None] * stride_xm + offs_ke[None, :] * stride_xk,
            mask=m_mask[:, None], other=0.0,
        )
        x_o = tl.load(
            x_ptr + offs_m[:, None] * stride_xm + offs_ko[None, :] * stride_xk,
            mask=m_mask[:, None], other=0.0,
        )
        wq = tl.load(
            wq_ptr + (kh0 + offs_kh)[:, None] * stride_wkh + offs_n[None, :] * stride_wn,
            mask=n_mask[None, :], other=0,
        )
        lo = (wq & 0xF).to(tl.float32)
        hi = ((wq >> 4) & 0xF).to(tl.float32)
        s = tl.load(s_ptr + kb * stride_sg + offs_n * stride_sn, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_ptr + kb * stride_zg + offs_n * stride_zn, mask=n_mask, other=0.0).to(tl.float32)

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

        acc += tl.dot(x_e, w_lo)
        acc += tl.dot(x_o, w_hi)

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


def _w4a16_gemm(x, wq, scales, zeros, N):
    M, K = x.shape
    x = x.contiguous()
    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)

    def grid(meta):
        return (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)

    _w4a16_direct[grid](
        x, wq, scales, zeros, out,
        M, N, K,
        x.stride(0), x.stride(1),
        wq.stride(0), wq.stride(1),
        scales.stride(0), scales.stride(1),
        zeros.stride(0), zeros.stride(1),
        out.stride(0), out.stride(1),
    )
    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
        assert K % 2 == 0
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size

        torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
        w_full = torch.randn(K, N, dtype=torch.float32) * 0.02
        w_g = w_full.view(n_groups, group_size, N)
        w_min = w_g.min(dim=1, keepdim=True).values
        w_max = w_g.max(dim=1, keepdim=True).values
        scales = (w_max - w_min).clamp_min(1e-8) / 15.0
        zeros = (-w_min / scales).round().clamp(0, 15)
        w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)

        self.register_buffer("w_q", _pack_int4(w_q))
        self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16))
        self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16))

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


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]