# -*- coding: utf-8 -*-
# Copyright (C) Duncan Macleod (2018-2019)
#
# This file is part of GWpy.
#
# GWpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GWpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GWpy.  If not, see <http://www.gnu.org/licenses/>.

"""Unit tests for gwpy.types.series
"""

import warnings

import numpy

import pytest

from astropy import units

from ...segments import Segment
from ...testing import utils
from .. import (Series, Array2D)
from .test_array import TestArray as _TestArray


class TestSeries(_TestArray):
    TEST_CLASS = Series

    def test_new(self):
        array = super(TestSeries, self).test_new()
        assert array.x0 == units.Quantity(0, self.TEST_CLASS._default_xunit)
        assert array.dx == units.Quantity(1, self.TEST_CLASS._default_xunit)
        return array

    # -- test properties ------------------------

    def test_x0(self, array):
        array.x0 = 5
        # test simple
        assert array.x0 == 5 * self.TEST_CLASS._default_xunit

        # test deleter
        del array.x0
        del array.x0
        assert array.x0 == 0 * self.TEST_CLASS._default_xunit

        # test quantity
        array.x0 = 5 * units.m
        assert array.x0 == units.Quantity(5, 'm')

    def test_dx(self, array):
        array.dx = 5 * self.TEST_CLASS._default_xunit
        # test simple
        assert array.dx == units.Quantity(5, self.TEST_CLASS._default_xunit)

        # test deleter
        del array.dx
        del array.dx
        assert array.dx == units.Quantity(1, self.TEST_CLASS._default_xunit)

        # test quantity
        array.dx = 5 * units.m
        assert array.dx == units.Quantity(5, 'm')

    def test_xindex(self):
        x = numpy.linspace(0, 100, num=self.data.shape[0])

        # test simple
        series = self.create(xindex=x)
        utils.assert_quantity_equal(
            series.xindex, units.Quantity(x, self.TEST_CLASS._default_xunit))

        # test deleter
        del series.xindex
        del series.xindex
        x1 = series.x0.value + series.shape[0] * series.dx.value
        x_default = numpy.linspace(series.x0.value, x1, num=series.shape[0],
                                   endpoint=False)
        utils.assert_quantity_equal(
            series.xindex,
            units.Quantity(x_default, self.TEST_CLASS._default_xunit))

        # test setting of x0 and dx
        series = self.create(xindex=units.Quantity(x, 'Farad'))
        assert series.x0 == units.Quantity(x[0], 'Farad')
        assert series.dx == units.Quantity(x[1] - x[0], 'Farad')
        assert series.xunit == units.Farad
        assert series.xspan == (x[0], x[-1] + x[1] - x[0])

        # test that setting xindex warns about ignoring dx or x0
        with pytest.warns(UserWarning):
            series = self.create(xindex=units.Quantity(x, 'Farad'), dx=1)
        with pytest.warns(UserWarning):
            series = self.create(xindex=units.Quantity(x, 'Farad'), x0=0)

        # test non-regular xindex
        x = numpy.logspace(0, 2, num=self.data.shape[0])
        series = self.create(xindex=units.Quantity(x, 'Mpc'))
        with pytest.raises(AttributeError):
            series.dx
        assert series.x0 == units.Quantity(1, 'Mpc')
        assert series.xspan == (x[0], x[-1] + x[-1] - x[-2])

    def test_xindex_dtype(self):
        x0 = numpy.longdouble(100)
        dx = numpy.float32(1e-4)
        series = self.create(x0=x0, dx=dx)
        assert series.xindex.dtype is x0.dtype

    def test_xunit(self, unit=None):
        if unit is None:
            unit = self.TEST_CLASS._default_xunit
        series = self.create(dx=4*unit)
        assert series.xunit == unit
        assert series.x0 == 0*unit
        assert series.dx == 4*unit
        # for series only, test arbitrary xunit
        if self.TEST_CLASS is (Series, Array2D):
            series = self.create(dx=4, xunit=units.m)
            assert series.x0 == 0*units.m
            assert series.dx == 4*units.m

    def test_xspan(self):
        # test normal
        series = self.create(x0=1, dx=1)
        assert series.xspan == (1, 1 + 1 * series.shape[0])
        assert isinstance(series.xspan, Segment)
        # test from irregular xindex
        x = numpy.logspace(0, 2, num=self.data.shape[0])
        series = self.create(xindex=x)
        assert series.xspan == (x[0], x[-1] + x[-1] - x[-2])

    # -- test methods ---------------------------

    def test_getitem(self, array):
        # item access
        utils.assert_quantity_equal(
            array[0], units.Quantity(array.value[0], array.unit))

        # slice
        utils.assert_quantity_equal(array[1::2], self.TEST_CLASS(
            array.value[1::2], x0=array.x0+array.dx, dx=array.dx*2,
            name=array.name, epoch=array.epoch, unit=array.unit),
        )

        # index array
        a = numpy.array([3, 4, 1, 2])
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore', message='xindex was given',
                                    category=UserWarning)
            utils.assert_quantity_equal(array[a], self.TEST_CLASS(
                array.value[a], xindex=array.xindex[a],
                name=array.name, epoch=array.epoch, unit=array.unit),
            )

    def test_empty_slice(self, array):
        """Check that we can slice a `Series` into nothing

        This tests against a bug in sliceutils.py.
        """
        a2 = array[:0]
        assert a2.x0 == array.x0
        assert a2.dx == array.dx

        idx = numpy.array([False]*array.shape[0])  # False slice array
        a3 = array[idx]
        utils.assert_quantity_sub_equal(a2, a3)

        a4 = array[idx[:0]]  # empty slice array
        utils.assert_quantity_sub_equal(a2, a4)

    def test_zip(self, array):
        z = array.zip()
        utils.assert_array_equal(
            z, numpy.column_stack((array.xindex.value, array.value)))

    def test_crop(self, array):
        a2 = array.crop(10, 20)
        utils.assert_quantity_equal(array[10:20], a2)
        # check that warnings are printed for out-of-bounds
        with pytest.warns(UserWarning):
            array.crop(array.xspan[0]-1, array.xspan[1])
        with pytest.warns(UserWarning):
            array.crop(array.xspan[0], array.xspan[1]+1)

    def test_is_compatible(self, array):
        """Test the `Series.is_compatible` method
        """
        a2 = self.create(name='TEST CASE 2')
        assert array.is_compatible(a2)

        a3 = self.create(dx=2)
        with pytest.raises(ValueError):
            array.is_compatible(a3)

        a4 = self.create(unit='m')
        with pytest.raises(ValueError):
            array.is_compatible(a4)

        x = numpy.logspace(0, 2, num=self.data.shape[0])
        a5 = self.create(xindex=x)
        with pytest.raises(ValueError):
            array.is_compatible(a5)

    def test_is_contiguous(self, array):
        a2 = self.create(x0=array.xspan[1])
        assert array.is_contiguous(a2) == 1
        assert array.is_contiguous(a2.value) == 1

        ts3 = self.create(x0=array.xspan[1]+1)
        assert array.is_contiguous(ts3) == 0

        ts4 = self.create(x0=-array.xspan[1])
        assert array.is_contiguous(ts4) == -1

    def test_append(self, array):
        a2 = self.create(x0=array.xspan[1])

        # test basic append
        a3 = array.append(a2, inplace=False)
        assert a3.epoch == array.epoch
        assert a3.x0 == array.x0
        assert a3.size == array.size+a2.size
        assert a3.xspan == array.xspan+a2.xspan
        utils.assert_array_equal(a3.value[:array.shape[0]], array.value)
        utils.assert_array_equal(a3.value[-a2.shape[0]:], a2.value)

        # check that appending again causes a problem
        with pytest.raises(ValueError):
            a3.append(array)

        # test appending with one xindex deletes it in the output
        array.xindex
        a3 = array.append(a2, inplace=False)
        assert hasattr(a3, '_xindex') is False

        # test appending with both xindex appends as well
        array.xindex
        a2.xindex
        a3 = array.append(a2, inplace=False)
        assert hasattr(a3, '_xindex')
        utils.assert_array_equal(
            a3.xindex.value,
            numpy.concatenate((array.xindex.value, a2.xindex.value)))

        # test appending with one only and not resize
        del a2.xindex
        a3 = array.append(a2, inplace=False, resize=False)
        assert a3.x0 == array.x0 + array.dx * a2.shape[0]

        # test discontiguous appends - gap='raise'
        a3 = self.create(x0=array.xspan[1] + 1)
        ts4 = array.copy()
        with pytest.raises(ValueError):
            array.append(a3)

        # gap='ignore'
        ts4.append(a3, gap='ignore')
        assert ts4.shape[0] == array.shape[0] + a3.shape[0]
        utils.assert_array_equal(
            ts4.value, numpy.concatenate((array.value, a3.value)))

        # gap='pad'
        ts4 = array.copy()
        ts4.append(a3, gap='pad', pad=0)
        assert ts4.shape[0] == array.shape[0] + 1 + a3.shape[0]
        z = numpy.zeros((1,) + array.shape[1:])
        utils.assert_array_equal(
            ts4.value, numpy.concatenate((array.value, z, a3.value)))

    def test_prepend(self, array):
        """Test the `Series.prepend` method
        """
        a2 = self.create(x0=array.xspan[1]) * 2
        a3 = a2.prepend(array, inplace=False)
        assert a3.x0 == array.x0
        assert a3.size == array.size + a2.size
        assert a3.xspan == array.xspan + a2.xspan
        with pytest.raises(ValueError):
            a3.prepend(array)
        utils.assert_array_equal(a3.value[:array.shape[0]], array.value)
        utils.assert_array_equal(a3.value[-a2.shape[0]:], a2.value)

    def test_update(self):
        """Test the `Series.update` method
        """
        ts1 = self.create()
        ts2 = self.create(x0=ts1.xspan[1])[:ts1.size//2]
        ts3 = ts1.update(ts2, inplace=False)
        assert ts3.x0 == ts1.x0 + abs(ts2.xspan)*ts1.x0.unit
        assert ts3.size == ts1.size
        with pytest.raises(ValueError):
            ts3.update(ts1)

    def test_pad(self):
        """Test the `Series.pad` method
        """
        ts1 = self.create()
        ts2 = ts1.pad(10)
        assert ts2.shape[0] == ts1.shape[0] + 20
        utils.assert_array_equal(
            ts2.value,
            numpy.concatenate((numpy.zeros(10), ts1.value, numpy.zeros(10))))
        assert ts2.x0 == ts1.x0 - 10*ts1.x0.unit
        # test pre-pad
        ts3 = ts1.pad((20, 10))
        assert ts3.size == ts1.size + 30
        utils.assert_array_equal(
            ts3.value,
            numpy.concatenate((numpy.zeros(20), ts1.value, numpy.zeros(10))))
        assert ts3.x0 == ts1.x0 - 20*ts1.x0.unit
        # test bogus input
        with pytest.raises(ValueError):
            ts1.pad(-1)

    def test_diff(self, array):
        """Test the `Series.diff` method

        This just ensures that the returned `Series` has the right length
        and the right x0
        """
        diff = array.diff(axis=0)
        assert isinstance(diff, type(array))
        assert array.shape[0] - 1 == diff.shape[0]
        assert diff.x0 == array.x0 + array.dx
        assert diff.xspan[1] == array.xspan[1]
        assert diff.channel == array.channel
        # test n=3
        diff = array.diff(n=3)
        assert array.shape[-1] - 3 == diff.shape[-1]
        assert diff.x0 == array.x0 + array.dx * 3

    def test_value_at(self):
        ts1 = self.create(dx=.5)
        assert ts1.value_at(1.5) == self.data[3] * ts1.unit
        assert ts1.value_at(1.5 * ts1.xunit) == self.data[3] * ts1.unit
        with pytest.raises(IndexError):
            ts1.value_at(1.6)
        # test TimeSeries unit conversion
        if ts1.xunit == units.s:
            assert ts1.value_at(1500 * units.millisecond) == (
                self.data[3] * ts1.unit)
        # test FrequencySeries unit conversion
        elif ts1.xunit == units.Hz:
            assert ts1.value_at(1500 * units.milliHertz) == (
                self.data[3] * ts1.unit)

    def test_shift(self):
        a = self.create(x0=0, dx=1, xunit='s')
        x0 = a.x0.copy()
        a.shift(5)
        assert a.x0 == x0 + 5 * x0.unit

        a.shift('1 hour')
        assert a.x0 == x0 + 3605 * x0.unit

        a.shift(-0.007)
        assert a.x0 == x0 + (3604.993) * x0.unit

        with pytest.raises(ValueError):
            a.shift('1 Hz')
