import decimal
import warnings
import functools
import contextlib
from decimal import Decimal
from datetime import datetime, timedelta

import pytest
from hypothesis import assume, example, given, target
from hypothesis.extra.numpy import array_shapes, arrays
from hypothesis.strategies import (composite, datetimes, floats, integers,
                                   one_of, sampled_from, timedeltas, tuples)

import numpy as np
import erfa
from erfa import ErfaError, ErfaWarning

import astropy.units as u
from astropy.tests.helper import assert_quantity_allclose
from astropy.time import STANDARD_TIME_SCALES, Time, TimeDelta
from astropy.time.utils import day_frac, two_sum
from astropy.utils import iers


allclose_jd = functools.partial(np.allclose, rtol=np.finfo(float).eps, atol=0)
allclose_jd2 = functools.partial(np.allclose, rtol=np.finfo(float).eps,
                                 atol=np.finfo(float).eps)  # 20 ps atol
allclose_sec = functools.partial(np.allclose, rtol=np.finfo(float).eps,
                                 atol=np.finfo(float).eps * 24 * 3600)

tiny = np.finfo(float).eps
dt_tiny = TimeDelta(tiny, format='jd')


def setup_module():
    # Pre-load leap seconds table to avoid flakiness in hypothesis runs.
    # See https://github.com/astropy/astropy/issues/11030
    Time('2020-01-01').ut1


@pytest.fixture(scope='module')
def iers_b():
    """This is an expensive operation, so we share it between tests using a
    module-scoped fixture instead of using the context manager form.  This
    is particularly important for Hypothesis, which invokes the decorated
    test function many times (100 by default; see conftest.py for details).
    """
    with iers.earth_orientation_table.set(iers.IERS_B.open(iers.IERS_B_FILE)):
        yield "<using IERS-B orientation table>"


@contextlib.contextmanager
def quiet_erfa():
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=ErfaWarning)
        yield


def assert_almost_equal(a, b, *, rtol=None, atol=None, label=''):
    """Assert numbers are almost equal.

    This version also lets hypothesis know how far apart the inputs are, so
    that it can work towards a failure and present the worst failure ever seen
    as well as the simplest, which often just barely exceeds the threshold.
    """
    __tracebackhide__ = True
    if rtol is None or rtol == 0:
        thresh = atol
    elif atol is None:
        thresh = rtol * (abs(a) + abs(b)) / 2
    else:
        thresh = atol + rtol * (abs(a) + abs(b)) / 2

    amb = (a - b)
    if isinstance(amb, TimeDelta):
        ambv = amb.to_value(u.s)
        target(ambv, label=label + " (a-b).to_value(u.s), from TimeDelta")
        target(-ambv, label=label + " (b-a).to_value(u.s), from TimeDelta")
        if isinstance(thresh, u.Quantity):
            amb = amb.to(thresh.unit)
    else:
        try:
            target_value = float(amb)
        except TypeError:
            pass
        else:
            target(target_value, label=label + " float(a-b)")
            target(-target_value, label=label + " float(b-a)")

    assert abs(amb) < thresh


# Days that end with leap seconds
# Some time scales use a so-called "leap smear" to cope with these, others
# have times they can't represent or can represent two different ways.
# In any case these days are liable to cause trouble in time conversions.
# Note that from_erfa includes some weird non-integer steps before 1970.
leap_second_table = iers.LeapSeconds.from_iers_leap_seconds()
# Days that contain leap_seconds
leap_second_days = leap_second_table["mjd"] - 1
leap_second_deltas = list(zip(leap_second_days[1:],
                              np.diff(leap_second_table["tai_utc"])))

today = Time.now()
mjd0 = Time(0, format="mjd")


def reasonable_ordinary_jd():
    return tuples(floats(2440000, 2470000), floats(-0.5, 0.5))


@composite
def leap_second_tricky(draw):
    mjd = draw(one_of(sampled_from(leap_second_days),
                      sampled_from(leap_second_days + 1),
                      sampled_from(leap_second_days - 1)))
    return mjd + mjd0.jd1 + mjd0.jd2, draw(floats(0, 1))


