"""W4A16 GEMM — fused int4 unpack + bf16 matmul in Triton.

Two kernels:
  * gemv_w4a16_kernel: M=1 decode path. Memory-bound on the int4 weight read.
    Programs are tiled over (N, K_group) and write partial sums to an fp32
    buffer via atomic add. A tiny finalize kernel casts fp32 -> bf16.
  * gemm_w4a16_kernel: M >= 16 prefill path. Standard matmul tile with the
    dequant fused into the K-loop. BLOCK_K = group_size (128) so each
    iteration loads exactly one scale/zero row.
"""
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


# ---------------------------------------------------------------------------
# Decoding kernel (M == 1) — memory-bound path
# ---------------------------------------------------------------------------
@triton.jit
def _gemv_w4a16_kernel(
    x_ptr,            # (K,) bf16 — x activation
    w_ptr,            # (K//2, N) uint8 packed
    s_ptr,            # (K//128, N) bf16 — scales
    z_ptr,            # (K//128, N) bf16 — zeros
    acc_ptr,          # (N,) fp32 — accumulator buffer (atomic add target)
    K, N,
    stride_wk, stride_wn,
    stride_sk, stride_sn,
    stride_zk, stride_zn,
    BLOCK_N: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
):
    """One program = one (N_BLOCK, K_group) tile. Computes the partial dot
    product of x[group] with the dequantized W[group, n_block] and atomic-adds
    it into the per-N fp32 accumulator. K_groups independent programs sum
    into the same N accumulators.
    """
    pid_n = tl.program_id(0)
    pid_g = tl.program_id(1)

    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    # Load 64 packed rows for this group, BLOCK_N cols each
    k_in_group = tl.arange(0, GROUP_SIZE)
    k_offs = pid_g * GROUP_SIZE + k_in_group
    kp_offs = k_offs // 2  # (GROUP_SIZE,)

    w_ptrs = w_ptr + kp_offs[:, None] * stride_wk + n_offs[None, :] * stride_wn
    packed = tl.load(w_ptrs, mask=n_mask[None, :], other=0)  # (GROUP_SIZE, BLOCK_N) uint8

    # Unpack int4: even K = lo nibble, odd K = hi nibble
    is_odd = (k_offs & 1)[:, None] == 1
    w_unp = tl.where(is_odd, (packed >> 4) & 0xF, packed & 0xF).to(tl.float32)

    # Load scale/zero for this group: (BLOCK_N,)
    s = tl.load(s_ptr + pid_g * stride_sk + n_offs * stride_sn, mask=n_mask, other=0.0).to(tl.float32)
    z = tl.load(z_ptr + pid_g * stride_zk + n_offs * stride_zn, mask=n_mask, other=0.0).to(tl.float32)

    # Dequant
    w_bf = (w_unp - z[None, :]) * s[None, :]

    # Load x for this group
    x_vals = tl.load(x_ptr + k_offs).to(tl.float32)  # (GROUP_SIZE,)

    # Dot: (GROUP_SIZE,) * (GROUP_SIZE, BLOCK_N) -> (BLOCK_N,)
    partial = tl.sum(x_vals[:, None] * w_bf, axis=0)

    # Atomic add to fp32 accumulator
    tl.atomic_add(acc_ptr + n_offs, partial, mask=n_mask)


@triton.jit
def _gemv_finalize_kernel(
    acc_ptr,    # (N,) fp32
    out_ptr,    # (N,) bf16
    N,
    BLOCK_N: tl.constexpr,
):
    pid = tl.program_id(0)
    offs = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    mask = offs < N
    v = tl.load(acc_ptr + offs, mask=mask, other=0.0)
    tl.store(out_ptr + offs, v.to(tl.bfloat16), mask=mask)


