import os
import shutil

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

if shutil.which("ninja") is None:
    import ninja
    os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ.get("PATH", "")


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

using bf16 = __nv_bfloat16;

__device__ __forceinline__ void copy_async(bf16* dst, const bf16* src, bool valid) {
    unsigned address = static_cast<unsigned>(__cvta_generic_to_shared(dst));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
                 [REDACTED: IP] "r"(address), "l"(src), "r"(valid ? 16 : 0));
}

__device__ __forceinline__ void commit_async() {
    asm volatile("cp.async.commit_group;");
}

template<int Pending>
__device__ __forceinline__ void wait_async() {
    asm volatile("cp.async.wait_group %0;" [REDACTED: IP] "n"(Pending));
}

__device__ __forceinline__ void load_a(unsigned (&value)[4], const bf16* ptr) {
    unsigned address = static_cast<unsigned>(__cvta_generic_to_shared(ptr));
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
                 : "=r"(value[0]), "=r"(value[1]), "=r"(value[2]), "=r"(value[3])
                 : "r"(address));
}

__device__ __forceinline__ void load_b(unsigned (&value)[2], const bf16* ptr) {
    unsigned address = static_cast<unsigned>(__cvta_generic_to_shared(ptr));
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
                 : "=r"(value[0]), "=r"(value[1]) : "r"(address));
}

__device__ __forceinline__ void mma(float (&acc)[4], const unsigned (&left)[4],
                                    const unsigned (&right)[2]) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3])
        : "r"(left[0]), "r"(left[1]), "r"(left[2]), "r"(left[3]),
          "r"(right[0]), "r"(right[1]));
}

__global__ void count_routes(const int64_t* ids, int* counts, int routes) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index < routes) atomicAdd(counts + ids[index], 1);
}

template<int TileM>
__global__ void make_offsets(const int* counts, int* offsets, int* tile_experts,
                             int* tile_count, int experts, int shared, int tokens) {
    using Scan = cub::BlockScan<int, 512>;
    __shared__ typename Scan::TempStorage storage;
    int expert = threadIdx.x;
    int rows = expert < experts ? counts[expert] : (expert < experts + shared ? tokens : 0);
    int tiles = (rows + TileM - 1) / TileM;
    int start, total;
    Scan(storage).ExclusiveSum(tiles, start, total);
    if (expert < experts + shared) {
        offsets[expert] = start * TileM;
        for (int tile = 0; tile < tiles; ++tile) tile_experts[start + tile] = expert;
    }
    if (expert == 0) *tile_count = total;
}

__global__ void scatter_routes(const int64_t* ids, const int* offsets, int* cursors,
                               int* packed_ids, int tokens, int experts, int topk, int shared) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int branches = topk + shared;
    if (index >= tokens * branches) return;
    int [REDACTED credential assignment] / branches;
    int branch = index % branches;
    int expert, position;
    if (branch < topk) {
        expert = ids[token * topk + branch];
        position = atomicAdd(cursors + expert, 1);
    } else {
        expert = experts + branch - topk;
        position = token;
    }
    packed_ids[offsets[expert] + position] = index;
}

template<int TileK>
__device__ __forceinline__ int swizzle(int row, int col) {
    return row * TileK + (col ^ ((row & (TileK / 8 - 1)) * 8));
}

template<int TileM, int TileN, int TileK, bool First, int Threads>
__device__ __forceinline__ void load_stage(
    bf16* shared_ptr, const bf16* input, const bf16* weight, const int* packed_ids,
    int first_row, int valid_rows, int branches, int hidden, int intermediate,
    int first_col, int first_k) {
    constexpr int WeightRows = First ? 2 * TileN : TileN;
    constexpr int Entries = (TileM + WeightRows) * TileK / 8;
    int reduction = First ? hidden : intermediate;
    #pragma unroll
    for (int vector = threadIdx.x; vector < Entries; vector += Threads) {
        int row = vector / (TileK / 8);
        int col = (vector % (TileK / 8)) * 8;
        const bf16* source;
        bool valid = first_k < reduction;
        if (row < TileM) {
            valid = valid && row < valid_rows;
            int source_row = first_row + row;
            if constexpr (First) source_row = row < valid_rows ? packed_ids[source_row] / branches : 0;
            source = input + static_cast<int64_t>(source_row) * reduction + first_k + col;
        } else {
            int weight_row = row - TileM;
            if constexpr (First) {
                weight_row = first_col + weight_row % TileN + (weight_row / TileN) * intermediate;
            } else {
                weight_row += first_col;
            }
            source = weight + static_cast<int64_t>(weight_row) * reduction + first_k + col;
        }
        if (first_k < reduction) copy_async(shared_ptr + swizzle<TileK>(row, col), source, valid);
    }
}

