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

__all__ = ['logger', 'GEFCom2012', 'GEFCom2012_L', 'GEFCom2012_W']

# Cell
import os
import logging

from dataclasses import dataclass
from typing import Tuple

import pandas as pd

from .utils import (
    download_file,
    create_calendar_variables,
    create_us_holiday_distance_variables,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Cell
class GEFCom2012:

    source_url = 'https://www.dropbox.com/s/epj9b57eivn79j7/GEFCom2012.zip?dl=1'

    @staticmethod
    def download(directory: str) -> None:
        """
        Downloads GEFCom2012 Dataset.

        Parameters
        ----------
        directory: str
            Directory path to download dataset.
        """
        path = f'{directory}/gefcom2012'
        if not os.path.exists(path):
            download_file(directory=path,
                          source_url=GEFCom2012.source_url,
                          decompress=True)

# Cell
class GEFCom2012_L:

    @staticmethod
    def load_Y(directory: str) -> pd.DataFrame:
        """Load target time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        Y_df: pd.DataFrame
            Target time series with columns ['unique_id', 'ds', 'y'].
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Load'
        filepath = f'{path}/Load_history.csv'
        Y_df = pd.read_csv(filepath, sep=',', thousands=',')

        # Parsing temperature data
        Y_df['ds'] = pd.to_datetime(dict(year=Y_df.year,
                                         month=Y_df.month,
                                         day=Y_df.day))
        del Y_df['year'], Y_df['month'], Y_df['day']
        Y_df = pd.wide_to_long(Y_df, ['h'], i=['zone_id', 'ds'], j="hour")
        Y_df.reset_index(inplace=True)
        Y_df['tdelta']   = pd.to_timedelta(Y_df.hour, unit="h")
        Y_df['ds']       = Y_df['ds'] + Y_df['tdelta']
        del Y_df['tdelta'], Y_df['hour']
        Y_df.rename(columns={'zone_id': 'unique_id', 'h': 'y'}, inplace=True)
        #Y_df['y'] = pd.to_numeric(Y_df['y'], errors='coerce')
        return Y_df

    @staticmethod
    def load_X(directory: str) -> pd.DataFrame:
        """Load exogenous time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        X_df: pd.DataFrame
            Exogenous time series with columns ['unique_id', 'ds', 'y'].
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Load'
        filepath = f'{path}/temperature_history.csv'
        X_df = pd.read_csv(filepath, sep=',')

        # Parsing temperature data
        X_df['ds'] = pd.to_datetime(dict(year=X_df.year,
                                         month=X_df.month,
                                         day=X_df.day))
        del X_df['year'], X_df['month'], X_df['day']
        X_df = pd.wide_to_long(X_df, ['h'], i=['station_id', 'ds'], j="hour")
        X_df.reset_index(inplace=True)
        X_df['tdelta']   = pd.to_timedelta(X_df.hour, unit="h")
        X_df['ds']       = X_df['ds'] + X_df['tdelta']
        del X_df['tdelta'], X_df['hour']
        X_df['station_id'] = 'station_' + X_df['station_id'].astype(str)

        X_df = X_df.pivot(index='ds', columns='station_id', values='h').reset_index('ds')
        X_df.reset_index(drop=True, inplace=True)
        X_df = create_calendar_variables(X_df=X_df)
        X_df = create_us_holiday_distance_variables(X_df=X_df)
        return X_df

    @staticmethod
    def load_benchmark(directory: str) -> pd.DataFrame:
        """Load benchmark time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        benchmark_df: pd.DataFrame
            Benchmark time series form gefcom2012 dataset.
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Load'
        filepath = f'{path}/Load_benchmark.csv'
        benchmark_df = pd.read_csv(filepath, sep=',')

        # Parsing benchmark data
        benchmark_df['ds'] = pd.to_datetime(dict(year=benchmark_df.year,
                                                 month=benchmark_df.month,
                                                 day=benchmark_df.day))
        del benchmark_df['year'], benchmark_df['month'], benchmark_df['day'], benchmark_df['id']
        benchmark_df.rename(columns={'zone_id': 'unique_id'}, inplace=True)

        benchmark_df = pd.wide_to_long(benchmark_df, ['h'], i=['unique_id', 'ds'], j="hour")
        benchmark_df.reset_index(inplace=True)

        benchmark_df['tdelta']   = pd.to_timedelta(benchmark_df.hour, unit="h")
        benchmark_df['ds']       = benchmark_df['ds'] + benchmark_df['tdelta']
        del benchmark_df['tdelta'], benchmark_df['hour']
        benchmark_df.rename(columns={'h': 'y_hat'}, inplace=True)
        return benchmark_df

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

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        Y_df: pd.DataFrame
            Target time series with columns ['unique_id', 'ds', 'y'].
        X_df: pd.DataFrame
            Exogenous time series with columns ['unique_id', 'ds', 'y'].
        benchmark_df: pd.DataFrame
            Benchmark time series form gefcom2012 dataset.
        """

        GEFCom2012.download(directory)

        Y_df = GEFCom2012_L.load_Y(directory)
        X_df = GEFCom2012_L.load_X(directory)
        benchmark_df = GEFCom2012_L.load_benchmark(directory)
        return Y_df, X_df, benchmark_df

# Cell
class GEFCom2012_W:

    train_start = '2009-07-01 00:00:00'
    train_end   = '2010-12-31 23:00:00'

    test_start  = '2011-01-01 01:00:00'
    test_end    = '2012-06-28 12:00:00'

    @staticmethod
    def load_benchmark(directory: str) -> pd.DataFrame:
        """Load benchmark time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        benchmark_df: pd.DataFrame
            Benchmark time series form gefcom2012 dataset.
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Wind'
        filepath = f'{path}/benchmark.csv'
        benchmark_df = pd.read_csv(filepath, sep=',')

        benchmark_df['ds'] = pd.to_datetime(benchmark_df.date, format='%Y%m%d%H')
        del benchmark_df['date']
        benchmark_df = pd.wide_to_long(benchmark_df, ['wp'], i='ds', j="unique_id")
        benchmark_df.reset_index(inplace=True)
        return benchmark_df

    @staticmethod
    def load_Y(directory: str) -> pd.DataFrame:
        """Load target time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        Y_df: pd.DataFrame
            Target time series with columns ['unique_id', 'ds', 'y'].
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Wind'
        yfilepath = f'{path}/train.csv'

        # Read and parse Y data
        Y_df = pd.read_csv(yfilepath, sep='\t')
        Y_df['ds'] = pd.date_range(start=GEFCom2012_W.train_start,
                                   end=GEFCom2012_W.train_end, freq='H')
        del Y_df['date']
        Y_df = pd.wide_to_long(Y_df, ['wp'], i='ds', j="unique_id")
        Y_df.reset_index(inplace=True)
        return Y_df

    @staticmethod
    def load_X_group(directory: str, group: str) -> pd.DataFrame:
        """Load exogenous time series.

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.
        group: str
            Group name.

        Returns
        -------
        X_df: pd.DataFrame
            Exogenous time series with columns ['unique_id', 'ds', 'y'].
        """

        # Meta data
        path = f'{directory}/gefcom2012/GEFCOM2012_Data/Wind'
        xfilepath = f'{path}/windforecasts_wf{group}.csv'
        X_df = pd.read_csv(xfilepath, sep=',')

        # Create ds associated to each forecast from forecast creation date
        X_df['date']     = X_df.date.astype(str)
        X_df['fcd']      = pd.to_datetime(X_df.date, format='%Y%m%d%H')
        X_df['tdelta']   = pd.to_timedelta(X_df.hors, unit="h")
        X_df['ds']       = X_df['fcd'] + X_df['tdelta']

        # Separate forecasts by lead time
        X_lead12_df = X_df[(X_df.hors>0)  & (X_df.hors<=12)].reset_index(drop=True)
        X_lead24_df = X_df[(X_df.hors>12) & (X_df.hors<=24)].reset_index(drop=True)
        X_lead36_df = X_df[(X_df.hors>24) & (X_df.hors<=36)].reset_index(drop=True)
        X_lead48_df = X_df[(X_df.hors>36) & (X_df.hors<=48)].reset_index(drop=True)
        del X_df

        # Cleaning auxiliary variables and reconstructing X_df
        X_df = pd.DataFrame({'ds': pd.date_range(start='2009-07-01 01:00:00',
                                                 end=GEFCom2012_W.test_end, freq='H')})
        for lead, df in zip(['_lead12', '_lead24', '_lead36', '_lead48'], \
                            [X_lead12_df, X_lead24_df, X_lead36_df, X_lead48_df]):
            df.drop(['fcd', 'tdelta', 'date', 'hors'], axis=1, inplace=True)
            df.columns = [f'u{lead}', f'v{lead}', f'ws{lead}', f'wd{lead}', 'ds']
            X_df = X_df.merge(df, on='ds', how='left')

        # Removing nans in hierarchical fashion (priority to shorter lead forecasts)
        for var in ['u', 'v', 'ws', 'wd']:
            X_df[var] = X_df[f'{var}_lead12']
            for lead in ['_lead24', '_lead36', '_lead48']:
                X_df[var].fillna(X_df[f'{var}{lead}'], inplace=True)

        for var in ['u', 'v', 'ws', 'wd']:
            for lead in ['_lead12', '_lead24', '_lead36', '_lead48']:
                X_df[f'{var}{lead}'].fillna(X_df[var], inplace=True)
            del X_df[var]
        del X_lead12_df, X_lead24_df, X_lead36_df, X_lead48_df
        X_df['unique_id'] = group
        return X_df

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

        Parameters
        ----------
        directory: str
            Directory where data will be downloaded.

        Returns
        -------
        Y_df: pd.DataFrame
            Target time series with columns ['unique_id', 'ds', 'y'].
        X_df: pd.DataFrame
            Exogenous time series with columns ['unique_id', 'ds', 'y'].
        benchmark_df: pd.DataFrame
            Benchmark time series form gefcom2012 dataset.
        """

        GEFCom2012.download(directory)

        Y_df = GEFCom2012_W.load_Y(directory)
        X_df_list = [GEFCom2012_W.load_X_group(directory, group) for group in range(1,8)]
        X_df = pd.concat(X_df_list)

        benchmark_df = GEFCom2012_W.load_benchmark(directory)
        return Y_df, X_df, benchmark_df