from hr_selection_function.gaia import (
    M10Estimator,
    MSubsampleEstimator,
    GaiaDensityEstimator,
)
import numpy as np
import healpy as hp


pixels_for_testing = [0, 57, 5893, 58932, 3, 382, 382, 159532, 9694, 196607]


def test_M10Estimator():
    model = M10Estimator()

    # Some values looked up manually from the GaiaUnlimited package
    test_values = [
        21.0278,
        20.9648,
        21.0118,
        20.9611,
        20.9353,
        20.9351,
        20.9351,
        21.0907,
        20.9455,
        21.0904,
    ]

    # Convert into ra/dec
    ra, dec = hp.pix2ang(2**7, pixels_for_testing, nest=True, lonlat=True)

    result = model(ra, dec)
    np.testing.assert_allclose(result, test_values)


def test_MSubsampleEstimator():
    model = MSubsampleEstimator()

    # Some values looked up manually from the dataframe
    test_values = [
        20.473399,
        20.460272,
        20.783859,
        20.615404,
        20.396479,
        20.648832,
        20.648832,
        20.727514,
        20.797851,
        20.266360,
    ]

    # Convert into ra/dec
    l, b = hp.pix2ang(2**7, pixels_for_testing, nest=True, lonlat=True)  # noqa: E741

    result = model(l, b)
    np.testing.assert_allclose(result, test_values)


def test_GaiaDensityEstimator():
    l, b, pmra, pmdec, parallax, result_expected = _get_test_density_data()

    model = GaiaDensityEstimator()
    result_actual = model(l, b, pmra, pmdec, parallax)

    np.testing.assert_allclose(result_actual, result_expected)


def _get_test_density_data(l_and_b=True):
    l = [
        162.9045963311823,
        222.2312721732947,
        128.44871235039162,
        160.9499288226896,
        213.06037341723444,
        199.36275026706875,
        128.60961396338558,
        136.28666583570796,
        125.97131795417525,
        213.97631324679773,
    ]
    b = [
        -0.9346927099010592,
        3.6108482027563507,
        1.9453330185828903,
        0.2056845242850199,
        1.1661362043217294,
        19.410708776697685,
        0.12021936533026821,
        1.3539080788454123,
        -1.75718947932212,
        1.1693380678294931,
    ]
    ra = [
        73.32579352199228,
        110.35512999194107,
        25.645086126893748,
        72.82866993029752,
        103.95590766540087,
        114.8592052923773,
        25.208684973667218,
        41.51797920137684,
        19.128741698462704,
        104.3764145121306,
    ]
    dec = [
        42.41642894773442,
        -6.5928011091731085,
        64.26790490403921,
        44.650500836239445,
        0.4174126564070922,
        20.39438531322255,
        62.44618488047631,
        61.2175017009223,
        60.966482841194676,
        -0.39608875827834383,
    ]

    x, y = ra, dec
    if l_and_b:
        x, y = l, b

    return [
        x,
        y,
        [
            -0.7766432588848509,
            -0.915231855730751,
            -1.6372068440755732,
            -0.2305765209103827,
            -0.4524758968610532,
            -0.9638938991196692,
            -0.5086517349842387,
            -1.0101726801170998,
            0.0820812064803136,
            0.6355885023846625,
        ],
        [
            -0.14081969572092248,
            0.9609431030745171,
            0.12285028500838725,
            -0.024477925088355806,
            -0.0185474630769482,
            -0.8602342854349014,
            -0.3938208999571014,
            0.5334403415731419,
            -0.1922518589849786,
            -0.6764031476209058,
        ],
        [
            0.08060393781080663,
            0.12747343843901612,
            0.1004291966661565,
            0.08204558753982645,
            0.11759992793618043,
            0.15105175627504655,
            0.07888521783700321,
            0.09479025225955634,
            0.11303189760716631,
            0.06938962146299331,
        ],
        [
            32393.884212237157,
            161238.37642755947,
            223081.47449782598,
            51176.42713227133,
            111278.63400948483,
            3853.1979999303608,
            123062.78494792822,
            76527.43449485092,
            142593.45652828,
            49079.85544204984,
        ],
    ]