template<int TileM, int TileN, int TileK, bool First, int Threads, int Stages, bool Preload = false>
__global__ __launch_bounds__(Threads) void expert_gemm(
    const bf16* input, const bf16* routed_weight, const bf16* shared_weight,
    bf16* activation, float* output, const bf16* routing_weights,
    const int* packed_ids, const int* counts, const int* offsets,
    const int* tile_experts, const int* tile_count,
    int tokens, int experts, int topk, int shared, int hidden, int intermediate) {
    int tile = blockIdx.y;
    if (tile >= *tile_count) return;
    int expert = tile_experts[tile];
    int first_row = tile * TileM;
    int rows = expert < experts ? counts[expert] : tokens;
    int valid_rows = min(TileM, rows - (first_row - offsets[expert]));
    int first_col = blockIdx.x * TileN;
    int branches = topk + shared;
    int reduction = First ? hidden : intermediate;
    int64_t weight_stride = static_cast<int64_t>(hidden) * intermediate * (First ? 2 : 1);
    const bf16* weight = expert < experts ? routed_weight + expert * weight_stride
                                        : shared_weight + (expert - experts) * weight_stride;
    constexpr int StageElements = (TileM + (First ? 2 * TileN : TileN)) * TileK;
    constexpr int WarpsM = TileM / 16 < Threads / 64 ? TileM / 16 : Threads / 64;
    constexpr int WarpsN = Threads / 32 / WarpsM;
    constexpr int MFragments = TileM / (16 * WarpsM);
    constexpr int NFragments = TileN / (8 * WarpsN);
    constexpr int Groups = First ? 2 : 1;
    extern __shared__ __align__(16) unsigned char smem_bytes[];
    bf16* smem = reinterpret_cast<bf16*>(smem_bytes);
    int lane = threadIdx.x % 32;
    int warp = threadIdx.x / 32;
    int warp_row = (warp / WarpsN) * (TileM / WarpsM);
    int warp_col = (warp % WarpsN) * (TileN / WarpsN);
    float accum[Groups][MFragments][NFragments][4] = {};

    #pragma unroll
    for (int stage = 0; stage < Stages - 1; ++stage) {
        load_stage<TileM, TileN, TileK, First, Threads>(smem + stage * StageElements,
            input, weight, packed_ids, first_row, valid_rows, branches,
            hidden, intermediate, first_col, stage * TileK);
        commit_async();
    }
    wait_async<Stages - 2>();
    __syncthreads();
    int steps = reduction / TileK;
    for (int step = 0; step < steps; ++step) {
        int next = step + Stages - 1;
        load_stage<TileM, TileN, TileK, First, Threads>(smem + (next % Stages) * StageElements,
            input, weight, packed_ids, first_row, valid_rows, branches,
            hidden, intermediate, first_col, next * TileK);
        commit_async();
        const bf16* stage_ptr = smem + (step % Stages) * StageElements;
        #pragma unroll
        for (int kfragment = 0; kfragment < TileK / 16; ++kfragment) {
            unsigned left[MFragments][4];
            unsigned prefetched[Groups][NFragments][2];
            #pragma unroll
            for (int mfragment = 0; mfragment < MFragments; ++mfragment) {
                int row = warp_row + mfragment * 16 + lane % 16;
                int col = kfragment * 16 + (lane / 16) * 8;
                load_a(left[mfragment], stage_ptr + swizzle<TileK>(row, col));
            }
            if constexpr (Preload) {
                #pragma unroll
                for (int group = 0; group < Groups; ++group) {
                    #pragma unroll
                    for (int nfragment = 0; nfragment < NFragments; ++nfragment) {
                        int row = TileM + group * TileN + warp_col + nfragment * 8 + lane % 8;
                        int col = kfragment * 16 + ((lane / 8) % 2) * 8;
                        load_b(prefetched[group][nfragment], stage_ptr + swizzle<TileK>(row, col));
                    }
                }
            }
            #pragma unroll
            for (int group = 0; group < Groups; ++group) {
                #pragma unroll
                for (int nfragment = 0; nfragment < NFragments; ++nfragment) {
                    if constexpr (Preload) {
                        #pragma unroll
                        for (int mfragment = 0; mfragment < MFragments; ++mfragment)
                            mma(accum[group][mfragment][nfragment], left[mfragment], prefetched[group][nfragment]);
                    } else {
                        unsigned right[2];
                        int row = TileM + group * TileN + warp_col + nfragment * 8 + lane % 8;
                        int col = kfragment * 16 + ((lane / 8) % 2) * 8;
                        load_b(right, stage_ptr + swizzle<TileK>(row, col));
                        #pragma unroll
                        for (int mfragment = 0; mfragment < MFragments; ++mfragment)
                            mma(accum[group][mfragment][nfragment], left[mfragment], right);
                    }
                }
            }
        }
        wait_async<Stages - 2>();
        __syncthreads();
    }
    wait_async<0>();
    #pragma unroll
    for (int mfragment = 0; mfragment < MFragments; ++mfragment) {
        #pragma unroll
        for (int half = 0; half < 2; ++half) {
            int row = warp_row + mfragment * 16 + lane / 4 + half * 8;
            if (row < valid_rows) {
                int packed_row = first_row + row;
                int route = packed_ids[packed_row];
                int [REDACTED credential assignment] / branches;
                int branch = route % branches;
                float probability = branch < topk ? __bfloat162float(routing_weights[token * topk + branch]) : 1.f;
                #pragma unroll
                for (int nfragment = 0; nfragment < NFragments; ++nfragment) {
                    int col = first_col + warp_col + nfragment * 8 + (lane % 4) * 2;
                    float first = accum[0][mfragment][nfragment][half * 2];
                    float second = accum[0][mfragment][nfragment][half * 2 + 1];
                    if constexpr (First) {
                        first = first / (1.f + __expf(-first)) * accum[1][mfragment][nfragment][half * 2];
                        second = second / (1.f + __expf(-second)) * accum[1][mfragment][nfragment][half * 2 + 1];
                        *reinterpret_cast<__nv_bfloat162*>(activation + static_cast<int64_t>(packed_row) * intermediate + col)
                            = __floats2bfloat162_rn(first, second);
                    } else {
                        atomicAdd(output + static_cast<int64_t>(token) * hidden + col, probability * first);
                        atomicAdd(output + static_cast<int64_t>(token) * hidden + col + 1, probability * second);
                    }
                }
            }
        }
    }
}

