import numpy as np

from darts.tests.base_test_class import DartsBaseTestClass
from darts.models.forecasting.fft import _find_relevant_timestamp_attributes
from darts.utils import timeseries_generation as tg


class FFTTestCase(DartsBaseTestClass):
    def helper_relevant_attributes(self, freq, length, period_attributes_tuples):

        # test random walk
        random_walk_ts = tg.random_walk_timeseries(freq=freq, length=length)
        self.assertEqual(_find_relevant_timestamp_attributes(random_walk_ts), set())

        for period, relevant_attributes in period_attributes_tuples:

            # test seasonal period with no noise
            seasonal_ts = tg.sine_timeseries(
                freq=freq, value_frequency=1 / period, length=length
            )
            self.assertEqual(
                _find_relevant_timestamp_attributes(seasonal_ts),
                relevant_attributes,
                "failed to recognize season in non-noisy timeseries",
            )

            # test seasonal period with no noise
            seasonal_noisy_ts = seasonal_ts + tg.gaussian_timeseries(
                freq=freq, length=length
            )
            self.assertEqual(
                _find_relevant_timestamp_attributes(seasonal_noisy_ts),
                relevant_attributes,
                "failed to recognize season in noisy timeseries",
            )

    def test_find_relevant_timestamp_attributes(self):

        np.random.seed(0)

        # monthly frequency
        self.helper_relevant_attributes("M", 150, [(12, {"month"})])

        # daily frequency
        self.helper_relevant_attributes(
            "D", 1000, [(365, {"month", "day"}), (30, {"day"}), (7, {"weekday"})]
        )

        # hourly frequency
        self.helper_relevant_attributes(
            "H",
            3000,
            [(730, {"day", "hour"}), (168, {"weekday", "hour"}), (24, {"hour"})],
        )

        # minutely frequency
        self.helper_relevant_attributes(
            "min",
            5000,
            [
                (1440, {"hour", "minute"}),
                (60, {"minute"}),
            ],
        )
