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):
    """
    Farthest Point Sampling (FPS) on 3D point coordinates (B, N, 3).

    Used in navigation / manipulation perception (PointNet++, PVCNN, voxelization
    pipelines) to subsample a dense cloud while spreading support across the shape.

    Greedy construction (reference fusion target): fix the first index to **0** for
    determinism, then repeatedly choose the point whose **minimum** squared distance
    to the set of already chosen points is **largest** (classical FPS).
    """

    def __init__(self, num_samples: int):
        super().__init__()
        assert num_samples >= 1
        self.num_samples = num_samples

    def forward(self, points: torch.Tensor) -> torch.Tensor:
        """
        Args:
            points: (B, N, 3) point coordinates.
        Returns:
            indices: (B, M) int64, each row is M chosen indices in [0, N).
        """
        B, N, _ = points.shape
        M = self.num_samples
        assert M <= N, "FPS cannot sample more points than exist"
        device = points.device
        dtype = points.dtype

        batch_idx = torch.arange(B, device=device)
        indices = torch.empty(B, M, dtype=torch.long, device=device)

        # min squared distance from each point to the current chosen set (large init)
        min_dist_sq = torch.full((B, N), float("inf"), device=device, dtype=dtype)
        # First center: point index 0 (deterministic)
        farthest = torch.zeros(B, dtype=torch.long, device=device)

        for i in range(M):
            indices[:, i] = farthest
            center = points[batch_idx, farthest]
            dist_sq = ((points - center.unsqueeze(1)) ** 2).sum(dim=-1)
            min_dist_sq = torch.minimum(min_dist_sq, dist_sq)
            # Exclude the just-chosen index from being selected again
            min_dist_sq[batch_idx, farthest] = -1.0
            farthest = min_dist_sq.argmax(dim=-1)

        return indices


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

    def __init__(self, num_samples: int):
        super().__init__()
        self._impl = Model(num_samples)

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


batch_size = 4
num_points = 1024
num_samples = 128


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


def get_init_inputs():
    return [num_samples]

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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>

__global__ void fps_kernel(
    const float* __restrict__ points,  // (B, N, 3)
    long* __restrict__ indices,        // (B, M)
    const int B,
    const int N,
    const int M
) {
    extern __shared__ float shared_mem[];
    float* min_dist_sq = shared_mem;  // size N

    const int batch_idx = blockIdx.x;
    if (batch_idx >= B) return;

    const float* batch_points = points + batch_idx * N * 3;

    // Initialize min_dist_sq to inf
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        min_dist_sq[i] = FLT_MAX;
    }
    __syncthreads();

    // First point: index 0
    int farthest = 0;
    if (threadIdx.x == 0) {
        indices[batch_idx * M] = 0;
    }
    __syncthreads();

    for (int i = 1; i < M; ++i) {
        // Compute distances to current farthest and update min_dist_sq
        float cx = batch_points[farthest * 3 + 0];
        float cy = batch_points[farthest * 3 + 1];
        float cz = batch_points[farthest * 3 + 2];

        for (int j = threadIdx.x; j < N; j += blockDim.x) {
            float dx = batch_points[j * 3 + 0] - cx;
            float dy = batch_points[j * 3 + 1] - cy;
            float dz = batch_points[j * 3 + 2] - cz;
            float dist_sq = dx*dx + dy*dy + dz*dz;
            if (dist_sq < min_dist_sq[j]) {
                min_dist_sq[j] = dist_sq;
            }
        }
        __syncthreads();

        // Mark chosen point
        if (threadIdx.x == 0) {
            min_dist_sq[farthest] = -1.0f;
        }
        __syncthreads();

        // Find argmax of min_dist_sq using warp shuffles
        float local_max = -FLT_MAX;
        int local_argmax = 0;
        for (int j = threadIdx.x; j < N; j += blockDim.x) {
            if (min_dist_sq[j] > local_max) {
                local_max = min_dist_sq[j];
                local_argmax = j;
            }
        }

        __shared__ float s_max[32];
        __shared__ int s_argmax[32];

        int warp_id = threadIdx.x / 32;
        int lane = threadIdx.x % 32;

        for (int offset = 16; offset > 0; offset /= 2) {
            float other = __shfl_down_sync(0xffffffff, local_max, offset);
            int other_idx = __shfl_down_sync(0xffffffff, local_argmax, offset);
            if (other > local_max) {
                local_max = other;
                local_argmax = other_idx;
            }
        }

        if (lane == 0) {
            s_max[warp_id] = local_max;
            s_argmax[warp_id] = local_argmax;
        }
        __syncthreads();

        if (warp_id == 0) {
            local_max = (threadIdx.x < (blockDim.x + 31) / 32) ? s_max[lane] : -FLT_MAX;
            local_argmax = (threadIdx.x < (blockDim.x + 31) / 32) ? s_argmax[lane] : 0;

            for (int offset = 16; offset > 0; offset /= 2) {
                float other = __shfl_down_sync(0xffffffff, local_max, offset);
                int other_idx = __shfl_down_sync(0xffffffff, local_argmax, offset);
                if (other > local_max) {
                    local_max = other;
                    local_argmax = other_idx;
                }
            }

            if (threadIdx.x == 0) {
                farthest = local_argmax;
                indices[batch_idx * M + i] = farthest;
            }
        }
        __syncthreads();
    }
}

torch::Tensor farthest_point_sampling_cuda(torch::Tensor points, int num_samples) {
    const int B = points.size(0);
    const int N = points.size(1);
    const int M = num_samples;

    auto indices = torch::empty({B, M}, torch::TensorOptions().dtype(torch::kInt64).device(points.device()));

    const int threads = 256;
    size_t shared_size = N * sizeof(float);

    fps_kernel<<<B, threads, shared_size>>>(
        points.data_ptr<float>(),
        indices.data_ptr<long>(),
        B, N, M
    );

    return indices;
}
'''

cpp_source = '''
torch::Tensor farthest_point_sampling_cuda(torch::Tensor points, int num_samples);
'''

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

class ModelNew(nn.Module):
    def __init__(self, num_samples: int):
        super().__init__()
        self.num_samples = num_samples

    def forward(self, points: torch.Tensor) -> torch.Tensor:
        return module.farthest_point_sampling_cuda(points, self.num_samples)
```