"""Fused W4A16 GEMM kernel using Triton.

AWQ/GPTQ-style asymmetric int4 with per-group bf16 scales and zero-points.
Fuses unpack + dequant + matmul in a single pass.
"""
from __future__ import annotations

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


# ---------------------------------------------------------------------------
# GEMV kernel for M=1 (decode / memory-bound), autotuned
# ---------------------------------------------------------------------------
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_N': 32, 'BLOCK_K': 32}, num_warps=4),
        triton.Config({'BLOCK_N': 32, 'BLOCK_K': 64}, num_warps=4),
        triton.Config({'BLOCK_N': 32, 'BLOCK_K': 128}, num_warps=4),
        triton.Config({'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=4),
        triton.Config({'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=4),
        triton.Config({'BLOCK_N': 64, 'BLOCK_K': 128}, num_warps=4),
        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=4),
        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=4),
        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=4),
        triton.Config({'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=8),
        triton.Config({'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=8),
        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=8),
        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=8),
        triton.Config({'BLOCK_N': 256, 'BLOCK_K': 128}, num_warps=8),
    ],
    key=['N', 'K'],
)
@triton.jit
def _w4a16_gemv_kernel(
    X_ptr,
    WQ_ptr,
    S_ptr,
    Z_ptr,
    Out_ptr,
    N: tl.constexpr,
    K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid = tl.program_id(0)
    off_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_n = off_n < N

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

    num_k_blocks = K // BLOCK_K
    for kb in range(num_k_blocks):
        off_k = kb * BLOCK_K + tl.arange(0, BLOCK_K)

        x_vals = tl.load(X_ptr + off_k, mask=off_k < K, other=0.0).to(tl.float32)

        off_kh = off_k // 2
        nibble = off_k % 2

        wq_ptrs = WQ_ptr + off_kh[:, None] * N + off_n[None, :]
        wq_vals = tl.load(wq_ptrs, mask=(off_kh[:, None] < K // 2) & mask_n[None, :], other=0)
        lo = wq_vals & 0xF
        hi = (wq_vals >> 4) & 0xF
        unpacked = tl.where(nibble[:, None] == 0, lo, hi).to(tl.float32)

        off_g = off_k // GROUP_SIZE
        s_ptrs = S_ptr + off_g[:, None] * N + off_n[None, :]
        z_ptrs = Z_ptr + off_g[:, None] * N + off_n[None, :]
        s_vals = tl.load(s_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)
        z_vals = tl.load(z_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)

        w_deq = (unpacked - z_vals) * s_vals
        acc += tl.sum(x_vals[:, None] * w_deq, axis=0)

    tl.store(Out_ptr + off_n, acc.to(tl.bfloat16), mask=mask_n)


# ---------------------------------------------------------------------------
# Batched GEMV: 1D grid over (M, N_tiles). Each program handles one row.
# ---------------------------------------------------------------------------
@triton.jit
def _w4a16_batched_gemv_kernel(
    X_ptr,
    WQ_ptr,
    S_ptr,
    Z_ptr,
    Out_ptr,
    N: tl.constexpr,
    K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    N_TILES: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid = tl.program_id(0)
    m = pid // N_TILES
    nt = pid % N_TILES

    off_n = nt * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_n = off_n < N

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

    num_k_blocks = K // BLOCK_K
    for kb in range(num_k_blocks):
        off_k = kb * BLOCK_K + tl.arange(0, BLOCK_K)

        x_ptrs = X_ptr + m * K + off_k
        x_vals = tl.load(x_ptrs, mask=off_k < K, other=0.0).to(tl.float32)

        off_kh = off_k // 2
        nibble = off_k % 2

        wq_ptrs = WQ_ptr + off_kh[:, None] * N + off_n[None, :]
        wq_vals = tl.load(wq_ptrs, mask=(off_kh[:, None] < K // 2) & mask_n[None, :], other=0)
        lo = wq_vals & 0xF
        hi = (wq_vals >> 4) & 0xF
        unpacked = tl.where(nibble[:, None] == 0, lo, hi).to(tl.float32)

        off_g = off_k // GROUP_SIZE
        s_ptrs = S_ptr + off_g[:, None] * N + off_n[None, :]
        z_ptrs = Z_ptr + off_g[:, None] * N + off_n[None, :]
        s_vals = tl.load(s_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)
        z_vals = tl.load(z_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)

        w_deq = (unpacked - z_vals) * s_vals
        acc += tl.sum(x_vals[:, None] * w_deq, axis=0)

    tl.store(Out_ptr + m * N + off_n, acc.to(tl.bfloat16), mask=mask_n)


# ---------------------------------------------------------------------------
# GEMM kernel using tl.dot for tensor cores, autotuned
# ---------------------------------------------------------------------------
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 32, 'BLOCK_K': 64}, num_stages=2, num_warps=2),
        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=2, num_warps=2),
        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'BLOCK_K': 64}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 128}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 128}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 128}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_stages=3, num_warps=8),
    ],
    key=['M_total', 'N', 'K'],
)
@triton.jit
def _w4a16_gemm_kernel(
    X_ptr,
    WQ_ptr,
    S_ptr,
    Z_ptr,
    Out_ptr,
    M_total: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    off_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    mask_m = off_m < M_total
    mask_n = off_n < N

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

    num_k_blocks = K // BLOCK_K
    for kb in range(num_k_blocks):
        off_k = kb * BLOCK_K + tl.arange(0, BLOCK_K)

        x_ptrs = X_ptr + off_m[:, None] * K + off_k[None, :]
        x_vals = tl.load(x_ptrs, mask=mask_m[:, None] & (off_k[None, :] < K), other=0.0)

        off_kh = off_k // 2
        nibble = off_k % 2

        wq_ptrs = WQ_ptr + off_kh[:, None] * N + off_n[None, :]
        wq_vals = tl.load(wq_ptrs, mask=(off_kh[:, None] < K // 2) & mask_n[None, :], other=0)
        lo = wq_vals & 0xF
        hi = (wq_vals >> 4) & 0xF
        unpacked = tl.where(nibble[:, None] == 0, lo, hi).to(tl.float32)

        off_g = off_k // GROUP_SIZE
        s_ptrs = S_ptr + off_g[:, None] * N + off_n[None, :]
        z_ptrs = Z_ptr + off_g[:, None] * N + off_n[None, :]
        s_vals = tl.load(s_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)
        z_vals = tl.load(z_ptrs, mask=(off_g[:, None] < K // GROUP_SIZE) & mask_n[None, :], other=0.0).to(tl.float32)

        w_deq = (unpacked - z_vals) * s_vals
        w_bf = w_deq.to(tl.bfloat16)

        acc = tl.dot(x_vals, w_bf, acc=acc)

    out_ptrs = Out_ptr + off_m[:, None] * N + off_n[None, :]
    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=mask_m[:, None] & mask_n[None, :])


def _w4a16_forward(x, w_q, scales, zeros, K, group_size):
    M_val = x.shape[0]
    N_val = w_q.shape[1]
    out = torch.empty((M_val, N_val), dtype=torch.bfloat16, device=x.device)

    if M_val == 1:
        # GEMV for M=1 decode
        grid = lambda META: (triton.cdiv(N_val, META['BLOCK_N']),)
        _w4a16_gemv_kernel[grid](
            x, w_q, scales, zeros, out,
            N=N_val, K=K, GROUP_SIZE=group_size,
        )
    else:
        # GEMM with tensor cores for M > 16
        grid = lambda META: (
            triton.cdiv(M_val, META['BLOCK_M']),
            triton.cdiv(N_val, META['BLOCK_N']),
        )
        _w4a16_gemm_kernel[grid](
            x, w_q, scales, zeros, out,
            M_total=M_val, N=N_val, K=K, GROUP_SIZE=group_size,
        )

    return out


class Model(nn.Module):
    """W4A16 GEMM with fused Triton kernel."""

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _w4a16_forward(x, self.w_q, self.scales, self.zeros, self.K, self.group_size)


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]