import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Row-wise softmax over edge scores stored in CSR format.

    This is a core primitive in graph attention networks and sparse attention on
    irregular neighborhoods. Each node owns a contiguous edge segment in
    `edge_scores[row_ptr[i]:row_ptr[i + 1]]`, and the kernel normalizes scores
    independently within each segment.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            row_ptr:      (num_nodes + 1,) CSR row offsets
            edge_scores:  (num_edges,) unnormalized attention logits
        Returns:
            (num_edges,) softmax-normalized edge weights per source node
        """
        num_nodes = row_ptr.numel() - 1
        out = torch.empty_like(edge_scores)

        for node in range(num_nodes):
            start = int(row_ptr[node].item())
            end = int(row_ptr[node + 1].item())
            if end > start:
                out[start:end] = torch.softmax(edge_scores[start:end], dim=0)

        return out


num_nodes = 512
avg_degree = 24
batchless_num_edges = num_nodes * avg_degree


def _make_row_ptr(num_nodes: int, avg_degree: int) -> torch.Tensor:
    base = torch.full((num_nodes,), avg_degree, dtype=torch.int32)
    return row_ptr


def get_inputs():
    return list(popcorn_pri.make_csr_scalar_edge_scores(num_nodes, avg_degree))



def get_init_inputs():
    return []

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

# Define the CUDA kernel
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/extension.h>

// Warp-level reduction for max
__inline__ __device__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset /= 2) {
        val = max(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

// Warp-level reduction for sum
__inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Each warp handles one row (node)
__global__ void row_softmax_kernel(
    torch::PackedTensorAccessor32<int, 1, torch::RestrictPtrTraits> row_ptr,
    torch::PackedTensorAccessor32<float, 1, torch::RestrictPtrTraits> edge_scores,
    torch::PackedTensorAccessor32<float, 1, torch::RestrictPtrTraits> out,
    int num_nodes) {
    
    int node = blockIdx.x * blockDim.y + threadIdx.y;  // One warp per row
    int lane = threadIdx.x;  // Lane within warp
    
    if (node >= num_nodes) return;
    
    int start = row_ptr[node];
    int end = row_ptr[node + 1];
    int row_len = end - start;
    
    // Find max in this row
    float max_val = -INFINITY;
    for (int i = lane; i < row_len; i += 32) {
        float val = edge_scores[start + i];
        max_val = max(max_val, val);
    }
    max_val = warp_reduce_max(max_val);
    max_val = __shfl_sync(0xffffffff, max_val, 0);  // Broadcast to all lanes
    
    if (row_len == 0) return;  // No edges for this node
    
    // Compute exp(x - max) and sum
    float sum = 0.0f;
    for (int i = lane; i < row_len; i += 32) {
        float val = expf(edge_scores[start + i] - max_val);
        out[start + i] = val;
        sum += val;
    }
    sum = warp_reduce_sum(sum);
    sum = __shfl_sync(0xffffffff, sum, 0);  // Broadcast to all lanes
    
    // Normalize
    float inv_sum = 1.0f / sum;
    for (int i = lane; i < row_len; i += 32) {
        out[start + i] *= inv_sum;
    }
}

// Host function to launch kernel
torch::Tensor row_softmax_cuda(
    torch::Tensor row_ptr,
    torch::Tensor edge_scores) {
    
    int num_nodes = row_ptr.numel() - 1;
    
    // Ensure contiguous and correct types
    row_ptr = row_ptr.contiguous().to(torch::kInt);
    edge_scores = edge_scores.contiguous();
    
    auto out = torch::empty_like(edge_scores);
    
    // Launch with warps per row
    int warps_per_block = 8;  // 8 warps = 256 threads per block
    int blocks = (num_nodes + warps_per_block - 1) / warps_per_block;
    
    dim3 block(32, warps_per_block);  // 32 threads per warp, 8 warps per block
    
    row_softmax_kernel<<<blocks, block>>>(
        row_ptr.packed_accessor32<int, 1, torch::RestrictPtrTraits>(),
        edge_scores.packed_accessor32<float, 1, torch::RestrictPtrTraits>(),
        out.packed_accessor32<float, 1, torch::RestrictPtrTraits>(),
        num_nodes);
    
    return out;
}
"""

cpp_source = "torch::Tensor row_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores);"

# Load the extension
module = load_inline(
    name='row_softmax_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['row_softmax_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.ext = module
    
    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        return self.ext.row_softmax_cuda(row_ptr, edge_scores)