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

cuda_source = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

// Optimized kernel with:
// 1. Vectorized loads (float4) for features when C is multiple of 4
// 2. ILP via loop unrolling
// 3. Cache points in registers for the distance check

__global__ void radius_ball_maxpool_kernel(
    const float* __restrict__ points,    // (B, P, 3)
    const float* __restrict__ queries,   // (B, Q, 3)
    const float* __restrict__ features,  // (B, P, C)
    float* __restrict__ output,           // (B, Q, C)
    const int B, const int P, const int Q, const int C,
    const float radius_sq
) {
    const int bq = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * Q;
    if (bq >= total) return;
    
    const int b = bq / Q;
    const int q = bq % Q;
    
    // Get query point coordinates
    const float qx = queries[b * Q * 3 + q * 3 + 0];
    const float qy = queries[b * Q * 3 + q * 3 + 1];
    const float qz = queries[b * Q * 3 + q * 3 + 2];
    
    // Track max values in registers
    // C=32 channels, 32 floats = 128 bytes per thread
    float m0 = -FLT_MAX, m1 = -FLT_MAX, m2 = -FLT_MAX, m3 = -FLT_MAX;
    float m4 = -FLT_MAX, m5 = -FLT_MAX, m6 = -FLT_MAX, m7 = -FLT_MAX;
    float m8 = -FLT_MAX, m9 = -FLT_MAX, m10 = -FLT_MAX, m11 = -FLT_MAX;
    float m12 = -FLT_MAX, m13 = -FLT_MAX, m14 = -FLT_MAX, m15 = -FLT_MAX;
    float m16 = -FLT_MAX, m17 = -FLT_MAX, m18 = -FLT_MAX, m19 = -FLT_MAX;
    float m20 = -FLT_MAX, m21 = -FLT_MAX, m22 = -FLT_MAX, m23 = -FLT_MAX;
    float m24 = -FLT_MAX, m25 = -FLT_MAX, m26 = -FLT_MAX, m27 = -FLT_MAX;
    float m28 = -FLT_MAX, m29 = -FLT_MAX, m30 = -FLT_MAX, m31 = -FLT_MAX;
    
    bool has_neighbor = false;
    
    // Base pointers for this batch
    const float* batch_points = points + b * P * 3;
    const float* batch_features = features + b * P * C;
    
    // Iterate over all points
    for (int p = 0; p < P; p++) {
        // Get point coordinates
        const float px = batch_points[p * 3 + 0];
        const float py = batch_points[p * 3 + 1];
        const float pz = batch_points[p * 3 + 2];
        
        // Compute squared distance
        const float dx = px - qx;
        const float dy = py - qy;
        const float dz = pz - qz;
        const float dist_sq = dx * dx + dy * dy + dz * dz;
        
        // Check if in ball
        if (dist_sq <= radius_sq) {
            has_neighbor = true;
            
            // Load features and update max using vectorized loads
            const float4* feat4 = (const float4*)(batch_features + p * C);
            
            float4 f0 = feat4[0];
            m0 = fmaxf(m0, f0.x); m1 = fmaxf(m1, f0.y); m2 = fmaxf(m2, f0.z); m3 = fmaxf(m3, f0.w);
            
            float4 f1 = feat4[1];
            m4 = fmaxf(m4, f1.x); m5 = fmaxf(m5, f1.y); m6 = fmaxf(m6, f1.z); m7 = fmaxf(m7, f1.w);
            
            float4 f2 = feat4[2];
            m8 = fmaxf(m8, f2.x); m9 = fmaxf(m9, f2.y); m10 = fmaxf(m10, f2.z); m11 = fmaxf(m11, f2.w);
            
            float4 f3 = feat4[3];
            m12 = fmaxf(m12, f3.x); m13 = fmaxf(m13, f3.y); m14 = fmaxf(m14, f3.z); m15 = fmaxf(m15, f3.w);
            
            float4 f4 = feat4[4];
            m16 = fmaxf(m16, f4.x); m17 = fmaxf(m17, f4.y); m18 = fmaxf(m18, f4.z); m19 = fmaxf(m19, f4.w);
            
            float4 f5 = feat4[5];
            m20 = fmaxf(m20, f5.x); m21 = fmaxf(m21, f5.y); m22 = fmaxf(m22, f5.z); m23 = fmaxf(m23, f5.w);
            
            float4 f6 = feat4[6];
            m24 = fmaxf(m24, f6.x); m25 = fmaxf(m25, f6.y); m26 = fmaxf(m26, f6.z); m27 = fmaxf(m27, f6.w);
            
            float4 f7 = feat4[7];
            m28 = fmaxf(m28, f7.x); m29 = fmaxf(m29, f7.y); m30 = fmaxf(m30, f7.z); m31 = fmaxf(m31, f7.w);
        }
    }
    
    // Write output
    float* out_ptr = output + b * Q * C + q * C;
    if (has_neighbor) {
        out_ptr[0] = m0; out_ptr[1] = m1; out_ptr[2] = m2; out_ptr[3] = m3;
        out_ptr[4] = m4; out_ptr[5] = m5; out_ptr[6] = m6; out_ptr[7] = m7;
        out_ptr[8] = m8; out_ptr[9] = m9; out_ptr[10] = m10; out_ptr[11] = m11;
        out_ptr[12] = m12; out_ptr[13] = m13; out_ptr[14] = m14; out_ptr[15] = m15;
        out_ptr[16] = m16; out_ptr[17] = m17; out_ptr[18] = m18; out_ptr[19] = m19;
        out_ptr[20] = m20; out_ptr[21] = m21; out_ptr[22] = m22; out_ptr[23] = m23;
        out_ptr[24] = m24; out_ptr[25] = m25; out_ptr[26] = m26; out_ptr[27] = m27;
        out_ptr[28] = m28; out_ptr[29] = m29; out_ptr[30] = m30; out_ptr[31] = m31;
    } else {
        for (int c = 0; c < C; c++) {
            out_ptr[c] = 0.0f;
        }
    }
}

