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 <cfloat>
#include <cstdint>

__global__ void fps_kernel_float(const float* __restrict__ pts,
                                 int64_t* __restrict__ out,
                                 int B, int N, int M) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int block_threads = blockDim.x;

    extern __shared__ unsigned char smem_raw[];
    float* min_d = reinterpret_cast<float*>(smem_raw);                  // N floats
    float* red_val = min_d + N;                                          // blockDim floats
    int* red_idx = reinterpret_cast<int*>(red_val + block_threads);      // blockDim ints

    for (int j = tid; j < N; j += block_threads) {
        min_d[j] = CUDART_INF_F;
    }
    __syncthreads();

    const float* batch_pts = pts + ((int64_t)b) * N * 3;
    int farthest = 0;

    for (int i = 0; i < M; ++i) {
        int selected = farthest;
        if (tid == 0) {
            out[((int64_t)b) * M + i] = (int64_t)selected;
        }

        // The next farthest point is not returned after the final sample.
        if (i == M - 1) break;

        float cx = batch_pts[selected * 3 + 0];
        float cy = batch_pts[selected * 3 + 1];
        float cz = batch_pts[selected * 3 + 2];

        float best = -FLT_MAX;
        int best_i = 0;

        for (int j = tid; j < N; j += block_threads) {
            float x = batch_pts[j * 3 + 0] - cx;
            float y = batch_pts[j * 3 + 1] - cy;
            float z = batch_pts[j * 3 + 2] - cz;
            float d = x * x + y * y + z * z;
            float md = min_d[j];
            if (d < md) md = d;
            if (j == selected) md = -1.0f;
            min_d[j] = md;

            // Match torch.argmax tie behavior: first (lowest) index wins.
            if ((md > best) || (md == best && j < best_i)) {
                best = md;
                best_i = j;
            }
        }

        red_val[tid] = best;
        red_idx[tid] = best_i;
        __syncthreads();

        // Power-of-two block size reduction, retaining lowest index on ties.
        for (int offset = block_threads >> 1; offset > 0; offset >>= 1) {
            if (tid < offset) {
                float ov = red_val[tid + offset];
                int oi = red_idx[tid + offset];
                float cv = red_val[tid];
                int ci = red_idx[tid];
                if ((ov > cv) || (ov == cv && oi < ci)) {
                    red_val[tid] = ov;
                    red_idx[tid] = oi;
                }
            }
            __syncthreads();
        }

        farthest = red_idx[0];
        __syncthreads();
    }
}

torch::Tensor fps_forward(torch::Tensor points, int64_t num_samples) {
    TORCH_CHECK(points.is_cuda(), "points must be a CUDA tensor");
    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");

    const int B = (int)points.size(0);
    const int N = (int)points.size(1);
    const int M = (int)num_samples;
    TORCH_CHECK(M >= 1, "num_samples must be >= 1");
    TORCH_CHECK(M <= N, "FPS cannot sample more points than exist");

    auto out = torch::empty({B, M}, points.options().dtype(torch::kInt64));
    if (B == 0) return out;

    constexpr int threads = 256;
    size_t shmem = ((size_t)N + threads) * sizeof(float) + threads * sizeof(int);
    fps_kernel_float<<<B, threads, shmem, at::cuda::getCurrentCUDAStream()>>>(
        points.data_ptr<float>(), out.data_ptr<int64_t>(), B, N, M);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fps_ext = load_inline(
    name="fps_kernelbench_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fps_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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.contiguous(), self.num_samples)