MAGIC_NUMBER = magic_num_reader() def gaia_default_dr(dr=None):
    """
    Check if dr argument is provided, if none then use default

    :param dr: Gaia DR
    :type dr: Union[NoneType, int]
    :return: Gaia DR
    :rtype: int
    :History: 2017-Oct-26 - Written - Henry Leung (University of Toronto)
    """
    if dr is None:
        dr = 2
        print(f'dr is not provided, using default dr={dr}')
    else:
        pass
    return dr def mag_to_absmag(mag, parallax, parallax_err=None):
    """
    To convert apparent magnitude to absolute magnitude, Magic Number will be preserved

    :param mag: apparent magnitude
    :type mag: Union[float, ndarray]
    :param parallax: parallax (mas) or with astropy (can be distance with units) so astroNN will convert to appropriate units
    :type parallax: Union[float, ndarray, astropy Quantity]
    :param parallax_err: parallax_error (mas) or with astropy so astroNN will convert to appropriate units
    :type parallax_err: Union[NoneType, float, ndarray, astropy Quantity]
    :return: absolute magnitude  (with additional return of propagated error if parallax_err is provided)
    :rtype: Union[float, ndarray]
    :History: 2017-Oct-14 - Written - Henry Leung (University of Toronto)
    """
    # Check unit if available
    if isinstance(parallax, u.Quantity):
        original_parallax_unit = parallax.unit
        parallax = parallax.to(default_parallax_unit, equivalencies=u.parallax())
        if parallax_err is not None:
            if not isinstance(parallax_err, u.Quantity):
                # assume parallax error carry the same original unit as parallax if no units detected
                parallax_err = (parallax_err * original_parallax_unit).to(default_parallax_unit).value
            if isinstance(parallax_err, u.Quantity):
                parallax_err = parallax_err.to(default_parallax_unit).value
    #     warnings.warn(f'Please be advised that astroNN mag_to_absmag() expects {default_parallax_unit.name}, '
    #           f'astroNN has corrected the unit according to astropy unit framework')
    # else:
    #     warnings.warn(f'Please be advised that astroNN mag_to_absmag expects parallax in {default_parallax_unit.name}')

    mag = np.array(mag)
    parallax_unitless = np.array(parallax)  # Take the value as we cant apply log10 to astropy unit

    magic_idx = ((parallax_unitless == MAGIC_NUMBER) | (mag == MAGIC_NUMBER) | (mag < -90.) | np.isnan(parallax_unitless) | np.isnan(mag))  # check for magic number

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        absmag = mag + 5. * (np.log10(parallax_unitless) - 2.)

    if parallax_unitless.shape != ():  # check if its only 1 element
        absmag[magic_idx] = MAGIC_NUMBER
    else:
        absmag = MAGIC_NUMBER if magic_idx == [1] else absmag
    if parallax_err is None:
        return absmag
    else:
        absmag_err = 5. * np.abs(parallax_err / (parallax_unitless * absmag))
        if parallax_unitless.shape != ():  # check if its only 1 element
            absmag_err[magic_idx] = MAGIC_NUMBER
        else:
            absmag_err = MAGIC_NUMBER if magic_idx == [1] else absmag_err
        return absmag, absmag_err def mag_to_fakemag(mag, parallax, parallax_err=None):
    """
    To convert apparent magnitude to astroNN fakemag, Magic Number will be preserved

    :param mag: apparent magnitude
    :type mag: Union[float, ndarray]
    :param parallax: parallax (mas) or with astropy(can be distance with units) so astroNN will convert to appropriate units
    :type parallax: Union[float, ndarray, astropy Quantity]
    :param parallax_err: parallax_error (mas) or with astropy so astroNN will convert to appropriate units
    :type parallax_err: Union[NoneType, float, ndarray, astropy Quantity]
    :return: astroNN fakemag, with addition (with additional return of propagated error if parallax_err is provided)
    :rtype: Union[float, ndarray]
    :History: 2017-Oct-14 - Written - Henry Leung (University of Toronto)
    """
    # Check unit if available
    if isinstance(parallax, u.Quantity):
        original_parallax_unit = parallax.unit
        parallax = parallax.to(default_parallax_unit, equivalencies=u.parallax())
        if parallax_err is not None:
            if not isinstance(parallax_err, u.Quantity):
                # assume parallax error carry the same original unit as parallax if no units detected
                parallax_err = (parallax_err * original_parallax_unit).to(default_parallax_unit).value
            if isinstance(parallax_err, u.Quantity):
                parallax_err = parallax_err.to(default_parallax_unit).value

    mag = np.array(mag)
    parallax_unitless = np.array(parallax)  # Take the value as we cant apply pow() to astropy unit

    magic_idx = ((parallax_unitless == MAGIC_NUMBER) | (mag == MAGIC_NUMBER) | (mag < -90.) | np.isnan(parallax_unitless) | np.isnan(mag))  # check for magic number

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        fakemag = parallax_unitless * (10. ** (0.2 * mag))
    if parallax_unitless.shape != ():  # check if its only 1 element
        fakemag[magic_idx] = MAGIC_NUMBER
    else:
        fakemag = MAGIC_NUMBER if magic_idx == [1] else fakemag

    if parallax_err is None:
        return fakemag
    else:
        fakemag_err = np.abs((parallax_err / parallax_unitless) * fakemag)
        if parallax_unitless.shape != ():  # check if its only 1 element
            fakemag_err[magic_idx] = MAGIC_NUMBER
        else:
            fakemag_err = MAGIC_NUMBER if magic_idx == [1] else fakemag_err
        return fakemag, fakemag_err def absmag_to_pc(absmag, mag):
    """
    To convert absolute magnitude to parsec, Magic Number will be preserved

    :param absmag: absolute magnitude
    :type absmag: Union[float, ndarray]
    :param mag: apparent magnitude
    :type mag: Union[float, ndarray]
    :return: parsec
    :rtype: astropy Quantity
    :History: 2017-Nov-16 - Written - Henry Leung (University of Toronto)
    """
    absmag = np.array(absmag)
    mag = np.array(mag)
    magic_idx = ((absmag == MAGIC_NUMBER) | (mag == MAGIC_NUMBER) | np.isnan(absmag) | np.isnan(mag))  # check for magic number

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        pc = (10. ** (((mag - absmag) / 5.) + 1.))

    if absmag.shape != ():  # check if its only 1 element
        pc[magic_idx] = MAGIC_NUMBER
        return pc * u.parsec
    else:
        return (MAGIC_NUMBER if magic_idx == [1] else pc) * u.parsec def fakemag_to_absmag(fakemag):
    """
    To convert fakemag to absmag, Magic Number will be preserved

    :param fakemag: eastroNN fakemag
    :type fakemag: Union[float, ndarray]
    :return: absolute magnitude
    :rtype: Union[float, ndarray]
    :History: 2018-Jan-31 - Written - Henry Leung (University of Toronto)
    """
    fakemag = np.array(fakemag)
    # treat non-negative fakemag as MAGIC_NUMBER
    magic_idx = ((fakemag == MAGIC_NUMBER) | (fakemag <= 0.) | np.isnan(fakemag))  # check for magic number and negative fakemag

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        absmag = 5. * (np.log10(fakemag) - 2.)

    if absmag.shape != ():  # check if its only 1 element
        absmag[magic_idx] = MAGIC_NUMBER
    else:  # for float
        absmag = MAGIC_NUMBER if magic_idx == [1] else absmag
    return absmag def absmag_to_fakemag(absmag):
    """
    To convert absmag to fakemag, Magic Number will be preserved

    :param absmag: absolute magnitude
    :type absmag: Union[float, ndarray]
    :return: astroNN fakemag
    :rtype: Union[float, ndarray]
    :History: 2018-Jan-31 - Written - Henry Leung (University of Toronto)
    """
    absmag = np.array(absmag)
    magic_idx = ((absmag == MAGIC_NUMBER) | np.isnan(absmag))  # check for magic number

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        fakemag = 10. ** (0.2 * absmag + 2.)
    if fakemag.shape != ():  # check if its only 1 element
        fakemag[magic_idx] = MAGIC_NUMBER
    else:  # for float
        fakemag = MAGIC_NUMBER if magic_idx == [1] else fakemag
    return fakemag def fakemag_to_pc(fakemag, mag, fakemag_err=None):
    """
    To convert fakemag to parsec, Magic Number will be preserved

    :param fakemag: astroNN fakemag
    :type fakemag: Union[float, ndarray]
    :param mag: apparent magnitude
    :type mag: Union[float, ndarray]
    :param fakemag_err: Optional, fakemag_err
    :type fakemag_err: Union[NoneType, float, ndarray]
    :return: array of pc with astropy Quantity (with additional return of propagated error if fakemag_err is provided)
    :rtype: astropy Quantity
    :History: 2018-Jan-31 - Written - Henry Leung (University of Toronto)
    """
    fakemag = np.array(fakemag)
    mag = np.array(mag)
    # treat non-positive fakemag as MAGIC_NUMBER, check for magic number and negative fakemag
    magic_idx = ((fakemag == MAGIC_NUMBER) | (mag == MAGIC_NUMBER) | (fakemag <= 0.) | np.isnan(fakemag) | np.isnan(mag))

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        pc = 1000. * (10. ** (0.2 * mag)) / fakemag
    if fakemag.shape != ():  # check if its only 1 element
        pc[magic_idx] = MAGIC_NUMBER
    else:  # for float
        pc = MAGIC_NUMBER if magic_idx == [1] else pc

    if fakemag_err is None:
        return pc * u.parsec
    else:
        with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
            warnings.simplefilter("ignore")
            pc_err = (fakemag_err / fakemag) * pc
        if fakemag.shape != ():  # check if its only 1 element
            pc_err[magic_idx] = MAGIC_NUMBER
        else:  # for float
            pc_err = MAGIC_NUMBER if magic_idx == [1] else pc_err
        return pc * u.parsec, pc_err * u.parsec def fakemag_to_parallax(fakemag, mag, fakemag_err=None):
    """
    To convert fakemag to parallax, Magic Number will be preserved

    :param fakemag: astroNN fakemag
    :type fakemag: Union[float, ndarray]
    :param mag: apparent magnitude
    :type mag: Union[float, ndarray]
    :param fakemag_err: Optional, fakemag_err
    :type fakemag_err: Union[NoneType, float, ndarray]
    :return: array of parallax in mas with astropy Quantity (with additional return of propagated error if fakemag_err is provided)
    :rtype: astropy Quantity
    :History: 2018-Aug-11 - Written - Henry Leung (University of Toronto)
    """
    fakemag = np.array(fakemag)
    mag = np.array(mag)
    # treat non-positive fakemag as MAGIC_NUMBER, check for magic number and negative fakemag
    magic_idx = ((fakemag == MAGIC_NUMBER) | (mag == MAGIC_NUMBER) | (fakemag <= 0.) | np.isnan(fakemag) | np.isnan(mag))

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        parallax = fakemag / (10. ** (0.2 * mag))
    if fakemag.shape != ():  # check if its only 1 element
        parallax[magic_idx] = MAGIC_NUMBER
    else:  # for float
        parallax = MAGIC_NUMBER if magic_idx == [1] else parallax

    if fakemag_err is None:
        return parallax * u.mas
    else:
        with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
            warnings.simplefilter("ignore")
            parallax_err = (fakemag_err / fakemag) * parallax
        if fakemag.shape != ():  # check if its only 1 element
            parallax_err[magic_idx] = MAGIC_NUMBER
        else:  # for float
            parallax_err = MAGIC_NUMBER if magic_idx == [1] else parallax_err
        return parallax * u.mas, parallax_err * u.mas def fakemag_to_mag(fakemag, pc, pc_err=None):
    """
    To convert apparent magnitude to astroNN fakemag, Magic Number will be preserved

    :param fakemag: fakemag
    :type fakemag: Union[float, ndarray]
    :param pc: parsec or with astropy (can be parallax with units) so astroNN will convert to appropriate units
    :type pc: Union[float, ndarray, astropy Quantity]
    :param pc_error: parsec uncertainty or with astropy so astroNN will convert to appropriate units
    :type pc_error: Union[NoneType, float, ndarray, astropy Quantity]
    :return: astroNN fakemag, with addition (with additional return of propagated error if parallax_err is provided)
    :rtype: Union[float, ndarray]
    :History: 2018-Aug-1 - Written - Henry Leung (University of Toronto)
    """
    # Check unit if available
    if isinstance(pc, u.Quantity):
        original_parallax_unit = pc.unit
        pc = pc.to(u.parsec, equivalencies=u.parallax())
        if pc_err is not None:
            if not isinstance(pc_err, u.Quantity):
                # assume parallax error carry the same original unit as parallax if no units detected
                pc_err = (pc_err * original_parallax_unit).to(u.parsec).value
            if isinstance(pc_err, u.Quantity):
                pc_err = pc_err.to(u.parsec).value

    fakemag = np.array(fakemag)
    pc_unitless = np.array(pc)  # Take the value as we cant apply pow() to astropy unit

    magic_idx = ((pc_unitless == MAGIC_NUMBER) | (fakemag == MAGIC_NUMBER) | (fakemag <= 0.) | np.isnan(pc_unitless) | np.isnan(fakemag))  # check for magic number

    with warnings.catch_warnings():  # suppress numpy Runtime warning caused by MAGIC_NUMBER
        warnings.simplefilter("ignore")
        mag = np.log10((pc_unitless / 1000) * fakemag) / 0.2
    if pc_unitless.shape != ():  # check if its only 1 element
        mag[magic_idx] = MAGIC_NUMBER
    else:
        fakemag = MAGIC_NUMBER if magic_idx == [1] else fakemag

    if pc_err is None:
        return mag
    else:
        mag_err = np.abs((pc_err / pc) * fakemag)
        if pc_unitless.shape != ():  # check if its only 1 element
            mag_err[magic_idx] = MAGIC_NUMBER
        else:
            mag_err = MAGIC_NUMBER if magic_idx == [1] else mag_err
        return mag, mag_err
