# Data management
# ==============================================================================
import numpy as np
import pandas as pd
from datetime import datetime
from skforecast.datasets import fetch_dataset
from tqdm.auto import tqdm

# Plots
# ==============================================================================
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from skforecast.plot import set_dark_theme

# Forecasting
# ==============================================================================
import lightgbm
from lightgbm import LGBMRegressor
import skforecast
from skforecast.recursive import ForecasterRecursive, ForecasterRecursiveMultiSeries
from skforecast.model_selection import (
    backtesting_forecaster,
    TimeSeriesFold,
    backtesting_forecaster_multiseries
)
from skforecast.preprocessing import (
    CalendarFeatures,
    reshape_series_long_to_dict,
    reshape_exog_long_to_dict
)

# Feature extraction
# ==============================================================================
import tsfresh
from tsfresh import extract_features, select_features
from tsfresh.feature_extraction.settings import ComprehensiveFCParameters, from_columns

# Clustering
# ==============================================================================
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sktime.clustering.k_means import TimeSeriesKMeans

# Warnings configuration
# ==============================================================================
import warnings
warnings.filterwarnings('once')
warnings.simplefilter('ignore', category=skforecast.exceptions.InputTypeWarning)

color = '\033[1m\033[38;5;208m' 
print(f"{color}Version skforecast: {skforecast.__version__}")
print(f"{color}Version lightgbm: {lightgbm.__version__}")

# Load preprocessed data
# ==============================================================================
data = fetch_dataset('ashrae_daily')
data.head(3)

# Ensure all time series indexes are complete without gaps
# ==============================================================================
data = (
    data
    .groupby('building_id')
    .apply(lambda group: group.asfreq('D', fill_value=np.nan), include_groups=False)
    .reset_index()
    .set_index('timestamp')
)

# Fill missing values of air_temperature and wind_speed using fowrard and backward fill
# ==============================================================================
# Imputation must be done separately for each building
data = data.sort_values(by=['building_id', 'timestamp'])
data['air_temperature'] = data.groupby('building_id')['air_temperature'].ffill().bfill()
data['wind_speed'] = data.groupby('building_id')['wind_speed'].ffill().bfill()
data = data.sort_index()

print(
    f"Range of dates available : {data.index.min()} --- {data.index.max()}  "
    f"(n_days={(data.index.max() - data.index.min()).days})"
)

# Number of buildings and type of buildings based on primary use
# ==============================================================================
n_building = data['building_id'].nunique()
n_type_building = data['primary_use'].nunique()

print(f"Number of buildings: {n_building}")
print(f"Number of building types: {n_type_building}")
display(data.drop_duplicates(subset=['building_id'])['primary_use'].value_counts())

# Types of buildings (primary use) with less than 100 samples are grouped as "Other".
# ==============================================================================
infrequent_categories = (
    data
    .drop_duplicates(subset=['building_id'])['primary_use']
    .value_counts()
    .loc[lambda x: x < 100]
    .index
    .tolist()
)
print("Infrequent categories:")
print("=====================")
print('\n'.join(infrequent_categories))

data['primary_use'] = np.where(
    data['primary_use'].isin(infrequent_categories),
    'Other',
    data['primary_use']
)

# Time series for 1 random selected building per group
# ==============================================================================
set_dark_theme()
fig, axs = plt.subplots(nrows=3, ncols=2, figsize=(8, 5.5), sharex=True, sharey=False)
sample_ids = (
    data
    .groupby('primary_use')['building_id']
    .apply(lambda x: x.sample(1, random_state=333))
    .tolist()
)
axs = axs.flatten()

for i, building_id in enumerate(sample_ids):
    data_sample = data[data['building_id'] == building_id]
    building_type = data_sample['primary_use'].unique()[0]
    data_sample.plot(
        y        = 'meter_reading',
        ax       = axs[i],
        legend   = False,
        title    = f"Building: {building_id}, type: {building_type}",
        fontsize = 8
    )
    axs[i].set_xlabel("")
    axs[i].set_ylabel("")
    # Scientific notation for y axis
    axs[i].ticklabel_format(axis='y', style='sci', scilimits=(0, 0))
    axs[i].title.set_size(9)

