import abc
from typing import Any
from numbers import Integral
from functools import reduce, partial
from collections.abc import Iterable

import numpy as np

import astropy.units as u
from astropy.coordinates import SkyCoord
from astropy.time import Time
from astropy.wcs import WCS
from astropy.wcs.wcsapi import BaseHighLevelWCS
from astropy.wcs.wcsapi.high_level_wcs_wrapper import HighLevelWCSWrapper
from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS, sanitize_slices

from ndcube.utils.wcs import convert_between_array_and_pixel_axes
from ndcube.wcs.wrappers import CompoundLowLevelWCS, ResampledLowLevelWCS
from .table_coord import (
    BaseTableCoordinate,
    MultipleTableCoordinate,
    QuantityTableCoordinate,
    SkyCoordTableCoordinate,
    TimeTableCoordinate,
)

__all__ = ['ExtraCoords', 'ExtraCoordsABC']


class ExtraCoordsABC(abc.ABC):
    """
    A representation of additional world coordinates associated with pixel axes.

    ExtraCoords can be initialised by either specifying a
    `~astropy.wcs.wcsapi.BaseLowLevelWCS` object and a ``mapping``, or it can be
    built up by specifying one or more lookup tables.

    Parameters
    ----------
    wcs
        The WCS specifying the extra coordinates.
    mapping
       The mapping between the array dimensions and pixel dimensions in the
       extra coords object. This is an iterable of ``(array_dimension, pixel_dimension)`` pairs
       of length equal to the number of pixel dimensions in the extra coords.

    """
    @abc.abstractmethod
    def add(self,
            name: str | Iterable[str],
            array_dimension: int | Iterable[int],
            lookup_table: Any,
            physical_types: str | Iterable[str] = None,
            **kwargs):
        """
        Add a coordinate to this `~ndcube.ExtraCoords` based on a lookup table.

        Parameters
        ----------
        name : `str` or sequence of `str`
            The name(s) for these world coordinate(s).
        array_dimension : `int` or `tuple` of `int`
            The array dimension(s), to which this lookup table corresponds.
        lookup_table : `object` or sequence of `object`
            The lookup table. A `~ndcube.extra_coords.BaseTableCoordinate` subclass or anything
            that can instantiate one, i.e. currently a `~astropy.time.Time`,
            `~astropy.coordinates.SkyCoord`, or a (sequence of) `~astropy.units.Quantity`.
        physical_types: `str` or iterable of `str`, optional
            Descriptor(s) of the :ref:`<dimensions and physical types <dimensions>`
            associated with each axis; length must match the number of dimensions in
            ``lookup_table``.
        """

    @abc.abstractmethod
    def keys(self) -> Iterable[str]:
        """
        The world axis names for all the coordinates in the extra coords.
        """

    @property
    @abc.abstractmethod
    def mapping(self) -> Iterable[tuple[int, int]]:
        """
        The mapping between the array dimensions and pixel dimensions.

        This is an iterable of ``(array_dimension, pixel_dimension)`` pairs
        of length equal to the number of pixel dimensions in the extra coords.
        """

    @property
    @abc.abstractmethod
    def wcs(self) -> BaseHighLevelWCS:
        """
        A WCS object representing the world coordinates described by this ``ExtraCoords``.

        .. note::
            This WCS object does not map to the pixel dimensions of the data array
            in the `~ndcube.NDCube` object. It only includes pixel dimensions associated
            with the extra coordinates. For example, if there is only one extra coordinate
            associated with a single pixel dimension, this WCS will only have 1 pixel dimension,
            even if the `~ndcube.NDCube` object has a data array of 2-D or greater.
            Therefore using this WCS directly might lead to some confusing results.

        """

    @property
    @abc.abstractproperty
    def is_empty(self):
        """Return True if no extra coords present, else return False."""

    @abc.abstractmethod
    def __getitem__(self, item: str | int | slice | Iterable[str | int | slice]) -> "ExtraCoordsABC":
        """
        ExtraCoords can be sliced with either a string, or a numpy like slice.

        When sliced with a string it should return a new ExtraCoords object
        with only those coordinates with the given names. When sliced with a
        numpy array like slice it should return a new ExtraCoords with the
        slice applied. Supporting step is not required and "fancy indexing" is
        not supported.
        """


