import unittest
import pandas as pd

from ..timeseries import TimeSeries
from ..logging import get_logger

logger = get_logger(__name__)

try:
    import torch
    from ..models.torch_forecasting_model import _TimeSeriesSequentialDataset
    TORCH_AVAILABLE = True
except ImportError:
    logger.warning('Torch is not available. Torch Dataset will not be tested')
    TORCH_AVAILABLE = False

if TORCH_AVAILABLE:
    class DatasetTestCase(unittest.TestCase):
        __test__ = True
        times = pd.date_range('20130101', '20130410')
        pd_series = pd.Series(range(100), index=times)
        series: TimeSeries = TimeSeries.from_series(pd_series)

        def test_creation(self):
            # Cannot have train window <= 0
            with self.assertRaises(ValueError):
                _TimeSeriesSequentialDataset(self.series, self.series, -1)
            # Cannot have label window <= 0
            with self.assertRaises(ValueError):
                _TimeSeriesSequentialDataset(self.series, self.series, 1, -1)

        def test_content(self):
            # Can have no transformation
            dataset = _TimeSeriesSequentialDataset(self.series, self.series, 12, 4)
            self.assertEqual(len(dataset), 100 - 12 - 4 + 1)
            # correct types
            self.assertEqual(type(dataset[0][0]), torch.Tensor)
            self.assertEqual(dataset[0][0].dtype, torch.float32)
            # take correct elements
            self.assertEqual(dataset[5][0].numpy()[0], self.pd_series.values[5])
            self.assertEqual(dataset[5][1].numpy()[0], self.pd_series.values[5 + 12])
            # content same size from beginning to end
            self.assertEqual(dataset[0][1].size(), dataset[len(dataset) - 1][1].size())
