import os
from glue.core import DataCollection
from glue.core.roi import XRangeROI
import numpy as np
from astropy.time import Time
from lightkurve import search_lightcurve, LightCurve


SMALL_LC_PATH = os.path.join(
    os.path.dirname(__file__), os.pardir, 'data',
    'HAT-P-11_Kepler_Q0_snippet.fits'
)


def test_translator(light_curve_like_kepler_quarter):
    light_curve = light_curve_like_kepler_quarter
    dc = DataCollection()
    dc['dummy-lc'] = light_curve
    translated_lc = dc['dummy-lc'].get_object()

    # check for equality within these attrs:
    attrs = ['time', 'flux', 'flux_err']

    for attr in attrs:
        expected = getattr(light_curve, attr)
        translated = getattr(translated_lc, attr)

        if attr == 'time':
            args = (
                expected.tdb.jd,
                translated.tdb.jd
            )
        else:
            args = (
                expected.value,
                translated.to_value(expected.unit)
            )

        # check that attribute's values match expectations:
        np.testing.assert_allclose(*args)

        if hasattr(expected, 'mask'):
            # check that mask is preserved
            assert np.all(expected.mask == translated.mask)


def _download_and_write_lc_snippet():
    """
    Download and write a local copy of the first 10 long cadences of the
    Kepler light curve of HAT-P-11 in Quarter 0.

    Run it with:
    python -c "from lcviz.tests.test_translator import _download_and_write_lc_snippet as d; d()"
    """
    lc = search_lightcurve(
        "HAT-P-11", mission="Kepler", cadence="long", quarter=0
    ).download()
    lc[0:10].write(SMALL_LC_PATH, overwrite=True)


def test_round_trip(helper):
    """
    The FITS file used in this test is generated by the
    function ``_download_and_write_lc_snippet``.
    """
    lc = LightCurve.read(SMALL_LC_PATH, astropy_native=True)
    helper.load_data(lc)

    # these times chosen by hand during Kepler Q0:
    near_transit = Time([
        '2009-05-02 02:51',
        '2009-05-02 03:52'
    ], format='iso')

    viewer = helper.default_time_viewer._obj
    viewer.apply_roi(XRangeROI(*near_transit))

    columns_to_check = ['time', 'flux', 'flux_err']
    subset_translated = helper.get_data(subset='Subset 1')[columns_to_check]
    row_mask = (lc.time > near_transit[0]) & (lc.time < near_transit[1])
    lc_subset = lc[row_mask][columns_to_check]

    for column in lc_subset.values_equal(subset_translated).itercols():
        assert np.all(column)