fig.suptitle('Energy consumption for 6 random buildings', fontsize=12)
fig.tight_layout()
plt.show()

# Sample 600 buildings
# ==============================================================================
rng = np.random.default_rng(12345)
buildings = data['building_id'].unique()
buildings_selected = rng.choice(
                         buildings,
                         size    = 600,
                         replace = False
                     )

data = data.query("building_id in @buildings_selected")

# Default features and settings created by tsfresh
# ==============================================================================
default_features = ComprehensiveFCParameters()
print("Name of the features extracted by tsfresh:")
print("=========================================")
print('\n'.join(list(default_features)))

# Default configuration for feature "partial_autocorrelation"
# ==============================================================================
default_features['partial_autocorrelation']

# Default configuration for all extracted feature
# ==============================================================================
# from pprint import pprint
# pprint(default_features)

# Feature extraction
# ==============================================================================
ts_features = extract_features(
    timeseries_container  = data[['building_id', 'meter_reading']].reset_index(),
    column_id             = "building_id",
    column_sort           = "timestamp",
    column_value          = "meter_reading",
    default_fc_parameters = default_features,
    impute_function       = tsfresh.utilities.dataframe_functions.impute,
    n_jobs                = 4
)

print("Shape of ts_features:", ts_features.shape)
ts_features.head(3)

# Select relevant features
# ==============================================================================
target = (
    data[['building_id', 'primary_use']]
    .drop_duplicates()
    .set_index('building_id')
    .loc[ts_features.index, :]
    ['primary_use']
)

assert ts_features.index.equals(target.index)

ts_features_selected = select_features(
    X = ts_features,
    y = target,
    fdr_level = 0.001 # Very restrictive filter
)

ts_features_selected.index.name = 'building_id'
print(f"Number of features before selection: {ts_features.shape[1]}")
print(f"Number of features after selection: {ts_features_selected.shape[1]}")
ts_features_selected.head()

# Remove features with missing values
# ==============================================================================
ts_features_selected = ts_features_selected.dropna(axis=1, how='any')
print(f"Number of features after selection and removing missing values: {ts_features_selected.shape[1]}")

# Dictionary with selected features and their configuration
# ==============================================================================
selected_features_info = from_columns(ts_features_selected)
# pprint(selected_features_info['meter_reading'])

# Scale features so all of them have mean 0 and std 1
# ==============================================================================
scaler = StandardScaler().set_output(transform="pandas")
ts_features_selected_scaled = scaler.fit_transform(ts_features_selected)
ts_features_selected_scaled.head(2)

# Plot variance reduction as a function of the number of PCA components
# ==============================================================================
pca = PCA()
pca.fit(ts_features_selected_scaled)
n_components = np.argmax(np.cumsum(pca.explained_variance_ratio_) > 0.85) + 1

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 2.5))
ax.plot(np.cumsum(pca.explained_variance_ratio_))
ax.set_title('Variance explained by PCA components')
ax.set_xlabel('Number of components')
ax.set_ylabel('Cumulative explained variance')
ax.axhline(y=0.85, color='red', linestyle='--')
ax.axvline(x=n_components, color='red', linestyle='--')
ax.text(
    x=n_components+10,
    y=0.4,
    s=f"n_components = {n_components}",
    color='red',
    verticalalignment='center'
)

plt.show();

# PCA with as many components as necessary to explain 85% of the variance
# ==============================================================================
pca = PCA(n_components=0.85)
pca_projections = pca.fit_transform(ts_features_selected_scaled)

# Create a data frame with the projections. Each column is a principal component
# and each row is the id of the building.
pca_projections = pd.DataFrame(
    pca_projections,
    index   = ts_features_selected_scaled.index,
    columns = [f"PC{i}" for i in range(1, pca.n_components_ + 1)]
)

print(f"Number of components selected: {pca.n_components_}")
print(f"Explained variance: {pca.explained_variance_ratio_.sum():.3f}")
pca_projections.head(3)

# Optimal number of clusters
# ==============================================================================
# Identify the optimal number of clusters using the elbow method
range_n_clusters = range(1, 20)
inertias = []
size_of_clusters = []

