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 <cmath>

__global__ void sparse_attention_kernel(
    const int* __restrict__ row_ptr,
    const int* __restrict__ col_idx,
    const float* __restrict__ edge_scores,
    const float* __restrict__ node_value,
    float* __restrict__ out,
    const int num_nodes,
    const int feat_dim) {
    
    // Each block handles one destination node
    const int dst = blockIdx.x;
    if (dst >= num_nodes) return;
    
    const int start = row_ptr[dst];
    const int end = row_ptr[dst + 1];
    const int num_edges = end - start;
    
    if (num_edges == 0) return;
    
    // Use shared memory for softmax weights and reduction
    extern __shared__ float smem[];
    float* edge_weights = smem;
    float* shared_reduce = smem + num_edges;  // 32 floats for reduction
    
    const int lane = threadIdx.x % 32;
    const int warp = threadIdx.x / 32;
    const int num_warps = blockDim.x / 32;
    
    // Find max score for numerical stability
    float local_max = -INFINITY;
    for (int i = threadIdx.x; i < num_edges; i += blockDim.x) {
        float score = edge_scores[start + i];
        if (score > local_max) local_max = score;
    }
    
    // Warp reduce max using shuffle
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, local_max, offset);
        if (other > local_max) local_max = other;
    }
    
    // Write warp result to shared memory
    if (lane == 0) {
        shared_reduce[warp] = local_max;
    }
    __syncthreads();
    
    // First warp reads all results and reduces
    float max_score = -INFINITY;
    if (warp == 0) {
        max_score = (lane < num_warps) ? shared_reduce[lane] : -INFINITY;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            float other = __shfl_down_sync(0xffffffff, max_score, offset);
            if (other > max_score) max_score = other;
        }
    }
    
    // Broadcast max_score to all threads
    max_score = __shfl_sync(0xffffffff, max_score, 0);
    
    // Compute exp(score - max) and store in shared memory
    float local_sum = 0.0f;
    for (int i = threadIdx.x; i < num_edges; i += blockDim.x) {
        float exp_val = expf(edge_scores[start + i] - max_score);
        edge_weights[i] = exp_val;
        local_sum += exp_val;
    }
    
    // Warp reduce sum using shuffle
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    
    // Write warp result to shared memory
    if (lane == 0) {
        shared_reduce[warp] = local_sum;
    }
    __syncthreads();
    
    // First warp reads all results and reduces
    float total_sum = 0.0f;
    if (warp == 0) {
        total_sum = (lane < num_warps) ? shared_reduce[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            total_sum += __shfl_down_sync(0xffffffff, total_sum, offset);
        }
    }
    
    // Broadcast total_sum to all threads
    total_sum = __shfl_sync(0xffffffff, total_sum, 0);
    
    // Normalize weights
    for (int i = threadIdx.x; i < num_edges; i += blockDim.x) {
        edge_weights[i] /= total_sum;
    }
    __syncthreads();
    
    // Compute weighted sum of node values
    // Each thread handles a subset of feature dimensions
    for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {
        float sum = 0.0f;
        for (int e = 0; e < num_edges; e++) {
            int src = col_idx[start + e];
            sum += node_value[src * feat_dim + f] * edge_weights[e];
        }
        out[dst * feat_dim + f] = sum;
    }
}

torch::Tensor sparse_attention_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor edge_scores,
    torch::Tensor node_value) {
    
    const int num_nodes = row_ptr.numel() - 1;
    const int feat_dim = node_value.size(1);
    
    // Ensure contiguous and correct types
    row_ptr = row_ptr.contiguous();
    col_idx = col_idx.contiguous();
    edge_scores = edge_scores.contiguous();
    node_value = node_value.contiguous();
    
    auto out = torch::zeros({num_nodes, feat_dim}, node_value.options());
    
    const int block_size = 256;
    // Shared memory for edge weights + 32 floats for reduction
    const size_t smem_size = edge_scores.size(0) * sizeof(float) + 32 * sizeof(float);
    
    sparse_attention_kernel<<<num_nodes, block_size, smem_size>>>(
        row_ptr.data_ptr<int>(),
        col_idx.data_ptr<int>(),
        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_attention_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_scores, torch::Tensor node_value);
"""

sparse_attention_extension = load_inline(
    name='sparse_attention_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sparse_attention_cuda'],
    verbose=False,
)

class ModelNew(nn.Module):
    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:
        return sparse_attention_extension.sparse_attention_cuda(row_ptr, col_idx, edge_scores, node_value)