"""test_transitfit.py - Luke Bouma (luke@astro.princeton.edu) - Aug 2019
License: MIT - see the LICENSE file for details.

NOTE: the tests in this module take a very long time. To enable and actually run
them, you must set the environmental variable RUN_LONG_TESTS=1 before running
pytest, like so: ``RUN_LONG_TESTS=1 pytest test_transitfit.py``, or set this
environmental variable in your CI runner.

When run, this file downloads two TESS light curves from the
astrobase-notebooks/nb-data directory.

Implemented tests include:

- fit an easy single-sector light curve (6 big HJ transits) using
fivetransitparam_fit_magseries, derive uncertainties, check them against
theoretical predictions.

- fit a trickier single-sector light curve (2 transits), do as above.

- test_multithread_speed: ensure that increasing nworkers speeds up the MCMC
sampling. 2019/08/15: this test fails.

"""

###########
# imports #
###########

from pytest import mark

import os
import multiprocessing
import time as _time
try:
    from urllib import urlretrieve
except Exception:
    from urllib.request import urlretrieve

import numpy as np
from numpy.testing import assert_allclose

from astropy.io import fits
import astrobase.imageutils as iu

try:
    import transitleastsquares
    from astrobase.lcfit.transits import fivetransitparam_fit_magseries
    test_ok = True
except Exception:
    test_ok = False


##########
# config #
##########

# this function is used to check progress of the download
def on_download_chunk(transferred,blocksize,totalsize):
    progress = transferred*blocksize/float(totalsize)*100.0
    print('downloading test LC: {progress:.1f}%'.format(progress=progress),
          end='\r')


# download the light curves used for tests if they do not exist. first is a
# nice easy hot jupiter, 6 transits.  second is a tricky two-transit warm
# jupiter. (TOI-450).
LCURLS = [
    ("https://github.com/waqasbhatti/astrobase-notebooks/raw/master/nb-data/"
     "hlsp_cdips_tess_ffi_gaiatwo0003007171311355035136-0006_tess_v01_llc.fits"),
    ("https://github.com/waqasbhatti/astrobase-notebooks/raw/master/nb-data/"
     "hlsp_cdips_tess_ffi_gaiatwo0004827527233363019776-0006_tess_v01_llc.fits")
]

modpath = os.path.abspath(__file__)
LCPATHS = [
    os.path.abspath(os.path.join(
        os.getcwd(),
        'hlsp_cdips_tess_ffi_gaiatwo0003007171311355035136-0006_tess_v01_llc.fits')
    ),
    os.path.abspath(os.path.join(
        os.getcwd(),
        'hlsp_cdips_tess_ffi_gaiatwo0004827527233363019776-0006_tess_v01_llc.fits')
    )
]

for LCPATH, LCURL in zip(LCPATHS, LCURLS):
    if not os.path.exists(LCPATH):
        localf, headerr = urlretrieve(
            LCURL,LCPATH,reporthook=on_download_chunk)


if not os.environ.get('RUN_LONG_TESTS'):
    test_ok = False


#########
# tests #
#########

if test_ok:

    def test_fivetransitparam_fit_magseries_easy():
        """
        Fit one TESS sector of data, with a HJ (candidate) in it, for a transit
        model (t0, period, incl, sma, rp/star) with believable error bars.
        """

        # path and identifier for GaiaDR2 3007171311355035136
        lcpath = LCPATHS[0]
        identifier = str(lcpath.split('gaiatwo')[1].split('-')[0].lstrip('0'))

        lc = iu.get_data_keyword_list(lcpath, ['TMID_BJD', 'TFA2'])
        hdr = iu.get_header_keyword_list(lcpath,
                                         ['TICTEFF', 'TICRAD', 'TICLOGG'])

        time = lc['TMID_BJD']
        mag = lc['TFA2']
        mag_0, f_0 = 12, 1e4