import inspect
from functools import wraps
from itertools import chain

import numpy as np

import astropy.nddata
from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS, HighLevelWCSWrapper, SlicedLowLevelWCS

from ndcube.utils import wcs as wcs_utils
from ndcube.utils.exceptions import warn_user

__all__ = [
    "get_crop_item_from_points",
    "propagate_rebin_uncertainties",
    "sanitize_crop_inputs",
    "sanitize_wcs",
]


def sanitize_wcs(func):
    """
    A wrapper for NDCube methods to sanitise the wcs argument.

    This decorator is only designed to be used on methods of NDCube.

    It will find the wcs argument, keyword or positional and if it is `None`, set
    it to ``self.wcs``.
    It will then verify that the WCS has a matching number of pixel dimensions
    to the dimensionality of the array. It will finally verify that the object
    passed is a HighLevelWCS object, or an ExtraCoords object.
    """
    # This needs to be here to prevent a circular import
    from ndcube.extra_coords.extra_coords import ExtraCoords  # noqa: PLC0415

    @wraps(func)
    def wcs_wrapper(*args, **kwargs):
        sig = inspect.signature(func)
        params = sig.bind(*args, **kwargs)
        wcs = params.arguments.get('wcs', None)
        self = params.arguments['self']

        if wcs is None:
            wcs = self.wcs

        if not isinstance(wcs, ExtraCoords):
            if not wcs.pixel_n_dim == self.data.ndim:
                raise ValueError("The supplied WCS must have the same number of "
                                 "pixel dimensions as the NDCube object. "
                                 "If you specified `cube.extra_coords.wcs` "
                                 "please just pass `cube.extra_coords`.")

        if not isinstance(wcs, (BaseHighLevelWCS, ExtraCoords)):
            raise TypeError("wcs argument must be a High Level WCS or an ExtraCoords object.")

        params.arguments['wcs'] = wcs

        return func(*params.args, **params.kwargs)

    return wcs_wrapper


def sanitize_crop_inputs(points, wcs):
    """Sanitize inputs to NDCube crop methods.

    First arg returned signifies whether the inputs imply that cropping
    should be performed or not.
    """
    points = list(points)
    n_points = len(points)
    n_coords = [None] * n_points
    values_are_none = [False] * n_points
    for i, point in enumerate(points):
        # Ensure each point is a list
        if isinstance(point, (tuple, list)):
            points[i] = list(point)
        else:
            points[i] = [point]
        # Record number of objects in each point.
        # Later we will ensure all points have same number of objects.
        n_coords[i] = len(points[i])
        # Confirm whether point contains at least one None entry.
        if all(coord is None for coord in points[i]):
            values_are_none[i] = True
        # Squeeze length-1 coordinate objects to scalars.
        points[i] = [coord.squeeze() if hasattr(coord, "squeeze") else coord for coord in points[i]]
    # If no points contain a coord, i.e. if all entries in all points are None,
    # set no-op flag to True and exit.
    if all(values_are_none):
        return True, points, wcs
    # Not not all points are of same length, error.
    if len(set(n_coords)) != 1:
        raise ValueError("All points must have same number of coordinate objects."
                         f"Number of objects in each point: {n_coords}")
    # Import must be here to avoid circular import.
    from ndcube.extra_coords.extra_coords import ExtraCoords  # noqa: PLC0415
    if isinstance(wcs, ExtraCoords):
        # Determine how many dummy axes are needed
        n_dummy_axes = len(wcs._cube_array_axes_without_extra_coords)
        if n_dummy_axes > 0:
            points = [point + [None] * n_dummy_axes for point in points]
        # Convert extra coords to WCS describing whole cube.
        wcs = wcs.cube_wcs
    # Ensure WCS is low level.
    if isinstance(wcs, BaseHighLevelWCS):
        wcs = wcs.low_level_wcs
    return False, points, wcs


