# 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 re
import logging
import zipfile

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union
from pandas.tseries.holiday import USFederalHolidayCalendar as calendar

import numpy as np
import pandas as pd

from .utils import (
    download_file,
    Info,
    TimeSeriesDataclass,
    create_calendar_variables,
    create_us_holiday_distance_variables,
)
from ..tsdataset import TimeSeriesDataset

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."""
        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) -> pd.DataFrame:
        # 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) -> pd.DataFrame:
        # 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) -> pd.DataFrame:
        # 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) -> Tuple[pd.DataFrame,
                                 pd.DataFrame,
                                 pd.DataFrame]:

        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):
        # 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):
        # 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, group):
        # 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]:
        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