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> csr_topk_forward(torch::Tensor row_ptr, torch::Tensor edge_scores, int64_t k);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <limits>

__global__ void csr_topk4_kernel(const int* __restrict__ row_ptr,
                                 const float* __restrict__ scores,
                                 float* __restrict__ out_vals,
                                 long long* __restrict__ out_idx,
                                 int n) {
    int node = blockIdx.x * blockDim.x + threadIdx.x;
    if (node >= n) return;

    int start = row_ptr[node];
    int end   = row_ptr[node + 1];

    float v0 = -CUDART_INF_F, v1 = -CUDART_INF_F, v2 = -CUDART_INF_F, v3 = -CUDART_INF_F;
    int   i0 = -1,            i1 = -1,            i2 = -1,            i3 = -1;

    #pragma unroll 16
    for (int e = start; e < end; ++e) {
        float v = scores[e];
        if (v > v0) {
            v3 = v2; i3 = i2; v2 = v1; i2 = i1; v1 = v0; i1 = i0; v0 = v; i0 = e;
        } else if (v > v1) {
            v3 = v2; i3 = i2; v2 = v1; i2 = i1; v1 = v; i1 = e;
        } else if (v > v2) {
            v3 = v2; i3 = i2; v2 = v; i2 = e;
        } else if (v > v3) {
            v3 = v; i3 = e;
        }
    }

    int o = node * 4;
    out_vals[o + 0] = v0; out_vals[o + 1] = v1; out_vals[o + 2] = v2; out_vals[o + 3] = v3;
    out_idx [o + 0] = (long long)i0; out_idx [o + 1] = (long long)i1;
    out_idx [o + 2] = (long long)i2; out_idx [o + 3] = (long long)i3;
}

template<int MAXK>
__global__ void csr_topk_generic_kernel(const int* __restrict__ row_ptr,
                                        const float* __restrict__ scores,
                                        float* __restrict__ out_vals,
                                        long long* __restrict__ out_idx,
                                        int n, int k) {
    int node = blockIdx.x * blockDim.x + threadIdx.x;
    if (node >= n) return;

    float vals[MAXK];
    int idxs[MAXK];
    #pragma unroll
    for (int j = 0; j < MAXK; ++j) { vals[j] = -CUDART_INF_F; idxs[j] = -1; }

    int kk = k < MAXK ? k : MAXK;
    int start = row_ptr[node];
    int end   = row_ptr[node + 1];
    for (int e = start; e < end; ++e) {
        float v = scores[e];
        if (v > vals[kk - 1]) {
            int pos = kk - 1;
            while (pos > 0 && v > vals[pos - 1]) {
                vals[pos] = vals[pos - 1];
                idxs[pos] = idxs[pos - 1];
                --pos;
            }
            vals[pos] = v;
            idxs[pos] = e;
        }
    }

    int base = node * k;
    for (int j = 0; j < k; ++j) {
        if (j < MAXK) {
            out_vals[base + j] = vals[j];
            out_idx[base + j] = (long long)idxs[j];
        } else {
            out_vals[base + j] = -CUDART_INF_F;
            out_idx[base + j] = -1LL;
        }
    }
}

std::vector<torch::Tensor> csr_topk_forward(torch::Tensor row_ptr, torch::Tensor edge_scores, int64_t k64) {
    const auto n = (int)(row_ptr.numel() - 1);
    auto vals = torch::empty({n, k64}, edge_scores.options());
    auto idx = torch::empty({n, k64}, row_ptr.options().dtype(torch::kInt64));
    if (n == 0 || k64 <= 0) return {vals, idx};

    const at::cuda::OptionalCUDAGuard device_guard(device_of(edge_scores));
    int threads = 256;
    int blocks = (n + threads - 1) / threads;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (k64 == 4) {
        csr_topk4_kernel<<<blocks, threads, 0, stream>>>(
            row_ptr.data_ptr<int>(), edge_scores.data_ptr<float>(), vals.data_ptr<float>(),
            reinterpret_cast<long long*>(idx.data_ptr<int64_t>()), n);
    } else {
        int k = (int)k64;
        csr_topk_generic_kernel<64><<<blocks, threads, 0, stream>>>(
            row_ptr.data_ptr<int>(), edge_scores.data_ptr<float>(), vals.data_ptr<float>(),
            reinterpret_cast<long long*>(idx.data_ptr<int64_t>()), n, k);
    }
    return {vals, idx};
}
"""

_ext = load_inline(
    name="csr_row_topk_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["csr_topk_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, k):
        super().__init__()
        self.k = k

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor):
        vals, idx = _ext.csr_topk_forward(row_ptr, edge_scores, int(self.k))
        return vals, idx