import inspect
import typing
from functools import wraps
from astropy import cosmology as cosmo

import autofit as af
from autolens.model import cosmology_util
from autolens import exc


def convert_units_to_input_units(func):
    """

    Parameters
    ----------
    func : (profiles, *args, **kwargs) -> Object
        A function that requries the units of its input parameters to be checked.

    Returns
    -------
        The original function
    """

    @wraps(func)
    def wrapper(profile, *args, **kwargs):
        """

        Parameters
        ----------
        profile : DimensionsProfile
            The profiles that owns the function

        Returns
        -------
            The DimensionsProfile
        """

        # Extract units of calculation, to convert the input variables and profile to use these units.

        unit_length = kwargs['unit_length'] if 'unit_length' in kwargs else profile.unit_length
        unit_luminosity = kwargs['unit_luminosity'] if 'unit_luminosity' in kwargs else profile.unit_luminosity
        unit_mass = kwargs['unit_mass'] if 'unit_mass' in kwargs else profile.unit_mass

        # Extract input values which are used for conversions

        cosmology = kwargs['cosmology'] if 'cosmology' in kwargs else cosmo.Planck15
        redshift_profile = kwargs['redshift_profile'] if 'redshift_profile' in kwargs else None
        redshift_source = kwargs['redshift_source'] if 'redshift_source' in kwargs else None
        exposure_time = kwargs['exposure_time'] if 'exposure_time' in kwargs else None
        redshift_of_cosmic_average_density = kwargs['redshift_of_cosmic_average_density'] if 'redshift_of_cosmic_average_density' in kwargs else 'profile'

        # Use cosmology and redshifts to compute conversion factors.

        if redshift_profile is not None and cosmology is not None:

            kpc_per_arcsec = cosmology_util.kpc_per_arcsec_from_redshift_and_cosmology(redshift=redshift_profile,
                                                                                       cosmology=cosmology)

        else:

            kpc_per_arcsec = None

        if redshift_profile is not None and redshift_source is not None and cosmology is not None and \
                unit_length is not None and unit_mass is not None:

            critical_surface_density = \
                cosmology_util.critical_surface_density_between_redshifts_from_redshifts_and_cosmology(
                    redshift_0=redshift_profile, redshift_1=redshift_source, cosmology=cosmology,
                    unit_length=unit_length, unit_mass=unit_mass)

        else:

            critical_surface_density = None

        if redshift_profile is not None and cosmology is not None and unit_length is not None and \
                unit_mass is not None and unit_mass is not 'angular':

            if redshift_of_cosmic_average_density is 'profile':
                redshift_calc = redshift_profile
            elif redshift_of_cosmic_average_density is 'local':
                redshift_calc = 0.0
            else:
                raise exc.UnitsException('The redshift of the cosmic average density haas been specified as an invalid '
                                         'string. Must be (local | profile)')

            cosmic_average_density = \
                cosmology_util.cosmic_average_density_from_redshift_and_cosmology(redshift=redshift_calc,
                    cosmology=cosmology, unit_length=unit_length, unit_mass=unit_mass)

        else:

            cosmic_average_density = None

        # Convert all input parameters to units of input units.

        for key, value in kwargs.items():

            if isinstance(value, Length):
                kwargs[key] = value.convert(unit_length=unit_length, kpc_per_arcsec=kpc_per_arcsec)

            if isinstance(value, Luminosity):
                kwargs[key] = value.convert(unit_luminosity=unit_luminosity, exposure_time=exposure_time)

            if isinstance(value, Mass):
                kwargs[key] = value.convert(unit_mass=unit_mass, critical_surface_density=critical_surface_density)

            if isinstance(value, MassOverLuminosity):
                kwargs[key] = value.convert(unit_luminosity=unit_luminosity, unit_mass=unit_mass,
                                            exposure_time=exposure_time,
                                            critical_surface_density=critical_surface_density)

            if isinstance(value, MassOverLength2):
                kwargs[key] = value.convert(unit_length=unit_length, unit_mass=unit_mass,
                                            kpc_per_arcsec=kpc_per_arcsec,
                                            critical_surface_density=critical_surface_density)

            if isinstance(value, MassOverLength3):

                kwargs[key] = value.convert(unit_length=unit_length, unit_mass=unit_mass,
                                            kpc_per_arcsec=kpc_per_arcsec,
                                            critical_surface_density=critical_surface_density)

        # Convert profile to input parameter units

        profile = profile.new_profile_with_units_converted(unit_length=unit_length, unit_luminosity=unit_luminosity,
                                                           unit_mass=unit_mass, kpc_per_arcsec=kpc_per_arcsec,
                                                           exposure_time=exposure_time,
                                                           critical_surface_density=critical_surface_density)

        if kpc_per_arcsec is not None:
            kwargs['kpc_per_arcsec'] = kpc_per_arcsec

        if critical_surface_density is not None:
            kwargs['critical_surface_density'] = critical_surface_density

        if cosmic_average_density is not None:
            kwargs['cosmic_average_density'] = cosmic_average_density

        return func(profile, *args, **kwargs)

    return wrapper