class ExtraCoords(ExtraCoordsABC):
    """
    A representation of additional world coordinates associated with pixel axes.

    ExtraCoords can be initialised by either specifying a
    `~astropy.wcs.wcsapi.BaseLowLevelWCS` object and a ``mapping``, or it can be
    built up by specifying one or more lookup tables.

    Parameters
    ----------
    wcs
        The WCS specifying the extra coordinates.
    mapping
       The mapping between the array dimensions and pixel dimensions in the
       extra coords object. This is an iterable of ``(array_dimension, pixel_dimension)`` pairs
       of length equal to the number of pixel dimensions in the extra coords.

    """

    def __init__(self, ndcube=None):
        super().__init__()

        # Setup private attributes
        self._wcs = None
        self._mapping = None

        # Lookup tables is a list of (pixel_dim, LookupTableCoord) to allow for
        # one pixel dimension having more than one lookup coord.
        self._lookup_tables = []
        self._dropped_tables = []

        # We need a reference to the parent NDCube
        self._ndcube = ndcube

    @classmethod
    def from_lookup_tables(cls, names, pixel_dimensions, lookup_tables, physical_types=None):
        """
        Construct a new ExtraCoords instance from lookup tables.

        This is a convenience wrapper around `ndcube.ExtraCoords.add` which does not
        expose all the options available in that method.

        Parameters
        ----------
        names : `tuple` of `str`
            The names of the world coordinates.
        pixel_dimensions : `tuple` of `int`
            The pixel dimensions (in the array) to which the ``lookup_tables``
            apply. Must be the same length as ``lookup_tables``.
        lookup_tables : iterable of `object`
            The lookup tables which specify the world coordinates for the ``pixel_dimensions``.
            Must be `~ndcube.extra_coords.BaseTableCoordinate` subclass instances or objects from
            which to instantiate them (see `ndcube.ExtraCoords.add`).
        physical_types: sequence of `str` or of sequences of `str`, optional
            Descriptors of the :ref:`dimensions`
            associated with each axis in the tables. Must be the same length as ``lookup_tables``;
            and length of each element must match the number of dimensions in corresponding
            ``lookup_tables[i]``.
            Defaults to `None`.

        Returns
        -------
        `ndcube.ExtraCoords`

        """
        if len(pixel_dimensions) != len(lookup_tables):
            raise ValueError(
                "The length of pixel_dimensions and lookup_tables must match."
            )

        if physical_types is None:
            physical_types = len(lookup_tables) * [physical_types]
        elif len(physical_types) != len(lookup_tables):
            raise ValueError("The number of physical types and lookup_tables must match.")

        extra_coords = cls()

        for name, pixel_dim, lookup_table, physical_type in zip(names, pixel_dimensions,
                                                                lookup_tables, physical_types):
            extra_coords.add(name, pixel_dim, lookup_table, physical_types=physical_type)

        return extra_coords

    def add(self, name, array_dimension, lookup_table, physical_types=None, **kwargs):
        # docstring in ABC

        if self._wcs is not None:
            raise ValueError(
                "Can not add a lookup_table to an ExtraCoords which was instantiated with a WCS object."
            )

        kwargs['names'] = [name] if not isinstance(name, (list, tuple)) else name

        if isinstance(lookup_table, BaseTableCoordinate):
            coord = lookup_table
        elif isinstance(lookup_table, Time):
            coord = TimeTableCoordinate(lookup_table, physical_types=physical_types, **kwargs)
        elif isinstance(lookup_table, SkyCoord):
            coord = SkyCoordTableCoordinate(lookup_table, physical_types=physical_types, **kwargs)
        elif isinstance(lookup_table, (list, tuple)):
            coord = QuantityTableCoordinate(*lookup_table, physical_types=physical_types, **kwargs)
        elif isinstance(lookup_table, u.Quantity):
            coord = QuantityTableCoordinate(lookup_table, physical_types=physical_types, **kwargs)
        else:
            raise TypeError(f"The input type {type(lookup_table)} isn't supported")

        self._lookup_tables.append((array_dimension, coord))

        # Sort the LUTs so that the mapping and the wcs are ordered in pixel dim order
        self._lookup_tables = sorted(self._lookup_tables,
                                          key=lambda x: x[0] if isinstance(x[0], Integral) else x[0][0])

    @property
    def _name_lut_map(self):
        """
        Map of world names to the corresponding `.LookupTableCoord`
        """
        return {lut[1].wcs.world_axis_names: lut for lut in self._lookup_tables}

    def keys(self):
        # docstring in ABC
        if not self.wcs:
            return ()

        return tuple(self.wcs.world_axis_names) if self.wcs.world_axis_names else None

    @property
    def mapping(self):
        # docstring in ABC
        if self._mapping:
            return self._mapping

        # If mapping is not set but lookup_tables is empty then the extra
        # coords is empty, so there is no mapping.
        if not self._lookup_tables:
            return ()

        # The mapping is from the array index (position in the list) to the
        # pixel dimensions (numbers in the list)
        lts = [list([lt[0]] if isinstance(lt[0], Integral) else lt[0]) for lt in self._lookup_tables]
        converter = partial(convert_between_array_and_pixel_axes, naxes=len(self._ndcube.shape))
        pixel_indicies = [list(converter(np.array(ids))) for ids in lts]
        return tuple(reduce(list.__add__, pixel_indicies))

    @mapping.setter
    def mapping(self, mapping):
        if self._mapping is not None:
            raise AttributeError("Can't set mapping if a mapping has already been specified.")

        if self._lookup_tables:
            raise AttributeError(
                "Can't set mapping manually when ExtraCoords is built from lookup tables."
            )

        if self._wcs is not None:
            if not max(mapping) <= self._wcs.pixel_n_dim - 1:
                raise ValueError(
                    "Values in the mapping can not be larger than the number of pixel dimensions in the WCS."
                )

        self._mapping = mapping

    @property
    def wcs(self):
        # docstring in ABC
        if self._wcs is not None:
            return self._wcs

        if not self._lookup_tables:
            return None

        tcoords = {lt[1] for lt in self._lookup_tables}
        # created a sorted list of unique items
        _tmp = set()  # a temporary set
        tcoords = [x[1] for x in self._lookup_tables if x[1] not in _tmp and _tmp.add(x[1]) is None]
        return MultipleTableCoordinate(*tcoords).wcs

    @wcs.setter
    def wcs(self, wcs):
        if self._wcs is not None:
            raise AttributeError(
                "Can't set wcs if a WCS has already been specified."
            )

        if self._lookup_tables:
            raise AttributeError(
                "Can't set wcs manually when ExtraCoords is built from lookup tables."
            )

        if self._mapping is not None:
            if not max(self._mapping) <= wcs.pixel_n_dim - 1:
                raise ValueError(
                    "Values in the mapping can not be larger than the number of pixel dimensions in the WCS."
                )

        self._wcs = wcs

    @property
    def is_empty(self):
        # docstring in ABC
        if not self._wcs and not self._lookup_tables:
            return True
        return False

    def _getitem_string(self, item):
        """
        Slice the Extracoords based on axis names.
        """

        for names, lut in self._name_lut_map.items():
            if item in names:
                new_ec = ExtraCoords(ndcube=self._ndcube)
                new_ec._lookup_tables = [lut]
                return new_ec

        raise KeyError(f"Can't find the world axis named {item} in this ExtraCoords object.")

    def _getitem_lookup_tables(self, item):
        """
        Apply an array slice to the lookup tables.

        Returns a new ExtraCoords object with modified lookup tables.
        """
        dropped_tables = set()
        new_lookup_tables = set()
        ndims = max([lut[0] if isinstance(lut[0], Integral) else max(lut[0])
                     for lut in self._lookup_tables]) + 1
        # Determine how many dimensions will be dropped by slicing below each dimension.
        if isinstance(item, Integral):
            n_dropped_dims = np.ones(ndims, dtype=int)
            item = tuple([item] + [slice(None)] * (ndims - 1))
        elif isinstance(item, slice):
            n_dropped_dims = np.zeros(ndims, dtype=int)
            item = tuple([item] + [slice(None)] * (ndims - 1))
        else:
            item = list(item) + [slice(None)] * (ndims - len(item))
            n_dropped_dims = np.cumsum([isinstance(i, Integral) for i in item])
        for lut_axis, lut in self._lookup_tables:
            lut_axes = (lut_axis,) if not isinstance(lut_axis, tuple) else lut_axis
            new_lut_axes = tuple(ax - n_dropped_dims[ax] for ax in lut_axes)
            lut_slice = tuple(item[i] for i in lut_axes)
            if isinstance(lut_slice, tuple) and len(lut_slice) == 1:
                lut_slice = lut_slice[0]

            sliced_lut = lut[lut_slice]

            if sliced_lut.is_scalar():
                dropped_tables.add(sliced_lut)
            else:
                new_lookup_tables.add((new_lut_axes, sliced_lut))
        new_extra_coords = type(self)()
        new_extra_coords._lookup_tables = list(new_lookup_tables)
        new_extra_coords._dropped_tables = list(dropped_tables)
        return new_extra_coords

    def _getitem_wcs(self, item):
        item = sanitize_slices(item, self.wcs.pixel_n_dim)

        # It's valid to slice down the EC such that there is nothing left,
        # which is not a valid way to slice the WCS
        if len(item) == self.wcs.pixel_n_dim and all(isinstance(i, Integral) for i in item):
            return type(self)()

        subwcs = self.wcs[item]

        new_mapping = [self.mapping[i] for i, subitem in enumerate(item) if not isinstance(subitem, Integral)]

        new_ec = type(self)()
        new_ec.wcs = subwcs
        new_ec.mapping = new_mapping
        return new_ec

    def __getitem__(self, item):
        # docstring in ABC
        if isinstance(item, str):
            return self._getitem_string(item)

        if self._wcs:
            return self._getitem_wcs(item)

        if self._lookup_tables:
            return self._getitem_lookup_tables(item)

        # If we get here this object is empty, so just return an empty extra coords
        # This is done to simplify the slicing in NDCube
        return self

    @property
    def dropped_world_dimensions(self):
        """
        Return an APE-14 like representation of any sliced out world dimensions.
        """

        if self._wcs:
            if isinstance(self._wcs, SlicedLowLevelWCS):
                return self._wcs.dropped_world_dimensions

        if self._lookup_tables or self._dropped_tables:
            mtc = MultipleTableCoordinate(*[lt[1] for lt in self._lookup_tables])
            mtc._dropped_coords = self._dropped_tables

            return mtc.dropped_world_dimensions

        return {}

    def resample(self, factor, offset=0, ndcube=None, **kwargs):
        """
        Resample all extra coords by given factors in array-index-space.

        One resample factor must be supplied for each array axis in array-axis order.

        Parameters
        ----------
        factor: `int`, `float`, or iterable thereof.
            The factor by which each array axis is resampled.
            If scalar, same factor is applied to all axes.
            Otherwise a factor for each axis must be provided.

        offset: `int`, `float`, or iterable therefore.
            The location on the underlying grid which corresponds
            to the zeroth element after resampling. If iterable, must have an entry
            for each dimension. If a scalar, the grid will be
            shifted by the same amount in all dimensions.

        ndcube: `~ndcube.NDCube`
            The NDCube instance with which the output ExtraCoords object is associated.

        kwargs
            All remaining kwargs are passed to `numpy.interp`.

        Returns
        -------
        new_ec: `~ndcube.ExtraCoords`
            A new ExtraCoords object holding the interpolated coords.
        """
        new_ec = type(self)(ndcube)
        if self.is_empty:
            return new_ec
        if self._ndcube is not None:
            cube_shape = self._ndcube.data.shape
            ndim = len(cube_shape)
        elif self._wcs is not None:
            ndim = self._wcs.pixel_n_dim
        else:
            raise NotImplementedError(
                "Resampling a lookup-table-based ExtraCoords not yet implemented. "
                "Please raise an issue at https://github.com/sunpy/ndcube/issues "
                "if you need this functionality")
        if np.isscalar(factor):
            factor = [factor] * ndim
        if len(factor) != ndim:
            raise ValueError(
                "factor must be scalar or an iterable with length equal to number of cube "
                f"dimensions: len(factor) = {len(factor)}; No. cube dimensions = {ndim}.")
        if np.isscalar(offset):
            offset = [offset] * ndim
        if len(offset) != ndim:
            raise ValueError(
                "offset must be scalar or an iterable with length equal to number of cube "
                f"dimensions: len(offset) = {len(offset)}; No. cube dimensions = {ndim}.")
        # If ExtraCoords object built on WCS, resample using WCS insfrastructure
        if self._wcs is not None:
            new_ec.wcs = HighLevelWCSWrapper(ResampledLowLevelWCS(self._wcs.low_level_wcs,
                                                                  factor, offset))
            return new_ec
        # Else interpolate the lookup table coordinates.
        factor = np.asarray(factor)
        new_grids = []
        for c, d, f in zip(offset, cube_shape, factor):
            x = np.arange(c, d+f, f)
            x = x[x <= d-1]
            new_grids.append(x)
        new_grids = np.array(new_grids, dtype=object)
        for array_axes, coord in self._lookup_tables:
            if np.isscalar(array_axes):
                new_coord = coord.interpolate(new_grids[array_axes], **kwargs)
            else:
                new_coord = coord.interpolate(*new_grids[np.asarray(array_axes)], **kwargs)
            new_ec.add(coord.names, array_axes, new_coord, physical_types=coord.physical_types)
        return new_ec

    @property
    def cube_wcs(self):
        """Produce a WCS that describes the associated NDCube with just the extra coords.

        For NDCube pixel axes without any extra coord, dummy axes are inserted.
        """
        wcses = [self.wcs]
        mapping = list(self.mapping)
        dummy_axes = self._cube_array_axes_without_extra_coords
        n_dummy_axes = len(dummy_axes)
        if n_dummy_axes > 0:
            dummy_wcs = WCS(naxis=n_dummy_axes)
            dummy_wcs.wcs.crpix = [1] * n_dummy_axes
            dummy_wcs.wcs.cdelt = [1] * n_dummy_axes
            dummy_wcs.wcs.crval = [0] * n_dummy_axes
            dummy_wcs.wcs.ctype = ["PIXEL"] * n_dummy_axes
            dummy_wcs.wcs.cunit = ["pix"] * n_dummy_axes
            wcses.append(dummy_wcs)
            mapping += list(dummy_axes)
        return CompoundLowLevelWCS(*wcses, mapping=mapping)

    @property
    def _cube_array_axes_without_extra_coords(self):
        """Return the array axes not associated with any extra coord."""
        return set(range(len(self._ndcube.shape))) - set(self.mapping)

    def __str__(self):
        classname = self.__class__.__name__
        elements = [f"{', '.join(table.names)} ({axes}) {table.physical_types}: {table}"
                    for axes, table in self._lookup_tables]
        length = len(classname) + 2 * len(elements) + sum(len(e) for e in elements)
        if length > np.get_printoptions()['linewidth']:
            joiner = ',\n ' + len(classname) * ' '
        else:
            joiner = ', '

        return f"{classname}({joiner.join(elements)})"

    def __repr__(self):
        return f"{object.__repr__(self)}\n{self}"