for n_clusters in range_n_clusters:
    kmeans = KMeans(
                 n_clusters   = n_clusters, 
                 n_init       = 20, 
                 random_state = 963852
             )
    kmeans.fit(pca_projections)
    inertias.append(kmeans.inertia_)
    sizes = [int(i) for i in np.unique(kmeans.labels_, return_counts=True)[1]]
    size_of_clusters.append(sizes)

inertias = pd.Series(inertias, index=range_n_clusters)
dicrease_in_inertia = inertias.pct_change() * -100

fig, axs = plt.subplots(1, 2, figsize=(9, 3.5))
inertias.plot(marker='o', ax=axs[0])
axs[0].xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
axs[0].set_title("Cluster inertia vs number of clusters", fontsize=10)
axs[0].set_xlabel('Number of clusters')
axs[0].set_ylabel('Intra-cluster variance (inertia)')

dicrease_in_inertia.plot(kind='bar', ax=axs[1])
axs[1].set_title("Decrease in inertia", fontsize=10)
axs[1].set_xlabel('Number of clusters')
axs[1].set_ylabel('Decrease in inertia (%)')
fig.tight_layout()
plt.show();

# Distribution of cluster sizes
# ==============================================================================
for n_clusters, sizes in zip(range_n_clusters, size_of_clusters):
    print(f"Size of clusters (n = {n_clusters}): {sizes}")

# Train clustering model with 9 clusters and assign each building to a cluster
# ==============================================================================
kmeans = KMeans(n_clusters=9, n_init=20, random_state=963852)
kmeans.fit(pca_projections)
clusters = kmeans.predict(pca_projections)
clusters = pd.DataFrame({
    'building_id': pca_projections.index,
    'cluster_base_on_features': clusters.astype(str)
})

# Merge cluster minor clusters into a single cluster
threshold = 20
cluster_size = clusters['cluster_base_on_features'].value_counts()
cluster_size = cluster_size[cluster_size < threshold].index.tolist()
clusters['cluster_base_on_features'] = np.where(
    clusters['cluster_base_on_features'].isin(cluster_size),
    'Other',
    clusters['cluster_base_on_features']
)

clusters['cluster_base_on_features'].value_counts()

# Add the predicted cluster to the data frame with the building information
# ==============================================================================
data = pd.merge(
           data.reset_index(), # To avoid losing the index
           clusters,
           on       ='building_id',
           how      ='left',
           validate = 'm:1'
       ).set_index('timestamp')

data.head(3)

# Percentaje of building types per cluster
# ==============================================================================
primary_usage_per_cluster = ( 
    data
    .groupby('cluster_base_on_features')['primary_use']
    .value_counts(normalize=True)
    .unstack()
    .fillna(0)
)

primary_usage_per_cluster = (100 * primary_usage_per_cluster).round(1)
primary_usage_per_cluster

# Convert data to long format with a multiindex: (building_id, timestamp)
# ==============================================================================
data_long = (
    data
    .reset_index()
    .loc[:, ['timestamp', 'building_id', 'meter_reading']]
    .set_index(['building_id', 'timestamp'])
    .sort_index(ascending=True)
)
data_long

# Fit clustering model
# ==============================================================================
model = TimeSeriesKMeans(
            n_clusters   = 4,
            metric       = "dtw",
            max_iter     = 10,
            random_state = 123456
        )

model.fit(data_long)

# Cluster prediction
# ==============================================================================
clusters = model.predict(data_long)
clusters = pd.DataFrame({
               'building_id': data_long.index.get_level_values('building_id').drop_duplicates(),
               'cluster_base_on_dtw': clusters.astype(str),
           })

# Size of each cluster
clusters['cluster_base_on_dtw'].value_counts().sort_index().to_frame(name='Number of series')

# Add the predicted cluster to the data frame with the building information
# ==============================================================================
data = pd.merge(
           data.reset_index(),  # To avoid losing the index
           clusters,
           on       = 'building_id',
           how      = 'left',
           validate = 'm:1'
       ).set_index('timestamp')

data.head(3)

# Save data for modelling to avoid repeating the previous steps
# ==============================================================================
data.to_parquet('data_modelling.parquet')

