__author__ = "giacomov"

import collections

from astropy import coordinates

from astromodels.core.parameter import Parameter
from astromodels.core.tree import Node


class WrongCoordinatePair(ValueError):
    pass


class IllegalCoordinateValue(ValueError):
    pass


class WrongCoordinateSystem(ValueError):
    pass


class SkyDirection(Node):
    """
    This is essentially a wrapper around astropy.coordinates.SkyCoord with a possibility for
    being serialized and deserialized with YAML.
    """

    def __init__(self, ra=None, dec=None, l=None, b=None, equinox="J2000"):
        """

        :param ra: Right Ascension in degrees
        :param dec: Declination in degrees
        :param l: Galactic latitude in degrees
        :param b: Galactic longitude in degrees
        :param equinox: string
        :return:
        """

        self._equinox = equinox

        # Create the node

        Node.__init__(self, "position")

        # Check that we have the right pairs of coordinates

        if ra is not None and dec is not None:

            # This goes against duck typing, but it is needed to provide a means of initiating this class
            # with either Parameter instances or just floats

            # Try to transform it to float, if it works than we transform it to a parameter

            ra = self._get_parameter_from_input(ra, 0, 360, "ra", "Right Ascension")

            dec = self._get_parameter_from_input(dec, -90, 90, "dec", "Declination")

            self._coord_type = "equatorial"

            self._add_child(ra)
            self._add_child(dec)

        elif l is not None and b is not None:

            # This goes against duck typing, but it is needed to provide a means of initiating this class
            # with either Parameter instances or just floats

            # Try to transform it to float, if it works than we transform it to a parameter

            l = self._get_parameter_from_input(l, 0, 360, "l", "Galactic longitude")

            b = self._get_parameter_from_input(b, -90, 90, "b", "Galactic latitude")

            self._coord_type = "galactic"
            self._add_child(l)
            self._add_child(b)

        else:

            raise WrongCoordinatePair("You have to specify either (ra, dec) or (l, b).")

    @staticmethod
    def _get_parameter_from_input(number_or_parameter, minimum, maximum, what, desc):

        # Try to transform it to float, if it works than we transform it to a parameter

        try:

            number_or_parameter = float(number_or_parameter)

        except TypeError:

            assert isinstance(number_or_parameter, Parameter), (
                "%s must be either a number or a " "parameter instance" % what
            )

            # So this is a Parameter instance already. Enforce that it has the right maximum and minimum

            parameter = number_or_parameter

            assert (
                parameter.min_value >= minimum
            ), "%s must have a minimum greater than or equal to %s" % (
                what,
                minimum,
            )
            assert (
                parameter.max_value <= maximum
            ), "%s must have a maximum less than or equal to %s" % (
                what,
                maximum,
            )

        else:

            # This was a float. Enforce that it has a legal value

            assert (
                minimum <= number_or_parameter <= maximum
            ), "%s cannot have a value of %s, " "it must be %s <= %s <= %s" % (
                what,
                number_or_parameter,
                minimum,
                what,
                maximum,
            )

            parameter = Parameter(
                what,
                number_or_parameter,
                desc=desc,
                min_value=minimum,
                max_value=maximum,
                unit="deg",
                free=False,
            )

        return parameter

    def get_ra(self):
        """
        Get R.A. corresponding to the current position (ICRS, J2000)

        :return: Right Ascension
        """

        try:

            return self.ra.value

        except AttributeError:

            # Transform from L,B to R.A., Dec

            return self.sky_coord.transform_to("icrs").ra.value

    def get_dec(self):
        """
        Get Dec. corresponding to the current position (ICRS, J2000)

        :return: Declination
        """

        try:

            return self.dec.value

        except AttributeError:

            # Transform from L,B to R.A., Dec

            return self.sky_coord.transform_to("icrs").dec.value

    def get_l(self):
        """
        Get Galactic Longitude (l) corresponding to the current position

        :return: Galactic Longitude
        """

        try:

            return self.l.value

        except AttributeError:

            # Transform from L,B to R.A., Dec

            return self.sky_coord.transform_to("galactic").l.value

    def get_b(self):
        """
        Get Galactic latitude (b) corresponding to the current position

        :return: Latitude
        """

        try:

            return self.b.value

        except AttributeError:

            # Transform from L,B to R.A., Dec

            return self.sky_coord.transform_to("galactic").b.value

    def _get_sky_coord(self):

        if self._coord_type == "galactic":

            l = self.l.value
            b = self.b.value

            return coordinates.SkyCoord(
                l=l, b=b, frame="galactic", equinox=self._equinox, unit="deg"
            )

        else:

            ra = self.ra.value
            dec = self.dec.value

            return coordinates.SkyCoord(
                ra=ra, dec=dec, frame="icrs", equinox=self._equinox, unit="deg"
            )

    @property
    def sky_coord(self):
        """
        Return an instance of astropy.coordinates.SkyCoord which can be used to make all transformations supported
        by it

        :return: astropy.coordinates.SkyCoord
        """
        return self._get_sky_coord()

    @property
    def parameters(self):
        """
        Get the dictionary of parameters (either ra,dec or l,b)

        :return: dictionary of parameters
        """

        if self._coord_type == "galactic":

            return collections.OrderedDict((("l", self.l), ("b", self.b)))

        else:

            return collections.OrderedDict((("ra", self.ra), ("dec", self.dec)))

    @property
    def equinox(self):
        """
        Returns the equinox for the coordinates.

        :return:
        """
        return self._equinox

    def to_dict(self, minimal=False):

        data = collections.OrderedDict()

        if self._coord_type == "equatorial":

            data["ra"] = self.ra.to_dict(minimal)
            data["dec"] = self.dec.to_dict(minimal)
            data["equinox"] = self._equinox

        else:

            data["l"] = self.l.to_dict(minimal)
            data["b"] = self.b.to_dict(minimal)
            data["equinox"] = self._equinox

        return data

    def fix(self):
        """
        Fix the parameters with the coordinates (either ra,dec or l,b depending on how the class
        has been instanced)

        """

        if self._coord_type == "equatorial":

            self.ra.fix = True
            self.dec.fix = True

        else:

            self.l.fix = True
            self.b.fix = True

    def free(self):
        """
        Free the parameters with the coordinates (either ra,dec or l,b depending on how the class
        has been instanced)

        """

        if self._coord_type == "equatorial":

            self.ra.fix = False
            self.dec.fix = False

        else:

            self.l.fix = False
            self.b.fix = False

    @classmethod
    def from_dict(cls, data):

        return cls(**data)

    def _repr__base(self, rich_output):

        if self._coord_type == "equatorial":

            representation = "Sky direction (R.A., Dec.) = (%.5f, %.5f) (%s)" % (
                self.ra.value,
                self.dec.value,
                self.equinox,
            )

        else:

            representation = "Sky direction (l, b) = (%.5f, %.5f) (%s)" % (
                self.l.value,
                self.b.value,
                self.equinox,
            )

        return representation
