# Libraries
# ==============================================================================
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.colors import ListedColormap, BoundaryNorm
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import accuracy_score, classification_report
from sklearn.calibration import CalibratedClassifierCV, CalibrationDisplay, calibration_curve
from skforecast.datasets import fetch_dataset
from skforecast.preprocessing import RollingFeaturesClassification
from skforecast.recursive import ForecasterRecursiveClassifier
from skforecast.model_selection import (
    TimeSeriesFold, OneStepAheadFold, backtesting_forecaster, bayesian_search_forecaster
)
from skforecast.plot import set_dark_theme

# Data download
# ==============================================================================
data = fetch_dataset(name='vic_electricity_classification')
data

# Distribution of classes in the series
# ==============================================================================
print(f"Distribution of classes:\n{data['Demand'].value_counts(normalize=True)}")

# Split train-test
# ==============================================================================
end_train = '2014-10-31 23:59:00'
data_train = data[:end_train].copy()
data_test  = data[end_train:].copy()

print(
    f"Train dates : {data_train.index.min()} --- "
    f"{data_train.index.max()}  (n={len(data_train)})"
)
print(
    f"Test dates  : {data_test.index.min()} --- "
    f"{data_test.index.max()}  (n={len(data_test)})"
)

# Data preprocessing
# ==============================================================================
# Encode categorical variable as numerical for plotting
encode_mapping = {'low': 0, 'medium': 1, 'high': 2}
print(f"Encoding map: {encode_mapping}")
data['encoded_value'] = data['Demand'].map(encode_mapping).astype(float)
data_train['encoded_value'] = data_train['Demand'].map(encode_mapping).astype(float)
data_test['encoded_value'] = data_test['Demand'].map(encode_mapping).astype(float)
data.head(3)

# Plot: line plot
# ==============================================================================
set_dark_theme()
fig, ax = plt.subplots(figsize=(8, 3))
data_train['encoded_value'].plot(ax=ax, label='train')
data_test['encoded_value'].plot(ax=ax, label='test')
y_ticks = list(encode_mapping.values())
y_labels = list(encode_mapping.keys())
ax.set_yticks(y_ticks)
ax.set_yticklabels(y_labels)
ax.set_title("Categorical Time Series (States)")
ax.set_xlabel("Date")
ax.legend();

# Plot: 1-row heatmap-style scatter
# ==================================================================================================
colors = ["skyblue", "gold", "tomato"]
cmap = ListedColormap(colors)
norm = BoundaryNorm([0, 1, 2, 3], cmap.N)
stripe_height = 20
state_array = np.tile(data["encoded_value"].values, (stripe_height, 1))
train_mask = data.index <= pd.to_datetime(end_train)
train_indices = np.where(train_mask)[0]
test_indices = np.where(~train_mask)[0]

fig, ax = plt.subplots(figsize=(8, 2.5))
ax.imshow(
    state_array,
    aspect="auto",
    cmap=cmap,
    norm=norm,
    extent=[0, len(data) - 1, 0, stripe_height],
)

ax.hlines(y=30, xmin=train_indices[0], xmax=train_indices[-1], color="#30a2da", linewidth=4)
ax.hlines(y=30, xmin=test_indices[0], xmax=test_indices[-1], color="#fc4f30", linewidth=4)
ax.set_yticks([])
ax.set_frame_on(False)
tick_indices = np.linspace(0, len(data) - 1, 10, dtype=int)
ax.set_xticks(tick_indices)
ax.set_xticklabels([data.index[i].strftime("%Y-%m") for i in tick_indices], rotation=0, ha="center")
ax.set_xlabel("Date")
ax.set_title("Categorical Daily States Over Time", pad=8)

cbar = fig.colorbar(
    plt.cm.ScalarMappable(cmap=cmap, norm=norm),
    ax=ax,
    orientation="horizontal",
    ticks=[0.5, 1.5, 2.5],
    pad=0.3,       # space between plot and colorbar
    fraction=0.05,  # thickness
    aspect=20       # length-to-height ratio
)
cbar.outline.set_visible(False)
cbar.ax.set_xticklabels(["low", "medium", "high"], fontsize=8)
cbar.set_label("", labelpad=2, fontsize=9)

