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

# Compile CUDA kernel at module load time
fps_cuda = load_inline(
    name='fps_cuda',
    cpp_sources='''
#include <torch/extension.h>

torch::Tensor fps_cuda_forward(torch::Tensor points, int M);
''',
    cuda_sources='''
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>
#include <limits>
#include <torch/extension.h>

template <typename scalar_t, int BLOCK_SIZE>
__global__ void fps_kernel(
    const scalar_t* __restrict__ points,
    int64_t* __restrict__ indices,
    int B, int N, int M
) {
    int b = blockIdx.x;
    if (b >= B) return;
    
    int tid = threadIdx.x;
    
    extern __shared__ char smem[];
    scalar_t* min_dist_sq = reinterpret_cast<scalar_t*>(smem);
    
    // Each thread handles a subset of points for initialization
    for (int j = tid; j < N; j += BLOCK_SIZE) {
        min_dist_sq[j] = std::numeric_limits<scalar_t>::infinity();
    }
    __syncthreads();
    
    const scalar_t* batch_points = points + static_cast<size_t>(b) * N * 3;
    int64_t* batch_indices = indices + static_cast<size_t>(b) * M;
    
    // Shared memory for reduction
    __shared__ scalar_t shared_max[BLOCK_SIZE];
    __shared__ int shared_farthest[BLOCK_SIZE];
    __shared__ int farthest;
    
    // Initialize: first center is index 0
    if (tid == 0) {
        farthest = 0;
    }
    __syncthreads();
    
    for (int i = 0; i < M; i++) {
        // Record current farthest
        if (tid == 0) {
            batch_indices[i] = farthest;
        }
        __syncthreads();
        
        // Get center coordinates
        scalar_t cx = batch_points[farthest * 3 + 0];
        scalar_t cy = batch_points[farthest * 3 + 1];
        scalar_t cz = batch_points[farthest * 3 + 2];
        
        // Update min_dist_sq: each thread handles a subset of points
        for (int j = tid; j < N; j += BLOCK_SIZE) {
            scalar_t px = batch_points[j * 3 + 0];
            scalar_t py = batch_points[j * 3 + 1];
            scalar_t pz = batch_points[j * 3 + 2];
            
            scalar_t dx = px - cx;
            scalar_t dy = py - cy;
            scalar_t dz = pz - cz;
            scalar_t dist_sq = dx * dx + dy * dy + dz * dz;
            
            if (dist_sq < min_dist_sq[j]) {
                min_dist_sq[j] = dist_sq;
            }
        }
        __syncthreads();
        
        // Mark farthest as excluded
        if (tid == 0) {
            min_dist_sq[farthest] = static_cast<scalar_t>(-1.0);
        }
        __syncthreads();
        
        // Find argmax using parallel reduction
        // Each thread finds local max in its assigned range
        scalar_t local_max = static_cast<scalar_t>(-2.0);
        int local_farthest = -1;
        
        for (int j = tid; j < N; j += BLOCK_SIZE) {
            scalar_t val = min_dist_sq[j];
            if (val > local_max) {
                local_max = val;
                local_farthest = j;
            }
        }
        
        shared_max[tid] = local_max;
        shared_farthest[tid] = local_farthest;
        __syncthreads();
        
        // Parallel reduction to find global max
        // Reduce from BLOCK_SIZE to 1
        for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
            if (tid < stride) {
                if (shared_max[tid + stride] > shared_max[tid]) {
                    shared_max[tid] = shared_max[tid + stride];
                    shared_farthest[tid] = shared_farthest[tid + stride];
                }
            }
            __syncthreads();
        }
        
        // Thread 0 writes the result
        if (tid == 0) {
            farthest = shared_farthest[0];
        }
        __syncthreads();
    }
}

torch::Tensor fps_cuda_forward(torch::Tensor points, int M) {
    TORCH_CHECK(points.dim() == 3, "points must be 3D");
    int B = points.size(0);
    int N = points.size(1);
    TORCH_CHECK(points.size(2) == 3, "points must have 3 coordinates");
    
    auto indices = torch::empty({B, M}, torch::TensorOptions().dtype(torch::kInt64).device(points.device()));
    
    const int BLOCK_SIZE = 256;
    
    dim3 grid(B);
    dim3 block(BLOCK_SIZE);
    
    AT_DISPATCH_FLOATING_TYPES(points.scalar_type(), "fps_kernel", ([&] {
        size_t shared_mem_size = N * sizeof(scalar_t);
        fps_kernel<scalar_t, BLOCK_SIZE><<<grid, block, shared_mem_size>>>(
            points.data_ptr<scalar_t>(),
            indices.data_ptr<int64_t>(),
            B, N, M
        );
    }));
    
    C10_CUDA_CHECK(cudaGetLastError());
    return indices;
}
''',
    functions=['fps_cuda_forward'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


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:
        B, N, _ = points.shape
        M = self.num_samples
        assert M <= N, "FPS cannot sample more points than exist"
        
        return fps_cuda.fps_cuda_forward(points, M)