class DimensionsProfile(object):

    def __init__(self):

        pass

    def new_profile_with_units_converted(self, unit_length=None, unit_luminosity=None, unit_mass=None,
                                         kpc_per_arcsec=None, exposure_time=None, critical_surface_density=None):

        constructor_args = inspect.getfullargspec(self.__init__).args

        def convert(value):
            if unit_length is not None:
                if isinstance(value, Length):
                    return value.convert(unit_length, kpc_per_arcsec)
                if isinstance(value, tuple):
                    return tuple(convert(item) for item in value)
            if unit_luminosity is not None and isinstance(value, Luminosity):
                return value.convert(unit_luminosity, exposure_time)
            if unit_mass is not None and isinstance(value, Mass):
                return value.convert(unit_mass, critical_surface_density)
            if (unit_mass is not None or unit_luminosity is not None) and isinstance(value, MassOverLuminosity):
                return value.convert(unit_luminosity, unit_mass, exposure_time, critical_surface_density)
            return value

        return self.__class__(
            **{key: convert(value) for key, value in self.__dict__.items() if key in constructor_args})

    @property
    def unit_length(self):

        for attr, value in self.__dict__.items():
            if hasattr(value, 'unit_length'):
                return value.unit_length

        return None

    @property
    def unit_luminosity(self):

        for attr, value in self.__dict__.items():
            if hasattr(value, 'unit_luminosity'):
                return value.unit_luminosity

        return None

    @property
    def unit_mass(self):

        for attr, value in self.__dict__.items():
            if hasattr(value, 'unit_mass'):
                return value.unit_mass

        return None


class Length(af.DimensionType):

    def __init__(self, value, unit_length="arcsec"):
        super().__init__(value)
        self.unit_length = unit_length
        self.unit_length_power = 1.0

    @property
    def unit(self):
        return self.unit_length

    def convert(self, unit_length, kpc_per_arcsec=None):
        value = self

        value = convert_length(value=value, unit_current=self.unit_length, unit_new=unit_length,
                               power=self.unit_length_power, kpc_per_arcsec=kpc_per_arcsec)

        return Length(value=value, unit_length=unit_length)


class Luminosity(af.DimensionType):

    def __init__(self, value, unit_luminosity='eps'):
        super().__init__(value)
        self.unit_luminosity = unit_luminosity
        self.unit_luminosity_power = 1.0

    @property
    def unit(self):
        return self.unit_luminosity

    def convert(self, unit_luminosity, exposure_time=None):
        value = self

        value = convert_luminosity(value=value, unit_current=self.unit_luminosity, unit_new=unit_luminosity,
                                   power=self.unit_luminosity_power, exposure_time=exposure_time)

        return Luminosity(value=value, unit_luminosity=unit_luminosity)


class Mass(af.DimensionType):

    def __init__(self, value, unit_mass="angular"):
        super().__init__(value)
        self.unit_mass = unit_mass
        self.unit_mass_power = 1.0

    @property
    def unit(self):
        return self.unit_mass

    def convert(self, unit_mass, critical_surface_density=None):
        value = self

        value = convert_mass(value=value, unit_current=self.unit_mass, unit_new=unit_mass,
                             critical_surface_density=critical_surface_density)

        return Mass(value=value, unit_mass=unit_mass)


class MassOverLuminosity(af.DimensionType):

    def __init__(self, value, unit_luminosity="eps", unit_mass="angular"):

        super().__init__(value)
        self.unit_luminosity = unit_luminosity
        self.unit_luminosity_power = -1.0
        self.unit_mass = unit_mass
        self.unit_mass_power = 1.0

    @property
    def unit(self):
        return self.unit_mass + ' / ' + self.unit_luminosity

    def convert(self, unit_luminosity, unit_mass, exposure_time=None, critical_surface_density=None):

        value = self
        if unit_luminosity is not None:
            value = convert_luminosity(value=value, unit_current=self.unit_luminosity, unit_new=unit_luminosity,
                                       power=self.unit_luminosity_power, exposure_time=exposure_time)
        else:
            unit_luminosity = self.unit_luminosity

        if unit_mass is not None:
            value = convert_mass(value=value, unit_current=self.unit_mass, unit_new=unit_mass,
                                 critical_surface_density=critical_surface_density)
        else:
            unit_mass = self.unit_mass

        return MassOverLuminosity(value=value, unit_mass=unit_mass, unit_luminosity=unit_luminosity)


