# adapted from triton_kernels package
# original code https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/matmul_ogs_details/_matmul_ogs.py

import torch
import triton
import triton.language as tl
from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid
from aiter.ops.triton._triton_kernels.moe.activations import _swiglu


def matmul_launch_metadata(grid, kernel, args):
    ret = dict()
    M, N, K = None, args["N"], args["K"]
    Y, X, W = args["Y"], args["X"], args["W"]
    hist = args["ExptHist"]
    if hist is not None:
        n_rows = int(hist.float().mean())
        n_tokens = float(hist.sum())
        n_w_bytes = (W.numel() * W.element_size() // hist.numel()) * (hist > 0).sum()
    else:
        n_tokens = None
        n_w_bytes = W.numel() * W.element_size()

    def repr(s, x):
        return f"{s}={x}" if x is not None else f"E_{len(hist)}({s})={n_rows}"

    nbits = X.dtype.itemsize * 8
    ret["name"] = f"{kernel.name} [{repr('M', M)}, {repr('N', N)}, {repr('K', K)}]"
    gindx = args.get("GatherIndx", None)
    if gindx is not None:
        ret["name"] += "_layer1"
    else:
        ret["name"] += "_layer2"
    if args["B"] is not None:
        ret["name"] += "_bias"
    alpha = args.get("alpha", 0)
    act_red_n = args.get("ACTIVATION_REDUCTION_N", 1)
    if alpha != 0:
        if act_red_n == 1:
            ret["name"] += "_silu"
        else:
            ret["name"] += "_swiglu"

    fM = n_tokens
    fK = K if K is not None else n_tokens
    ret[f"flops{nbits}"] = 2.0 * fM * N * fK

    n_x_bytes = X.numel() * X.element_size()
    n_y_bytes = Y.numel() * Y.element_size()
    if hist is not None:
        assert n_tokens is not None
        n_expts_act = args["N_EXPTS_ACT"]

        if gindx is not None:
            # recreate inverse GatherIndx.
            dst = torch.full_like(gindx, -1)
            idx = torch.arange(len(gindx), device=gindx.device, dtype=torch.int32)
            mask = gindx != -1
            dst[gindx[mask]] = idx[mask]
            n_read_rows = (dst.view((-1, n_expts_act)) != -1).any(dim=1).sum()
        else:
            n_read_rows = n_tokens
        n_x_bytes = n_read_rows * X.shape[-1] * X.element_size()
        n_y_bytes = n_tokens * Y.shape[-1] * Y.element_size()
    ret["bytes"] = int(n_x_bytes + n_y_bytes + n_w_bytes)

    return ret


@triton.jit
def unshuffle_weights(w, BLOCK_N, BLOCK_K):
    w = w.trans()
    w = w.reshape(1, BLOCK_N // 16, BLOCK_K // 32, 2, 16, 16)
    w = w.permute(0, 1, 4, 2, 3, 5)
    w = w.reshape(BLOCK_N, BLOCK_K)
    w = w.trans()
    return w


@triton.jit(launch_metadata=matmul_launch_metadata)
def _moe_gemm_int8_smoothquant(
    Y,
    stride_y_k,
    stride_y_m,
    stride_y_n,
    X,
    stride_x_m,
    stride_x_k,
    XScale,
    stride_x_scale,
    W,
    stride_w_e,
    stride_w_k,
    stride_w_n,
    WScale,
    stride_w_scale_e,
    stride_w_scale_n,
    B,
    stride_b_e,  # Bias
    Gammas,
    N,
    K,
    # expt data
    GatherIndx,
    ExptHist,
    ExptOffs,
    ExptOffsSum,
    ExptData,
    # true grid size
    grid_m,
    grid_n,
    alpha,
    limit,
    ACTIVATION_REDUCTION_N: tl.constexpr,
    APPLY_ACTIVATION: tl.constexpr,
    ADD_RESIDUAL: tl.constexpr,
    # MoE config
    N_EXPTS_ACT: tl.constexpr,
    # optimization config
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
    PRESHUFFLED: tl.constexpr,
    EVEN_K: tl.constexpr,
    MASK_K_LIMIT: tl.constexpr,
    SPLIT_K: tl.constexpr,
    W_CACHE_MODIFIER: tl.constexpr,
    UPCAST_INDICES: tl.constexpr = False,
):
    """
    Int8 MoE GEMM with SmoothQuant support and per-token per-channel scaling.

    SmoothQuant formula:
        Y = (X * diag(s)^-1) @ (diag(s) * W)

    Where s is the smoothing factor

    Implementation:
        Y = (X @ W) * x_scale * w_scale

    Key parameters:
    - X is int8 activations [M, K] (quantized X * diag(s)^-1)
    - W is int8 weights [E, K, N] (quantized diag(s) * W)
    - x_scale is fp32 per-token scale [M] (dequant scale for X)
    - w_scale is fp32 per-output-channel scale [E, N] (dequant scale for W)

    Activation functions:
    - alpha=0: No activation
    - alpha==1, ADD_RESIDUAL=False: SiLU
    - alpha!=0: SwiGLU
    """
    # Assume positive strides for compiler hints
    tl.assume(stride_y_k >= 0)
    tl.assume(stride_y_m >= 0)
    tl.assume(stride_y_n >= 0)
    tl.assume(stride_x_m >= 0)
    tl.assume(stride_x_k >= 0)
    tl.assume(stride_w_e >= 0)
    tl.assume(stride_w_k >= 0)
    tl.assume(stride_w_n >= 0)
    if B is not None:
        tl.assume(stride_b_e >= 0)
    tl.assume(grid_m >= 0)
    tl.assume(grid_n >= 0)

    OUT_BLOCK_N: tl.constexpr = BLOCK_N // ACTIVATION_REDUCTION_N
    yN = N // ACTIVATION_REDUCTION_N

    pid = tl.program_id(0)
    padding_m: tl.constexpr = 0

    index_type: tl.constexpr = tl.int64 if UPCAST_INDICES else tl.int32

    unpadded_m = grid_m - padding_m
    tl.assume(unpadded_m >= 0)
    total_actual_tiles = unpadded_m * grid_n * SPLIT_K
    if padding_m > 0 and pid >= total_actual_tiles:
        return

    pid_emnk = pid
    pid_mnk = pid_emnk % (unpadded_m * grid_n * SPLIT_K)
    pid_k = pid_mnk % SPLIT_K
    pid_mn = pid_mnk // SPLIT_K
    pid_m, pid_n = pid_grid(pid_mn, unpadded_m, grid_n, GROUP_M)
    # For split-k, advance to the output k slice
    if SPLIT_K > 1:
        Y += pid_k.to(index_type) * stride_y_k
    # unpack expert data
    expt_data = tl.load(ExptData + pid_m)
    if expt_data == -1:
        return
    expt_id = expt_data & 0x0000FFFF
    block_id = expt_data >> 16
    M = tl.load(ExptHist + expt_id)
    start_m = tl.load(ExptOffs + expt_id)
    expt_id, block_id = expt_id.to(index_type), block_id.to(index_type)
    start_m = start_m.to(index_type)
    pid_n, pid_k = pid_n.to(index_type), pid_k.to(index_type)

    # A pointers
    offs_x_m = BLOCK_M * block_id + tl.arange(0, BLOCK_M)
    offs_x_m = tl.max_contiguous(tl.multiple_of(offs_x_m % M, BLOCK_M), BLOCK_M)
    if GatherIndx is None:
        X += start_m * stride_x_m
        XScale += start_m * stride_x_scale
    else:
        GatherIndx += start_m
        # no needs to bounds-check here because `offs_x_m` wraps around M dim
        offs_x_m = tl.load(GatherIndx + offs_x_m) // N_EXPTS_ACT
    offs_x_k = BLOCK_K * pid_k + tl.arange(0, BLOCK_K)
    XPtrs = (
        X
        + offs_x_m.to(index_type)[:, None] * stride_x_m
        + offs_x_k.to(index_type)[None, :] * stride_x_k
    )

    # B pointers
    if PRESHUFFLED:
        PACKED_BLOCK_N: tl.constexpr = BLOCK_N // 16
        PACKED_BLOCK_K: tl.constexpr = BLOCK_K * 16
        PACKED_N = N // 16
    else:
        PACKED_BLOCK_N: tl.constexpr = BLOCK_N
        PACKED_BLOCK_K: tl.constexpr = BLOCK_K
        PACKED_N = N

    offs_w_n = pid_n * PACKED_BLOCK_N + tl.arange(0, PACKED_BLOCK_N)
    offs_w_n = tl.max_contiguous(
        tl.multiple_of(offs_w_n % PACKED_N, PACKED_BLOCK_N),
        PACKED_BLOCK_N,
    )
    offs_w_k = pid_k * PACKED_BLOCK_K + tl.arange(0, PACKED_BLOCK_K)
    W += expt_id * stride_w_e
    WPtrs = W + (
        offs_w_k.to(index_type)[:, None] * stride_w_k
        + offs_w_n.to(index_type)[None, :] * stride_w_n
    )

    num_k_iter = tl.cdiv(K, BLOCK_K * SPLIT_K)
    if not EVEN_K:
        num_k_iter -= 1

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32)
    for k in range(num_k_iter):
        x = tl.load(XPtrs)
        w = tl.load(WPtrs, cache_modifier=W_CACHE_MODIFIER)
        if PRESHUFFLED:
            w = unshuffle_weights(w, BLOCK_N, BLOCK_K)

        acc += tl.dot(x, w, input_precision="ieee")

        XPtrs += (BLOCK_K * SPLIT_K) * stride_x_k
        WPtrs += (PACKED_BLOCK_K * SPLIT_K) * stride_w_k

    if not EVEN_K:
        mask_x_k = offs_x_k < MASK_K_LIMIT
        if PRESHUFFLED:
            mask_w_k = offs_w_k < MASK_K_LIMIT * 16
        else:
            mask_w_k = offs_w_k < MASK_K_LIMIT

        x = tl.load(XPtrs, mask=mask_x_k[None, :], other=0)
        w = tl.load(
            WPtrs, mask=mask_w_k[:, None], other=0, cache_modifier=W_CACHE_MODIFIER
        )
        if PRESHUFFLED:
            w = unshuffle_weights(w, BLOCK_N, BLOCK_K)

        acc += tl.dot(x, w, input_precision="ieee")

    # per-token activation scale
    offs_m = BLOCK_M * block_id + tl.arange(0, BLOCK_M)
    mask_m = offs_m < M
    x_scale_ptrs = XScale + offs_x_m * stride_x_scale

    # per-channel weight scale (load using full BLOCK_N)
    offs_y_n = BLOCK_N * pid_n + tl.arange(0, BLOCK_N)
    mask_n = offs_y_n < N
    w_scale_ptrs = WScale + expt_id * stride_w_scale_e + offs_y_n * stride_w_scale_n

    x_scale = tl.load(x_scale_ptrs, mask=mask_m, other=1.0)
    w_scale = tl.load(w_scale_ptrs, mask=mask_n, other=1.0)
    acc = acc * x_scale[:, None] * w_scale[None, :]

    # bias
    if B is not None:
        BPtrs = B + expt_id * stride_b_e + offs_y_n
        if pid_k == 0:
            bias = tl.load(BPtrs, mask=mask_n, other=0, cache_modifier=W_CACHE_MODIFIER)
        else:
            bias = tl.full([BLOCK_N], 0, dtype=tl.float32)
        acc = acc + bias[None, :]

    if APPLY_ACTIVATION and SPLIT_K == 1:
        out = _swiglu(acc, alpha, limit, ADD_RESIDUAL=ADD_RESIDUAL)
        tl.static_assert(
            out.shape[1] == OUT_BLOCK_N,
            f"Activation fn out.shape[1] ({out.shape[1]}) doesn't match computed OUT_BLOCK_N ({OUT_BLOCK_N})",
        )
        offs_y_n = OUT_BLOCK_N * pid_n + tl.arange(0, OUT_BLOCK_N)
        mask_n = offs_y_n < yN
    else:
        tl.static_assert(
            ACTIVATION_REDUCTION_N == 1,
            "Activation reduction must be 1 if no activation fn is provided",
        )
        out = acc

    # Apply gammas if provided
    if Gammas is not None:
        gammas = tl.load(Gammas + start_m + offs_m, mask=mask_m, other=0.0)
        out *= gammas[:, None]

    # Write back
    Y += start_m * stride_y_m
    offs_y_m = offs_m
    YPtrs = (
        Y
        + offs_y_m.to(index_type)[:, None] * stride_y_m
        + offs_y_n.to(index_type)[None, :] * stride_y_n
    )
    mask = mask_m[:, None] & mask_n[None, :]
    tl.store(YPtrs, out, mask=mask)
