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


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

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

torch::Tensor sparse_agg_forward(torch::Tensor row_ptr,
                                 torch::Tensor col_idx,
                                 torch::Tensor edge_weight,
                                 torch::Tensor node_feat) {
    return sparse_agg_forward_cuda(row_ptr, col_idx, edge_weight, node_feat);
}
'''

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

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

    int d = tid % head_dim;
    int t = tid / head_dim;
    int h = t % num_heads;
    int dst = t / num_heads;

    int start = row_ptr[dst];
    int end = row_ptr[dst + 1];
    float sum = 0.0f;

    // node_feat layout: [node][head][dim], edge_weight: [edge][head]
    int hd = num_heads * head_dim;
    for (int e = start; e < end; ++e) {
        int src = col_idx[e];
        float w = edge_weight[e * num_heads + h];
        sum += node_feat[src * hd + h * head_dim + d] * w;
    }
    out[tid] = sum;
}

torch::Tensor sparse_agg_forward_cuda(torch::Tensor row_ptr,
                                      torch::Tensor col_idx,
                                      torch::Tensor edge_weight,
                                      torch::Tensor node_feat) {
    const int num_nodes = row_ptr.size(0) - 1;
    const int num_heads = node_feat.size(1);
    const int head_dim = node_feat.size(2);
    auto out = torch::empty({num_nodes, num_heads, head_dim}, node_feat.options());
    int total = num_nodes * num_heads * head_dim;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    sparse_agg_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);
    return out;
}
'''

_ext = load_inline(
    name='sparse_agg_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['sparse_agg_forward'],
    extra_cflags=['-O3'],
    extra_cuda_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.sparse_agg_forward(row_ptr, col_idx, edge_weight, node_feat)