import torch
import torch.nn as nn

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


class Model(nn.Module):
    """
    Fused radius (ball) neighborhood query over a 3D point map with per-query feature
    aggregation—common in navigation / manipulation perception (PointNet++-style set
    abstraction, local maps around end-effector or goal queries on LiDAR / depth).

    Single forward fusion target:
      1) pairwise Euclidean distances between each query center and every point,
      2) membership in the closed ball of radius ``r``,
      3) channel-wise **max** pooling of point features over in-ball neighbors.

    Points with no neighbors inside the radius contribute a **zero** feature vector
    for that query (instead of ``-inf`` max).
    """

    def __init__(self, in_channels: int, radius: float):
        super().__init__()
        self.in_channels = in_channels
        self.register_buffer("radius_sq", torch.tensor(float(radius) ** 2))

    def forward(
        self,
        points: torch.Tensor,
        queries: torch.Tensor,
        features: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            points: (B, P, 3) point coordinates (e.g. world or sensor frame).
            queries: (B, Q, 3) ball centers (e.g. navigation waypoints, grasp samples).
            features: (B, P, C) per-point descriptors (geometry or lifted RGB-D).
        Returns:
            (B, Q, C) max-pooled features over neighbors with ||p - q||_2 <= r.
        """
        B, P, _ = points.shape
        _, Q, _ = queries.shape
        _, P2, C = features.shape
        assert P == P2 and C == self.in_channels

        # Squared distances (B, Q, P)
        dist_sq = torch.cdist(queries, points, p=2.0) ** 2
        in_ball = dist_sq <= self.radius_sq

        # (B, 1, P, C) broadcasts with (B, Q, P) for masked max over P
        feats = features.unsqueeze(1)
        neg_large = torch.finfo(features.dtype).min
        masked = torch.where(
            in_ball.unsqueeze(-1),
            feats,
            torch.full_like(feats, neg_large),
        )
        pooled, _ = masked.max(dim=2)

        empty = in_ball.sum(dim=2) == 0
        if empty.any():
            pooled = pooled.masked_fill(empty.unsqueeze(-1), 0.0)

        return pooled


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self, in_channels: int, radius: float):
        super().__init__()
        self._impl = Model(in_channels, radius)

    def forward(
        self,
        points: torch.Tensor,
        queries: torch.Tensor,
        features: torch.Tensor,
    ) -> torch.Tensor:
        return self._impl(points, queries, features)


batch_size = 4
num_points = 512
num_queries = 64
in_channels = 32
radius = 0.12


def get_inputs():
    p = popcorn_pri
    points = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_points), 3)
    queries = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_queries), 3)
    features = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_points), in_channels)
    return [points, queries, features]


def get_init_inputs():
    return [in_channels, radius]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


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

torch::Tensor ball_max_forward(torch::Tensor points, torch::Tensor queries, torch::Tensor features, double radius_sq);
"""

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

__device__ __forceinline__ void atomicMaxFloatShared(float* addr, float value) {
    int* addr_i = reinterpret_cast<int*>(addr);
    int old = *addr_i;
    int assumed;
    // CAS loop using floating point comparison, works for negative values too.
    while (value > __int_as_float(old)) {
        assumed = old;
        old = atomicCAS(addr_i, assumed, __float_as_int(value));
        if (old == assumed) break;
    }
}

__global__ void ball_max_sparse_kernel(
    const float* __restrict__ points,
    const float* __restrict__ queries,
    const float* __restrict__ features,
    float* __restrict__ out,
    int B, int P, int Q, int C,
    float radius_sq)
{
    int b = blockIdx.x;
    int q = blockIdx.y;
    int tid = threadIdx.x;

    extern __shared__ unsigned char smem[];
    float* smax = reinterpret_cast<float*>(smem);
    int* sfound = reinterpret_cast<int*>(smax + C);

    for (int c = tid; c < C; c += blockDim.x) {
        smax[c] = -FLT_MAX;
    }
    if (tid == 0) *sfound = 0;
    __syncthreads();

    const float qx = queries[(b * Q + q) * 3 + 0];
    const float qy = queries[(b * Q + q) * 3 + 1];
    const float qz = queries[(b * Q + q) * 3 + 2];

    const int point_base = b * P * 3;
    const int feat_base = b * P * C;

    for (int p = tid; p < P; p += blockDim.x) {
        float dx = points[point_base + p * 3 + 0] - qx;
        float dy = points[point_base + p * 3 + 1] - qy;
        float dz = points[point_base + p * 3 + 2] - qz;
        float d2 = dx * dx + dy * dy + dz * dz;
        if (d2 <= radius_sq) {
            atomicExch(sfound, 1);
            const float* fptr = features + feat_base + p * C;
            // In the benchmark C=32, so this is a short, rare loop because the radius is small.
            for (int c = 0; c < C; ++c) {
                atomicMaxFloatShared(smax + c, fptr[c]);
            }
        }
    }
    __syncthreads();

    const int out_base = (b * Q + q) * C;
    if (*sfound) {
        for (int c = tid; c < C; c += blockDim.x) {
            out[out_base + c] = smax[c];
        }
    } else {
        for (int c = tid; c < C; c += blockDim.x) {
            out[out_base + c] = 0.0f;
        }
    }
}

torch::Tensor ball_max_forward(torch::Tensor points, torch::Tensor queries, torch::Tensor features, double radius_sq) {
    const int B = (int)points.size(0);
    const int P = (int)points.size(1);
    const int Q = (int)queries.size(1);
    const int C = (int)features.size(2);

    auto out = torch::empty({B, Q, C}, features.options());

    dim3 grid(B, Q);
    const int threads = 256;
    size_t shmem = (size_t)C * sizeof(float) + sizeof(int);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    ball_max_sparse_kernel<<<grid, threads, shmem, stream>>>(
        points.data_ptr<float>(), queries.data_ptr<float>(), features.data_ptr<float>(), out.data_ptr<float>(),
        B, P, Q, C, (float)radius_sq);
    return out;
}
"""

_ext = load_inline(
    name="ball_max_pool_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["ball_max_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, in_channels: int, radius: float):
        super().__init__()
        self.in_channels = in_channels
        self.radius_sq = float(radius) * float(radius)

    def forward(self, points: torch.Tensor, queries: torch.Tensor, features: torch.Tensor) -> torch.Tensor:
        return _ext.ball_max_forward(points, queries, features, self.radius_sq)