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

__all__ = ['NP', 'PJM', 'BE', 'FR', 'DE', 'EPFInfo', 'EPF', 'epf_naive_forecast']

# Cell
import os
from datetime import timedelta
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

# Cell
@dataclass
class NP:
    test_date: str = '2016-12-27'
    name: str = 'NP'

@dataclass
class PJM:
    test_date: str = '2016-12-27'
    name: str = 'PJM'

@dataclass
class BE:
    test_date: str = '2015-01-04'
    name: str = 'BE'

@dataclass
class FR:
    test_date: str = '2015-01-04'
    name: str = 'FR'

@dataclass
class DE:
    test_date: str = '2016-01-04'
    name: str = 'DE'

# Cell
EPFInfo = Info(groups=('NP', 'PJM', 'BE', 'FR', 'DE'),
               class_groups=(NP, PJM, BE, FR, DE))

# Cell
class EPF:

    source_url = 'https://sandbox.zenodo.org/api/files/da5b2c6f-8418-4550-a7d0-7f2497b40f1b/'

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

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.
        group: str
            Group name.
            Allowed groups: 'NP', 'PJM', 'BE', 'FR', 'DE'.
        """
        EPF.download(directory)
        class_group = EPFInfo.get_group(group)

        path = f'{directory}/epf/datasets'
        file = f'{path}/{group}.csv'

        df = pd.read_csv(file)

        df.columns = ['ds', 'y'] + \
                     [f'Exogenous{i}' for i in range(1, len(df.columns) - 1)]

        df['unique_id'] = group
        df['ds'] = pd.to_datetime(df['ds'])
        df['week_day'] = df['ds'].dt.dayofweek

        dummies = pd.get_dummies(df['week_day'], prefix='day')
        df = pd.concat([df, dummies], axis=1)

        dummies_cols = [col for col in df if col.startswith('day')]

        Y = df.filter(items=['unique_id', 'ds', 'y'])
        X = df.filter(items=['unique_id', 'ds', 'Exogenous1', 'Exogenous2', 'week_day'] + \
                      dummies_cols)

        return Y, X, None

    @staticmethod
    def load_groups(directory: str,
                    groups: List[str]) -> Tuple[pd.DataFrame,
                                                Optional[pd.DataFrame],
                                                Optional[pd.DataFrame]]:
        """
        Downloads and loads panel of EPF data
        according of groups.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.
        groups: List[str]
            Group names.
            Allowed groups: 'NP', 'PJM', 'BE', 'FR', 'DE'.
        """
        Y = []
        X = []
        for group in groups:
            Y_df, X_df, S_df = EPF.load(directory=directory, group=group)
            Y.append(Y_df)
            X.append(X_df)

        Y = pd.concat(Y).sort_values(['unique_id', 'ds']).reset_index(drop=True)
        X = pd.concat(X).sort_values(['unique_id', 'ds']).reset_index(drop=True)

        S = Y[['unique_id']].drop_duplicates().reset_index(drop=True)
        dummies = pd.get_dummies(S['unique_id'], prefix='static')
        S = pd.concat([S, dummies], axis=1)

        return Y, X, S

    @staticmethod
    def download(directory: str) -> None:
        """Downloads EPF Dataset."""
        path = f'{directory}/epf/datasets'
        if not os.path.exists(path):
            for group in EPFInfo.groups:
                download_file(path, EPF.source_url + f'{group}.csv')

# Cell
# TODO: extend this to group_by unique_id application
def epf_naive_forecast(Y_df):
    """Function to build the naive forecast for electricity price forecasting

    The function is used to compute the accuracy metrics MASE and RMAE, the function
    assumes that the number of prices per day is 24. And computes naive forecast for
    days of the week and seasonal Naive forecast for weekends.

    Parameters
    ----------
    Y_df : pandas.DataFrame
        Dataframe containing the real prices in long format
        that contains variables ['ds', 'unique_id', 'y']

    Returns
    -------
    Y_hat_df : pandas.DataFrame
        Dataframe containing the predictions of the epf naive forecast.
    """
    assert type(Y_df) == pd.core.frame.DataFrame
    assert all([(col in Y_df) for col in ['unique_id', 'ds', 'y']])

    # Init the naive forecast
    Y_hat_df = Y_df[24 * 7:].copy()
    Y_hat_df['dayofweek'] = Y_df['ds'].dt.dayofweek

    # Monday, Saturday and Sunday
    # we have a naive forecast using weekly seasonality
    weekend_indicator = Y_hat_df['dayofweek'].isin([0,5,6])

    # Tuesday, Wednesday, Thursday, Friday
    # we have a naive forecast using daily seasonality
    week_indicator = Y_hat_df['dayofweek'].isin([1,2,3,4])

    naive = Y_df['y'].shift(24).values[24 * 7:]
    seasonal_naive = Y_df['y'].shift(24*7).values[24 * 7:]

    Y_hat_df['y_hat'] = naive * week_indicator + seasonal_naive * weekend_indicator
    return Y_hat_df