# Load data for modelling
# ==============================================================================
data = pd.read_parquet('data_modelling.parquet')

# Transform series and exog to dictionaries
# ==============================================================================
series_dict = reshape_series_long_to_dict(
    data      = data.reset_index(),
    series_id = 'building_id',
    index     = 'timestamp',
    values    = 'meter_reading',
    freq      = 'D'
)

exog_dict = reshape_exog_long_to_dict(
    data      = data[['building_id', 'primary_use', 'air_temperature', 'wind_speed']].reset_index(),
    series_id = 'building_id',
    index     = 'timestamp',
    freq      = 'D'
)

# Partition data in train and test
# ==============================================================================
# The train/test split used for modelling is handled internally by TimeSeriesFold
# (see `cv` below), which receives the full, undivided series_dict/exog_dict.
# The splits are created in case they need to be used for additional analysis.
end_train = '2016-07-31 23:59:00'
series_dict_train = {k: v.loc[: end_train,] for k, v in series_dict.items()}
exog_dict_train   = {k: v.loc[: end_train,] for k, v in exog_dict.items()}
series_dict_test  = {k: v.loc[end_train:,] for k, v in series_dict.items()}
exog_dict_test    = {k: v.loc[end_train:,] for k, v in exog_dict.items()}

# Forecaster definitions for single and global models
# ==============================================================================
calendar_features = CalendarFeatures(features = ['day_of_week'],  encoding='cyclical')

params_lgbm_single = {
    'n_estimators': 500,
    'learning_rate': 0.01,
    'max_depth': 4,
    'random_state': 8520,
    'verbose': -1
}
forecaster_single = ForecasterRecursive(
                        estimator = LGBMRegressor(**params_lgbm_single),
                        lags      = 31,
                    )

params_lgbm_global = {
    'n_estimators': 500,
    'learning_rate': 0.01,
    'max_depth': 10,
    'random_state': 8520,
    'verbose': -1
}
forecaster_global = ForecasterRecursiveMultiSeries(
                        estimator          = LGBMRegressor(**params_lgbm_global),
                        lags               = 31,
                        calendar_features  = calendar_features,
                        encoding           = "ordinal_category"
                    )

# Backtesting definition
# ==============================================================================
cv = TimeSeriesFold(
        steps              = 7,
        initial_train_size = end_train,
        refit              = False
     )

# Exogenous variables included in the model
# ==============================================================================
exog_features = ['primary_use', 'air_temperature', 'wind_speed']

# Table of results for all models
# ==============================================================================
table_results = pd.DataFrame(columns=['model', 'mae', 'abs_error', 'bias', 'elapsed_time'])
table_results = table_results.set_index('model')
table_results = table_results.astype({'mae': float, 'abs_error': float, 'bias': float, 'elapsed_time': object})

# Train and test a model for each building
# ==============================================================================
predictions_all_buildings = {}
metrics_all_buildings = {}
errors_all_buildings = {}

# Train and predict for each building
start = datetime.now()

for building in tqdm(data['building_id'].unique(), desc='Modelling buildings'):

    # Get data for the building
    data_building = data[data['building_id'] == building]
    data_building = data_building.asfreq('D').sort_index()

    # Backtesting
    try:
        metric, predictions = backtesting_forecaster(
                                  forecaster    = forecaster_single,
                                  y             = data_building['meter_reading'],
                                  exog          = data_building[exog_features],
                                  cv            = cv,
                                  metric        = 'mean_absolute_error',
                                  show_progress = False
                              )
        predictions_all_buildings[building] = predictions['pred']
        metrics_all_buildings[building] = metric.at[0, 'mean_absolute_error']
        errors_all_buildings[building] = (
            predictions['pred'] - data_building.loc[predictions.index, 'meter_reading']
        )
    except Exception as e:
        print(f"Error modelling building {building}: {e}")

end = datetime.now()

predictions_all_buildings = pd.DataFrame(predictions_all_buildings)
errors_all_buildings = pd.DataFrame(errors_all_buildings)
mean_metric_all_buildings = pd.Series(metrics_all_buildings).mean()
sum_abs_errors_all_buildings = errors_all_buildings.abs().sum(axis=1).sum()
sum_bias_all_buildings = errors_all_buildings.sum().sum()

