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>
torch::Tensor fps_forward(torch::Tensor points, int64_t num_samples);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math_constants.h>

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

__global__ void fps_kernel512(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 center[3];
    __shared__ float redv[16];
    __shared__ int redi[16];

    const float* p = points + (long long)b * N * 3;
    int j0 = tid, j1 = tid + 512;
    float x0 = 0.f, y0 = 0.f, z0 = 0.f, m0 = -1.f;
    float x1 = 0.f, y1 = 0.f, z1 = 0.f, m1 = -1.f;

    if (j0 < N) { x0 = p[j0*3]; y0 = p[j0*3+1]; z0 = p[j0*3+2]; m0 = CUDART_INF_F; }
    if (j1 < N) { x1 = p[j1*3]; y1 = p[j1*3+1]; z1 = p[j1*3+2]; m1 = CUDART_INF_F; }

    int far = 0;
    for (int i = 0; i < M; ++i) {
        if (tid == 0) out[(long long)b * M + i] = (long long)far;

        if (far == j0) { center[0] = x0; center[1] = y0; center[2] = z0; }
        if (far == j1) { center[0] = x1; center[1] = y1; center[2] = z1; }
        __syncthreads();

        float cx = center[0], cy = center[1], cz = center[2];
        if (j0 < N) {
            float dx = x0 - cx, dy = y0 - cy, dz = z0 - cz;
            float d = dx*dx + dy*dy + dz*dz;
            m0 = d < m0 ? d : m0;
            if (far == j0) m0 = -1.f;
        }
        if (j1 < N) {
            float dx = x1 - cx, dy = y1 - cy, dz = z1 - cz;
            float d = dx*dx + dy*dy + dz*dz;
            m1 = d < m1 ? d : m1;
            if (far == j1) m1 = -1.f;
        }

        float bv = m0; int bi = j0;
        if (m1 > bv || (m1 == bv && j1 < bi)) { bv = m1; bi = j1; }
        wrmax(bv, bi);
        if (lane == 0) { redv[warp] = bv; redi[warp] = bi; }
        __syncthreads();

        if (warp == 0) {
            bv = lane < 16 ? redv[lane] : -1.f;
            bi = lane < 16 ? redi[lane] : 0;
            wrmax(bv, bi);
            if (lane == 0) redi[0] = bi;
        }
        __syncthreads();
        far = redi[0];
    }
}

torch::Tensor fps_forward(torch::Tensor points, int64_t num_samples) {
    TORCH_CHECK(points.is_cuda(), "points must be CUDA");
    TORCH_CHECK(points.scalar_type() == torch::kFloat32, "points must be float32");
    TORCH_CHECK(points.dim() == 3 && points.size(2) == 3, "points shape must be (B,N,3)");
    auto pc = points.contiguous();
    int B = (int)pc.size(0), N = (int)pc.size(1), M = (int)num_samples;
    TORCH_CHECK(N <= 1024 && M <= N, "bad FPS sizes");
    auto out = torch::empty({B, M}, torch::TensorOptions().dtype(torch::kInt64).device(pc.device()));
    fps_kernel512<<<B, 512, 0, at::cuda::getCurrentCUDAStream()>>>(
        pc.data_ptr<float>(), (long long*)out.data_ptr<int64_t>(), N, M);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_fps_ext = load_inline(
    name='fps_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fps_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

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

    def forward(self, points: torch.Tensor) -> torch.Tensor:
        return _fps_ext.fps_forward(points, self.num_samples)