def reasonable_jd():
    """Pick a reasonable JD.

    These should be not too far in the past or future (so that date conversion
    routines don't have to deal with anything too exotic), but they should
    include leap second days as a special case, and they should include several
    particularly simple cases (today, the beginning of the MJD scale, a
    reasonable date) so that hypothesis' example simplification produces
    obviously simple examples when they trigger problems.
    """
    moments = [(2455000., 0.), (mjd0.jd1, mjd0.jd2), (today.jd1, today.jd2)]
    return one_of(sampled_from(moments),
                  reasonable_ordinary_jd(),
                  leap_second_tricky())


def unreasonable_ordinary_jd():
    """JD pair that might be unordered or far away"""
    return tuples(floats(-1e7, 1e7), floats(-1e7, 1e7))


def ordered_jd():
    """JD pair that is ordered but not necessarily near now"""
    return tuples(floats(-1e7, 1e7), floats(-0.5, 0.5))


def unreasonable_jd():
    return one_of(reasonable_jd(), ordered_jd(), unreasonable_ordinary_jd())


@composite
def jd_arrays(draw, jd_values):
    s = draw(array_shapes())
    d = np.dtype([("jd1", float), ("jd2", float)])
    jdv = jd_values.map(lambda x: np.array(x, dtype=d))
    a = draw(arrays(d, s, elements=jdv))
    return a["jd1"], a["jd2"]


def unreasonable_delta():
    return tuples(floats(-1e7, 1e7), floats(-1e7, 1e7))


def reasonable_delta():
    return tuples(floats(-1e4, 1e4), floats(-0.5, 0.5))


# redundant?
def test_abs_jd2_always_less_than_half():
    """Make jd2 approach +/-0.5, and check that it doesn't go over."""
    t1 = Time(2400000.5, [-tiny, +tiny], format='jd')
    assert np.all(t1.jd1 % 1 == 0)
    assert np.all(abs(t1.jd2) < 0.5)
    t2 = Time(2400000., [[0.5 - tiny, 0.5 + tiny],
                         [-0.5 - tiny, -0.5 + tiny]], format='jd')
    assert np.all(t2.jd1 % 1 == 0)
    assert np.all(abs(t2.jd2) < 0.5)


@given(jd_arrays(unreasonable_jd()))
def test_abs_jd2_always_less_than_half_on_construction(jds):
    jd1, jd2 = jds
    t = Time(jd1, jd2, format="jd")
    target(np.amax(np.abs(t.jd2)))
    assert np.all(t.jd1 % 1 == 0)
    assert np.all(abs(t.jd2) <= 0.5)
    assert np.all((abs(t.jd2) < 0.5) | (t.jd1 % 2 == 0))


@given(integers(-10**8, 10**8), sampled_from([-0.5, 0.5]))
def test_round_to_even(jd1, jd2):
    t = Time(jd1, jd2, format="jd")
    assert (abs(t.jd2) == 0.5) and (t.jd1 % 2 == 0)


def test_addition():
    """Check that an addition at the limit of precision (2^-52) is seen"""
    t = Time(2455555., 0.5, format='jd', scale='utc')

    t_dt = t + dt_tiny
    assert t_dt.jd1 == t.jd1 and t_dt.jd2 != t.jd2

    # Check that the addition is exactly reversed by the corresponding
    # subtraction
    t2 = t_dt - dt_tiny
    assert t2.jd1 == t.jd1 and t2.jd2 == t.jd2


def test_mult_div():
    """Test precision with multiply and divide"""
    dt_small = 6 * dt_tiny
    # pick a number that will leave remainder if divided by 6.
    dt_big = TimeDelta(20000., format='jd')
    dt_big_small_by_6 = (dt_big + dt_small) / 6.
    dt_frac = dt_big_small_by_6 - TimeDelta(3333., format='jd')
    assert allclose_jd2(dt_frac.jd2, 0.33333333333333354)


def test_init_variations():
    """Check that 3 ways of specifying a time + small offset are equivalent"""
    dt_tiny_sec = dt_tiny.jd2 * 86400.
    t1 = Time(1e11, format='cxcsec') + dt_tiny
    t2 = Time(1e11, dt_tiny_sec, format='cxcsec')
    t3 = Time(dt_tiny_sec, 1e11, format='cxcsec')
    assert t1.jd1 == t2.jd1
    assert t1.jd2 == t3.jd2
    assert t1.jd1 == t2.jd1
    assert t1.jd2 == t3.jd2


