"""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
from scipy.ndimage import zoom as _ndi_zoom

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

_KERNEL = "thin_plate_spline"
_SMOOTHING = 0.0

# Sensor density (≤244 points) carries no more spatial information than this;
# evaluating on a bigger grid than the widget's pixel size just burns CPU that
# a GPU-bilinear QImage upscale reproduces visually. See alpha-release-polish-plan.md §2.2.
GRID_RES = 100

# Bilinear-scaling the raw 100x100 field straight to widget size shows visible
# facets on real (higher-contrast) sensor data. A cheap cubic pre-smooth of the
# field (RBF grid is the expensive part; resampling a small array is not) removes
# the facets before the final QImage upscale does the rest.
REFINE_FACTOR = 2


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 float64 array of interpolated values, capped at GRID_RES per side.

    Args:
        xs, ys:      sensor positions in mm (1-D arrays, length N)
        vs:          sensor values (length N)
        width/height: desired output dimensions in pixels; the grid is solved at
                      min(dim, GRID_RES) and the caller upscales the resulting image
        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)

    grid_w = min(width, GRID_RES)
    grid_h = min(height, GRID_RES)

    xmin, xmax, ymin, ymax = extent
    gx = np.linspace(xmin, xmax, grid_w)
    gy = np.linspace(ymin, ymax, grid_h)
    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(grid_h, grid_w)

    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)

    result: np.ndarray = field.astype(np.float64)
    return result


def refine_field(field: np.ndarray, factor: int = REFINE_FACTOR) -> np.ndarray:
    """Cubic-upsample a coarse (NaN-free) field by `factor`.

    Callers that need a circular boundary should pass an *unclipped* field
    (round_clip=False) and mask with ellipse_alpha() afterwards — cubic
    resampling across a NaN/zero-filled hole rings at the boundary, and any
    mask derived from the coarse grid keeps its staircase when upscaled.
    """
    refined: np.ndarray = _ndi_zoom(field, factor, order=3)
    return refined


def ellipse_alpha(height: int, width: int, feather_px: float = 1.2) -> np.ndarray:
    """Anti-aliased coverage mask (0..1) of the ellipse inscribed in the grid.

    Computed analytically at the target resolution — unlike upsampling a
    binary clip mask from the coarse RBF grid, this cannot show grid steps.
    `feather_px` is the half-width of the linear edge ramp, in grid pixels.
    """
    y = (np.arange(height) + 0.5) / height * 2.0 - 1.0
    x = (np.arange(width) + 0.5) / width * 2.0 - 1.0
    yy, xx = np.meshgrid(y, x, indexing="ij")
    r = np.sqrt(xx * xx + yy * yy)
    f = feather_px * 2.0 / min(width, height)
    return np.clip((1.0 - r) / f + 0.5, 0.0, 1.0)