legend_elements = [
    Patch(facecolor="#30a2da", label="Train"),
    Patch(facecolor="#fc4f30", label="Test"),
]
ax.legend(handles=legend_elements, loc="upper right", frameon=False, ncol=2, borderaxespad=0.5);

# Plot: states timeline
# ==============================================================================
fig, ax = plt.subplots(figsize=(8, 3))
color_map = {"low": "skyblue", "medium": "gold", "high": "tomato"}
for state, color in color_map.items():
    mask = data["Demand"] == state
    ax.plot(
        data.index[mask], [state] * mask.sum(), '|', color=color, markersize=8, label=state
    )
ax.set_yticks(y_ticks)
ax.set_title("Categorical States Timeline")
ax.set_xlabel("Date")
ax.legend()
plt.tight_layout()
plt.show()

# Create and fit forecaster with rolling window features
# ==============================================================================
forecaster = ForecasterRecursiveClassifier(
                 estimator         = HistGradientBoostingClassifier(random_state=123),
                 lags              = 24,
                 window_features   = RollingFeaturesClassification(stats=['proportion'], window_sizes=72),
                 features_encoding = 'auto', 
             )

forecaster.fit(y=data_train['Demand'], exog=data_train[['Temperature', 'Holiday', 'Hour_of_day']])
forecaster

# Predict
# ==============================================================================
predictions = forecaster.predict(steps=24, exog=data_test[['Temperature', 'Holiday', 'Hour_of_day']])
predictions

# Plot predictions
# ==============================================================================
fig, ax = plt.subplots(figsize=(7, 3))
data_test.loc[predictions.index, 'encoded_value'].plot(ax=ax, label='test')
(
    predictions
    .to_frame()
    .assign(encoded_value=predictions.map(encode_mapping))['encoded_value']
    .plot(ax=ax, label='predictions')
)
y_ticks = list(encode_mapping.values())
y_labels = list(encode_mapping.keys())
ax.set_yticks(y_ticks)
ax.set_yticklabels(y_labels)
ax.set_title("Categorical Time Series (States)")
ax.set_xlabel("Date")
ax.set_ylabel("State")
ax.legend();

# Prediction metrics
# ==============================================================================
error_mse = accuracy_score(
                y_true = data_test.loc[predictions.index, 'Demand'],
                y_pred = predictions
            )
print(f"Test error (accuracy): {error_mse}")
print("Classification Report \n---------------------")
print(
    classification_report(
        data.loc[predictions.index, 'Demand'],
        predictions
    )
)

# Predict probabilities
# ==============================================================================
predictions = forecaster.predict_proba(steps=24, exog=data_test[['Temperature', 'Holiday', 'Hour_of_day']])
predictions

# Native feature importance of the estimator used in the forecaster
# ==============================================================================
forecaster.get_feature_importances()

# Backtesting forecaster
# ==============================================================================
forecaster = ForecasterRecursiveClassifier(
                estimator         = HistGradientBoostingClassifier(random_state=123),
                lags              = 24,
                window_features   = RollingFeaturesClassification(stats=['proportion'], window_sizes=[72]),
                features_encoding = 'auto',
             )

cv = TimeSeriesFold(steps = 24, initial_train_size = end_train)

metric, predictions = backtesting_forecaster(
                          forecaster    = forecaster,
                          y             = data['Demand'],
                          exog          = data[['Temperature', 'Holiday', 'Hour_of_day']],
                          cv            = cv,
                          metric        = 'accuracy_score',
                          n_jobs        = 'auto',
                          verbose       = False,
                          show_progress = True
                      )

display(metric)
display(predictions)

# Summary of classification metrics on test set
# ==============================================================================
print("Classification Report \n---------------------")
print(
    classification_report(
        data.loc[predictions.index, 'Demand'],
        predictions['pred']
    )
)