def test_precision_exceeds_64bit():
    """
    Check that Time object really holds more precision than float64 by looking
    at the (naively) summed 64-bit result and asserting equality at the
    bit level.
    """
    t1 = Time(1.23456789e11, format='cxcsec')
    t2 = t1 + dt_tiny
    assert t1.jd == t2.jd


def test_through_scale_change():
    """Check that precision holds through scale change (cxcsec is TT)"""
    t0 = Time(1.0, format='cxcsec')
    t1 = Time(1.23456789e11, format='cxcsec')
    dt_tt = t1 - t0
    dt_tai = t1.tai - t0.tai
    assert allclose_jd(dt_tt.jd1, dt_tai.jd1)
    assert allclose_jd2(dt_tt.jd2, dt_tai.jd2)


def test_iso_init():
    """Check when initializing from ISO date"""
    t1 = Time('2000:001:00:00:00.00000001', scale='tai')
    t2 = Time('3000:001:13:00:00.00000002', scale='tai')
    dt = t2 - t1
    assert allclose_jd2(dt.jd2, 13. / 24. + 1e-8 / 86400. - 1.0)


def test_jd1_is_mult_of_one():
    """
    Check that jd1 is a multiple of 1.
    """
    t1 = Time('2000:001:00:00:00.00000001', scale='tai')
    assert np.round(t1.jd1) == t1.jd1
    t1 = Time(1.23456789, 12345678.90123456, format='jd', scale='tai')
    assert np.round(t1.jd1) == t1.jd1


def test_precision_neg():
    """
    Check precision when jd1 is negative.  This used to fail because ERFA
    routines use a test like jd1 > jd2 to decide which component to update.
    It was updated to abs(jd1) > abs(jd2) in erfa 1.6 (sofa 20190722).
    """
    t1 = Time(-100000.123456, format='jd', scale='tt')
    assert np.round(t1.jd1) == t1.jd1
    t1_tai = t1.tai
    assert np.round(t1_tai.jd1) == t1_tai.jd1


def test_precision_epoch():
    """
    Check that input via epoch also has full precision, i.e., against
    regression on https://github.com/astropy/astropy/pull/366
    """
    t_utc = Time(range(1980, 2001), format='jyear', scale='utc')
    t_tai = Time(range(1980, 2001), format='jyear', scale='tai')
    dt = t_utc - t_tai
    assert allclose_sec(dt.sec, np.round(dt.sec))


def test_leap_seconds_rounded_correctly():
    """Regression tests against #2083, where a leap second was rounded
    incorrectly by the underlying ERFA routine."""
    with iers.conf.set_temp('auto_download', False):
        t = Time(['2012-06-30 23:59:59.413',
                  '2012-07-01 00:00:00.413'], scale='ut1', precision=3).utc
        assert np.all(t.iso == np.array(['2012-06-30 23:59:60.000',
                                         '2012-07-01 00:00:00.000']))
    # with the bug, both yielded '2012-06-30 23:59:60.000'


@given(integers(-2**52+2, 2**52-2), floats(-1, 1))
@example(i=65536, f=3.637978807091714e-12)
def test_two_sum(i, f):
    with decimal.localcontext(decimal.Context(prec=40)):
        a = Decimal(i) + Decimal(f)
        s, r = two_sum(i, f)
        b = Decimal(s) + Decimal(r)
        assert_almost_equal(a, b, atol=Decimal(tiny), rtol=Decimal(0))


@given(floats(), floats())
def test_two_sum_symmetric(f1, f2):
    np.testing.assert_equal(two_sum(f1, f2), two_sum(f2, f1))


@given(floats(allow_nan=False, allow_infinity=False),
       floats(allow_nan=False, allow_infinity=False))
@example(f1=8.988465674311579e+307, f2=8.98846567431158e+307)
@example(f1=8.988465674311579e+307, f2=-8.98846567431158e+307)
@example(f1=-8.988465674311579e+307, f2=-8.98846567431158e+307)
def test_two_sum_size(f1, f2):
    r1, r2 = two_sum(f1, f2)
    assert (abs(r1) > abs(r2) / np.finfo(float).eps
            or r1 == r2 == 0
            or not np.isfinite(f1 + f2))


@given(integers(-2**52+2, 2**52-2), floats(-1, 1))
@example(i=65536, f=3.637978807091714e-12)
def test_day_frac_harmless(i, f):
    with decimal.localcontext(decimal.Context(prec=40)):
        a = Decimal(i) + Decimal(f)
        i_d, f_d = day_frac(i, f)
        a_d = Decimal(i_d) + Decimal(f_d)
        assert_almost_equal(a, a_d, atol=Decimal(tiny), rtol=Decimal(0))


