# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/data_datasets__tourism.ipynb (unless otherwise specified).

__all__ = ['Yearly', 'Quarterly', 'Monthly', 'TourismInfo', 'Tourism']

# Cell
import os
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
from pandas.tseries.frequencies import to_offset

from .utils import download_file, Info, TimeSeriesDataclass

# Cell
@dataclass
class Yearly:
    seasonality: int = 1
    horizon: int = 4
    freq: str = 'Y'
    rows: int = 2
    name: str = 'Yearly'
    n_ts: int = 518

@dataclass
class Quarterly:
    seasonality: int = 4
    horizon: int = 8
    freq: str = 'Q'
    rows: int = 3
    name: str = 'Quarterly'
    n_ts: int = 427

@dataclass
class Monthly:
    seasonality: int = 12
    horizon: int = 24
    freq: str = 'M'
    rows: int = 3
    name: str = 'Monthly'
    n_ts: int = 366

# Cell
TourismInfo = Info(groups=('Yearly', 'Quarterly', 'Monthly'),
                   class_groups=(Yearly, Quarterly, Monthly))

# Cell
class Tourism(TimeSeriesDataclass):

    source_url = 'https://robjhyndman.com/data/27-3-Athanasopoulos1.zip'

    @staticmethod
    def load(directory: str,
             group: str) -> Tuple[pd.DataFrame,
                                  Optional[pd.DataFrame],
                                  Optional[pd.DataFrame]]:
        """
        Downloads and loads Tourism data.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.
        group: str
            Group name.
            Allowed groups: 'Yearly', 'Quarterly', 'Monthly'.

        Notes
        -----
        [1] Returns train+test sets.
        """
        Tourism.download(directory)

        path = f'{directory}/tourism/datasets'

        class_group = TourismInfo.get_group(group)
        train_file = f'{path}/{class_group.name.lower()}_in.csv'
        test_file  = f'{path}/{class_group.name.lower()}_oos.csv'

        train, test = pd.read_csv(train_file), pd.read_csv(test_file)

        dfs = []
        freq = to_offset(class_group.freq)
        for col in train.columns:
            df_appended = []
            for df, training in zip([train, test], [True, False]):
                df_col = df[col]
                length, year = df_col[:2].astype(int)
                skip_rows = class_group.rows
                start_date = pd.to_datetime(f'{year}-01-01')
                if group != 'Yearly':
                    n_offsets = df_col[2].astype(int)
                    start_date += n_offsets * freq
                elif col == 'Y18' and not training: # viene mal en el archivo esta serie
                    start_date += 2 * freq
                df_col = df_col[skip_rows:length + skip_rows]
                df_col = df_col.rename('y').to_frame()
                df_col['unique_id'] = col
                df_col['ds'] = pd.date_range(start_date, periods=length, freq=freq)
                df_appended.append(df_col)
            df_appended = pd.concat(df_appended)
            dfs.append(df_appended)

        df = pd.concat(dfs)

        df = df.reset_index().filter(items=['unique_id', 'ds', 'y'])
        df = df.sort_values(['unique_id', 'ds'])

        return df, None, None

    @staticmethod
    def download(directory: str) -> None:
        """Downloads Tourism Dataset."""
        path = f'{directory}/tourism/datasets'

        if not os.path.exists(path):
            download_file(path, Tourism.source_url, decompress=True)