"""RBF (thin-plate spline) heatmap field.

Uses CuPy for GPU acceleration when available, falls back to NumPy + SciPy.
"""
from __future__ import annotations

import numpy as np
from scipy.interpolate import RBFInterpolator

try:
    import cupy as _cupy  # type: ignore
    BACKEND = "cupy"
except Exception:
    _cupy = None
    BACKEND = "numpy"

_KERNEL = "thin_plate_spline"
_SMOOTHING = 0.0


def interpolate_field(
    xs: np.ndarray,
    ys: np.ndarray,
    vs: np.ndarray,
    *,
    width: int,
    height: int,
    extent: tuple[float, float, float, float],  # (xmin, xmax, ymin, ymax) in mm
    round_clip: bool = False,
) -> np.ndarray:
    """Return a (height, width) float64 array of interpolated values.

    Args:
        xs, ys:      sensor positions in mm (1-D arrays, length N)
        vs:          sensor values (length N)
        width/height: output grid dimensions in pixels
        extent:      (xmin, xmax, ymin, ymax) in the same mm space as xs/ys
        round_clip:  if True, pixels outside the inscribed ellipse become NaN
    """
    coords = np.column_stack([xs, ys])
    rbf = RBFInterpolator(coords, vs, kernel=_KERNEL, smoothing=_SMOOTHING)

    xmin, xmax, ymin, ymax = extent
    gx = np.linspace(xmin, xmax, width)
    gy = np.linspace(ymin, ymax, height)
    grid_x, grid_y = np.meshgrid(gx, gy)
    flat = np.column_stack([grid_x.ravel(), grid_y.ravel()])

    # RBFInterpolator always runs on CPU; CuPy only accelerates other ops if added later
    field = rbf(flat).reshape(height, width)

    if round_clip:
        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2
        rx = (xmax - xmin) / 2
        ry = (ymax - ymin) / 2
        dist = ((grid_x - cx) / rx) ** 2 + ((grid_y - cy) / ry) ** 2
        field = np.where(dist <= 1.0, field, np.nan)

    return field.astype(np.float64)