import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Fused sparse attention-value aggregation in CSR format.

    For each destination row, compute a softmax over edge logits and then use
    the normalized weights to aggregate source node values. This is a common
    fused graph-attention pattern.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        edge_scores: torch.Tensor,
        node_value: torch.Tensor,
    ) -> torch.Tensor:
        num_nodes = row_ptr.numel() - 1
        feat_dim = node_value.shape[1]
        out = torch.zeros(num_nodes, feat_dim, dtype=node_value.dtype, device=node_value.device)

        for dst in range(num_nodes):
            start = int(row_ptr[dst].item())
            end = int(row_ptr[dst + 1].item())
            if end > start:
                weights = torch.softmax(edge_scores[start:end], dim=0)
                src = col_idx[start:end].long()
                out[dst] = (node_value[src] * weights.unsqueeze(-1)).sum(dim=0)

        return out


num_nodes = 384
avg_degree = 18
feat_dim = 64


def get_inputs():
    return list(
        popcorn_pri.make_csr_fused_attention_value_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>
#include <cstdint>

__global__ void sparse_attention_aggregation(
    const int32_t* __restrict__ row_ptr,
    const int32_t* __restrict__ col_idx,
    const float* __restrict__ edge_scores,
    const float* __restrict__ node_value,
    float* __restrict__ out,
    int num_nodes,
    int feat_dim
) {
    int dst = blockIdx.x;
    if (dst >= num_nodes) return;

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

    if (degree <= 0) return;

    // Compute max
    float thread_max = -INFINITY;
    for (int i = threadIdx.x; i < degree; i += blockDim.x) {
        float s = edge_scores[start + i];
        thread_max = fmaxf(thread_max, s);
    }

    // Block reduce max
    __shared__ float shared_max[32];
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;
    thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 16));
    thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 8));
    thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 4));
    thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 2));
    thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 1));
    if (lane == 0) shared_max[wid] = thread_max;
    __syncthreads();
    if (wid == 0) {
        thread_max = (threadIdx.x < (blockDim.x + 31)/32) ? shared_max[lane] : -INFINITY;
        thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 16));
        thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 8));
        thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 4));
        thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 2));
        thread_max = fmaxf(thread_max, __shfl_down_sync(0xffffffff, thread_max, 1));
        if (lane == 0) shared_max[0] = thread_max;
    }
    __syncthreads();
    float max_val = shared_max[0];

    // Compute sum_exp
    float thread_sum = 0.0f;
    for (int i = threadIdx.x; i < degree; i += blockDim.x) {
        float s = expf(edge_scores[start + i] - max_val);
        thread_sum += s;
    }

    // Block reduce sum
    __shared__ float shared_sum[32];
    thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 16);
    thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 8);
    thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 4);
    thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 2);
    thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 1);
    if (lane == 0) shared_sum[wid] = thread_sum;
    __syncthreads();
    if (wid == 0) {
        thread_sum = (threadIdx.x < (blockDim.x + 31)/32) ? shared_sum[lane] : 0.0f;
        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 16);
        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 8);
        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 4);
        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 2);
        thread_sum += __shfl_down_sync(0xffffffff, thread_sum, 1);
        if (lane == 0) shared_sum[0] = thread_sum;
    }
    __syncthreads();
    float sum_exp = shared_sum[0];
    float inv_sum = 1.0f / sum_exp;

    // Now compute weighted sum for each feature
    for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {
        float acc = 0.0f;
        for (int i = 0; i < degree; i++) {
            int src = col_idx[start + i];
            float w = expf(edge_scores[start + i] - max_val) * inv_sum;
            acc += node_value[src * feat_dim + f] * w;
        }
        out[dst * feat_dim + f] = acc;
    }
}

at::Tensor sparse_agg_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor edge_scores,
    torch::Tensor node_value
) {
    int num_nodes = row_ptr.size(0) - 1;
    int feat_dim = node_value.size(1);
    auto out = torch::zeros({num_nodes, feat_dim}, node_value.options());

    dim3 blocks(num_nodes);
    dim3 threads(256);
    sparse_attention_aggregation<<<blocks, threads>>>(
        row_ptr.data_ptr<int32_t>(),
        col_idx.data_ptr<int32_t>(),
        edge_scores.data_ptr<float>(),
        node_value.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes,
        feat_dim
    );
    return out;
}
"""

cpp_source = """
torch::Tensor sparse_agg_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_scores, torch::Tensor node_value);
"""

module = load_inline(
    name="sparse_agg",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["sparse_agg_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3"]
)

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

    def forward(self, row_ptr, col_idx, edge_scores, node_value):
        return self.module.sparse_agg_cuda(row_ptr, col_idx, edge_scores, node_value)