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

import pytest

from numpy.testing import assert_equal

from astropy import units as u
from astropy.table import Table, QTable, vstack
from astropy.time import Time

from astropy.timeseries.sampled import TimeSeries
from astropy.timeseries.binned import BinnedTimeSeries


INPUT_TIME = Time(['2016-03-22T12:30:31', '2015-01-21T12:30:32', '2016-03-22T12:30:40'])
PLAIN_TABLE = Table([[1., 2., 11.], [3, 4, 1], ['x', 'y', 'z']], names=['a', 'b', 'c'])


class CommonTimeSeriesTests:

    def test_stacking(self):
        ts = vstack([self.series, self.series])
        assert isinstance(ts, self.series.__class__)

    def test_row_slicing(self):
        ts = self.series[:2]
        assert isinstance(ts, self.series.__class__)

    def test_row_indexing(self):
        self.series[0][self.time_attr] == Time('2015-01-21T12:30:32')
        self.series[self.time_attr][0] == Time('2015-01-21T12:30:32')

    def test_column_indexing(self):
        assert_equal(self.series['a'], [1, 2, 11])

    def test_column_slicing_notime(self):
        tab = self.series['a', 'b']
        assert not isinstance(tab, self.series.__class__)
        assert isinstance(tab, QTable)

    def test_add_column(self):
        self.series['d'] = [1, 2, 3]

    def test_add_row(self):
        self.series.add_row(self._row)

    def test_required_after_stacking(self):
        # When stacking, we have to temporarily relax the checking of the
        # columns in the time series, but we need to make sure that the
        # checking works again afterwards
        ts = vstack([self.series, self.series])
        with pytest.raises(ValueError) as exc:
            ts.remove_columns(ts.colnames)
        assert 'TimeSeries object is invalid' in exc.value.args[0]


class TestTimeSeries(CommonTimeSeriesTests):

    _row = {'time': '2016-03-22T12:30:40', 'a': 1., 'b': 2, 'c': 'a'}

    def setup_method(self, method):
        self.series = TimeSeries(time=INPUT_TIME, data=PLAIN_TABLE)
        self.time_attr = 'time'

    def test_column_slicing(self):
        ts = self.series['time', 'a']
        assert isinstance(ts, TimeSeries)


class TestBinnedTimeSeries(CommonTimeSeriesTests):

    _row = {'time_bin_start': '2016-03-22T12:30:40',
            'time_bin_size': 2 * u.s, 'a': 1., 'b': 2, 'c': 'a'}

    def setup_method(self, method):
        self.series = BinnedTimeSeries(time_bin_start=INPUT_TIME,
                                       time_bin_size=3 * u.s,
                                       data=PLAIN_TABLE)
        self.time_attr = 'time_bin_start'

    def test_column_slicing(self):
        ts = self.series['time_bin_start', 'time_bin_size', 'a']
        assert isinstance(ts, BinnedTimeSeries)
