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 os
_tmpdir = os.path.join(os.getcwd(), 'tmp_nvcc_csr')
os.makedirs(_tmpdir, exist_ok=True)
os.environ['TMPDIR'] = _tmpdir
os.environ.setdefault('MAX_JOBS', '1')

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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor csr_forward_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_weight, torch::Tensor node_feat);

torch::Tensor csr_forward(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_weight, torch::Tensor node_feat) {
    return csr_forward_cuda(row_ptr, col_idx, edge_weight, node_feat);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

template <typename row_t, typename col_t>
__global__ __launch_bounds__(96, 8) void csr96_kernel_t(
        const row_t* __restrict__ row_ptr,
        const col_t* __restrict__ col_idx,
        const float* __restrict__ edge_weight,
        const float* __restrict__ node_feat,
        float* __restrict__ out) {
    const int dst = blockIdx.x;
    const int f = threadIdx.x;
    const int start = (int)row_ptr[dst];
    const int end = (int)row_ptr[dst + 1];
    float acc = 0.0f;
    int e = start;
    for (; e + 3 < end; e += 4) {
        int s0 = (int)col_idx[e];
        float w0 = edge_weight[e];
        acc = fmaf(node_feat[s0 * 96 + f], w0, acc);
        int s1 = (int)col_idx[e + 1];
        float w1 = edge_weight[e + 1];
        acc = fmaf(node_feat[s1 * 96 + f], w1, acc);
        int s2 = (int)col_idx[e + 2];
        float w2 = edge_weight[e + 2];
        acc = fmaf(node_feat[s2 * 96 + f], w2, acc);
        int s3 = (int)col_idx[e + 3];
        float w3 = edge_weight[e + 3];
        acc = fmaf(node_feat[s3 * 96 + f], w3, acc);
    }
    for (; e < end; ++e) {
        int src = (int)col_idx[e];
        float w = edge_weight[e];
        acc = fmaf(node_feat[src * 96 + f], w, acc);
    }
    out[dst * 96 + f] = acc;
}

template <typename row_t, typename col_t>
static inline void launch_csr(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_weight,
                              torch::Tensor node_feat, torch::Tensor out, int num_nodes, cudaStream_t stream) {
    csr96_kernel_t<row_t, col_t><<<num_nodes, 96, 0, stream>>>(
        row_ptr.data_ptr<row_t>(), col_idx.data_ptr<col_t>(), edge_weight.data_ptr<float>(),
        node_feat.data_ptr<float>(), out.data_ptr<float>());
}

torch::Tensor csr_forward_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_weight, torch::Tensor node_feat) {
    const int num_nodes = row_ptr.numel() - 1;
    const int feat_dim = node_feat.size(1);
    auto out = torch::empty({num_nodes, feat_dim}, node_feat.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (row_ptr.scalar_type() == at::kInt && col_idx.scalar_type() == at::kInt) {
        launch_csr<int,int>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kFloat && col_idx.scalar_type() == at::kFloat) {
        launch_csr<float,float>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kLong && col_idx.scalar_type() == at::kLong) {
        launch_csr<long,long>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kInt && col_idx.scalar_type() == at::kFloat) {
        launch_csr<int,float>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kFloat && col_idx.scalar_type() == at::kInt) {
        launch_csr<float,int>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kInt && col_idx.scalar_type() == at::kLong) {
        launch_csr<int,long>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else if (row_ptr.scalar_type() == at::kLong && col_idx.scalar_type() == at::kInt) {
        launch_csr<long,int>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    } else {
        launch_csr<float,float>(row_ptr, col_idx, edge_weight, node_feat, out, num_nodes, stream);
    }
    return out;
}
'''

_ext = load_inline(
    name='csr_mp_ext_submit_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['csr_forward'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    verbose=False,
)

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 _ext.csr_forward(row_ptr, col_idx, edge_weight, node_feat)