# Licensed under a 3-clause BSD style license - see LICENSE.rst

import functools
import numpy as np

from ...utils.compat import NUMPY_LT_1_14
from ...tests.helper import pytest
from .. import Time
from ...table import Table

try:
    import h5py  # pylint: disable=W0611
except ImportError:
    HAS_H5PY = False
else:
    HAS_H5PY = True

try:
    import yaml  # pylint: disable=W0611
    HAS_YAML = True
except ImportError:
    HAS_YAML = False

allclose_sec = functools.partial(np.allclose, rtol=2. ** -52,
                                 atol=2. ** -52 * 24 * 3600)  # 20 ps atol
is_masked = np.ma.is_masked


def test_simple():
    t = Time([1, 2, 3], format='cxcsec')
    assert t.masked is False
    assert np.all(t.mask == [False, False, False])

    # Before masking, format output is not a masked array (it is an ndarray
    # like always)
    assert not isinstance(t.value, np.ma.MaskedArray)
    assert not isinstance(t.unix, np.ma.MaskedArray)

    t[2] = np.ma.masked
    assert t.masked is True
    assert np.all(t.mask == [False, False, True])
    assert allclose_sec(t.value[:2], [1, 2])
    assert is_masked(t.value[2])
    assert is_masked(t[2].value)

    # After masking format output is a masked array
    assert isinstance(t.value, np.ma.MaskedArray)
    assert isinstance(t.unix, np.ma.MaskedArray)
    # Todo : test all formats


def test_scalar_init():
    t = Time('2000:001')
    assert t.masked is False
    assert t.mask == np.array(False)


def test_mask_not_writeable():
    t = Time('2000:001')
    with pytest.raises(AttributeError) as err:
        t.mask = True
    assert "can't set attribute" in str(err)

    t = Time(['2000:001'])
    with pytest.raises(ValueError) as err:
        t.mask[0] = True
    assert "assignment destination is read-only" in str(err)


def test_str():
    t = Time(['2000:001', '2000:002'])
    t[1] = np.ma.masked
    assert str(t) == "['2000:001:00:00:00.000' --]"
    assert repr(t) == "<Time object: scale='utc' format='yday' value=['2000:001:00:00:00.000' --]>"

    if NUMPY_LT_1_14:
        expected = ["masked_array(data = ['2000-01-01 00:00:00.000' --],",
                    "             mask = [False  True],",
                    "       fill_value = N/A)"]
    else:
        expected = ["masked_array(data=['2000-01-01 00:00:00.000', --],",
                    '             mask=[False,  True],',
                    "       fill_value='N/A',",
                    "            dtype='<U23')"]
    assert repr(t.iso).splitlines() == expected

    # Assign value to unmask
    t[1] = '2000:111'
    assert str(t) == "['2000:001:00:00:00.000' '2000:111:00:00:00.000']"
    assert t.masked is False


def test_transform():
    t = Time(['2000:001', '2000:002'])
    t[1] = np.ma.masked

    # Change scale (this tests the ERFA machinery with masking as well)
    t_ut1 = t.ut1
    assert is_masked(t_ut1.value[1])
    assert not is_masked(t_ut1.value[0])
    assert np.all(t_ut1.mask == [False, True])

    # Change format
    t_unix = t.unix
    assert is_masked(t_unix[1])
    assert not is_masked(t_unix[0])
    assert np.all(t_unix.mask == [False, True])


def test_masked_input():
    v0 = np.ma.MaskedArray([[1, 2], [3, 4]])  # No masked elements
    v1 = np.ma.MaskedArray([[1, 2], [3, 4]], mask=[[True, False], [False, False]])
    v2 = np.ma.MaskedArray([[10, 20], [30, 40]], mask=[[False, False], [False, True]])

    # Init from various combinations of masked arrays
    t = Time(v0, format='cxcsec')
    assert np.ma.allclose(t.value, v0)
    assert np.all(t.mask == [[False, False], [False, False]])
    assert t.masked is False

    t = Time(v1, format='cxcsec')
    assert np.ma.allclose(t.value, v1)
    assert np.all(t.mask == v1.mask)
    assert np.all(t.value.mask == v1.mask)
    assert t.masked is True

    t = Time(v1, v2, format='cxcsec')
    assert np.ma.allclose(t.value, v1 + v2)
    assert np.all(t.mask == (v1 + v2).mask)
    assert t.masked is True

    t = Time(v0, v1, format='cxcsec')
    assert np.ma.allclose(t.value, v0 + v1)
    assert np.all(t.mask == (v0 + v1).mask)
    assert t.masked is True

    t = Time(0, v2, format='cxcsec')
    assert np.ma.allclose(t.value, v2)
    assert np.all(t.mask == v2.mask)
    assert t.masked is True

    # Init from a string masked array
    t_iso = t.iso
    t2 = Time(t_iso)
    assert np.all(t2.value == t_iso)
    assert np.all(t2.mask == v2.mask)
    assert t2.masked is True


def test_serialize_fits_masked(tmpdir):
    tm = Time([1, 2, 3], format='cxcsec')
    tm[1] = np.ma.masked

    fn = str(tmpdir.join('tempfile.fits'))
    t = Table([tm])
    t.write(fn)
    t2 = Table.read(fn, astropy_native=True)

    # Time FITS handling does not current round-trip format in FITS
    t2['col0'].format = tm.format

    assert t2['col0'].masked
    assert np.all(t2['col0'].mask == [False, True, False])
    assert np.all(t2['col0'].value == t['col0'].value)


@pytest.mark.skipif('not HAS_H5PY')
def test_serialize_hdf5_masked(tmpdir):
    tm = Time([1, 2, 3], format='cxcsec')
    tm[1] = np.ma.masked

    fn = str(tmpdir.join('tempfile.hdf5'))
    t = Table([tm])
    t.write(fn, path='root', serialize_meta=True)
    t2 = Table.read(fn)

    assert t2['col0'].masked
    assert np.all(t2['col0'].mask == [False, True, False])
    assert np.all(t2['col0'].value == t['col0'].value)


@pytest.mark.skipif('not HAS_YAML')
def test_serialize_ecsv_masked(tmpdir):
    tm = Time([1, 2, 3], format='cxcsec')
    tm[1] = np.ma.masked

    # Serializing in the default way for ECSV fails to round-trip
    # because it writes out a "nan" instead of "".  But for jd1/jd2
    # this works OK.
    tm.info.serialize_method['ecsv'] = 'jd1_jd2'

    fn = str(tmpdir.join('tempfile.ecsv'))
    t = Table([tm])
    t.write(fn)
    t2 = Table.read(fn)

    assert t2['col0'].masked
    assert np.all(t2['col0'].mask == [False, True, False])
    # Serializing floats to ASCII loses some precision so use allclose
    # and 1e-7 seconds tolerance.
    assert np.allclose(t2['col0'].value, t['col0'].value, rtol=0, atol=1e-7)
