import scipy.spatial

from interp.grid import grid as basegrid


def get_simplex_extra_points(X, points, triangulation, kdtree, extra_points=8):
    simplex_id = triangulation.find_simplex(X)
    simplex_verts_ids = set(triangulation.vertices[simplex_id])

    distances, kdt_ids = kdtree.query(X, extra_points + len(simplex_verts_ids))
    kdt_ids = set(kdt_ids)

    simplex_ids = list(simplex_verts_ids)
    extra_points_ids = list(kdt_ids - simplex_verts_ids)

    return simplex_ids, extra_points_ids

class dgrid(basegrid):
    def __init__(self, points, values):
        self.points = points
        self.values = values
        self.triangulation = scipy.spatial.Delaunay(points)
        self.kdtree = scipy.spatial.KDTree(points)

    def get_simplex_extra_points(self, X, extra_points=8):
        return get_simplex_extra_points(X, self.points, self.triangulation,
                                self.kdtree, extra_points=extra_points)