import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Sparse message passing with CSR adjacency.

    For each destination node, aggregate weighted source-node features from the
    incoming edges in its CSR row. This is a core primitive in sparse graph
    neural network layers and graph convolution operators.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        edge_weight: torch.Tensor,
        node_feat: torch.Tensor,
    ) -> torch.Tensor:
        num_nodes = row_ptr.numel() - 1
        feat_dim = node_feat.shape[1]
        out = torch.zeros(num_nodes, feat_dim, dtype=node_feat.dtype, device=node_feat.device)

        for dst in range(num_nodes):
            start = int(row_ptr[dst].item())
            end = int(row_ptr[dst + 1].item())
            if end > start:
                src = col_idx[start:end].long()
                weights = edge_weight[start:end].unsqueeze(-1)
                out[dst] = (node_feat[src] * weights).sum(dim=0)

        return out


num_nodes = 768
avg_degree = 20
feat_dim = 96




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



def get_init_inputs():
    return []

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

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void sparse_message_pass_kernel(
    const int* __restrict__ row_ptr,
    const int* __restrict__ col_idx,
    const float* __restrict__ edge_weight,
    const float* __restrict__ node_feat,
    float* __restrict__ out,
    int num_nodes,
    int feat_dim
) {
    int dst = blockIdx.x * blockDim.y + threadIdx.y;
    
    if (dst >= num_nodes) return;
    
    int tid = threadIdx.x;
    int start = row_ptr[dst];
    int end = row_ptr[dst + 1];
    
    // Each thread handles a subset of feature dimensions
    int feats_per_thread = (feat_dim + blockDim.x - 1) / blockDim.x;
    
    // Initialize output to zero
    for (int i = 0; i < feats_per_thread; i++) {
        int feat_idx = tid * feats_per_thread + i;
        if (feat_idx < feat_dim) {
            out[dst * feat_dim + feat_idx] = 0.0f;
        }
    }
    
    // Iterate over incoming edges
    for (int edge = start; edge < end; edge++) {
        int src = col_idx[edge];
        float weight = edge_weight[edge];
        
        for (int i = 0; i < feats_per_thread; i++) {
            int feat_idx = tid * feats_per_thread + i;
            if (feat_idx < feat_dim) {
                float feat_val = node_feat[src * feat_dim + feat_idx];
                out[dst * feat_dim + feat_idx] += weight * feat_val;
            }
        }
    }
}

torch::Tensor sparse_message_pass_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor edge_weight,
    torch::Tensor node_feat
) {
    int num_nodes = row_ptr.size(0) - 1;
    int feat_dim = node_feat.size(1);
    
    auto out = torch::zeros({num_nodes, feat_dim}, node_feat.options());
    
    int threads_per_block_x = 32;
    int threads_per_block_y = 8;
    int blocks = (num_nodes + threads_per_block_y - 1) / threads_per_block_y;
    
    dim3 block_dim(threads_per_block_x, threads_per_block_y);
    
    // Use data_ptr<int>() - PyTorch's int32 tensors use C++ int type
    sparse_message_pass_kernel<<<blocks, block_dim>>>(
        row_ptr.data_ptr<int>(),
        col_idx.data_ptr<int>(),
        edge_weight.data_ptr<float>(),
        node_feat.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes,
        feat_dim
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor sparse_message_pass_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_weight, torch::Tensor node_feat);
"""

sparse_ops = load_inline(
    name='sparse_message_pass_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sparse_message_pass_cuda'],
    verbose=True,
    extra_cuda_cflags=['-O3']
)


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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        edge_weight: torch.Tensor,
        node_feat: torch.Tensor,
    ) -> torch.Tensor:
        return sparse_ops.sparse_message_pass_cuda(
            row_ptr.int(), col_idx.int(), edge_weight, node_feat
        )