from __future__ import annotations

import math

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_gemm(
    xe_ptr, xo_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
    M, N, K,
    sxe, sN_w, sN_s,
    s_out_m, s_out_n,
    HALF: tl.constexpr, N_GROUPS: tl.constexpr,
    TM: tl.constexpr, TN: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    m0 = pid_m * TM
    n0 = pid_n * TN
    ms = m0 + tl.arange(0, TM)
    ns = n0 + tl.arange(0, TN)
    m_mask = ms < M
    n_mask = ns < N

    acc = tl.zeros((TM, TN), dtype=tl.float32)

    for g in range(N_GROUPS):
        sg = tl.load(s_ptr + g * sN_s + ns, mask=n_mask, other=0.0)
        zg = tl.load(z_ptr + g * sN_s + ns, mask=n_mask, other=0.0)

        k0 = g * HALF
        tk = tl.arange(0, HALF)

        w_bytes = tl.load(wq_ptr + (k0 + tk[:, None]) * sN_w + ns[None, :],
                          mask=n_mask[None, :], other=0)
        ev = (w_bytes & 0x0F).to(tl.bfloat16)
        od = ((w_bytes.to(tl.uint8) >> 4) & 0x0F).to(tl.bfloat16)
        we = (ev - zg) * sg
        wo = (od - zg) * sg

        k_idx = k0 + tk[None, :]
        rxe = tl.load(xe_ptr + ms[:, None] * sxe + k_idx,
                      mask=m_mask[:, None], other=0.0)
        rxo = tl.load(xo_ptr + ms[:, None] * sxe + k_idx,
                      mask=m_mask[:, None], other=0.0)

        acc += tl.dot(rxe, we)
        acc += tl.dot(rxo, wo)

    tl.store(out_ptr + ms[:, None] * s_out_m + ns[None, :],
             acc.to(tl.bfloat16), mask=m_mask[:, None] & n_mask[None, :])


# Per-(M,N) best tile configuration from sweep.
_TILE = {
    (1, 4096): (8, 32),
    (1, 12288): (1, 64),
    (16, 14336): (16, 64),
    (32, 12288): (32, 64),
    (256, 12288): (64, 64),
}


def _pick_tile(M: int, N: int):
    return _TILE.get((M, N), (64, 64))


@torch.no_grad()
def w4a16_gemm_reuse(x, w_q, scales, zeros, xe, xo, out):
    M, K = x.shape
    N = w_q.shape[1]
    TM, TN = _pick_tile(M, N)
    xe[:, :] = x[:, 0::2]
    xo[:, :] = x[:, 1::2]

    HALF = GROUP_SIZE // 2
    N_GROUPS = K // GROUP_SIZE
    grid = (math.ceil(M / TM), math.ceil(N / TN))
    _w4a16_gemm[grid](
        xe, xo, w_q, scales, zeros, out,
        M, N, K,
        xe.stride(0), w_q.stride(0), scales.stride(0),
        out.stride(0), out.stride(1),
        HALF=HALF, N_GROUPS=N_GROUPS, TM=TM, TN=TN,
    )
    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_g = (-w_min / scales).round().clamp(0, 15)
        w_q = ((w_g / scales) + zeros_g).round().clamp(0, 15).to(torch.uint8).view(K, N)

        scales_2d = scales.squeeze(1).to(torch.bfloat16)
        zeros_2d = zeros_g.squeeze(1).to(torch.bfloat16)
        Kh = K // 2
        lo = w_q[0::2].to(torch.uint8) & 0xF
        hi = w_q[1::2].to(torch.uint8) & 0xF
        w_packed = (lo | (hi << 4)).contiguous()

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

        dev = torch.device('cuda')
        self._buf_xe = torch.empty(M, Kh, dtype=torch.bfloat16, device=dev)
        self._buf_xo = torch.empty(M, Kh, dtype=torch.bfloat16, device=dev)
        self._buf_out = torch.empty(M, N, dtype=torch.bfloat16, device=dev)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return w4a16_gemm_reuse(x, self.w_q, self.scales, self.zeros,
                                self._buf_xe, self._buf_xo, self._buf_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]