class MassOverLength2(af.DimensionType):

    def __init__(self, value, unit_length="arcsec", unit_mass="angular"):
        super().__init__(value)
        self.unit_length = unit_length
        self.unit_length_power = -2.0
        self.unit_mass = unit_mass
        self.unit_mass_power = 1.0

    @property
    def unit(self):
        return self.unit_mass + ' / ' + self.unit_length + '^2'

    def convert(self, unit_length, unit_mass, kpc_per_arcsec=None, critical_surface_density=None, ):

        value = self

        if unit_length is not None:
            value = convert_length(value=value, unit_current=self.unit_length, unit_new=unit_length,
                                   power=self.unit_length_power, kpc_per_arcsec=kpc_per_arcsec)
        else:
            unit_length = value.unit_length

        if unit_mass is not None:
            value = convert_mass(value=value, unit_current=self.unit_mass, unit_new=unit_mass,
                                 critical_surface_density=critical_surface_density)
        else:
            unit_mass = value.unit_mass

        return MassOverLength2(value=value, unit_mass=unit_mass, unit_length=unit_length)


class MassOverLength3(af.DimensionType):

    def __init__(self, value, unit_length="arcsec", unit_mass="angular"):
        super().__init__(value)
        self.unit_length = unit_length
        self.unit_length_power = -3.0
        self.unit_mass = unit_mass
        self.unit_mass_power = 1.0

    @property
    def unit(self):
        return self.unit_mass + ' / ' + self.unit_length + '^3'

    def convert(self, unit_length, unit_mass, kpc_per_arcsec=None, critical_surface_density=None):

        value = self

        if unit_length is not None:
            value = convert_length(value=value, unit_current=self.unit_length, unit_new=unit_length,
                                   power=self.unit_length_power, kpc_per_arcsec=kpc_per_arcsec)
        else:
            unit_length = value.unit_length

        if unit_mass is not None:
            value = convert_mass(value=value, unit_current=self.unit_mass, unit_new=unit_mass,
                                 critical_surface_density=critical_surface_density)
        else:
            unit_mass = value.unit_mass

        return MassOverLength3(value=value, unit_mass=unit_mass, unit_length=unit_length)


Position = typing.Tuple[Length, Length]


def convert_length(value, unit_current, unit_new, power, kpc_per_arcsec):

    if unit_current is not unit_new and kpc_per_arcsec is None:
        raise exc.UnitsException('The length for a value has been requested in new units without a '
                                 'kpc_per_arcsec conversion factor.')

    if unit_current is unit_new:
        return value
    elif unit_current is 'arcsec' and unit_new is 'kpc':
        return (kpc_per_arcsec ** power) * value
    elif unit_current is 'kpc' and unit_new is 'arcsec':
        return value / (kpc_per_arcsec ** power)
    else:
        raise exc.UnitsException('The unit specified for the length of a value was an invalid string, you '
                                 'must use (arcsec | kpc)')


def convert_luminosity(value, unit_current, unit_new, power, exposure_time):
    if unit_current is not unit_new and exposure_time is None:
        raise exc.UnitsException('The luminosity for a value has been requested in new units '
                                 'without an  exposure time conversion factor.')

    if unit_current is unit_new:
        return value
    elif unit_current is 'eps' and unit_new is 'counts':
        return (exposure_time ** power) * value
    elif unit_current is 'counts' and unit_new is 'eps':
        return value / (exposure_time ** power)
    else:
        raise exc.UnitsException('The unit specified for the luminosity of a value was an invalid string, you '
                                 'must use (electrons per second | counts)')


def convert_mass(value, unit_current, unit_new, critical_surface_density):
    if unit_current is not unit_new and critical_surface_density is None:
        raise exc.UnitsException('The mass for a value has been requested in new units '
                                 'without a critical surface mass density conversion factor.')

    if unit_current is unit_new:
        return value
    elif unit_current is 'angular' and unit_new is 'solMass':
        return critical_surface_density * value
    elif unit_current is 'solMass' and unit_new is 'angular':
        return value / critical_surface_density
    else:
        raise exc.UnitsException('The unit specified for the mass of a value was an invalid string, you '
                                 'must use (angular | solMass)')