"""FP8 e4m3 GEMM solution — Triton kernel.

y = ((x @ w.T) * weight_scale).to(bf16)  where x, w are fp8_e4m3.

This implements a real fp8 x fp8 tensor-core MMA via Triton tl.dot
(fp32 accumulate) and applies the per-output-channel dequant scale.
"""
import torch
import torch.nn as nn
import triton
import triton.language as tl


# --- Triton FP8 GEMM kernel ----------------------------------------------------
def _gen_configs():
    cfgs = []
    # Standard compute-bound large tiles.
    for bm, bn, bk, ns, nw in [
        (128, 128, 128, 4, 8),
        (128, 256, 128, 3, 8),
        (256, 128, 128, 3, 8),
        (128, 128, 64,  4, 4),
        (256, 256, 64,  3, 8),
        (128, 256, 64,  3, 8),
        (256, 128, 64,  3, 8),
        (128, 128, 128, 3, 4),
        (64,  128, 128, 4, 4),
        (128, 64,  128, 4, 4),
    ]:
        cfgs.append(triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8},
                                  num_stages=ns, num_warps=nw))
    # Skinny-M (decode): BLOCK_M small or padded.
    for bm, bn, bk, ns, nw in [
        (16,  256, 128, 4, 4),
        (16,  128, 128, 4, 4),
        (32,  256, 128, 4, 8),
        (32,  128, 128, 4, 4),
        (64,  256, 128, 4, 8),
    ]:
        cfgs.append(triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": 8},
                                  num_stages=ns, num_warps=nw))
    # Wider GROUP_M variants for better L2 reuse.
    for bm, bn, bk, ns, nw in [
        (128, 128, 128, 4, 8),
        (256, 128, 128, 3, 8),
    ]:
        for gm in (4, 16):
            cfgs.append(triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_M": gm},
                                      num_stages=ns, num_warps=nw))
    return cfgs


@triton.autotune(
    configs=_gen_configs(),
    key=["M", "N", "K"],
)
@triton.jit
def fp8_gemm_kernel(
    A_ptr, B_ptr, C_ptr, S_ptr,
    M, N, K,
    sa, sb,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    # Group order for L2 reuse: walk BLOCK_M groups of BLOCK_M*GROUP_M rows,
    # with all BLOCK_N tiles of each row-band, before moving to the next.
    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 % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B_ptr + offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        # Masked loads for the final partial tile.
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
        # fp8 x fp8 -> fp32 accumulate.
        acc = tl.dot(a, b.T, acc, out_dtype=tl.float32)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    # Per-output-channel dequant: scale (M, N) by weight_scale (N,).
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    scale = tl.load(S_ptr + offs_n, mask=offs_n < N, other=0.0)
    acc = acc * scale[None, :]

    # Write as bf16.
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = C_ptr + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, acc.to(tl.bfloat16), mask=c_mask)


def fp8_gemm_triton(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor:
    """x: (M, K) fp8_e4m3; weight: (N, K) fp8_e4m3; weight_scale: (N,) f32.
    Returns bf16 (M, N).
    """
    assert x.is_cuda and weight.is_cuda
    M, K = x.shape
    N = weight.shape[0]
    assert weight.shape == (N, K)
    assert weight_scale.shape == (N,)

    # Make sure layouts are right (x is row-major, weight row-major, output row-major).
    y = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)

    grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
    fp8_gemm_kernel[grid](
        x, weight, y, weight_scale,
        M, N, K,
        1.0, 1.0,  # scales (unused — kernel takes fp8 directly and applies weight_scale)
        x.stride(0), x.stride(1),
        weight.stride(0), weight.stride(1),
        y.stride(0), y.stride(1),
    )
    return y


# --- Module: mirrors reference.Model interface -------------------------------
class Model(nn.Module):
    """y = ((x @ w.T) * weight_scale).to(bf16).

    Same buffers as reference.Model so the graded state_dict loads.
    """

    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.register_buffer("weight", torch.empty(N, K, dtype=torch.float8_e4m3fn))
        self.register_buffer("weight_scale", torch.empty(N, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return fp8_gemm_triton(x, self.weight, self.weight_scale)


# Local module-level shims (reference.py style — overwritten by benchmark/check).
M = 4096
N = 4096
K = 4096


def get_inputs():
    x = (torch.rand(M, K) * 8 - 4).to(torch.float8_e4m3fn)
    return [x]


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