import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Degree-normalized neighbor aggregation.

    Computes a GCN-style sparse aggregation where each incoming message is
    scaled by 1 / sqrt(deg(dst) * deg(src)). This captures a common spectral
    graph normalization pattern over CSR adjacency.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        node_feat: torch.Tensor,
        degrees: 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:
                continue
            src = col_idx[start:end].long()
            norm = 1.0 / torch.sqrt(degrees[dst] * degrees[src]).unsqueeze(-1)
            out[dst] = (node_feat[src] * norm).sum(dim=0)

        return out


num_nodes = 640
avg_degree = 16
feat_dim = 80




def get_inputs():
    return list(popcorn_pri.make_csr_degree_normalized_graph(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 gcn_aggregate_kernel(
    const int* __restrict__ row_ptr,
    const int* __restrict__ col_idx,
    const float* __restrict__ node_feat,
    const float* __restrict__ degrees,
    float* __restrict__ out,
    int num_nodes,
    int feat_dim
) {
    int dst = blockIdx.x * blockDim.x + threadIdx.x;
    if (dst >= num_nodes) return;

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

    if (end <= start) return;

    float deg_dst = degrees[dst];
    float inv_sqrt_dst = rsqrtf(deg_dst);

    // Accumulate features for this destination node
    float* local_out = out + dst * feat_dim;

    // Initialize local output
    for (int f = 0; f < feat_dim; f++) {
        local_out[f] = 0.0f;
    }

    // Iterate over neighbors
    for (int edge = start; edge < end; edge++) {
        int src = col_idx[edge];
        float deg_src = degrees[src];
        float norm = inv_sqrt_dst * rsqrtf(deg_src);

        const float* src_feat = node_feat + src * feat_dim;
        for (int f = 0; f < feat_dim; f++) {
            local_out[f] += src_feat[f] * norm;
        }
    }
}

torch::Tensor gcn_aggregate_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor node_feat,
    torch::Tensor degrees
) {
    int num_nodes = row_ptr.numel() - 1;
    int feat_dim = node_feat.size(1);

    // Convert to correct types
    torch::Tensor row_ptr_int = row_ptr.to(torch::kInt32).contiguous();
    torch::Tensor col_idx_int = col_idx.to(torch::kInt32).contiguous();
    torch::Tensor node_feat_float = node_feat.to(torch::kFloat32).contiguous();
    torch::Tensor degrees_float = degrees.to(torch::kFloat32).contiguous();

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(node_feat.device());
    auto out = torch::zeros({num_nodes, feat_dim}, options);

    int threads = 256;
    int blocks = (num_nodes + threads - 1) / threads;

    gcn_aggregate_kernel<<<blocks, threads>>>(
        row_ptr_int.data_ptr<int>(),
        col_idx_int.data_ptr<int>(),
        node_feat_float.data_ptr<float>(),
        degrees_float.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes,
        feat_dim
    );

    return out;
}
"""

cpp_source = "torch::Tensor gcn_aggregate_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat, torch::Tensor degrees);"

_gcn_module = load_inline(
    name='gcn_aggregate',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gcn_aggregate_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.gcn_aggregate = _gcn_module.gcn_aggregate_cuda

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        node_feat: torch.Tensor,
        degrees: torch.Tensor,
    ) -> torch.Tensor:
        return self.gcn_aggregate(row_ptr, col_idx, node_feat, degrees)