__global__ void convert_output(const float* input, bf16* output, int elements) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index < elements / 2)
        reinterpret_cast<__nv_bfloat162*>(output)[index] = __floats2bfloat162_rn(input[index * 2], input[index * 2 + 1]);
}

__device__ __forceinline__ float warp_sum(float value) {
    #pragma unroll
    for (int distance = 16; distance > 0; distance /= 2)
        value += __shfl_down_sync(0xffffffff, value, distance);
    return value;
}

template<bool First>
__global__ __launch_bounds__(256) void decode_gemv(
    const bf16* input, const bf16* routed_weight, const bf16* shared_weight,
    const int64_t* ids, bf16* activation, float* partial,
    int topk, int shared, int hidden, int intermediate) {
    int branch = blockIdx.y;
    int [REDACTED credential assignment];
    int branches = topk + shared;
    int expert = branch < topk ? ids[token * topk + branch] : branch - topk;
    int reduction = First ? hidden : intermediate;
    int width = First ? intermediate : hidden;
    int64_t weight_stride = static_cast<int64_t>(hidden) * intermediate * (First ? 2 : 1);
    const bf16* weight = (branch < topk ? routed_weight : shared_weight) + expert * weight_stride;
    const bf16* source = input + static_cast<int64_t>(First ? [REDACTED credential assignment] * branches + branch) * reduction;
    int row = blockIdx.x * 8 + threadIdx.x / 32;
    int lane = threadIdx.x % 32;
    float gate = 0.f;
    float up = 0.f;
    if (row >= width) return;
    for (int col = lane * 8; col < reduction; col += 256) {
        uint4 inputs = *reinterpret_cast<const uint4*>(source + col);
        uint4 gates = *reinterpret_cast<const uint4*>(weight + static_cast<int64_t>(row) * reduction + col);
        const bf16* in_values = reinterpret_cast<const bf16*>(&inputs);
        const bf16* gate_values = reinterpret_cast<const bf16*>(&gates);
        #pragma unroll
        for (int element = 0; element < 8; ++element)
            gate = fmaf(__bfloat162float(in_values[element]), __bfloat162float(gate_values[element]), gate);
        if constexpr (First) {
            uint4 ups = *reinterpret_cast<const uint4*>(weight + static_cast<int64_t>(row + intermediate) * reduction + col);
            const bf16* up_values = reinterpret_cast<const bf16*>(&ups);
            #pragma unroll
            for (int element = 0; element < 8; ++element)
                up = fmaf(__bfloat162float(in_values[element]), __bfloat162float(up_values[element]), up);
        }
    }
    gate = warp_sum(gate);
    if constexpr (First) up = warp_sum(up);
    if (lane == 0) {
        if constexpr (First)
            activation[(token * branches + branch) * intermediate + row] = __float2bfloat16_rn(gate / (1.f + __expf(-gate)) * up);
        else
            partial[(token * branches + branch) * hidden + row] = gate;
    }
}