def get_crop_item_from_points(points, wcs, crop_by_values, keepdims, original_shape):
    """
    Find slice item that crops to minimum cube in array-space containing specified world points.

    Parameters
    ----------
    points : iterable of iterables
        Each iterable represents a point in real world space.
        Each element in a point gives the real world coordinate value of the point
        in high-level coordinate objects or quantities.
        (Must be consistently high or low level within and across points.)
        Objects must be in the order required by
        wcs.world_to_array_index/world_to_array_index_values.

    wcs : `~astropy.wcs.wcsapi.BaseHighLevelWCS`, `~astropy.wcs.wcsapi.BaseLowLevelWCS`
        The WCS to use to convert the world coordinates to array indices.

    crop_by_values : `bool`
        Denotes whether cropping is done using high-level objects or "values",
        i.e. low-level objects.

    keepdims : `bool`
        If `False`, returned item will drop length-1 dimensions otherwise, item will keep length-1 dimensions.

    original_shape: `tuple` of `int`
        The shape of the data cube before cropping.

    Returns
    -------
    item : `tuple` of `slice`
        The slice item for each axis of the cube which, when applied to the cube,
        will return the minimum cube in array-index-space that contains all the
        input world points.
    """
    # Define a list of lists to hold the pixel coordinates of the points
    # where each inner list gives the pixel coordinates of all points for that pixel axis.
    # Recall that pixel axis ordering is reversed compared to array axis ordering.
    combined_points_pixel_idx = [[]] * wcs.pixel_n_dim
    high_level_wcs = HighLevelWCSWrapper(wcs) if isinstance(wcs, BaseLowLevelWCS) else wcs
    low_level_wcs = high_level_wcs.low_level_wcs
    # For each point compute the corresponding array indices.
    for point in points:
        # Get the pixel axes associated with each element in point.
        point_inputs_pixel_axes = (
            tuple(wcs_utils.world_axis_to_pixel_axes(i, low_level_wcs.axis_correlation_matrix)
                  for i in range(low_level_wcs.world_n_dim)) if crop_by_values
            else wcs_utils.pixel_indices_for_world_objects(high_level_wcs))
        # Get indices of pixel axes which correspond to only None inputs in point
        # as well as those that correspond to a coord.
        point_indices_with_inputs = []
        pixel_axes_with_input = []
        for i, coord in enumerate(point):
            if coord is not None:
                point_indices_with_inputs.append(i)
                pixel_axes_with_input.append(point_inputs_pixel_axes[i])
        pixel_axes_with_input = set(chain.from_iterable(pixel_axes_with_input))
        pixel_axes_without_input = set(range(low_level_wcs.pixel_n_dim)) - pixel_axes_with_input
        pixel_axes_with_input = np.array(list(pixel_axes_with_input))
        pixel_axes_without_input = np.array(list(pixel_axes_without_input))
        # Slice out the axes that do not correspond to a coord
        # from the WCS and the input point.
        if len(pixel_axes_without_input) > 0:
            array_axes_without_input = wcs_utils.convert_between_array_and_pixel_axes(
                pixel_axes_without_input, low_level_wcs.pixel_n_dim)
            wcs_slice = np.array([slice(None)] * low_level_wcs.pixel_n_dim)
            wcs_slice[array_axes_without_input] = 0
            sliced_wcs = SlicedLowLevelWCS(low_level_wcs, slices=tuple(wcs_slice))
            sliced_point = np.array(point, dtype=object)[np.array(point_indices_with_inputs)]
        else:
            # Else, if all axes have at least one crop input, no need to slice the WCS.
            sliced_wcs, sliced_point = low_level_wcs, np.array(point, dtype=object)
        # Derive the pixel indices of the input point and place each index
        # in the list corresponding to its axis.
        # Use the to_pixel methods to preserve fractional indices for future rounding.
        if crop_by_values:
            # world_to_pixel_values is APE14 low-level API and expects plain
            # floats, not Quantity objects. So we need to strip units here; the values are
            # already in the correct units because _get_crop_by_values_item called
            # .to(wcs.world_axis_units[j]) before reaching this point.
            #
            # Passing Quantity objects raises TypeError in gWCS when the WCS's
            # declared high-level type is itself Quantity (e.g., a WCS built from
            # QuantityTableCoordinate), because gWCS cannot distinguish such inputs
            # from an accidental high-level API call.
            stripped_point = [p.value if hasattr(p, "value") else p for p in sliced_point]
            point_pixel_indices = sliced_wcs.world_to_pixel_values(*stripped_point)
        else:
            point_pixel_indices = HighLevelWCSWrapper(sliced_wcs).world_to_pixel(*sliced_point)
        # For each pixel axis associated with this point, place the pixel coords for
        # that pixel axis into the corresponding list within combined_points_pixel_idx.
        if sliced_wcs.pixel_n_dim == 1:
            point_pixel_indices = (point_pixel_indices,)
        for axis, index in zip(pixel_axes_with_input, point_pixel_indices):
            combined_points_pixel_idx[axis] = combined_points_pixel_idx[axis] + [index]

    # Iterate through each array axis to determine the min and max pixel coords
    # and then convert to array indices. Note that combined_points_pixel_idx holds the
    # pixel coords for each pixel axis. Therefore, to iterate in array axis order,
    # combined_points_pixel_idx must be reversed.
    item = []
    ambiguous = False
    message = ""
    result_is_scalar = True
    for array_axis, pixel_coords in enumerate(combined_points_pixel_idx[::-1]):
        if pixel_coords == []:
            result_is_scalar = False
            item.append(slice(None))
        else:
            # Calculate the index of the array element containing the pixel coordinate.
            # Note that integer pixel coordinates correspond to the pixel center,
            # while integer array indices correspond to lower edge of desired array element.
            # Therefore a shift of 0.5 is required in the conversion.
            # The max idx conversion below will discard right-ward array element if
            # max pixel coord corresponds to a pixel edge.
            min_array_idx = int(np.floor(min(pixel_coords) + 0.5))
            max_array_idx = int(np.ceil(max(pixel_coords) - 0.5)) + 1
            # Raise error if indices all lie below or all lie above array axis's extent.
            # Exception: min_array_idx == max_array_idx == 0 is allowed because max_array_idx
            # will be later changed to 1.
            if (min_array_idx < 0 and max_array_idx <= 0) or min_array_idx >= original_shape[array_axis]:
                raise ValueError(f"All world points associated with array axis {array_axis}"
                                 " are outside the range of the NDCube being cropped.")
            # world_to_array_index uses negative indices to represent locations to the left
            # of the 0th pixel, while python slicing uses them to count backwards from the
            # last element in the array. Therefore, set negative indices to 0.
            # Note that we've already checked that the max pixel_coord is >= 0.
            # Also note that there's no need to clip the max array idx, as values above
            # the array extent does not cause ambiguity in the slicing so long as the
            # min array idx is below that upper extent, which has also already been checked
            # by the above error.
            if min_array_idx < 0:
                min_array_idx = 0
            # Due to the above calculation, the above min and max array indices can only be
            # same if the original pixel coords correspond to the same pixel edge.
            # If this is the case, increment the max array index by 1 so the rightward array
            # element is kept. Also, build a warning message about this to be raised later.
            if min_array_idx == max_array_idx:
                ambiguous = True
                max_array_idx += 1
                if min_array_idx == 0:
                    message += (f"All input points corresponding to array axis {array_axis} lie on "
                                "the lower boundary of array element 0 (the first element). "
                                "The cropped NDCube will only include array element 0.\n")
                else:
                    message += (f"All input points corresponding to array axis {array_axis} lie on "
                                f"the boundary between array elements {min_array_idx - 1} and "
                                f"{min_array_idx}. The cropped NDCube will only include array "
                                f"element {min_array_idx}.\n")
            if max_array_idx - min_array_idx == 1 and not keepdims:
                item.append(min_array_idx)
            else:
                item.append(slice(min_array_idx, max_array_idx))
                result_is_scalar = False
    # Raise warning if all world values for any array axes correspond to a pixel edge.
    if ambiguous:
        warn_user(message)
    # If item will result in a scalar cube, raise an error as this is not currently supported.
    if result_is_scalar:
        raise ValueError("Input points causes cube to be cropped to a single pixel. "
                         "This is not supported when keepdims=False.")
    return tuple(item)


