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.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <limits>
#include <cmath>

template <typename scalar_t>
__global__ void topk_csr_kernel(
    const int32_t* __restrict__ row_ptr,
    const scalar_t* __restrict__ edge_scores,
    scalar_t* __restrict__ topk_vals,
    int64_t* __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];
    int seg_len = end - start;
    
    // Use registers for top-k values (small k is typical)
    // Initialize with -inf
    scalar_t v0 = -INFINITY, v1 = -INFINITY, v2 = -INFINITY, v3 = -INFINITY;
    int64_t i0 = -1, i1 = -1, i2 = -1, i3 = -1;
    
    // Process segment
    for (int i = 0; i < seg_len; i++) {
        scalar_t score = edge_scores[start + i];
        int64_t edge_idx = start + i;
        
        // Insert into sorted top-k (handles k=1,2,3,4)
        if (score > v3) {
            if (score > v2) {
                if (score > v1) {
                    if (score > v0) {
                        v3 = v2; i3 = i2;
                        v2 = v1; i2 = i1;
                        v1 = v0; i1 = i0;
                        v0 = score; i0 = edge_idx;
                    } else {
                        v3 = v2; i3 = i2;
                        v2 = v1; i2 = i1;
                        v1 = score; i1 = edge_idx;
                    }
                } else {
                    v3 = v2; i3 = i2;
                    v2 = score; i2 = edge_idx;
                }
            } else {
                v3 = score; i3 = edge_idx;
            }
        }
    }
    
    // Write results based on k
    if (k > 0) {
        if (seg_len > 0) {
            topk_vals[node * k + 0] = v0;
            topk_idx[node * k + 0] = i0;
        } else {
            topk_vals[node * k + 0] = -INFINITY;
            topk_idx[node * k + 0] = -1;
        }
    }
    if (k > 1) {
        if (seg_len > 1) {
            topk_vals[node * k + 1] = v1;
            topk_idx[node * k + 1] = i1;
        } else {
            topk_vals[node * k + 1] = -INFINITY;
            topk_idx[node * k + 1] = -1;
        }
    }
    if (k > 2) {
        if (seg_len > 2) {
            topk_vals[node * k + 2] = v2;
            topk_idx[node * k + 2] = i2;
        } else {
            topk_vals[node * k + 2] = -INFINITY;
            topk_idx[node * k + 2] = -1;
        }
    }
    if (k > 3) {
        if (seg_len > 3) {
            topk_vals[node * k + 3] = v3;
            topk_idx[node * k + 3] = i3;
        } else {
            topk_vals[node * k + 3] = -INFINITY;
            topk_idx[node * k + 3] = -1;
        }
    }
}

void topk_csr_forward_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores,
    torch::Tensor topk_vals,
    torch::Tensor topk_idx,
    int k) {
    
    int num_nodes = row_ptr.numel() - 1;
    const int block_size = 256;
    const int num_blocks = (num_nodes + block_size - 1) / block_size;
    
    AT_DISPATCH_FLOATING_TYPES(edge_scores.scalar_type(), "topk_csr_kernel", ([&] {
        topk_csr_kernel<scalar_t><<<num_blocks, block_size>>>(
            row_ptr.data_ptr<int32_t>(),
            edge_scores.data_ptr<scalar_t>(),
            topk_vals.data_ptr<scalar_t>(),
            topk_idx.data_ptr<int64_t>(),
            num_nodes,
            k);
    }));
}
"""

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

cuda_extension = load_inline(
    name='topk_csr_extension_v7',
    cuda_sources=cuda_source,
    cpp_sources=cpp_source,
    functions=['topk_csr_forward_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True,
)


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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor):
        num_nodes = row_ptr.numel() - 1
        
        # Create output tensors on Python side to avoid type issues
        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 CUDA kernel
        cuda_extension.topk_csr_forward_cuda(
            row_ptr.contiguous(),
            edge_scores.contiguous(),
            topk_vals,
            topk_idx,
            self.k
        )
        
        return topk_vals, topk_idx