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

__all__ = ['Yearly', 'Quarterly', 'Monthly', 'Other', 'M3Info', 'M3']

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

import numpy as np
import pandas as pd

from .utils import download_file, Info, TimeSeriesDataclass

# Cell
@dataclass
class Yearly:
    seasonality: int = 1
    horizon: int = 6
    freq: str = 'Y'
    sheet_name: str = 'M3Year'
    name: str = 'Yearly'
    n_ts: int = 645

@dataclass
class Quarterly:
    seasonality: int = 4
    horizon: int = 8
    freq: str = 'Q'
    sheet_name: str = 'M3Quart'
    name: str = 'Quarterly'
    n_ts: int = 756

@dataclass
class Monthly:
    seasonality: int = 12
    horizon: int = 18
    freq: str = 'M'
    sheet_name: str = 'M3Month'
    name: str = 'Monthly'
    n_ts: int = 1428

@dataclass
class Other:
    seasonality: int = 1
    horizon: int = 8
    freq: str = 'D'
    sheet_name: str = 'M3Other'
    name: str = 'Other'
    n_ts: int = 174

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

# Internal Cell
def _return_year(ts):
    year = ts.iloc[0]
    year = year if year != 0 else 1970

    return year

# Cell
@dataclass
class M3(TimeSeriesDataclass):

    source_url = 'https://forecasters.org/data/m3comp/M3C.xls'

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

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

        Notes
        -----
        [1] Returns train+test sets.
        [2] There are monthly time series without start year.
            This time series will start with 1970.
        [3] Other time series have no start date.
            This time series will start with 1970.
        """
        M3.download(directory)

        path = f'{directory}/m3/datasets/'
        file = f'{path}/M3C.xls'

        class_group = M3Info.get_group(group)

        df = pd.read_excel(file, sheet_name=class_group.sheet_name)
        df = df.rename(columns={'Series': 'unique_id',
                                'Category': 'category',
                                'Starting Year': 'year',
                                'Starting Month': 'month'})
        df['unique_id'] = [class_group.name[0] + str(i + 1) for i in range(len(df))]
        S = df.filter(items=['unique_id', 'category'])

        id_vars = list(df.columns[:6])
        df = pd.melt(df, id_vars=id_vars, var_name='ds', value_name='y')
        df = df.dropna().sort_values(['unique_id', 'ds']).reset_index(drop=True)

        freq = pd.tseries.frequencies.to_offset(class_group.freq)

        if group == 'Other':
            df['year'] = 1970

        df['ds'] = df.groupby('unique_id')['year'] \
                     .transform(lambda df: pd.date_range(f'{_return_year(df)}-01-01',
                                                         periods=df.shape[0],
                                                         freq=freq))
        df = df.filter(items=['unique_id', 'ds', 'y'])

        return df, None, None

    @staticmethod
    def download(directory: str) -> None:
        """Download M3 Dataset."""
        path = f'{directory}/m3/datasets/'
        if not os.path.exists(path):
            download_file(path, M3.source_url)