__global__ void finish_decode(const float* partial, const bf16* weights, bf16* output,
                              int tokens, int hidden, int topk, int shared) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= tokens * hidden) return;
    int [REDACTED credential assignment] / hidden;
    int col = index % hidden;
    int branches = topk + shared;
    float result = 0.f;
    for (int branch = 0; branch < branches; ++branch) {
        float probability = branch < topk ? __bfloat162float(weights[token * topk + branch]) : 1.f;
        result += partial[(token * branches + branch) * hidden + col] * probability;
    }
    output[index] = __float2bfloat16_rn(result);
}

template<int TileM, int TileN, int TileK, int Threads, int Stages, bool Preload = false>
void launch_prefill(torch::Tensor x, torch::Tensor ids, torch::Tensor weights,
                    torch::Tensor w1, torch::Tensor w2, torch::Tensor shared1,
                    torch::Tensor shared2, torch::Tensor result, cudaStream_t stream) {
    int tokens = x.size(0), hidden = x.size(1), experts = w1.size(0);
    int intermediate = w2.size(2), topk = ids.size(1), shared = shared1.size(0);
    int max_tiles = (tokens * (topk + shared) + TileM - 1) / TileM + experts + shared;
    auto integers = x.options().dtype(torch::kInt32);
    auto counts = torch::empty({experts}, integers);
    auto cursors = torch::empty_like(counts);
    auto offsets = torch::empty({experts + shared}, integers);
    auto tile_experts = torch::empty({max_tiles}, integers);
    auto tile_count = torch::empty({1}, integers);
    auto packed_ids = torch::empty({max_tiles * TileM}, integers);
    auto activation = torch::empty({max_tiles * TileM, intermediate}, x.options());
    auto accum = torch::empty({tokens, hidden}, x.options().dtype(torch::kFloat32));
    cudaMemsetAsync(counts.data_ptr(), 0, experts * sizeof(int), stream);
    cudaMemsetAsync(cursors.data_ptr(), 0, experts * sizeof(int), stream);
    cudaMemsetAsync(accum.data_ptr(), 0, tokens * hidden * sizeof(float), stream);
    count_routes<<<(tokens * topk + 255) / 256, 256, 0, stream>>>(ids.data_ptr<int64_t>(), counts.data_ptr<int>(), tokens * topk);
    make_offsets<TileM><<<1, 512, 0, stream>>>(counts.data_ptr<int>(), offsets.data_ptr<int>(), tile_experts.data_ptr<int>(),
        tile_count.data_ptr<int>(), experts, shared, tokens);
    scatter_routes<<<(tokens * (topk + shared) + 255) / 256, 256, 0, stream>>>(ids.data_ptr<int64_t>(), offsets.data_ptr<int>(),
        cursors.data_ptr<int>(), packed_ids.data_ptr<int>(), tokens, experts, topk, shared);
    constexpr int SharedBytes = Stages * (TileM + 2 * TileN) * TileK * sizeof(bf16);
    if constexpr (SharedBytes >= 48 * 1024) {
        C10_CUDA_CHECK(cudaFuncSetAttribute(expert_gemm<TileM, TileN, TileK, true, Threads, Stages, Preload>, cudaFuncAttributeMaxDynamicSharedMemorySize, SharedBytes));
        C10_CUDA_CHECK(cudaFuncSetAttribute(expert_gemm<TileM, 2 * TileN, TileK, false, Threads, Stages, Preload>, cudaFuncAttributeMaxDynamicSharedMemorySize, SharedBytes));
    }
    expert_gemm<TileM, TileN, TileK, true, Threads, Stages, Preload><<<dim3(intermediate / TileN, max_tiles), Threads, SharedBytes, stream>>>(
        reinterpret_cast<bf16*>(x.data_ptr()), reinterpret_cast<bf16*>(w1.data_ptr()), reinterpret_cast<bf16*>(shared1.data_ptr()),
        reinterpret_cast<bf16*>(activation.data_ptr()), accum.data_ptr<float>(), reinterpret_cast<bf16*>(weights.data_ptr()),
        packed_ids.data_ptr<int>(), counts.data_ptr<int>(), offsets.data_ptr<int>(), tile_experts.data_ptr<int>(), tile_count.data_ptr<int>(),
        tokens, experts, topk, shared, hidden, intermediate);
    expert_gemm<TileM, 2 * TileN, TileK, false, Threads, Stages, Preload><<<dim3(hidden / (2 * TileN), max_tiles), Threads, SharedBytes, stream>>>(
        reinterpret_cast<bf16*>(activation.data_ptr()), reinterpret_cast<bf16*>(w2.data_ptr()), reinterpret_cast<bf16*>(shared2.data_ptr()),
        reinterpret_cast<bf16*>(activation.data_ptr()), accum.data_ptr<float>(), reinterpret_cast<bf16*>(weights.data_ptr()),
        packed_ids.data_ptr<int>(), counts.data_ptr<int>(), offsets.data_ptr<int>(), tile_experts.data_ptr<int>(), tile_count.data_ptr<int>(),
        tokens, experts, topk, shared, hidden, intermediate);
    convert_output<<<(tokens * hidden / 2 + 255) / 256, 256, 0, stream>>>(accum.data_ptr<float>(), reinterpret_cast<bf16*>(result.data_ptr()), tokens * hidden);
}

torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor weights,
                          torch::Tensor w1, torch::Tensor w2, torch::Tensor shared1,
                          torch::Tensor shared2) {
    c10::cuda::CUDAGuard guard(x.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int tokens = x.size(0), hidden = x.size(1), intermediate = w2.size(2);
    int topk = ids.size(1), shared = shared1.size(0);
    auto output = torch::empty_like(x);
    if (tokens <= 4) {
        auto activation = torch::empty({tokens * (topk + shared), intermediate}, x.options());
        auto partial = torch::empty({tokens * (topk + shared), hidden}, x.options().dtype(torch::kFloat32));
        decode_gemv<true><<<dim3((intermediate + 7) / 8, topk + shared, tokens), 256, 0, stream>>>(
            reinterpret_cast<bf16*>(x.data_ptr()), reinterpret_cast<bf16*>(w1.data_ptr()), reinterpret_cast<bf16*>(shared1.data_ptr()),
            ids.data_ptr<int64_t>(), reinterpret_cast<bf16*>(activation.data_ptr()), partial.data_ptr<float>(), topk, shared, hidden, intermediate);
        decode_gemv<false><<<dim3((hidden + 7) / 8, topk + shared, tokens), 256, 0, stream>>>(
            reinterpret_cast<bf16*>(activation.data_ptr()), reinterpret_cast<bf16*>(w2.data_ptr()), reinterpret_cast<bf16*>(shared2.data_ptr()),
            ids.data_ptr<int64_t>(), reinterpret_cast<bf16*>(activation.data_ptr()), partial.data_ptr<float>(), topk, shared, hidden, intermediate);
        finish_decode<<<(tokens * hidden + 255) / 256, 256, 0, stream>>>(partial.data_ptr<float>(),
            reinterpret_cast<bf16*>(weights.data_ptr()), reinterpret_cast<bf16*>(output.data_ptr()), tokens, hidden, topk, shared);
    } else if (static_cast<int64_t>(tokens) * topk >= 192 * w1.size(0)) {
        launch_prefill<128, 128, 64, 512, 2, true>(x, ids, weights, w1, w2, shared1, shared2, output, stream);
    } else {
        launch_prefill<64, 64, 64, 256, 2>(x, ids, weights, w1, w2, shared1, shared2, output, stream);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}

"""

_extension = load_inline(
    name="glm52_moe_sm120",
    cpp_sources="torch::Tensor moe_forward(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);",
    cuda_sources=_CUDA,
    functions=["moe_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_120,code=sm_120"],
    extra_cflags=["-O3"],
    verbose=False,
)


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 parameter in self.parameters():
            nn.init.normal_(parameter, std=0.02)

    def forward(self, x: torch.Tensor, expert_ids: torch.Tensor, expert_weights: torch.Tensor):
        return _extension.moe_forward(
            x, expert_ids, expert_weights,
            self.w1_routed, self.w2_routed, self.w1_shared, self.w2_shared,
        )