import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Multi-head row-wise softmax over CSR edge scores.

    This extends scalar graph edge softmax to multiple attention heads while
    preserving sparse CSR layout. Each row is normalized independently for each
    head.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> torch.Tensor:
        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 = 384
avg_degree = 18
num_heads = 8




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



def get_init_inputs():
    return []

import torch
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void csr_softmax_kernel(const int32_t* __restrict__ row_ptr, const float* __restrict__ edge_scores, float* __restrict__ out, int num_nodes, int num_heads) {
    int node = blockIdx.x;
    if (node >= num_nodes) return;
    int32_t start = row_ptr[node];
    int32_t end = row_ptr[node + 1];
    int degree = end - start;
    if (degree <= 0) return;

    for(int h = 0; h < num_heads; h++) {
        float max_val = -INFINITY;
        for(int i = 0; i < degree; i++) {
            float v = edge_scores[(start + i) * num_heads + h];
            max_val = fmaxf(max_val, v);
        }
        float sum = 0.0f;
        for(int i = 0; i < degree; i++) {
            float v = edge_scores[(start + i) * num_heads + h];
            sum += expf(v - max_val);
        }
        for(int i = 0; i < degree; i++) {
            float v = edge_scores[(start + i) * num_heads + h];
            out[(start + i) * num_heads + h] = expf(v - max_val) / sum;
        }
    }
}

torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    TORCH_CHECK(row_ptr.is_cuda(), "row_ptr must be CUDA tensor");
    TORCH_CHECK(edge_scores.is_cuda(), "edge_scores must be CUDA tensor");
    TORCH_CHECK(row_ptr.dtype() == torch::kInt32, "row_ptr must be int32");
    int num_nodes = row_ptr.size(0) - 1;
    int num_heads = edge_scores.size(1);
    auto out = torch::empty_like(edge_scores);
    dim3 grid(num_nodes);
    csr_softmax_kernel<<<grid, 1>>>(row_ptr.data_ptr<int32_t>(), edge_scores.data_ptr<float>(), out.data_ptr<float>(), num_nodes, num_heads);
    cudaDeviceSynchronize();
    return out;
}
'''

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

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

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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        orig_device = row_ptr.device
        row_ptr = row_ptr.cuda().to(torch.int32)
        edge_scores = edge_scores.cuda()
        out = self.module.csr_softmax_cuda(row_ptr, edge_scores)
        return out.to(orig_device)