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 os
os.makedirs('/scratch/abaner/tmp', exist_ok=True)
os.environ['TMPDIR'] = '/scratch/abaner/tmp'
os.environ['TEMP'] = '/scratch/abaner/tmp'
os.environ['TMP'] = '/scratch/abaner/tmp'
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = """
#include <torch/extension.h>
torch::Tensor gcn_forward_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat, torch::Tensor degrees);
"""

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

#define BLOCK_THREADS 96
#define CHUNK 128

__global__ __launch_bounds__(BLOCK_THREADS, 8)
void gcn_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;
    int tid = threadIdx.x;
    int start = row_ptr[dst];
    int end = row_ptr[dst + 1];
    int deg = end - start;

    __shared__ int sh_src[CHUNK];
    __shared__ float sh_norm[CHUNK];

    float acc = 0.0f;
    bool active = tid < feat_dim;
    float ddeg = degrees[dst];

    for (int base = 0; base < deg; base += CHUNK) {
        int n = deg - base;
        if (n > CHUNK) n = CHUNK;
        if (tid < n) {
            int src = col_idx[start + base + tid];
            sh_src[tid] = src;
            sh_norm[tid] = 1.0f / sqrtf(ddeg * degrees[src]);
        }
        __syncthreads();
        if (active) {
            for (int j = 0; j < n; ++j) {
                int src = sh_src[j];
                acc += node_feat[src * feat_dim + tid] * sh_norm[j];
            }
        }
        __syncthreads();
    }
    if (active) out[dst * feat_dim + tid] = acc;
}

torch::Tensor gcn_forward_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);
    auto out = torch::empty({num_nodes, feat_dim}, node_feat.options());
    gcn_kernel<<<num_nodes, BLOCK_THREADS>>>(row_ptr.data_ptr<int>(), col_idx.data_ptr<int>(), node_feat.data_ptr<float>(), degrees.data_ptr<float>(), out.data_ptr<float>(), num_nodes, feat_dim);
    return out;
}
"""

_ext = load_inline(
    name="gcn_norm_final_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["gcn_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    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:
        if row_ptr.dtype is not torch.int32:
            row_ptr = row_ptr.to(torch.int32)
        if col_idx.dtype is not torch.int32:
            col_idx = col_idx.to(torch.int32)
        return _ext.gcn_forward_cuda(row_ptr, col_idx, node_feat, degrees)