from __future__ import annotations

import os

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_CPP = r"""
#include <torch/extension.h>

void pack_cuda(torch::Tensor x, torch::Tensor ids, torch::Tensor expert_map,
               torch::Tensor cursor, torch::Tensor packed, torch::Tensor inverse);
void reduce_cuda(torch::Tensor shared, torch::Tensor routed, torch::Tensor inverse,
                 torch::Tensor weights, torch::Tensor output, int64_t topk);

void pack(torch::Tensor x, torch::Tensor ids, torch::Tensor expert_map,
          torch::Tensor cursor, torch::Tensor packed, torch::Tensor inverse) {
  pack_cuda(x, ids, expert_map, cursor, packed, inverse);
}

void reduce(torch::Tensor shared, torch::Tensor routed, torch::Tensor inverse,
            torch::Tensor weights, torch::Tensor output, int64_t topk) {
  reduce_cuda(shared, routed, inverse, weights, output, topk);
}
"""

_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>

__global__ void pack_kernel(const __nv_bfloat16* x, const int64_t* ids,
                            const int32_t* expert_map, int32_t* cursor,
                            __nv_bfloat16* packed, int32_t* inverse,
                            int nassign, int topk, int hidden, int capacity) {
  int a = blockIdx.x;
  if (a >= nassign) return;
  int expert = static_cast<int>(ids[a]);
  int batch = expert_map[expert];
  __shared__ int packed_row_shared;
  if (threadIdx.x == 0) {
    int row = atomicAdd(cursor + expert, 1);
    packed_row_shared = batch * capacity + row;
    inverse[a] = packed_row_shared;
  }
  __syncthreads();
  int packed_row = packed_row_shared;
  int [REDACTED credential assignment] / topk;
  int64_t src = static_cast<int64_t>(token) * hidden;
  int64_t dst = static_cast<int64_t>(packed_row) * hidden;
  for (int h = threadIdx.x; h < hidden; h += blockDim.x) {
    packed[dst + h] = x[src + h];
  }
}

__global__ void reduce_kernel(const __nv_bfloat16* shared,
                              const __nv_bfloat16* routed,
                              const int32_t* inverse,
                              const __nv_bfloat16* weights,
                              __nv_bfloat16* output, int64_t total,
                              int hidden, int topk) {
  int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (i >= total) return;
  int [REDACTED credential assignment] / hidden);
  int h = static_cast<int>(i - static_cast<int64_t>(token) * hidden);
  float acc = __bfloat162float(shared[i]);
  int base = token * topk;
  for (int k = 0; k < topk; ++k) {
    int packed_row = inverse[base + k];
    float y = __bfloat162float(routed[static_cast<int64_t>(packed_row) * hidden + h]);
    float w = __bfloat162float(weights[base + k]);
    acc = fmaf(y, w, acc);
  }
  output[i] = __float2bfloat16_rn(acc);
}