def gemv_w4a16(x: torch.Tensor, w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor) -> torch.Tensor:
    assert x.dim() == 2 and x.shape[0] == 1
    K = x.shape[1]
    N = w_q.shape[1]
    out = torch.empty((1, N), dtype=torch.bfloat16, device=x.device)
    acc = torch.zeros((1, N), dtype=torch.float32, device=x.device)

    BLOCK_N = 32
    grid = (triton.cdiv(N, BLOCK_N), K // GROUP_SIZE)
    _gemv_w4a16_kernel[grid](
        x, w_q, scales, zeros, acc,
        K, N,
        w_q.stride(0), w_q.stride(1),
        scales.stride(0), scales.stride(1),
        zeros.stride(0), zeros.stride(1),
        BLOCK_N=BLOCK_N, GROUP_SIZE=GROUP_SIZE,
        num_warps=2, num_stages=2,
    )

    grid_f = (triton.cdiv(N, 256),)
    _gemv_finalize_kernel[grid_f](acc, out, N, BLOCK_N=256, num_warps=2)

    return out


# ---------------------------------------------------------------------------
# Prefill kernel (M >= 16)
# ---------------------------------------------------------------------------
@triton.jit
def _gemm_w4a16_kernel(
    x_ptr,            # (M, K) bf16
    w_ptr,            # (K//2, N) uint8
    s_ptr,            # (K//128, N) bf16
    z_ptr,            # (K//128, N) bf16
    o_ptr,            # (M, N) bf16
    M, N, K,
    stride_xm, stride_xk,
    stride_wk, stride_wn,
    stride_sk, stride_sn,
    stride_zk, stride_zn,
    stride_om, stride_on,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
):
    """Standard matmul tile with dequant in the inner loop. BLOCK_K is fixed at
    GROUP_SIZE so each iteration loads exactly one row of scales/zeros.
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    m_offs = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    m_mask = m_offs < M
    n_mask = n_offs < N

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

    for g in tl.static_range(0, NUM_GROUPS):
        k_in_group = tl.arange(0, GROUP_SIZE)
        k_offs = g * GROUP_SIZE + k_in_group
        kp_offs = k_offs // 2

        # Load packed weights: (GROUP_SIZE, BLOCK_N)
        w_offs = kp_offs[:, None] * stride_wk + n_offs[None, :] * stride_wn
        packed = tl.load(w_ptr + w_offs, mask=n_mask[None, :], other=0)
        is_odd = (k_offs & 1)[:, None] == 1
        w_unp = tl.where(is_odd, (packed >> 4) & 0xF, packed & 0xF).to(tl.float32)

        # Load scales/zeros for this group: (BLOCK_N,)
        s_offs = g * stride_sk + n_offs * stride_sn
        z_offs = g * stride_zk + n_offs * stride_zn
        s = tl.load(s_ptr + s_offs, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_ptr + z_offs, mask=n_mask, other=0.0).to(tl.float32)

        # Dequant: (GROUP_SIZE, BLOCK_N)
        w_bf = (w_unp - z[None, :]) * s[None, :]

        # Load x: (BLOCK_M, GROUP_SIZE)
        x_offs = m_offs[:, None] * stride_xm + k_offs[None, :] * stride_xk
        x_vals = tl.load(x_ptr + x_offs, mask=m_mask[:, None], other=0.0).to(tl.float32)

        # MMA accumulate
        acc += tl.dot(x_vals, w_bf, out_dtype=tl.float32)

    out_offs = m_offs[:, None] * stride_om + n_offs[None, :] * stride_on
    tl.store(o_ptr + out_offs, acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


def gemm_w4a16(x: torch.Tensor, w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor) -> torch.Tensor:
    M, K = x.shape
    N = w_q.shape[1]
    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)

    BLOCK_M = 128 if M >= 128 else triton.next_power_of_2(M)
    BLOCK_N = 128

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    _gemm_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),
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
        GROUP_SIZE=GROUP_SIZE, NUM_GROUPS=K // GROUP_SIZE,
        num_warps=4, num_stages=3,
    )
    return out


# ---------------------------------------------------------------------------
# Reference-aligned interface
# ---------------------------------------------------------------------------
def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    K, N = w_q.shape
    lo = w_q[0::2].to(torch.uint8) & 0xF
    hi = w_q[1::2].to(torch.uint8) & 0xF
    return (lo | (hi << 4)).contiguous()


def _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor:
    Kh, N = w_packed.shape
    out = torch.empty((K, N), dtype=torch.uint8, device=w_packed.device)
    out[0::2] = w_packed & 0xF
    out[1::2] = (w_packed >> 4) & 0xF
    return out


class Model(nn.Module):
    """W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""

    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
        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)
        w_q = w_q.view(K, N)

        scales_2d = scales.squeeze(1).to(torch.bfloat16)
        zeros_2d = zeros.squeeze(1).to(torch.bfloat16)
        w_packed = _pack_int4(w_q)

        self.register_buffer("w_q", w_packed)
        self.register_buffer("scales", scales_2d)
        self.register_buffer("zeros", zeros_2d)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        M = x.shape[0]
        if M == 1:
            return gemv_w4a16(x, self.w_q, self.scales, self.zeros)
        return gemm_w4a16(x, self.w_q, self.scales, self.zeros)


# Defaults are required by the harness; get_init_inputs overrides per shape.
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]