# Plot predictions
# ==============================================================================
fig, ax = plt.subplots(figsize=(7, 3))
data.loc[predictions.index, 'encoded_value'].plot(ax=ax, label='test')
predictions.assign(encoded_value=predictions['pred'].map(encode_mapping))['encoded_value'].plot(ax=ax, label='predictions')
y_ticks = list(encode_mapping.values())
y_labels = list(encode_mapping.keys())
ax.set_yticks(y_ticks)
ax.set_yticklabels(y_labels)
ax.set_title("Categorical Time Series (States)")
ax.set_xlabel("Date")
ax.set_ylabel("State")
ax.legend();

# Train-validation-test split
# ==============================================================================
end_train = '2014-08-31 23:59:00'
end_val ='2014-10-31 23:59:00'
print(
    f"Train dates      : {data.index.min()} --- {data.loc[:end_train].index.max()}"
    f"  (n={len(data.loc[:end_train])})"
)
print(
    f"Validation dates : {data.loc[end_train:].index.min()} --- {data.loc[:end_val].index.max()}"
    f"  (n={len(data.loc[end_train:end_val])})"
)
print(
    f"Test dates       : {data.loc[end_val:].index.min()} --- {data.index.max()}"
    f"  (n={len(data.loc[end_val:])})"
)
print("")

# Bayesian search hyperparameters and lags with Optuna
# ==============================================================================
estimator = HistGradientBoostingClassifier(
                early_stopping= True,
                random_state=123,
                verbose=0
            )
forecaster = ForecasterRecursiveClassifier(
                estimator         = estimator,
                lags              = 14,  # Placeholder, the value will be overwritten
                window_features   = RollingFeaturesClassification(stats=['proportion'], window_sizes=[72]),
                features_encoding = 'auto',
             )

# Lags used as predictors
lags_grid = [24, (1, 2, 3, 23, 24, 25, 47, 48, 49)]

# Search space
def search_space(trial):
    search_space  = {
        'max_iter' : trial.suggest_int('max_iter', 100, 500, step=100),
        'max_depth'    : trial.suggest_int('max_depth', 3, 10),
        'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.5),
        'l2_regularization': trial.suggest_float('l2_regularization', 0.0, 1.0),
        'lags'         : trial.suggest_categorical('lags', lags_grid)
    } 
    return search_space

# Folds
cv_search = OneStepAheadFold(initial_train_size = end_train)

results, best_trial = bayesian_search_forecaster(
                          forecaster          = forecaster,
                          y                   = data.loc[:end_val, 'Demand'],
                          exog                = data.loc[:end_val, ['Temperature', 'Holiday', 'Hour_of_day']],
                          search_space        = search_space,
                          cv                  = cv_search,
                          metric              = 'accuracy_score',
                          kwargs_create_study = {'direction': 'maximize'},
                          n_trials            = 20,
                          random_state        = 123,
                          return_best         = True,
                          n_jobs              = 'auto',
                          verbose             = False,
                          show_progress       = True,
                      )
results.head(4)

# Best model
# ==============================================================================
best_params = results.loc[0, 'params']
best_lags   = results.loc[0, 'lags']
forecaster = ForecasterRecursiveClassifier(
               estimator       = HistGradientBoostingClassifier(random_state=123, **best_params),
               lags            = best_lags,
               window_features = RollingFeaturesClassification(stats=['proportion'], window_sizes=[72]),
               features_encoding = 'auto',
             )

# Backtest final model on test data
# ==============================================================================
cv = TimeSeriesFold(steps = 24, initial_train_size = end_val)
metric, predictions = backtesting_forecaster(
                          forecaster = forecaster,
                          y          = data['Demand'],
                          exog       = data[['Temperature', 'Holiday', 'Hour_of_day']],
                          cv         = cv,
                          metric     = 'accuracy_score'
                      )
display(metric)
predictions

# Summary of classification metrics on test set
# ==============================================================================
print("Classification Report \n---------------------")
print(
    classification_report(
        data.loc[predictions.index, 'Demand'],
        predictions['pred']
    )
)

import session_info
session_info.show(html=False)