import numpy as np

from astropy.units import Unit, si
from astropy.units import equivalencies as eq
from astropy.units.decorators import quantity_input
from astropy.units.quantity import Quantity, SpecificTypeQuantity

__all__ = ["SpectralQuantity"]

# We don't want to run doctests in the docstrings we inherit from Quantity
__doctest_skip__ = ["SpectralQuantity.*"]

KMS = si.km / si.s

SPECTRAL_UNITS = (si.Hz, si.m, si.J, si.m**-1, KMS)

DOPPLER_CONVENTIONS = {
    "radio": eq.doppler_radio,
    "optical": eq.doppler_optical,
    "relativistic": eq.doppler_relativistic,
}


class SpectralQuantity(SpecificTypeQuantity):
    """
    One or more value(s) with spectral units.

    The spectral units should be those for frequencies, wavelengths, energies,
    wavenumbers, or velocities (interpreted as Doppler velocities relative to a
    rest spectral value). The advantage of using this class over the regular
    `~astropy.units.Quantity` class is that in `SpectralQuantity`, the
    ``u.spectral`` equivalency is enabled by default (allowing automatic
    conversion between spectral units), and a preferred Doppler rest value and
    convention can be stored for easy conversion to/from velocities.

    Parameters
    ----------
    value : ndarray or `~astropy.units.Quantity` or `SpectralQuantity`
        Spectral axis data values.
    unit : unit-like
        Unit for the given data.
    doppler_rest : `~astropy.units.Quantity` ['speed'], optional
        The rest value to use for conversions from/to velocities
    doppler_convention : str, optional
        The convention to use when converting the spectral data to/from
        velocities.
    """

    _equivalent_unit = SPECTRAL_UNITS

    _include_easy_conversion_members = True

    def __new__(
        cls, value, unit=None, doppler_rest=None, doppler_convention=None, **kwargs
    ):
        obj = super().__new__(cls, value, unit=unit, **kwargs)

        # If we're initializing from an existing SpectralQuantity, keep any
        # parameters that aren't being overridden
        if doppler_rest is None:
            doppler_rest = getattr(value, "doppler_rest", None)
        if doppler_convention is None:
            doppler_convention = getattr(value, "doppler_convention", None)

        obj._doppler_rest = doppler_rest
        obj._doppler_convention = doppler_convention

        return obj

    def __array_finalize__(self, obj):
        super().__array_finalize__(obj)
        self._doppler_rest = getattr(obj, "_doppler_rest", None)
        self._doppler_convention = getattr(obj, "_doppler_convention", None)

    def __quantity_subclass__(self, unit):
        # Always default to just returning a Quantity, unless we explicitly
        # choose to return a SpectralQuantity - even if the units match, we
        # want to avoid doing things like adding two SpectralQuantity instances
        # together and getting a SpectralQuantity back
        if unit is self.unit:
            return SpectralQuantity, True
        else:
            return Quantity, False

    def __array_ufunc__(self, function, method, *inputs, **kwargs):
        # We always return Quantity except in a few specific cases
        result = super().__array_ufunc__(function, method, *inputs, **kwargs)
        if (
            (
                function is np.multiply
                or (function is np.true_divide and inputs[0] is self)
            )
            and result.unit == self.unit
        ) or (
            function in (np.minimum, np.maximum, np.fmax, np.fmin)
            and method in ("reduce", "reduceat")
        ):
            result = result.view(self.__class__)
            result.__array_finalize__(self)
        else:
            if result is self:
                raise TypeError(
                    "Cannot store the result of this operation in"
                    f" {self.__class__.__name__}"
                )
            if result.dtype.kind == "b":
                result = result.view(np.ndarray)
            else:
                result = result.view(Quantity)
        return result

    @property
    def doppler_rest(self):
        """
        The rest value of the spectrum used for transformations to/from
        velocity space.

        Returns
        -------
        `~astropy.units.Quantity` ['speed']
            Rest value as an astropy `~astropy.units.Quantity` object.
        """
        return self._doppler_rest

    @doppler_rest.setter
    @quantity_input(value=SPECTRAL_UNITS)
    def doppler_rest(self, value):
        """
        New rest value needed for velocity-space conversions.

        Parameters
        ----------
        value : `~astropy.units.Quantity` ['speed']
            Rest value.
        """
        if self._doppler_rest is not None:
            raise AttributeError(
                "doppler_rest has already been set, and cannot be changed. Use the"
                " ``to`` method to convert the spectral values(s) to use a different"
                " rest value"
            )
        self._doppler_rest = value

    @property
    def doppler_convention(self):
        """
        The defined convention for conversions to/from velocity space.

        Returns
        -------
        str
            One of 'optical', 'radio', or 'relativistic' representing the
            equivalency used in the unit conversions.
        """
        return self._doppler_convention

    @doppler_convention.setter
    def doppler_convention(self, value):
        """
        New velocity convention used for velocity space conversions.

        Parameters
        ----------
        value

        Notes
        -----
        More information on the equations dictating the transformations can be
        found in the astropy documentation [1]_.

        References
        ----------
        .. [1] Astropy documentation: https://docs.astropy.org/en/stable/units/equivalencies.html#spectral-doppler-equivalencies

        """
        if self._doppler_convention is not None:
            raise AttributeError(
                "doppler_convention has already been set, and cannot be changed. Use"
                " the ``to`` method to convert the spectral values(s) to use a"
                " different convention"
            )

        if value is not None and value not in DOPPLER_CONVENTIONS:
            raise ValueError(
                "doppler_convention should be one of"
                f" {'/'.join(sorted(DOPPLER_CONVENTIONS))}"
            )

        self._doppler_convention = value

    @quantity_input(doppler_rest=SPECTRAL_UNITS)
    def to(self, unit, equivalencies=[], doppler_rest=None, doppler_convention=None):
        """
        Return a new `~astropy.coordinates.SpectralQuantity` object with the specified unit.

        By default, the ``spectral`` equivalency will be enabled, as well as
        one of the Doppler equivalencies if converting to/from velocities.

        Parameters
        ----------
        unit : unit-like
            An object that represents the unit to convert to. Must be
            an `~astropy.units.UnitBase` object or a string parseable
            by the `~astropy.units` package, and should be a spectral unit.
        equivalencies : list of `~astropy.units.equivalencies.Equivalency`, optional
            A list of equivalence pairs to try if the units are not
            directly convertible (along with spectral).
            See :ref:`astropy:unit_equivalencies`.
            If not provided or ``[]``, spectral equivalencies will be used.
            If `None`, no equivalencies will be applied at all, not even any
            set globally or within a context.
        doppler_rest : `~astropy.units.Quantity` ['speed'], optional
            The rest value used when converting to/from velocities. This will
            also be set at an attribute on the output
            `~astropy.coordinates.SpectralQuantity`.
        doppler_convention : {'relativistic', 'optical', 'radio'}, optional
            The Doppler convention used when converting to/from velocities.
            This will also be set at an attribute on the output
            `~astropy.coordinates.SpectralQuantity`.

        Returns
        -------
        `SpectralQuantity`
            New spectral coordinate object with data converted to the new unit.
        """
        # Make sure units can be passed as strings
        unit = Unit(unit)

        # If equivalencies is explicitly set to None, we should just use the
        # default Quantity.to with equivalencies also set to None
        if equivalencies is None:
            result = super().to(unit, equivalencies=None)
            result = result.view(self.__class__)
            result.__array_finalize__(self)
            return result

        # FIXME: need to consider case where doppler equivalency is passed in
        # equivalencies list, or is u.spectral equivalency is already passed

        if doppler_rest is None:
            doppler_rest = self._doppler_rest

        if doppler_convention is None:
            doppler_convention = self._doppler_convention
        elif doppler_convention not in DOPPLER_CONVENTIONS:
            raise ValueError(
                "doppler_convention should be one of"
                f" {'/'.join(sorted(DOPPLER_CONVENTIONS))}"
            )

        if self.unit.is_equivalent(KMS) and unit.is_equivalent(KMS):
            # Special case: if the current and final units are both velocity,
            # and either the rest value or the convention are different, we
            # need to convert back to frequency temporarily.

            if doppler_convention is not None and self._doppler_convention is None:
                raise ValueError("Original doppler_convention not set")

            if doppler_rest is not None and self._doppler_rest is None:
                raise ValueError("Original doppler_rest not set")

            if doppler_rest is None and doppler_convention is None:
                result = super().to(unit, equivalencies=equivalencies)
                result = result.view(self.__class__)
                result.__array_finalize__(self)
                return result

            elif (doppler_rest is None) is not (doppler_convention is None):
                raise ValueError(
                    "Either both or neither doppler_rest and doppler_convention should"
                    " be defined for velocity conversions"
                )

            vel_equiv1 = DOPPLER_CONVENTIONS[self._doppler_convention](
                self._doppler_rest
            )

            freq = super().to(si.Hz, equivalencies=equivalencies + vel_equiv1)

            vel_equiv2 = DOPPLER_CONVENTIONS[doppler_convention](doppler_rest)

            result = freq.to(unit, equivalencies=equivalencies + vel_equiv2)

        else:
            additional_equivalencies = eq.spectral()

            if self.unit.is_equivalent(KMS) or unit.is_equivalent(KMS):
                if doppler_convention is None:
                    raise ValueError(
                        "doppler_convention not set, cannot convert to/from velocities"
                    )

                if doppler_rest is None:
                    raise ValueError(
                        "doppler_rest not set, cannot convert to/from velocities"
                    )

                additional_equivalencies = (
                    additional_equivalencies
                    + DOPPLER_CONVENTIONS[doppler_convention](doppler_rest)
                )

            result = super().to(
                unit, equivalencies=equivalencies + additional_equivalencies
            )

        # Since we have to explicitly specify when we want to keep this as a
        # SpectralQuantity, we need to convert it back from a Quantity to
        # a SpectralQuantity here. Note that we don't use __array_finalize__
        # here since we might need to set the output doppler convention and
        # rest based on the parameters passed to 'to'
        result = result.view(self.__class__)
        result.__array_finalize__(self)
        result._doppler_convention = doppler_convention
        result._doppler_rest = doppler_rest

        return result

    def to_value(self, unit=None, *args, **kwargs):
        if unit is None:
            return self.view(np.ndarray)

        return self.to(unit, *args, **kwargs).value
