# pykdtree, Fast kd-tree implementation with OpenMP-enabled queries
#
# Copyright (C) 2013 - present  Esben S. Nielsen
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with this program.  If not, see <http://www.gnu.org/licenses/>.

from typing import Any
import numpy as np

class KDTree:
    """kd-tree for fast nearest-neighbour lookup.
    The interface is made to resemble the scipy.spatial kd-tree except
    only Euclidean distance measure is supported.

    :Parameters:
    data_pts : numpy array
        Data points with shape (n , dims)
    leafsize : int, optional
        Maximum number of data points in tree leaf
    """

    @property
    def data_pts(self) -> np.ndarray:
        """Data points used to construct the kd-tree."""

    @property
    def data(self) -> np.ndarray:
        """Data points used to construct the kd-tree."""

    @property
    def n(self) -> int:
        """Number of data points."""

    @property
    def ndim(self) -> int:
        """Number of dimensions."""

    @property
    def leafsize(self) -> int:
        """Maximum number of data points in tree leaf."""

    def __init__(self, data_pts: np.ndarray, leafsize: int = 16): ...
    def query(
        self,
        query_pts: np.ndarray,
        k: int = 1,
        eps: float = 0,
        distance_upper_bound: float | None = None,
        sqr_dists: bool = False,
        mask: np.ndarray | None = None,
    ):
        """Query the kd-tree for nearest neighbors

        :Parameters:
        query_pts : numpy array
            Query points with shape (m, dims)
        k : int
            The number of nearest neighbours to return
        eps : non-negative float
            Return approximate nearest neighbours; the k-th returned value
            is guaranteed to be no further than (1 + eps) times the distance
            to the real k-th nearest neighbour
        distance_upper_bound : non-negative float
            Return only neighbors within this distance.
            This is used to prune tree searches.
        sqr_dists : bool, optional
            Internally pykdtree works with squared distances.
            Determines if the squared or Euclidean distances are returned.
        mask : numpy array, optional
            Array of booleans where neighbors are considered invalid and
            should not be returned. A mask value of True represents an
            invalid pixel. Mask should have shape (n,) to match data points.
            By default all points are considered valid.

        """
        ...

# These are generated by Cython.
# Just in here to avoid errors in mypy tests.
__reduce_cython__: Any
__setstate_cython__: Any
__test__: Any