table_results.loc['One model per building', ['mae', 'abs_error', 'bias', 'elapsed_time']] = [
    mean_metric_all_buildings,
    sum_abs_errors_all_buildings,
    sum_bias_all_buildings,
    end - start
]

print(
    f"\nAverage mean absolute error for all buildings: {mean_metric_all_buildings:.0f}\n"
    f"Sum of absolute errors for all buildings (x 10,000): {sum_abs_errors_all_buildings/10000:.0f}\n"
    f"Bias (x 10,000): {sum_bias_all_buildings/10000:.0f}"
)

# Plot predictions vs real value for 2 random buildings
# ==============================================================================
rng = np.random.default_rng(147)
selected_buildings = rng.choice(data['building_id'].unique(), size=2, replace=False)

fig, axs = plt.subplots(2, 1, figsize=(6, 4), sharex=True)
axs = axs.flatten()

for i, building in enumerate(selected_buildings):
    series_dict_test[building].plot(ax=axs[i], label='test')
    predictions_all_buildings[building].plot(ax=axs[i], label='One model per building')
    axs[i].set_title(f"Building {building}", fontsize=10)
    axs[i].set_xlabel("")
    axs[i].legend()

fig.tight_layout()
plt.show();

# Forecaster multi-series to model all buildings at the same time
# ==============================================================================
start = datetime.now()

metric, predictions = backtesting_forecaster_multiseries(
                          forecaster            = forecaster_global,
                          series                = series_dict,
                          exog                  = exog_dict,
                          cv                    = cv,
                          metric                = 'mean_absolute_error',
                          add_aggregated_metric = False,
                          verbose               = False,
                          show_progress         = True
                      )

end = datetime.now()

# Combine predictions with real values
results = predictions.reset_index(names=['timestamp']).merge(
    data[['building_id', 'meter_reading']].reset_index(names=['timestamp']),
    left_on=['timestamp', 'level'],
    right_on=['timestamp', 'building_id'],
    how='inner'
)
results['error'] = results['pred'] - results['meter_reading']
mean_metric_all_buildings = metric['mean_absolute_error'].mean()
sum_abs_errors_all_buildings = results['error'].abs().sum()
sum_bias_all_buildings = results['error'].sum()
table_results.loc['Global model', ['mae', 'abs_error', 'bias', 'elapsed_time']] = [
    mean_metric_all_buildings,
    sum_abs_errors_all_buildings,
    sum_bias_all_buildings,
    end - start
]

print(
    f"\nAverage mean absolute error for all buildings: {mean_metric_all_buildings:.0f}\n"
    f"Sum of absolute errors for all buildings (x 10,000): {sum_abs_errors_all_buildings/10000:.0f}\n"
    f"Bias (x 10,000): {sum_bias_all_buildings/10000:.0f}"
)

# Add predictions to the already existing plot (not showing the plot yet)
# ==============================================================================
for i, building in enumerate(selected_buildings):
    predictions.query("level==@building")['pred'].plot(ax=axs[i], label='Global model')
    axs[i].legend()

# Forecaster multi-series models for buildings grouped by primary usage
# ==============================================================================
predictions_all_buildings = []
metrics_all_buildings = []
start = datetime.now()

for primary_usage in data['primary_use'].unique():

    print(
        f"Training and testing model for primary usage: {primary_usage} "
        f"(n = {data[data['primary_use'] == primary_usage]['building_id'].nunique()})"
    )

    # Create a subset based on primary use clusters
    building_id_subset = set(data.loc[data['primary_use'] == primary_usage, 'building_id'].unique())
    series_dict_subset = {k: v for k, v in series_dict.items() if k in building_id_subset}
    exog_dict_subset   = {k: v for k, v in exog_dict.items() if k in building_id_subset}

    metric, predictions = backtesting_forecaster_multiseries(
                            forecaster            = forecaster_global,
                            series                = series_dict_subset,
                            exog                  = exog_dict_subset,
                            cv                    = cv,
                            metric                = 'mean_absolute_error',
                            add_aggregated_metric = False,
                            verbose               = False,
                            show_progress         = True,
                            suppress_warnings     = True
                        )
    predictions_all_buildings.append(predictions)
    metrics_all_buildings.append(metric)

