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

_cpp_src = r"""
#include <torch/extension.h>
void fps_launcher(torch::Tensor points, torch::Tensor out, int M);
torch::Tensor fps_forward(torch::Tensor points, int64_t M) {
    TORCH_CHECK(points.is_cuda(), "points must be CUDA");
    TORCH_CHECK(points.scalar_type() == at::kFloat, "points must be float32");
    TORCH_CHECK(points.dim() == 3 && points.size(2) == 3, "points must have shape (B,N,3)");
    TORCH_CHECK(points.is_contiguous(), "points must be contiguous");
    int64_t B = points.size(0), N = points.size(1);
    TORCH_CHECK(M <= N, "FPS cannot sample more points than exist");
    auto out = torch::empty({B, M}, torch::TensorOptions().device(points.device()).dtype(torch::kInt64));
    fps_launcher(points, out, (int)M);
    return out;
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
#define BLOCK_SIZE 1024

__device__ __forceinline__ void warp_argmax(float &v, int &idx) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        float v2 = __shfl_down_sync(mask, v, off);
        int i2 = __shfl_down_sync(mask, idx, off);
        if ((v2 > v) || (v2 == v && i2 < idx)) { v = v2; idx = i2; }
    }
}

__global__ void fps_kernel_1024(const float* __restrict__ points, long long* __restrict__ out,
                                int N, int M) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    __shared__ float sx[BLOCK_SIZE];
    __shared__ float sy[BLOCK_SIZE];
    __shared__ float sz[BLOCK_SIZE];
    __shared__ float min_dist[BLOCK_SIZE];
    __shared__ float warp_val[32];
    __shared__ int warp_idx[32];
    __shared__ int farthest_s;
    __shared__ float cx, cy, cz;
    const float* pts = points + ((size_t)b * N * 3);
    long long* out_b = out + ((size_t)b * M);

    if (tid < N) {
        sx[tid] = pts[tid * 3 + 0];
        sy[tid] = pts[tid * 3 + 1];
        sz[tid] = pts[tid * 3 + 2];
        min_dist[tid] = FLT_MAX;
    } else {
        sx[tid] = sy[tid] = sz[tid] = 0.0f;
        min_dist[tid] = -1.0f;
    }
    if (tid == 0) farthest_s = 0;
    __syncthreads();

    for (int i = 0; i < M; ++i) {
        int farthest = farthest_s;
        if (tid == 0) {
            out_b[i] = (long long)farthest;
            cx = sx[farthest]; cy = sy[farthest]; cz = sz[farthest];
        }
        __syncthreads();

        float v = -1.0f;
        int idx = tid;
        if (tid < N) {
            float dx = sx[tid] - cx;
            float dy = sy[tid] - cy;
            float dz = sz[tid] - cz;
            float d = dx * dx + dy * dy + dz * dz;
            float old = min_dist[tid];
            v = d < old ? d : old;
            if (tid == farthest) v = -1.0f;
            min_dist[tid] = v;
        }
        warp_argmax(v, idx);
        if (lane == 0) { warp_val[warp] = v; warp_idx[warp] = idx; }
        __syncthreads();
        if (warp == 0) {
            float vv = warp_val[lane];
            int ii = warp_idx[lane];
            warp_argmax(vv, ii);
            if (lane == 0) farthest_s = ii;
        }
        __syncthreads();
    }
}

void fps_launcher(torch::Tensor points, torch::Tensor out, int M) {
    int B = (int)points.size(0);
    int N = (int)points.size(1);
    TORCH_CHECK(N <= BLOCK_SIZE, "This FPS kernel supports N <= 1024");
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fps_kernel_1024<<<B, BLOCK_SIZE, 0, stream>>>(points.data_ptr<float>(), (long long*)out.data_ptr<int64_t>(), N, M);
}
"""

_fps_ext = load_inline(name="fps_ext_kernel_v3", cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=["fps_forward"], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math"], with_cuda=True, verbose=False)

class ModelNew(nn.Module):
    def __init__(self, num_samples: int):
        super().__init__(); assert num_samples >= 1; self.num_samples = int(num_samples)
    def forward(self, points: torch.Tensor) -> torch.Tensor:
        return _fps_ext.fps_forward(points, self.num_samples)