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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cstdint>

__global__ void csr_softmax_kernel(const int64_t* __restrict__ row_ptr, const float* __restrict__ edge_scores, float* __restrict__ out, int num_nodes) {
    int node = blockIdx.x;
    if (node >= num_nodes) return;
    
    int64_t start = row_ptr[node];
    int64_t end = row_ptr[node + 1];
    int64_t len = end - start;
    if (len <= 0) return;
    
    float max_val = -FLT_MAX;
    for (int64_t i = start; i < end; ++i) {
        max_val = fmaxf(max_val, edge_scores[i]);
    }
    
    float sum = 0.0f;
    for (int64_t i = start; i < end; ++i) {
        sum += expf(edge_scores[i] - max_val);
    }
    
    for (int64_t i = start; i < end; ++i) {
        out[i] = expf(edge_scores[i] - max_val) / sum;
    }
}

torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    int num_nodes = row_ptr.size(0) - 1;
    auto out = torch::empty_like(edge_scores);
    
    dim3 grid(num_nodes);
    dim3 block(1);
    
    csr_softmax_kernel<<<grid, block>>>(
        row_ptr.data_ptr<int64_t>(),
        edge_scores.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes
    );
    
    return out;
}
'''

cpp_source = '''
torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores);
'''

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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module
    
    def forward(self, row_ptr, edge_scores):
        row_ptr = row_ptr.to(torch.int64)
        return self.module.csr_softmax_cuda(row_ptr, edge_scores)