end = datetime.now()

predictions_all_buildings = pd.concat(predictions_all_buildings, axis=0)
metrics_all_buildings = pd.concat(metrics_all_buildings, axis=0)
results = predictions_all_buildings.reset_index(names=['timestamp']).merge(
    data[['building_id', 'meter_reading']].reset_index(names=['timestamp']),
    left_on=['timestamp', 'level'],
    right_on=['timestamp', 'building_id'],
    how='inner'
)
results['error'] = results['pred'] - results['meter_reading']

mean_metric_all_buildings = metrics_all_buildings['mean_absolute_error'].mean()
sum_abs_errors_all_buildings = results['error'].abs().sum()
sum_bias_all_buildings = results['error'].sum()

table_results.loc['Global model per primary usage', ['mae', 'abs_error', 'bias', 'elapsed_time']] = [
    mean_metric_all_buildings,
    sum_abs_errors_all_buildings,
    sum_bias_all_buildings,
    end - start
]

print(
    f"\nAverage mean absolute error for all buildings: {mean_metric_all_buildings:.0f}\n"
    f"Sum of absolute errors for all buildings (x 10,000): {sum_abs_errors_all_buildings/10000:.0f}\n"
    f"Bias (x 10,000): {sum_bias_all_buildings/10000:.0f}"
)

# Add predictions to the already existing plot (not showing the plot yet)
# ==============================================================================
for i, building in enumerate(selected_buildings):
    predictions_all_buildings.query("level==@building")['pred'].plot(ax=axs[i], label='Global model per primary usage')
    axs[i].legend()

# Forecaster multi-series models for buildings grouped by time series features
# ==============================================================================
predictions_all_buildings = []
metrics_all_buildings = []
start = datetime.now()

for cluster in data['cluster_base_on_features'].unique():

    print(
        f"Training and testing model for cluster: {cluster} "
        f"(n = {data[data['cluster_base_on_features'] == cluster]['building_id'].nunique()})"
    )

    # Create subset based on DTW clusters
    building_id_subset = set(data.loc[data['cluster_base_on_features'] == cluster, 'building_id'].unique())
    series_dict_subset = {k: v for k, v in series_dict.items() if k in building_id_subset}
    exog_dict_subset   = {k: v for k, v in exog_dict.items() if k in building_id_subset}

    metric, predictions = backtesting_forecaster_multiseries(
                            forecaster            = forecaster_global,
                            series                = series_dict_subset,
                            exog                  = exog_dict_subset,
                            cv                    = cv,
                            metric                = 'mean_absolute_error',
                            add_aggregated_metric = False,
                            verbose               = False,
                            show_progress         = True,
                            suppress_warnings     = True
                        )
    predictions_all_buildings.append(predictions)
    metrics_all_buildings.append(metric)

end = datetime.now()

predictions_all_buildings = pd.concat(predictions_all_buildings, axis=0)
metrics_all_buildings = pd.concat(metrics_all_buildings, axis=0)
results = predictions_all_buildings.reset_index(names=['timestamp']).merge(
    data[['building_id', 'meter_reading']].reset_index(names=['timestamp']),
    left_on=['timestamp', 'level'],
    right_on=['timestamp', 'building_id'],
    how='inner'
)
results['error'] = results['pred'] - results['meter_reading']

mean_metric_all_buildings = metrics_all_buildings['mean_absolute_error'].mean()
sum_abs_errors_all_buildings = results['error'].abs().sum()
sum_bias_all_buildings = results['error'].sum()

table_results.loc['Global model per cluster (features)', ['mae', 'abs_error', 'bias', 'elapsed_time']] = [
    mean_metric_all_buildings,
    sum_abs_errors_all_buildings,
    sum_bias_all_buildings,
    end - start
]

