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

__all__ = ['get_latest_lolpdrm_df', 'clean_fcst_df', 'AxTransformer', 'set_date_ticks', 'create_fcst_htmp',
           'create_fcst_delta_htmp', 'save_lolpdrm_imgs', 'construct_drm_md_txt', 'generate_lolpdrm_imgs_text']

# Cell
import numpy as np
import pandas as pd

import seaborn as sns
import matplotlib.pyplot as plt

from collections.abc import Iterable
from sklearn import linear_model

from ..api import Client

# Cell
def get_latest_lolpdrm_df(api_key: str=None):
    start_date = pd.Timestamp.now(tz='Europe/London').strftime('%Y-%m-%d')
    end_date = pd.Timestamp.now(tz='Europe/London') + pd.Timedelta(days=3)

    client = Client(api_key=api_key)
    df_lolpdrm = client.get_LOLPDRM(start_date, end_date)

    return df_lolpdrm

# Cell
def clean_fcst_df(
    df_lolpdrm: pd.DataFrame,
    fcst_type: str='lolp',
    fcst_horizons: list=[8, 4, 2, 1]
):
    cols_renaming_map = {f'{fcst_type}12Forecast': '12:00'}

    cols_renaming_map.update({
        f'{fcst_type}{fcst_horizon}HourForecast': f"{fcst_horizon} Hour{'s' if fcst_horizon>1 else ''} Ahead"
        for fcst_horizon
        in fcst_horizons
    })

    df_clean_fcst = (df_lolpdrm
                     .set_index('local_datetime')
                     [cols_renaming_map.keys()]
                     .rename(columns=cols_renaming_map)
                     .astype(float)
                    )

    return df_clean_fcst

# Cell
class AxTransformer:
    def __init__(self, datetime_vals=False):
        self.datetime_vals = datetime_vals
        self.lr = linear_model.LinearRegression()

        return

    def process_tick_vals(self, tick_vals):
        if not isinstance(tick_vals, Iterable) or isinstance(tick_vals, str):
            tick_vals = [tick_vals]

        if self.datetime_vals == True:
            tick_vals = pd.to_datetime(tick_vals).astype(int).values

        tick_vals = np.array(tick_vals)

        return tick_vals

    def fit(self, ax, axis='x'):
        axis = getattr(ax, f'get_{axis}axis')()

        tick_locs = axis.get_ticklocs()
        tick_vals = self.process_tick_vals([label._text for label in axis.get_ticklabels()])

        self.lr.fit(tick_vals.reshape(-1, 1), tick_locs)

        return

    def transform(self, tick_vals):
        tick_vals = self.process_tick_vals(tick_vals)
        tick_locs = self.lr.predict(np.array(tick_vals).reshape(-1, 1))

        return tick_locs

def set_date_ticks(ax, start_date, end_date, axis='y', date_format='%Y-%m-%d', **date_range_kwargs):
    dt_rng = pd.date_range(start_date, end_date, **date_range_kwargs)

    ax_transformer = AxTransformer(datetime_vals=True)
    ax_transformer.fit(ax, axis=axis)

    getattr(ax, f'set_{axis}ticks')(ax_transformer.transform(dt_rng))
    getattr(ax, f'set_{axis}ticklabels')(dt_rng.strftime(date_format))

    ax.tick_params(axis=axis, which='both', bottom=True, top=False, labelbottom=True)

    return ax

# Cell
def create_fcst_htmp(
    df_clean: pd.DataFrame,
    img_fp: str='docs/img/vis/drm_fcst_htmp.png'
):
    # Plotting
    fig, ax = plt.subplots(dpi=250)

    htmp = sns.heatmap(df_clean, vmin=0, ax=ax)

    set_date_ticks(ax, df_clean.index.min(), df_clean.index.max(), freq='6H', date_format='%Y-%m-%d %H:%M')
    ax.set_xticklabels(ax.get_xticklabels(), rotation=25)
    ax.xaxis.tick_top()
    ax.set_ylabel('')
    cbar = htmp.collections[0].colorbar
    cbar.set_label('De-Rated Margin (MW)', labelpad=20, rotation=270)

    # Saving
    fig.tight_layout()
    fig.savefig(img_fp)