def propagate_rebin_uncertainties(uncertainty, data, mask, operation, operation_ignores_mask=False,
                                  propagation_operation=None, correlation=0, **kwargs):
    """
    Default algorithm for uncertainty propagation in :meth:`~ndcube.NDCube.rebin`.

    First dimension of uncertainty, data and mask inputs represent the pixels
    in the bin being aggregated by the rebin process while the latter dimensions
    must have the same shape as the rebinned data. The operation input is the
    function used to aggregate elements in the first dimension, e.g. `numpy.sum`.

    Parameters
    ----------
    uncertainty: `astropy.nddata.NDUncertainty`
        Cannot be instance of `astropy.nddata.UnknownUncertainty`.
        The uncertainties associated with the data. The first dimension represents
        pixels in each bin being aggregated while trailing dimensions must have
        the same shape as the rebinned data.
    data: array-like or `None`
        The data associated with the above uncertainties.
        Must have same shape as above.
    mask: array-like of `bool` or `None`
        Indicates whether any uncertainty elements should be ignored in propagation.
        If True, corresponding uncertainty element is ignored. If False, it is used.
        Must have same shape as above.
    operation: function
        The function used to aggregate the data for which the uncertainties are being
        propagated here.
    operation_ignores_mask: `bool`
        Determines whether masked values are used or excluded from calculation.
        Default is False causing masked data and uncertainty to be excluded.
    propagation_operation: function
        The operation which defines how the uncertainties are propagated.
        This can differ from operation, e.g. if operation is sum, then
        propagation_operation should be add.
    correlation: `int`
        Passed to `astropy.nddata.NDUncertainty.propagate`. See that method's docstring.
        Default=0.

    Returns
    -------
    new_uncertainty: `astropy.nddata.NDUncertainty`
        The propagated uncertainty.
        Same shape as input uncertainty without its first dimension.
    """
    flat_axis = 0
    operation_is_mean = True if operation in {np.mean, np.nanmean} else False
    operation_is_nantype = True if operation in {np.nansum, np.nanmean, np.nanprod} else False
    # If propagation_operation kwarg not set manually, try to set it based on operation kwarg.
    if not propagation_operation:
        if operation in {np.sum, np.nansum, np.mean, np.nanmean}:
            propagation_operation = np.add
        # TODO: product was renamed to prod for numpy 2.0
        elif operation in {np.prod, np.nanprod}:
            propagation_operation = np.multiply
        else:
            raise ValueError("propagation_operation not recognized.")
    # Build mask if not provided.
    new_uncertainty = uncertainty[0]  # Define uncertainty for initial iteration step.
    if operation_ignores_mask or mask is None:
        mask = False
    if mask is False:
        if operation_is_nantype:
            nan_mask = np.isnan(data)
            if nan_mask.any():
                mask = nan_mask
                idx = np.logical_not(mask)
                mask1 = mask[1:]
        else:
            # If there is no mask and operation is not nan-type, build generator
            # so non-mask can still be iterated.
            n_pix_per_bin = data.shape[flat_axis]
            mask1 = (False for i in range(1, n_pix_per_bin))
    else:
        # Mask uncertainties corresponding to nan data if operation is nantype.
        if operation_is_nantype:
            mask[np.isnan(data)] = True
        # Set masked uncertainties in first mask to 0
        # as they shouldn't count towards final uncertainty.
        mask1 = mask[1:]
        idx = np.logical_not(mask)
        uncertainty.array[mask] = 0
        new_uncertainty.array[mask[0]] = 0
    # Propagate uncertainties.
    # Note uncertainty must be associated with a parent nddata for some propagations.
    cumul_data = data[0]
    if mask is not False and operation_ignores_mask is False:
        cumul_data[idx[0]] = 0
    parent_nddata = astropy.nddata.NDData(cumul_data, uncertainty=new_uncertainty)
    new_uncertainty.parent_nddata = parent_nddata
    for j, mask_slice in enumerate(mask1):
        i = j + 1
        cumul_data = operation(data[:i+1]) if mask is False else operation(data[:i+1][idx[:i+1]])
        data_slice = astropy.nddata.NDData(data=data[i], mask=mask_slice,
                                           uncertainty=uncertainty[i])
        new_uncertainty = new_uncertainty.propagate(propagation_operation, data_slice,
                                                    cumul_data, correlation)
        parent_nddata = astropy.nddata.NDData(cumul_data, uncertainty=new_uncertainty)
        new_uncertainty.parent_nddata = parent_nddata
    # If aggregation operation is mean, uncertainties must be divided by
    # number of unmasked pixels in each bin.
    if operation_is_mean and propagation_operation is np.add:
        if mask is False:
            new_uncertainty.array /= n_pix_per_bin
        else:
            unmasked_per_bin = np.logical_not(mask).astype(int).sum(axis=flat_axis)
            new_uncertainty.array /= np.clip(unmasked_per_bin, 1, None)
    return new_uncertainty
