import numpy as np


__all__ = ["clip", "clipping_mask"]


def clip(orbit_instance, window_dimensions, **kwargs):
    """ Create Orbit instance whose state array is a subdomain of the provided Orbit.

    Parameters
    ----------
    orbit_instance : Orbit
        The orbit whose state the subdomain is extracted from.
    window_dimensions : tuple of tuples.
        Contains one tuple for each continuous dimension, each defining the interval of the dimension to slice out.

    kwargs : dict
        Keyword arguments for Orbit instantiate.

    Returns
    -------
    Orbit :
        Orbit whose state and parameters reflect the subdomain defined by the provided dimensions.

    Notes
    -----
    The intervals provided refer to the :meth:`Orbit._plotting_dimensions()` method. The motivation
    here is to allow for clipping using visualization techniques as a direct guide.
    If a dimension has zero extent; i.e. equilibrium in that dimension, then the corresponding window_dimension
    tuple must be passed as (None, None).




    Examples
    --------

    Extract subdomain from an Orbit

    >>> orb = Orbit(state=np.ones([128, 128, 128, 128]), basis='physical',
    ...                              parameters=(100, 100, 100, 100))
    >>> one_sixteeth_subdomain_orbit = clip(orb, ((0, 50), (0, 50), (0, 50), (0, 50)))

    It is 1/16th the size because it takes half of the points in 4 different dimensions.



    """
    clipping_type = kwargs.get("clipping_type", orbit_instance.__class__)
    slices, dimensions = _slices_from_window(orbit_instance, window_dimensions)

    # It of course is better to get the dimensions/parameters from the clipping directly, but if the user wants to
    # this gives them the ability to override.
    parameters = kwargs.pop(
        "parameters",
        tuple(
            dimensions[i] if i < len(dimensions) else p
            for i, p in enumerate(orbit_instance.parameters)
        ),
    )

    clipped_orbit = clipping_type(
        state=orbit_instance.transform(to=orbit_instance.bases_labels()[0]).state[
            slices
        ],
        basis=orbit_instance.bases_labels()[0],
        parameters=parameters,
        **kwargs
    )
    return clipped_orbit


def clipping_mask(orbit_instance, *windows, invert=True):
    """
    Produce an array mask which shows the clipped regions corresponding to windows upon plotting.

    Parameters
    ----------
    orbit_instance : Orbit
        An instance whose state is to be masked.
    windows : list or tuple
        An iterable of window tuples; see Notes below.
    invert : bool
        Whether to logically invert the boolean mask; equivalent to showing the "interior" or "exterior" of the clipping
        if True or False, respectively.

    Returns
    -------
    Orbit :
        Orbit instance whose state is a numpy masked array.

    """
    # Create boolean mask to manipulate for numpy masked arrays.
    mask = np.zeros(orbit_instance.shapes()[0]).astype(bool)
    if isinstance(*windows, tuple) and len(windows) == 1:
        windows = tuple(*windows)

    if type(windows) in [list, tuple]:
        for window in windows:
            # Do not need dimensions, as we are not clipping technically.
            window_slices, _ = _slices_from_window(orbit_instance, window)
            mask[window_slices] = True
    else:
        # Do not need dimensions, as we are not clipping technically.
        window_slices, _ = _slices_from_window(orbit_instance, windows)
        mask[window_slices] = True

    if invert:
        mask = np.invert(mask)
    masked_field = np.ma.masked_array(
        orbit_instance.transform(to=orbit_instance.bases_labels()[0]).state, mask=mask
    )
    return orbit_instance.__class__(
        state=masked_field,
        basis=orbit_instance.bases_labels()[0],
        parameters=orbit_instance.parameters,
    )


def _slices_from_window(orbit_instance, window_dimensions):
    """
    Slices for Orbit state which represents the subdomain defined by provided window_dimensions

    Parameters
    ----------
    orbit_instance : Orbit
        The orbit instance whose state will be sliced.
    window_dimensions : tuple of tuples
        A tuple containing the intervals which define the dimensions of the new slice.

    Returns
    -------
    tuple, tuple :
        Tuples containing the slices for the Orbit state and the new corresponding dimensions

    """
    shape = orbit_instance.shapes()[0]
    # Returns the dimensions which would be shown on a plot (easier to eye-ball clipping then), including units.
    # Should be a tuple of tuples (d_min, d_max), one for each dimension.
    plot_dimensions = orbit_instance.plotting_dimensions()
    # Returns a tuple of the "length" > 0 of each axis.
    actual_dimensions = orbit_instance.dimensions()

    clipping_slices = []
    clipping_dimensions = []
    for i, (d_min, d_max) in enumerate(window_dimensions):
        # the division and multiplication by 2's keeps the sizes even.
        if d_min is None:
            slice_start = 0
        else:
            # Clipping out of bounds does not make sense.
            assert (
                d_min >= plot_dimensions[i][0]
            ), "Trying to clip out of bounds. Please revise clipping domain."
            # Some coordinate axis range from, for example, -1 to 1. Account for this by rescaling the interval.
            # An example clipping in this case could be from -1 to -0.5. To handle this, rescale the plotting dimensions
            # to [0, plotting_dimension_max - plotting_dimension_min], by subtracting the minimum.
            # This makes the d_min value equal to the fraction of the domain.

            rescaled_domain_min = (d_min - plot_dimensions[i][0]) / (
                plot_dimensions[i][1] - plot_dimensions[i][0]
            )
            slice_start = int(shape[i] * rescaled_domain_min)

        if d_max is None:
            slice_end = shape[i]
        else:
            assert (
                d_max <= plot_dimensions[i][1]
            ), "Trying to clip out of bounds. Please revise clipping domain."
            rescaled_domain_max = (d_max - plot_dimensions[i][0]) / (
                plot_dimensions[i][1] - plot_dimensions[i][0]
            )
            slice_end = int(shape[i] * rescaled_domain_max)

        # Apply a transformation if increasing the index corresponds to the negative dimension direction.
        if not orbit_instance.positive_indexing()[i]:
            slice_start = shape[i] - slice_start
            slice_end = shape[i] - slice_end
            slice_start, slice_end = slice_end, slice_start

        if np.mod(slice_end - slice_start, 2):
            # If the difference is odd, then floor dividing and multiplying by two switches whichever is odd to even.
            # By definition, only one can be odd if the difference is odd; hence only once number is changing.
            slice_start, slice_end = 2 * (slice_start // 2), 2 * (slice_end // 2)
        clipping_slices.append(slice(slice_start, slice_end))

        # Find the correct fraction of the length>0 then subtract the minimum to rescale back to original plot units.
        ith_clipping_dim = (
            int(np.abs(slice_end - slice_start)) / shape[i]
        ) * actual_dimensions[i]
        clipping_dimensions.append(ith_clipping_dim)

    return tuple(clipping_slices), tuple(clipping_dimensions)