print(
    f"\nAverage mean absolute error for all buildings: {mean_metric_all_buildings:.0f}\n"
    f"Sum of absolute errors for all buildings (x 10,000): {sum_abs_errors_all_buildings/10000:.0f}\n"
    f"Bias (x 10,000): {sum_bias_all_buildings/10000:.0f}"
)

# Add predictions to the already existing plot (not showing the plot yet)
# ==============================================================================
for i, building in enumerate(selected_buildings):
    predictions_all_buildings.query("level==@building")['pred'].plot(ax=axs[i], label='Global model per cluster (features)')
    axs[i].legend()

# Forecaster multi-series models for buildings grouped by DTW
# ==============================================================================
predictions_all_buildings = []
metrics_all_buildings = []
start = datetime.now()

for cluster in data['cluster_base_on_dtw'].unique():

    print(
        f"Training and testing model for cluster: {cluster} "
        f"(n = {data[data['cluster_base_on_dtw'] == cluster]['building_id'].nunique()})"
    )

    # Create subset based on DTW clusters
    building_id_subset = set(data.loc[data['cluster_base_on_dtw'] == cluster, 'building_id'].unique())
    series_dict_subset = {k: v for k, v in series_dict.items() if k in building_id_subset}
    exog_dict_subset   = {k: v for k, v in exog_dict.items() if k in building_id_subset}

    metric, predictions = backtesting_forecaster_multiseries(
                            forecaster            = forecaster_global,
                            series                = series_dict_subset,
                            exog                  = exog_dict_subset,
                            cv                    = cv,
                            metric                = 'mean_absolute_error',
                            add_aggregated_metric = False,
                            verbose               = False,
                            show_progress         = True,
                            suppress_warnings     = True
                        )
    predictions_all_buildings.append(predictions)
    metrics_all_buildings.append(metric)

end = datetime.now()

predictions_all_buildings = pd.concat(predictions_all_buildings, axis=0)
metrics_all_buildings = pd.concat(metrics_all_buildings, axis=0)
results = predictions_all_buildings.reset_index(names=['timestamp']).merge(
    data[['building_id', 'meter_reading']].reset_index(names=['timestamp']),
    left_on=['timestamp', 'level'],
    right_on=['timestamp', 'building_id'],
    how='inner'
)
results['error'] = results['pred'] - results['meter_reading']

mean_metric_all_buildings = metrics_all_buildings['mean_absolute_error'].mean()
sum_abs_errors_all_buildings = results['error'].abs().sum()
sum_bias_all_buildings = results['error'].sum()

table_results.loc['Global model per cluster (DTW)', ['mae', 'abs_error', 'bias', 'elapsed_time']] = [
    mean_metric_all_buildings,
    sum_abs_errors_all_buildings,
    sum_bias_all_buildings,
    end - start
]

print(
    f"\nAverage mean absolute error for all buildings: {mean_metric_all_buildings:.0f}\n"
    f"Sum of absolute errors for all buildings (x 10,000): {sum_abs_errors_all_buildings/10000:.0f}\n"
    f"Bias (x 10,000): {sum_bias_all_buildings/10000:.0f}"
)

# Add predictions to the already existing plot (not showing the plot yet)
# ==============================================================================
for i, building in enumerate(selected_buildings):
    predictions_all_buildings.query("level==@building")['pred'].plot(ax=axs[i], label='Global model per cluster (DTW)')
    axs[i].legend()

# Table of results
# ==============================================================================
def highlight_best(column):
    # Bias is best when closest to zero, not when most negative;
    # mae and abs_error are best when smallest.
    best_idx = column.abs().idxmin() if column.name == 'bias' else column.idxmin()
    return ['background-color: green' if idx == best_idx else '' for idx in column.index]

table_results['elapsed_time'] = table_results['elapsed_time'].astype(str).str[:7]
table_results.style.apply(highlight_best, subset=['mae', 'abs_error', 'bias'], axis=0).format(precision=0)

# Plot predictions vs real value for 2 random buildings for all models
# ==============================================================================
handles, labels = axs[0].get_legend_handles_labels()
fig.legend(handles, labels, loc='upper center', bbox_to_anchor=(0.5, 0.05), ncol=2)
for ax in axs:
    ax.legend().remove()
fig

import session_info
session_info.show(html=False)