@given(integers(-2**52+2, 2**52-2), floats(-0.5, 0.5))
@example(i=65536, f=3.637978807091714e-12)
def test_day_frac_exact(i, f):
    assume(abs(f) < 0.5 or i % 2 == 0)
    i_d, f_d = day_frac(i, f)
    assert i == i_d
    assert f == f_d


@given(integers(-2**52+2, 2**52-2), floats(-1, 1))
@example(i=65536, f=3.637978807091714e-12)
def test_day_frac_idempotent(i, f):
    i_d, f_d = day_frac(i, f)
    assert (i_d, f_d) == day_frac(i_d, f_d)


@given(integers(-2**52+2, 2**52-2), floats(-1, 1))
@example(i=65536, f=3.637978807091714e-12)
def test_mjd_initialization_precise(i, f):
    t = Time(val=i, val2=f, format="mjd", scale="tai")
    jd1, jd2 = day_frac(i + erfa.DJM0, f)
    jd1_t, jd2_t = day_frac(t.jd1, t.jd2)
    assert (abs((jd1 - jd1_t) + (jd2 - jd2_t)) * u.day).to(u.ns) < 1 * u.ns


@given(jd_arrays(unreasonable_jd()))
def test_day_frac_always_less_than_half(jds):
    jd1, jd2 = jds
    t_jd1, t_jd2 = day_frac(jd1, jd2)
    assert np.all(t_jd1 % 1 == 0)
    assert np.all(abs(t_jd2) <= 0.5)
    assert np.all((abs(t_jd2) < 0.5) | (t_jd1 % 2 == 0))


@given(integers(-10**8, 10**8), sampled_from([-0.5, 0.5]))
def test_day_frac_round_to_even(jd1, jd2):
    t_jd1, t_jd2 = day_frac(jd1, jd2)
    assert (abs(t_jd2) == 0.5) and (t_jd1 % 2 == 0)


@given(scale=sampled_from(STANDARD_TIME_SCALES), jds=unreasonable_jd())
@example(scale="tai", jds=(0.0, 0.0))
@example(scale="tai", jds=(0.0, -31738.500000000346))
def test_resolution_never_decreases(scale, jds):
    jd1, jd2 = jds
    assume(not scale == 'utc' or 2440000 < jd1 + jd2 < 2460000)
    t = Time(jd1, jd2, format="jd", scale=scale)
    with quiet_erfa():
        assert t != t + dt_tiny


@given(reasonable_jd())
def test_resolution_never_decreases_utc(jds):
    """UTC is very unhappy with unreasonable times"""
    jd1, jd2 = jds
    t = Time(jd1, jd2, format="jd", scale="utc")
    with quiet_erfa():
        assert t != t + dt_tiny


@given(scale1=sampled_from(STANDARD_TIME_SCALES),
       scale2=sampled_from(STANDARD_TIME_SCALES),
       jds=unreasonable_jd())
@example(scale1='tcg', scale2='ut1', jds=(2445149.5, 0.47187700984387526))
@example(scale1='tai', scale2='tcb', jds=(2441316.5, 0.0))
@example(scale1='tai', scale2='tcb', jds=(0.0, 0.0))
def test_conversion_preserves_jd1_jd2_invariant(iers_b, scale1, scale2, jds):
    jd1, jd2 = jds
    t = Time(jd1, jd2, scale=scale1, format="jd")
    try:
        with quiet_erfa():
            t2 = getattr(t, scale2)
    except iers.IERSRangeError:  # UT1 conversion needs IERS data
        assume(False)
    except ErfaError:
        assume(False)
    assert t2.jd1 % 1 == 0
    assert abs(t2.jd2) <= 0.5
    assert abs(t2.jd2) < 0.5 or t2.jd1 % 2 == 0


@given(scale1=sampled_from(STANDARD_TIME_SCALES),
       scale2=sampled_from(STANDARD_TIME_SCALES),
       jds=unreasonable_jd())
@example(scale1='tai', scale2='utc', jds=(0.0, 0.0))
def test_conversion_never_loses_precision(iers_b, scale1, scale2, jds):
    """Check that time ordering remains if we convert to another scale.