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

# Compile CUDA kernel at module load time
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__device__ inline void update_topk(float* vals, long long* idxs, int k, float new_val, long long new_idx) {
    // Find position of minimum value
    int min_pos = 0;
    float min_val = vals[0];
    for (int i = 1; i < k; i++) {
        if (vals[i] < min_val) {
            min_val = vals[i];
            min_pos = i;
        }
    }
    // Replace if new value is larger
    if (new_val > min_val) {
        vals[min_pos] = new_val;
        idxs[min_pos] = new_idx;
    }
}

__global__ void topk_per_row_kernel(
    const int* __restrict__ row_ptr,
    const float* __restrict__ edge_scores,
    float* __restrict__ topk_vals,
    long long* __restrict__ topk_idx,
    int num_nodes,
    int k
) {
    int node = blockIdx.x * blockDim.x + threadIdx.x;
    if (node >= num_nodes) return;
    
    int start = row_ptr[node];
    int end = row_ptr[node + 1];
    
    // Initialize with -inf (register arrays for efficiency)
    float local_vals[16];
    long long local_idx[16];
    
    for (int i = 0; i < k; i++) {
        local_vals[i] = -FLT_MAX;
        local_idx[i] = -1;
    }
    
    // Process segment sequentially
    for (int i = start; i < end; i++) {
        update_topk(local_vals, local_idx, k, edge_scores[i], (long long)i);
    }
    
    // Write results (output is row-major [num_nodes, k])
    for (int i = 0; i < k; i++) {
        topk_vals[node * k + i] = local_vals[i];
        topk_idx[node * k + i] = local_idx[i];
    }
}

torch::Tensor topk_per_row_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores,
    torch::Tensor topk_vals,
    torch::Tensor topk_idx,
    int k
) {
    int num_nodes = row_ptr.size(0) - 1;
    
    const int threads = 256;
    const int blocks = (num_nodes + threads - 1) / threads;
    
    topk_per_row_kernel<<<blocks, threads>>>(
        row_ptr.data_ptr<int>(),
        edge_scores.data_ptr<float>(),
        topk_vals.data_ptr<float>(),
        topk_idx.data_ptr<long long>(),
        num_nodes,
        k
    );
    
    return topk_vals;
}
"""

cpp_source = """
torch::Tensor topk_per_row_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores,
    torch::Tensor topk_vals,
    torch::Tensor topk_idx,
    int k);
"""

# Load the CUDA extension
topk_extension = load_inline(
    name='topk_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['topk_per_row_cuda'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Per-row top-k selection over CSR edge segments using custom CUDA kernel.
    
    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
        
        # Preallocate output tensors with default values
        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
        )
        
        # Call custom CUDA kernel
        topk_extension.topk_per_row_cuda(
            row_ptr.int(),
            edge_scores,
            topk_vals,
            topk_idx,
            self.k
        )
        
        return topk_vals, topk_idx