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 = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void sparse_row_softmax_kernel(
    const int* __restrict__ row_ptr,
    const float* __restrict__ edge_scores,
    float* __restrict__ out,
    int num_nodes
) {
    int node = blockIdx.x;
    if (node >= num_nodes) return;

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

    if (degree <= 0) return;

    // Each block handles one row; use shared memory for reduction
    extern __shared__ float sdata[];

    float thread_max = -FLT_MAX;
    for (int i = threadIdx.x; i < degree; i += blockDim.x) {
        float val = edge_scores[start + i];
        thread_max = fmaxf(thread_max, val);
    }

    // Block reduce max
    sdata[threadIdx.x] = thread_max;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] = fmaxf(sdata[threadIdx.x], sdata[threadIdx.x + s]);
        }
        __syncthreads();
    }
    float row_max = sdata[0];
    __syncthreads();

    float thread_sum = 0.0f;
    for (int i = threadIdx.x; i < degree; i += blockDim.x) {
        float val = expf(edge_scores[start + i] - row_max);
        out[start + i] = val;  // temporarily store exp
        thread_sum += val;
    }

    sdata[threadIdx.x] = thread_sum;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }
    float row_sum = sdata[0];

    float inv_sum = 1.0f / row_sum;
    for (int i = threadIdx.x; i < degree; i += blockDim.x) {
        out[start + i] *= inv_sum;
    }
}

torch::Tensor sparse_row_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);

    const int threads = 32;  // one warp per row
    dim3 blocks(num_nodes);

    size_t shared_mem = threads * sizeof(float);

    sparse_row_softmax_kernel<<<blocks, threads, shared_mem>>>(
        row_ptr.data_ptr<int>(),
        edge_scores.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes
    );

    return out;
}
'''

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

sparse_softmax = load_inline(
    name='sparse_softmax',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['sparse_row_softmax_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        return sparse_softmax.sparse_row_softmax_cuda(row_ptr, edge_scores)
```