import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Per-destination max aggregation over neighbor features in CSR format.

    This pattern appears in max-pooling GNNs and neighborhood feature
    extraction where reductions are not sums but feature-wise maxima.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        node_feat: torch.Tensor,
    ) -> torch.Tensor:
        num_nodes = row_ptr.numel() - 1
        feat_dim = node_feat.shape[1]
        out = torch.full((num_nodes, feat_dim), float("-inf"), dtype=node_feat.dtype, device=node_feat.device)

        for dst in range(num_nodes):
            start = int(row_ptr[dst].item())
            end = int(row_ptr[dst + 1].item())
            if end > start:
                out[dst] = node_feat[col_idx[start:end].long()].amax(dim=0)

        return out


num_nodes = 640
avg_degree = 16
feat_dim = 80


def get_inputs():
    return list(popcorn_pri.make_csr_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_runtime.h>
#include <cstdint>

__global__ void max_aggregation_kernel(
    const int32_t* __restrict__ row_ptr,
    const int32_t* __restrict__ col_idx,
    const float* __restrict__ node_feat,
    float* __restrict__ out,
    int num_nodes,
    int feat_dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = num_nodes * feat_dim;
    if (idx >= total) return;

    int dst = idx / feat_dim;
    int f = idx % feat_dim;

    int32_t start = row_ptr[dst];
    int32_t end = row_ptr[dst + 1];

    float max_val = -INFINITY;
    for (int32_t i = start; i < end; ++i) {
        int32_t src = col_idx[i];
        float val = node_feat[src * feat_dim + f];
        if (val > max_val) max_val = val;
    }
    out[dst * feat_dim + f] = max_val;
}

torch::Tensor max_aggregation_cuda(
    torch::Tensor row_ptr,
    torch::Tensor col_idx,
    torch::Tensor node_feat
) {
    int num_nodes = row_ptr.size(0) - 1;
    int feat_dim = node_feat.size(1);

    auto out = torch::full({num_nodes, feat_dim}, -INFINITY, node_feat.options());

    int total_threads = num_nodes * feat_dim;
    int block_size = 256;
    int grid_size = (total_threads + block_size - 1) / block_size;

    max_aggregation_kernel<<<grid_size, block_size>>>(
        row_ptr.data_ptr<int32_t>(),
        col_idx.data_ptr<int32_t>(),
        node_feat.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes,
        feat_dim
    );

    return out;
}
"""

cpp_source = """
torch::Tensor max_aggregation_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat);
"""

module = load_inline(
    name="max_aggregation",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["max_aggregation_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

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

    def forward(self, row_ptr, col_idx, node_feat):
        if not row_ptr.is_cuda:
            row_ptr = row_ptr.cuda()
            col_idx = col_idx.cuda()
            node_feat = node_feat.cuda()
        out = self.module.max_aggregation_cuda(row_ptr, col_idx, node_feat)
        return out.cpu()
```