import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Multi-head sparse aggregation with per-edge attention weights.

    Each edge carries one weight per head, and source node features are arranged
    as `(num_nodes, num_heads, head_dim)`. The output aggregates weighted
    source features into `(num_nodes, num_heads, head_dim)`.
    """

    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
        num_heads = node_feat.shape[1]
        head_dim = node_feat.shape[2]
        out = torch.zeros(num_nodes, num_heads, head_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 = 320
avg_degree = 14
num_heads = 8
head_dim = 32


def get_inputs():
    return list(
        popcorn_pri.make_csr_multihead_spmm_graph(
            num_nodes, avg_degree, num_heads, head_dim
        )
    )



def get_init_inputs():
    return []

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>

__global__ void sparse_aggregate_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 num_heads,
    int head_dim,
    int num_edges
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = num_nodes * num_heads * head_dim;
    if (idx >= total) return;

    int d = idx % head_dim;
    int h = (idx / head_dim) % num_heads;
    int dst = idx / (num_heads * head_dim);

    float sum = 0.0f;
    int start = row_ptr[dst];
    int end = row_ptr[dst + 1];
    for (int e = start; e < end; ++e) {
        int src = col_idx[e];
        float w = edge_weight[e * num_heads + h];
        float feat = node_feat[(src * num_heads + h) * head_dim + d];
        sum += w * feat;
    }
    out[(dst * num_heads + h) * head_dim + d] = sum;
}

torch::Tensor sparse_aggregate_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 num_heads = node_feat.size(1);
    int head_dim = node_feat.size(2);
    int num_edges = col_idx.size(0);

    auto out = torch::zeros({num_nodes, num_heads, head_dim}, node_feat.options());

    int total = num_nodes * num_heads * head_dim;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    sparse_aggregate_kernel<<<blocks, threads>>>(
        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, num_heads, head_dim, num_edges
    );

    return out;
}
'''

cpp_source = r'''
torch::Tensor sparse_aggregate_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor edge_weight,
    torch::Tensor node_feat
);
'''

module = load_inline(
    name='sparse_agg',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['sparse_aggregate_cuda'],
    verbose=False
)

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        edge_weight: torch.Tensor,
        node_feat: torch.Tensor,
    ) -> torch.Tensor:
        return self.module.sparse_aggregate_cuda(row_ptr, col_idx, edge_weight, node_feat)