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

__all__ = ['load_EI_df', 'clean_df_for_plot', 'rgb_2_plt_tuple', 'convert_fuel_colour_dict_to_plt_tuple',
           'stacked_fuel_plot']

# Cell
import json
import pandas as pd

import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.transforms as mtf

# Cell
def load_EI_df(EI_fp):
    df = pd.read_csv(EI_fp)

    df['local_datetime'] = pd.to_datetime(df['local_datetime'], utc=True)
    df = df.set_index('local_datetime')

    return df

# Cell
def clean_df_for_plot(df, freq='7D'):
    fuel_order = ['Imports & Storage', 'nuclear', 'biomass', 'gas', 'coal', 'hydro', 'wind', 'solar']
    interconnectors = ['french', 'irish', 'dutch', 'belgian', 'ireland', 'northern_ireland']

    df = (df
          .copy()
          .assign(imports_storage=df[interconnectors+['pumped_storage']].sum(axis=1))
          .rename(columns={'imports_storage':'Imports & Storage'})
          .drop(columns=interconnectors+['demand', 'pumped_storage'])
          [fuel_order]
         )

    df_resampled = df.astype('float').resample(freq).mean()
    return df_resampled

# Cell
def rgb_2_plt_tuple(rgb_tuple):
    plt_tuple = tuple([x/255 for x in rgb_tuple])
    return plt_tuple

def convert_fuel_colour_dict_to_plt_tuple(fuel_colour_dict_rgb):
    fuel_colour_dict_plt = fuel_colour_dict_rgb.copy()

    fuel_colour_dict_plt = {
        fuel: rgb_2_plt_tuple(rgb_tuple)
        for fuel, rgb_tuple
        in fuel_colour_dict_plt.items()
    }

    return fuel_colour_dict_plt

# Cell
def stacked_fuel_plot(df, fuel_colour_dict, ax=None, save_path=None, dpi=150):
    df = df[fuel_colour_dict.keys()]

    if ax == None:
        fig = plt.figure(figsize=(10, 5), dpi=dpi)
        ax = plt.subplot()

    ax.stackplot(df.index.values, df.values.T, labels=df.columns.str.capitalize(), linewidth=0.25, edgecolor='white', colors=list(fuel_colour_dict.values()))

    plt.rcParams['axes.ymargin'] = 0
    ax.spines['bottom'].set_position('zero')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    ax.set_xlim(df.index.min(), df.index.max())
    ax.legend(ncol=4, bbox_to_anchor=(0.85, 1.15), frameon=False)
    ax.set_ylabel('Generation (GW)')

    if save_path:
        fig.savefig(save_path)

    return ax