from __future__ import annotations

from dataclasses import dataclass
import numpy as np

from .tkcoords import convert_axes_coordinates


@dataclass
class tkHitResult:
    dataset_index: int
    artist_index: int
    data_index: int
    distance2: float
    x: float
    y: float
    dataset: dict


class tkNearestFinder:
    def __init__(self, distance="r"):
        if distance not in ("r", "x", "y"):
            raise ValueError("distance must be 'r', 'x', or 'y'")
        self.distance = distance

    def find(self, event, dataset_indices, datasets):
        if event.xdata is None or event.ydata is None:
            return None
        best = None
        for local_index, dataset in enumerate(datasets):
            axis = dataset["axis"]
            scale_axis = dataset.get("axis_scale") or axis
            x0, y0 = convert_axes_coordinates(event.xdata, event.ydata, event.inaxes, scale_axis)
            if x0 is None:
                continue
            xlim = scale_axis.get_xlim()
            ylim = scale_axis.get_ylim()
            xr = abs(xlim[1] - xlim[0]) or 1.0
            yr = abs(ylim[1] - ylim[0]) or 1.0
            for artist_index, artist in enumerate(dataset.get("data", [])):
                xdata = np.asarray(getattr(artist, "xdata", []), dtype=float)
                ydata = np.asarray(getattr(artist, "ydata", []), dtype=float)
                if xdata.size == 0 or ydata.size == 0:
                    continue
                valid = np.isfinite(xdata) & np.isfinite(ydata)
                if not np.any(valid):
                    continue
                xv = xdata[valid]
                yv = ydata[valid]
                original_indices = np.flatnonzero(valid)
                dx2 = ((xv - x0) / xr) ** 2
                dy2 = ((yv - y0) / yr) ** 2
                if self.distance == "x":
                    d2 = dx2
                elif self.distance == "y":
                    d2 = dy2
                else:
                    d2 = dx2 + dy2
                j = int(np.argmin(d2))
                result = tkHitResult(
                    dataset_index=dataset_indices[local_index],
                    artist_index=artist_index,
                    data_index=int(original_indices[j]),
                    distance2=float(d2[j]),
                    x=float(xv[j]),
                    y=float(yv[j]),
                    dataset=dataset,
                )
                if best is None or result.distance2 < best.distance2:
                    best = result
        return best