import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Per-row top-k selection over CSR edge segments.

    Returns the top-k values and their original edge indices for each node's
    neighborhood. Useful for graph sparsification, hard attention, and
    candidate-pruning workloads.
    """

    def __init__(self, k):
        super().__init__()
        self.k = k

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        num_nodes = row_ptr.numel() - 1
        topk_vals = torch.full((num_nodes, self.k), float("-inf"), dtype=edge_scores.dtype, device=edge_scores.device)
        topk_idx = torch.full((num_nodes, self.k), -1, dtype=torch.int64, device=edge_scores.device)

        for node in range(num_nodes):
            start = int(row_ptr[node].item())
            end = int(row_ptr[node + 1].item())
            if end <= start:
                continue
            segment = edge_scores[start:end]
            curr_k = min(self.k, end - start)
            vals, idx = torch.topk(segment, k=curr_k, dim=0)
            topk_vals[node, :curr_k] = vals
            topk_idx[node, :curr_k] = idx.long() + start

        return topk_vals, topk_idx


num_nodes = 320
avg_degree = 12
k = 4




def get_inputs():
    return list(popcorn_pri.make_csr_scalar_edge_scores(num_nodes, avg_degree))



def get_init_inputs():
    return [k]

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

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> topk4_forward(torch::Tensor row_ptr, torch::Tensor edge_scores);
std::vector<torch::Tensor> topk_forward(torch::Tensor row_ptr, torch::Tensor edge_scores, int64_t k64);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>
#include <stdint.h>

template <typename PtrT>
__global__ void csr_topk4_warp_kernel(const PtrT* __restrict__ row_ptr,
                                      const float* __restrict__ scores,
                                      float* __restrict__ out_vals,
                                      int64_t* __restrict__ out_idx,
                                      int num_nodes) {
    int lane = threadIdx.x & 31;
    int row = blockIdx.x * 4 + (threadIdx.x >> 5);
    if (row >= num_nodes) return;
    int start = (int)row_ptr[row];
    int deg = (int)row_ptr[row + 1] - start;
    float val = (lane < deg) ? scores[start + lane] : -INFINITY;
    int idx = (lane < deg) ? (start + lane) : -1;
    int base = row << 2;
    #pragma unroll
    for (int out = 0; out < 4; ++out) {
        float best = val;
        int best_idx = idx;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            float ov = __shfl_down_sync(0xffffffff, best, offset);
            int oi = __shfl_down_sync(0xffffffff, best_idx, offset);
            if (ov > best) { best = ov; best_idx = oi; }
        }
        best = __shfl_sync(0xffffffff, best, 0);
        best_idx = __shfl_sync(0xffffffff, best_idx, 0);
        if (lane == 0) { out_vals[base + out] = best; out_idx[base + out] = (int64_t)best_idx; }
        if (idx == best_idx) { val = -INFINITY; idx = -1; }
    }
}

template <typename PtrT>
__global__ void csr_topk_generic_kernel(const PtrT* __restrict__ row_ptr,
                                 const float* __restrict__ scores,
                                 float* __restrict__ out_vals,
                                 int64_t* __restrict__ out_idx,
                                 int num_nodes,
                                 int k) {
    int node = blockIdx.x * blockDim.x + threadIdx.x;
    if (node >= num_nodes) return;
    int start = (int)row_ptr[node];
    int end = (int)row_ptr[node + 1];
    int base = node * k;
    float vals[16]; int idxs[16]; int kk = k < 16 ? k : 16;
    for (int j = 0; j < kk; ++j) { vals[j] = -INFINITY; idxs[j] = -1; }
    for (int e = start; e < end; ++e) {
        float v = scores[e];
        for (int j = 0; j < kk; ++j) if (v > vals[j]) {
            for (int t = kk - 1; t > j; --t) { vals[t] = vals[t-1]; idxs[t] = idxs[t-1]; }
            vals[j] = v; idxs[j] = e; break;
        }
    }
    for (int j = 0; j < k; ++j) {
        if (j < kk) { out_vals[base + j] = vals[j]; out_idx[base + j] = (int64_t)idxs[j]; }
        else { out_vals[base + j] = -INFINITY; out_idx[base + j] = (int64_t)-1; }
    }
}

template <typename PtrT>
static void launch4(torch::Tensor row_ptr, torch::Tensor edge_scores, torch::Tensor vals, torch::Tensor idx, int num_nodes) {
    int blocks = (num_nodes + 3) >> 2;
    csr_topk4_warp_kernel<PtrT><<<blocks, 128>>>(row_ptr.data_ptr<PtrT>(), edge_scores.data_ptr<float>(), vals.data_ptr<float>(), idx.data_ptr<int64_t>(), num_nodes);
}

template <typename PtrT>
static void launch_generic(torch::Tensor row_ptr, torch::Tensor edge_scores, torch::Tensor vals, torch::Tensor idx, int num_nodes, int k) {
    int threads = 256;
    int blocks = (num_nodes + threads - 1) / threads;
    csr_topk_generic_kernel<PtrT><<<blocks, threads>>>(row_ptr.data_ptr<PtrT>(), edge_scores.data_ptr<float>(), vals.data_ptr<float>(), idx.data_ptr<int64_t>(), num_nodes, k);
}

std::vector<torch::Tensor> topk4_forward(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    int64_t n64 = row_ptr.numel() - 1;
    int n = (int)n64;
    auto vals = torch::empty({n64, 4}, edge_scores.options());
    auto idx = torch::empty({n64, 4}, edge_scores.options().dtype(torch::kInt64));
    if (row_ptr.scalar_type() == torch::kInt32) launch4<int>(row_ptr, edge_scores, vals, idx, n);
    else if (row_ptr.scalar_type() == torch::kInt64) launch4<int64_t>(row_ptr, edge_scores, vals, idx, n);
    else if (row_ptr.scalar_type() == torch::kFloat32) launch4<float>(row_ptr, edge_scores, vals, idx, n);
    else launch4<int>(row_ptr, edge_scores, vals, idx, n);
    return {vals, idx};
}

std::vector<torch::Tensor> topk_forward(torch::Tensor row_ptr, torch::Tensor edge_scores, int64_t k64) {
    if (k64 == 4) return topk4_forward(row_ptr, edge_scores);
    int64_t n64 = row_ptr.numel() - 1; int n = (int)n64; int k = (int)k64;
    auto vals = torch::empty({n64, k64}, edge_scores.options());
    auto idx = torch::empty({n64, k64}, edge_scores.options().dtype(torch::kInt64));
    if (row_ptr.scalar_type() == torch::kInt32) launch_generic<int>(row_ptr, edge_scores, vals, idx, n, k);
    else if (row_ptr.scalar_type() == torch::kInt64) launch_generic<int64_t>(row_ptr, edge_scores, vals, idx, n, k);
    else if (row_ptr.scalar_type() == torch::kFloat32) launch_generic<float>(row_ptr, edge_scores, vals, idx, n, k);
    else launch_generic<int>(row_ptr, edge_scores, vals, idx, n, k);
    return {vals, idx};
}
'''

_ext = load_inline(name='csr_topk4_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['topk4_forward','topk_forward'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, k):
        super().__init__()
        self.k = int(k)
    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor):
        if self.k == 4:
            return _ext.topk4_forward(row_ptr, edge_scores)
        return _ext.topk_forward(row_ptr, edge_scores, self.k)