# Cell
def create_fcst_delta_htmp(
    df_clean_drm: pd.DataFrame,
    img_fp: str='docs/img/vis/drm_fcst_delta_htmp.png'
):
    # Preparing data
    df_drm_delta = df_clean_drm.iloc[:, 2:].apply(lambda col: col-df_clean_drm.iloc[:, 1]).dropna(how='all', axis=1).dropna(how='all')

    # Plotting
    fig, ax = plt.subplots(dpi=150)

    htmp = sns.heatmap(df_drm_delta, center=0, cmap='bwr_r', ax=ax)

    set_date_ticks(ax, df_drm_delta.index.min(), df_drm_delta.index.max(), freq='3H', date_format='%Y-%m-%d %H:%M')
    ax.set_xticklabels(ax.get_xticklabels(), rotation=25)
    ax.xaxis.tick_top()
    ax.set_ylabel('')
    cbar = htmp.collections[0].colorbar
    cbar.set_label('Delta to the 8 Hour-\nAhead Forecast Horizon', labelpad=25, rotation=270)

    # Saving
    fig.tight_layout()
    fig.savefig(img_fp)

# Cell
def save_lolpdrm_imgs(
    docs_dir: str='docs',
    api_key: str=None,
    fcst_type: str='lolp',
    fcst_horizons: list=[8, 4, 2, 1]
):
    df_lolpdrm = get_latest_lolpdrm_df(api_key=api_key)
    df_clean = clean_fcst_df(df_lolpdrm, fcst_type=fcst_type, fcst_horizons=fcst_horizons)

    create_fcst_htmp(df_clean, img_fp=f'{docs_dir}/img/vis/{fcst_type}_fcst_htmp.png')
    plt.close()

    create_fcst_delta_htmp(df_clean, img_fp=f'{docs_dir}/img/vis/{fcst_type}_fcst_delta_htmp.png')
    plt.close()

# Cell
def construct_drm_md_txt(
    update_time: str=None
):
    if update_time is None:
        update_time = pd.Timestamp.now().round('5min').strftime('%Y-%m-%d %H:%M')

    md_txt = f"""### De-Rated Margin

In each settlement period the system operator publishes the de-rated margin forecast calculated in accordance with the [Loss of Load Probability Calculation Statement](https://www.elexon.co.uk/documents/bsc-codes/lolp/loss-of-load-probability-calculation-statement/) at the following times:

* At 1200 hours on each calendar day for all Settlement Periods for which Gate Closure has not yet passed and which occur within the current Operational Day or the following Operational Day; and
* At eight, four, two and one hour(s) prior to the beginning of the Settlement Period to which the De-Rated Margin Forecast relates.

These figures will be updated on an hourly basis, the last update was at: {update_time}

<br>

#### Forecasts

The following heatmap shows the evolving de-rated margin forecast across the different forecast horizons.

![](img/vis/drm_fcst_htmp.png)

<br>

#### Forecast Deltas

The following heatmap shows how the more recent de-rated margin forecasts deviate from the 8 hours ahead forecast.

![](img/vis/drm_fcst_delta_htmp.png)
"""

    return md_txt

# Cell
def generate_lolpdrm_imgs_text(
    docs_dir: str='docs',
    api_key: str=None,
    fcst_horizons: list=[8, 4, 2, 1],
    update_time: str=None
):
    for fcst_type in ['drm', 'lolp']:
        save_lolpdrm_imgs(docs_dir=docs_dir, api_key=api_key, fcst_type=fcst_type, fcst_horizons=fcst_horizons)

    md_txt = construct_drm_md_txt(update_time=update_time)

    return md_txt