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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor maxagg_forward(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ __launch_bounds__(32) void maxagg80_kernel(
    const int* __restrict__ row_ptr,
    const int* __restrict__ col_idx,
    const float* __restrict__ node_feat,
    float* __restrict__ out,
    int num_nodes)
{
    int dst = blockIdx.x;
    int lane = threadIdx.x;
    if (dst >= num_nodes) return;

    if (lane < 20) {
        const int start = row_ptr[dst];
        const int end   = row_ptr[dst + 1];

        float4 m;
        m.x = -CUDART_INF_F;
        m.y = -CUDART_INF_F;
        m.z = -CUDART_INF_F;
        m.w = -CUDART_INF_F;

        const float4* __restrict__ feat4 = reinterpret_cast<const float4*>(node_feat);
        #pragma unroll 1
        for (int e = start; e < end; ++e) {
            int src = col_idx[e];
            float4 v = feat4[src * 20 + lane];
            m.x = fmaxf(m.x, v.x);
            m.y = fmaxf(m.y, v.y);
            m.z = fmaxf(m.z, v.z);
            m.w = fmaxf(m.w, v.w);
        }

        float4* __restrict__ out4 = reinterpret_cast<float4*>(out);
        out4[dst * 20 + lane] = m;
    }
}

__global__ __launch_bounds__(128) void maxagg_generic_kernel(
    const int* __restrict__ row_ptr,
    const int* __restrict__ col_idx,
    const float* __restrict__ node_feat,
    float* __restrict__ out,
    int num_nodes,
    int feat_dim)
{
    int dst = blockIdx.x;
    int tid = threadIdx.x;
    if (dst >= num_nodes) return;

    const int start = row_ptr[dst];
    const int end   = row_ptr[dst + 1];

    for (int f = tid; f < feat_dim; f += blockDim.x) {
        float m = -CUDART_INF_F;
        #pragma unroll 1
        for (int e = start; e < end; ++e) {
            int src = col_idx[e];
            float v = node_feat[src * feat_dim + f];
            m = fmaxf(m, v);
        }
        out[dst * feat_dim + f] = m;
    }
}

torch::Tensor maxagg_forward(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat) {
    const int num_nodes = (int)row_ptr.numel() - 1;
    const int feat_dim = (int)node_feat.size(1);
    auto out = torch::empty({num_nodes, feat_dim}, node_feat.options());

    c10::cuda::CUDAGuard device_guard(node_feat.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (feat_dim == 80) {
        maxagg80_kernel<<<num_nodes, 32, 0, stream>>>(
            row_ptr.data_ptr<int>(),
            col_idx.data_ptr<int>(),
            node_feat.data_ptr<float>(),
            out.data_ptr<float>(),
            num_nodes);
    } else {
        maxagg_generic_kernel<<<num_nodes, 128, 0, stream>>>(
            row_ptr.data_ptr<int>(),
            col_idx.data_ptr<int>(),
            node_feat.data_ptr<float>(),
            out.data_ptr<float>(),
            num_nodes,
            feat_dim);
    }
    return out;
}
"""

_ext = load_inline(
    name="maxagg_csr_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["maxagg_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


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

    def forward(self, row_ptr: torch.Tensor, col_idx: torch.Tensor, node_feat: torch.Tensor) -> torch.Tensor:
        return _ext.maxagg_forward(row_ptr, col_idx, node_feat)