from scipy.interpolate import RegularGridInterpolator, NearestNDInterpolator
import numpy as np
from ..field import Field

def make_nearest_interpolator_separated(field, grid=None):
    '''Make a nearest interpolator for a field on a separated grid.

    Parameters
    ----------
    field : Field or ndarray
        The field to interpolate.
    grid : Grid or None
        The grid of the field. If it is given, the grid of `field` is replaced by this grid.

    Returns
    -------
    Field generator
        The interpolator, as a Field generator. The grid on which this field generator will evaluated, does
        not have to be separated.
    '''
    if grid is None:
        grid = field.grid
    else:
        field = Field(field, grid)

    axes_reversed = np.array(grid.separated_coords)
    interp = RegularGridInterpolator(axes_reversed, field.shaped, 'nearest', False)

    def interpolator(evaluated_grid):
        evaluated_coords = np.flip(np.array(evaluated_grid.coords), 0)
        res = interp(evaluated_coords.T)
        return Field(res.ravel(), evaluated_grid)

    return interpolator

def make_nearest_interpolator_unstructured(field, grid=None):
    '''Make a nearest interpolator for an unstructured grid.

    Parameters
    ----------
    field : Field or array_like
        The field to interpolate.
    grid : Grid or None
        The grid of the field. If it is given, the grid of `field` is replaced by this grid.

    Returns
    -------
    Field generator
        The interpolator as a Field generator. The grid on which this field generator will be evaluated does
        not need to have any structure.
    '''
    if grid is None:
        grid = field.grid
    else:
        field = Field(field, grid)

    interp = NearestNDInterpolator(grid.points, field)

    def interpolator(evaluated_grid):
        res = interp(grid.points)
        return Field(res, evaluated_grid)

    return interpolator

def make_nearest_interpolator(field, grid=None):
    '''Make a nearest interpolator for any type of grid.

    Parameters
    ----------
    field : Field or array_like
        The field to interpolate.
    grid : Grid or None
        The grid of the field. If it is given, the grid of `field` is replaced by this grid.

    Returns
    -------
    Field generator
        The interpolator as a Field generator. The grid on which this field generator will be evaluated does
        not need to have any structure.
    '''
    if grid is None:
        grid = field.grid

    if grid.is_unstructured:
        return make_nearest_interpolator_unstructured(field, grid)
    else:
        return make_nearest_interpolator_separated(field, grid)
