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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void csr_topk_kernel(
    const int* __restrict__ row_ptr,
    const float* __restrict__ edge_scores,
    float* __restrict__ topk_vals,
    int64_t* __restrict__ topk_idx,
    const int num_nodes,
    const int k,
    const int max_deg
) {
    int node = blockIdx.x * blockDim.x + threadIdx.x;
    if (node >= num_nodes) return;

    int start = row_ptr[node];
    int end = row_ptr[node + 1];
    int len = end - start;
    if (len <= 0) {
        for (int i = 0; i < k; i++) {
            topk_vals[node * k + i] = -INFINITY;
            topk_idx[node * k + i] = -1;
        }
        return;
    }

    int curr_k = min(k, len);

    // Local buffers - assume max_deg <= 64 for this workload
    float local_scores[64];
    int local_indices[64];
    bool selected[64] = {false};

    for (int i = 0; i < len; i++) {
        local_scores[i] = edge_scores[start + i];
        local_indices[i] = start + i;
    }

    for (int pos = 0; pos < curr_k; pos++) {
        float max_val = -INFINITY;
        int max_pos = -1;
        for (int j = 0; j < len; j++) {
            if (!selected[j] && local_scores[j] > max_val) {
                max_val = local_scores[j];
                max_pos = j;
            }
        }
        if (max_pos != -1) {
            selected[max_pos] = true;
            topk_vals[node * k + pos] = max_val;
            topk_idx[node * k + pos] = local_indices[max_pos];
        }
    }

    // Fill remaining with -inf if curr_k < k
    for (int pos = curr_k; pos < k; pos++) {
        topk_vals[node * k + pos] = -INFINITY;
        topk_idx[node * k + pos] = -1;
    }
}

std::tuple<torch::Tensor, torch::Tensor> csr_topk_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores,
    int k
) {
    int num_nodes = row_ptr.size(0) - 1;
    auto options_val = torch::TensorOptions().dtype(edge_scores.dtype()).device(edge_scores.device());
    auto options_idx = torch::TensorOptions().dtype(torch::kInt64).device(edge_scores.device());

    auto topk_vals = torch::full({num_nodes, k}, -INFINITY, options_val);
    auto topk_idx = torch::full({num_nodes, k}, -1LL, options_idx);

    // Estimate max degree
    int max_deg = 0;
    auto row_ptr_cpu = row_ptr.to(torch::kCPU);
    for (int i = 0; i < num_nodes; i++) {
        int d = row_ptr_cpu[i+1].item<int>() - row_ptr_cpu[i].item<int>();
        if (d > max_deg) max_deg = d;
    }
    if (max_deg > 64) max_deg = 64; // safety

    dim3 block(256);
    dim3 grid((num_nodes + block.x - 1) / block.x);

    csr_topk_kernel<<<grid, block>>>(
        row_ptr.data_ptr<int>(),
        edge_scores.data_ptr<float>(),
        topk_vals.data_ptr<float>(),
        topk_idx.data_ptr<int64_t>(),
        num_nodes,
        k,
        max_deg
    );

    cudaDeviceSynchronize();
    return std::make_tuple(topk_vals, topk_idx);
}
'''

cpp_source = r'''
std::tuple<torch::Tensor, torch::Tensor> csr_topk_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores,
    int k
);
'''

module = load_inline(
    name='csr_topk_module',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['csr_topk_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor):
        return module.csr_topk_cuda(row_ptr, edge_scores, self.k)
```