import typing

import numpy
import numpy.typing
from scipy.interpolate import interp1d

from goats.core import numerical
from goats.core import variable


def restrict_coordinate(
    array: numpy.ndarray,
    target: float,
    axis: typing.SupportsIndex,
) -> range:
    """Restrict `array` to `target` along `axis`."""
    bounds = numerical.get_bounding_indices(array, target, axis=axis)
    if bounds.ndim == 1:
        return range(bounds[0], bounds[1]+1)
    if bounds.ndim == 2:
        lower = min(bounds[:, 0])
        upper = max(bounds[:, 1])
        return range(lower, upper+1)
    # NOTE: We can get `lower` and `upper` for N-D arrays by replacing `:` with
    # `...`. The problem is that they will be (N-1)-D arrays that we will then
    # need to split into individual index vectors before returning.
    raise NotImplementedError(f"Can't operate on {bounds.ndim}-D array")


class Restriction:
    """A subset of indices with which to create subsets of arrays."""

    def __init__(
        self,
        rule: typing.Callable[..., typing.Iterable[typing.SupportsIndex]],
        *args,
        **kwargs,
    ) -> None:
        """
        Create a new restriction.

        Parameters
        ----------
        rule : callable
            The rule by which to create a subset of array indices. This must be
            a callable object that returns an iterable of objects to be used as
            indices (e.g., `range`).

        *args
            Positional arguments to pass to `rule`.

        **kwargs
            Keyword arguments to pass to `rule`.
        """
        self.indices = rule(*args, **kwargs)

    def apply(self, array: numpy.typing.ArrayLike, axis: int=None):
        """Restrict `array` to a subset along `axis`.
        
        Parameters
        ----------
        array : array-like
            The array from which to extract a subarray corresponding to the
            restricted indices.

        axis : int, default=-1
            The axis of `array` to restrict.

        Returns
        -------
        array-like
            The subarray.
        """
        if axis is None:
            axis = -1
        permuted = numpy.moveaxis(array, axis, 0)
        result = permuted[self.indices, ...]
        return numpy.moveaxis(result, 0, axis)


class Quantity(variable.Quantity):
    """A variable quantity, interpolated over one or more axes."""

    def interpolate(
        self,
        coordinates: typing.Iterable[typing.Dict[str, typing.Any]],
    ) -> variable.Quantity:
        """Internal interpolation logic."""
        for coordinate in coordinates:
            self._array = self._interpolate_coordinate(
                coordinate['targets'],
                coordinate['reference'],
                axis=coordinate.get('axis'),
            )
        return self

    def _interpolate_coordinate(
        self,
        targets: numpy.ndarray,
        reference: variable.Quantity,
        axis: int=None,
    ) -> numpy.ndarray:
        """Interpolate a variable array based on a known coordinate."""
        array = numpy.array(self)
        indices = (self.axes.index(d) for d in reference.axes)
        dst, src = zip(*enumerate(indices))
        reordered = numpy.moveaxis(array, src, dst)
        interpolated = _apply_interpolation(
            reordered,
            numpy.array(reference),
            targets,
            axis=axis,
        )
        return numpy.moveaxis(interpolated, dst, src)


def _apply_interpolation(
    array: numpy.ndarray,
    reference: numpy.ndarray,
    targets: typing.Iterable[float],
    axis: int=None,
) -> numpy.ndarray:
    """Interpolate `array` to target values over `coordinate`."""
    interpolated = [
        _apply_interp1d(array, reference, target, axis=axis)
        for target in targets
    ]
    if reference.ndim == 2:
        return numpy.swapaxes(interpolated, 0, 1)
    return numpy.array(interpolated)


def _apply_interp1d(
    array: numpy.ndarray,
    reference: numpy.ndarray,
    target: float,
    axis: int=None,
) -> typing.List[float]:
    """Interpolate data to `target` along the leading axis."""
    if target in reference:
        idx = numerical.find_nearest(reference, target).index
        return array[idx]
    if axis is not None:
        restriction = Restriction(
            restrict_coordinate,
            reference,
            target,
            axis=axis,
        )
        ref = restriction.apply(reference, axis=axis)
        arr = restriction.apply(array, axis=axis)
    else:
        ref = reference
        arr = array
    if reference.ndim == 2:
        interps = [interp1d(x, y, axis=0) for x, y in zip(ref, arr)]
        return [interp(target) for interp in interps]
    interp = interp1d(ref, arr, axis=0)
    return interp(target)