void pack_cuda(torch::Tensor x, torch::Tensor ids, torch::Tensor expert_map,
               torch::Tensor cursor, torch::Tensor packed, torch::Tensor inverse) {
  c10::cuda::CUDAGuard guard(x.device());
  int nassign = static_cast<int>(ids.numel());
  int hidden = static_cast<int>(x.size(1));
  int topk = static_cast<int>(ids.size(1));
  int capacity = static_cast<int>(packed.size(1));
  auto stream = at::cuda::getCurrentCUDAStream();
  pack_kernel<<<nassign, 256, 0, stream>>>(
      reinterpret_cast<const __nv_bfloat16*>(x.data_ptr<at::BFloat16>()),
      ids.data_ptr<int64_t>(), expert_map.data_ptr<int32_t>(),
      cursor.data_ptr<int32_t>(),
      reinterpret_cast<__nv_bfloat16*>(packed.data_ptr<at::BFloat16>()),
      inverse.data_ptr<int32_t>(), nassign, topk, hidden, capacity);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void reduce_cuda(torch::Tensor shared, torch::Tensor routed, torch::Tensor inverse,
                 torch::Tensor weights, torch::Tensor output, int64_t topk) {
  c10::cuda::CUDAGuard guard(shared.device());
  int hidden = static_cast<int>(shared.size(1));
  int64_t total = shared.numel();
  int blocks = static_cast<int>((total + 255) / 256);
  auto stream = at::cuda::getCurrentCUDAStream();
  reduce_kernel<<<blocks, 256, 0, stream>>>(
      reinterpret_cast<const __nv_bfloat16*>(shared.data_ptr<at::BFloat16>()),
      reinterpret_cast<const __nv_bfloat16*>(routed.data_ptr<at::BFloat16>()),
      inverse.data_ptr<int32_t>(),
      reinterpret_cast<const __nv_bfloat16*>(weights.data_ptr<at::BFloat16>()),
      reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
      total, hidden, static_cast<int>(topk));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}
"""

_EXT = None


def _extension():
    global _EXT
    if _EXT is None:
        os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
        _EXT = load_inline(
            name="glm52_fused_moe_cuda",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["pack", "reduce"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            with_cuda=True,
            verbose=False,
        )
    return _EXT


class Model(nn.Module):
    def __init__(self, T: int, E: int, top_k: int, n_shared: int, H: int, I: int):
        super().__init__()
        self.T, self.E, self.top_k = T, E, top_k
        self.n_shared, self.H, self.I = n_shared, H, I
        self.w1_routed = nn.Parameter(torch.empty(E, 2 * I, H, dtype=torch.bfloat16))
        self.w2_routed = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        self.w1_shared = nn.Parameter(torch.empty(n_shared, 2 * I, H, dtype=torch.bfloat16))
        self.w2_shared = nn.Parameter(torch.empty(n_shared, H, I, dtype=torch.bfloat16))
        for p in self.parameters():
            nn.init.normal_(p, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        expert_ids: torch.Tensor,
        expert_weights: torch.Tensor,
    ) -> torch.Tensor:
        T, H = x.shape
        E, topk, I = self.E, self.top_k, self.I
        x = x.contiguous()
        expert_ids = expert_ids.contiguous()
        expert_weights = expert_weights.contiguous()

        # Shared expert: fused gate/up GEMM followed by activation and down GEMM.
        shared_gate_up = x @ self.w1_shared[0].transpose(0, 1)
        shared_hidden = F.silu(shared_gate_up[:, :I]) * shared_gate_up[:, I:]
        shared_out = shared_hidden @ self.w2_shared[0].transpose(0, 1)

        # Group routed rows by expert. Tiny decode batches compact the expert
        # dimension; larger batches retain the full contiguous weight table.
        counts = torch.bincount(expert_ids.reshape(-1), minlength=E)
        capacity = int(counts.max().item())
        if T < 32:
            active = torch.nonzero(counts, as_tuple=False).flatten()
            nactive = active.numel()
            expert_map = (counts != 0).to(torch.int32).cumsum(0, dtype=torch.int32) - 1
            w1 = self.w1_routed.index_select(0, active)
            w2 = self.w2_routed.index_select(0, active)
        else:
            nactive = E
            expert_map = torch.arange(E, device=x.device, dtype=torch.int32)
            w1 = self.w1_routed
            w2 = self.w2_routed

        packed_x = torch.zeros((nactive, capacity, H), device=x.device, dtype=x.dtype)
        cursor = torch.zeros((E,), device=x.device, dtype=torch.int32)
        inverse = torch.empty((T * topk,), device=x.device, dtype=torch.int32)
        ext = _extension()
        ext.pack(x, expert_ids, expert_map, cursor, packed_x, inverse)

        gate_up = torch.bmm(packed_x, w1.transpose(1, 2))
        hidden = F.silu(gate_up[..., :I]) * gate_up[..., I:]
        routed_out = torch.bmm(hidden, w2.transpose(1, 2))

        output = torch.empty_like(x)
        ext.reduce(shared_out, routed_out, inverse, expert_weights, output, topk)
        return output