import unittest
import astropy.units as u
import numpy as np
import numpy.testing as npt
from astroNN.config import MAGIC_NUMBER
from astroNN.gaia import absmag_to_pc, mag_to_absmag, fakemag_to_absmag, absmag_to_fakemag, fakemag_to_pc, \
    mag_to_fakemag, gaia_default_dr, fakemag_to_parallax, fakemag_to_mag
        from astroNN.gaia import fakemag_to_logsol, absmag_to_logsol, logsol_to_absmag, logsol_to_fakemag
        from astroNN.gaia import extinction_correction
        from astroNN.gaia import anderson_2017_parallax
        from astroNN.gaia import gaiadr2_parallax



# noinspection PyUnresolvedReferences
class GaiaToolsCase(unittest.TestCase):
    def test_astrometry_conversion(self):
        # Example data of [Vega, Sirius]
        absmag = np.array([0.582, 1.42])
        mag = np.array([0.026, -1.46])
        parallax = np.array([130.23, 379.21])
        parallax_err = np.array([0.36, 1.58])

        # Absmag related test
        pc = absmag_to_pc(absmag, mag)
        npt.assert_almost_equal(pc.value, 1000 / parallax, decimal=1)

        npt.assert_almost_equal(mag_to_absmag(mag, parallax * u.mas), absmag, decimal=1)