void radius_ball_maxpool_cuda(
    const float* points, const float* queries, const float* features,
    float* output, int B, int P, int Q, int C, float radius_sq
) {
    const int total = B * Q;
    const int threads_per_block = 256;
    const int num_blocks = (total + threads_per_block - 1) / threads_per_block;
    
    radius_ball_maxpool_kernel<<<num_blocks, threads_per_block>>>(
        points, queries, features, output, B, P, Q, C, radius_sq
    );
}
"""

cpp_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

void radius_ball_maxpool_cuda(
    const float* points, const float* queries, const float* features,
    float* output, int B, int P, int Q, int C, float radius_sq
);

void radius_ball_maxpool_wrapper(
    const torch::Tensor& points,
    const torch::Tensor& queries,
    const torch::Tensor& features,
    torch::Tensor& output,
    float radius_sq
) {
    const int B = points.size(0);
    const int P = points.size(1);
    const int Q = queries.size(1);
    const int C = features.size(2);
    
    radius_ball_maxpool_cuda(
        points.data_ptr<float>(),
        queries.data_ptr<float>(),
        features.data_ptr<float>(),
        output.data_ptr<float>(),
        B, P, Q, C, radius_sq
    );
}
"""

module = load_inline(
    name='radius_ball_maxpool_ext_v4',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['radius_ball_maxpool_wrapper'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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) ** 2

    def forward(
        self,
        points: torch.Tensor,
        queries: torch.Tensor,
        features: torch.Tensor,
    ) -> torch.Tensor:
        B, P, _ = points.shape
        _, Q, _ = queries.shape
        _, _, C = features.shape
        
        output = torch.empty(B, Q, C, dtype=features.dtype, device=features.device)
        
        module.radius_ball_maxpool_wrapper(points, queries, features, output, self.radius_sq)
        
        return output