# 数据处理
# ==============================================================================
import numpy as np
import pandas as pd
from astral.sun import sun
from astral import LocationInfo
from skforecast.datasets import fetch_dataset
from feature_engine.datetime import DatetimeFeatures
from feature_engine.creation import CyclicalFeatures
from feature_engine.timeseries.forecasting import WindowFeatures

# 绘图
# ==============================================================================
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from skforecast.plot import plot_residuals
import plotly.graph_objects as go
import plotly.io as pio
import plotly.offline as poff
pio.templates.default = "seaborn"
poff.init_notebook_mode(connected=True)
plt.style.use('seaborn-v0_8-darkgrid')
plt.rcParams.update({'font.size': 8})

# 建模与预测
# ==============================================================================
import skforecast
import lightgbm
import sklearn
from lightgbm import LGBMRegressor
from sklearn.preprocessing import PolynomialFeatures
from sklearn.feature_selection import RFECV
from skforecast.recursive import ForecasterEquivalentDate, ForecasterRecursive
from skforecast.direct import ForecasterDirect
from skforecast.model_selection import TimeSeriesFold, bayesian_search_forecaster, backtesting_forecaster
from skforecast.feature_selection import select_features
from skforecast.preprocessing import RollingFeatures
from skforecast.plot import calculate_lag_autocorrelation, plot_residuals
from skforecast.metrics import calculate_coverage
import shap


# 警告配置
# ==============================================================================
import warnings
warnings.filterwarnings('once')

color = '\033[1m\033[38;5;208m' 
print(f"{color}Version skforecast: {skforecast.__version__}")
print(f"{color}Version scikit-learn: {sklearn.__version__}")
print(f"{color}Version lightgbm: {lightgbm.__version__}")
print(f"{color}Version pandas: {pd.__version__}")
print(f"{color}Version numpy: {np.__version__}")

# 数据下载
# ==============================================================================
data = fetch_dataset(name='vic_electricity', raw=True)
data.info()

# 数据预处理
# ==============================================================================
data = data.copy()
data['Time'] = pd.to_datetime(data['Time'], format='%Y-%m-%dT%H:%M:%SZ')
data = data.set_index('Time')
data = data.asfreq('30min')
data = data.sort_index()
data.head(2)

# 验证时间索引是否完整
# ==============================================================================
start_date = data.index.min()
end_date = data.index.max()
complete_date_range = pd.date_range(start=start_date, end=end_date, freq=data.index.freq)
is_index_complete = (data.index == complete_date_range).all()
print(f"Index complete: {is_index_complete}")
print(f"Number of rows with missing values: {data.isnull().any(axis=1).mean()}")

# 填补时间索引中的缺口
# ==============================================================================
# data.asfreq(freq='30min', fill_value=np.nan)

# 以 1 小时为间隔进行聚合
# ==============================================================================
# 删除 Date 列，避免聚合时报错。
data = data.drop(columns="Date")
data = (
    data
    .resample(rule="h", closed="left", label="right")
    .agg({
        "Demand": "mean",
        "Temperature": "mean",
        "Holiday": "mean",
    })
)
data

# 划分数据为 train-val-test
# ==============================================================================
data = data.loc['2012-01-01 00:00:00':'2014-12-30 23:00:00', :].copy()
end_train = '2013-12-31 23:59:00'
end_validation = '2014-11-30 23:59:00'
data_train = data.loc[: end_train, :].copy()
data_val   = data.loc[end_train:end_validation, :].copy()
data_test  = data.loc[end_validation:, :].copy()

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

# 创建基线模型：前一天同一小时的值
# ==============================================================================
forecaster = ForecasterEquivalentDate(
                 offset    = pd.DateOffset(days=1),
                 n_offsets = 1
)

# 训练预测器
# ==============================================================================
forecaster.fit(y=data.loc[:end_validation, 'Demand'])
forecaster

# 回测
# ==============================================================================
cv = TimeSeriesFold(
        steps              = 24,
        initial_train_size = len(data.loc[:end_validation]),
        refit              = False
)
metric, predictions = backtesting_forecaster(
                          forecaster = forecaster,
                          y          = data['Demand'],
                          cv         = cv,
                          metric     = 'mean_absolute_error'
)
metric

# 创建预测器
# ==============================================================================
window_features = RollingFeatures(stats=["mean"], window_sizes=24 * 3)
forecaster = ForecasterRecursive(
                 regressor       = LGBMRegressor(random_state=15926, verbose=-1),
                 lags            = 24,
                 window_features = window_features
)

# 训练预测器
# ==============================================================================
forecaster.fit(y=data.loc[:end_validation, 'Demand'])
forecaster

# 回测
# ==============================================================================
metric, predictions = backtesting_forecaster(
                          forecaster    = forecaster,
                          y             = data['Demand'],
                          cv            = cv,
                          metric        = 'mean_absolute_error',
                          verbose       = True, # 设为 False 可避免打印
)

# 绘制预测值与真实值对比图
# ======================================================================================
fig = go.Figure()
trace1 = go.Scatter(x=data_test.index, y=data_test['Demand'], name="test", mode="lines")
trace2 = go.Scatter(x=predictions.index, y=predictions['pred'], name="prediction", mode="lines")
fig.add_trace(trace1)
fig.add_trace(trace2)
fig.update_layout(
    title="Real value vs predicted in test data",
    xaxis_title="Date time",
    yaxis_title="Demand",
    width=800,
    height=400,
    margin=dict(l=20, r=20, t=35, b=20),
    legend=dict(orientation="h", yanchor="top", y=1.01, xanchor="left", x=0)
)
fig.show()

# 回测误差
# ==============================================================================
metric

# 日历特征
# ==============================================================================
features_to_extract = [
    'month',
    'week',
    'day_of_week',
    'hour'
    ]
calendar_transformer = DatetimeFeatures(
    variables           = 'index',
    features_to_extract = features_to_extract,
    drop_original       = True,
    )
calendar_features = calendar_transformer.fit_transform(data)[features_to_extract]

# 日照特征
# ==============================================================================
location = LocationInfo(
    latitude  = -37.8,
    longitude = 144.95,
    timezone  = 'Australia/Melbourne'
    )
sunrise_hour = [
    sun(location.observer, date=date, tzinfo=location.timezone)['sunrise']
    for date in data.index
]
sunset_hour = [
    sun(location.observer, date=date, tzinfo=location.timezone)['sunset']
    for date in data.index
]
sunrise_hour = pd.Series(sunrise_hour, index=data.index).dt.round("h").dt.hour
sunset_hour = pd.Series(sunset_hour, index=data.index).dt.round("h").dt.hour
sun_light_features = pd.DataFrame({
                        'sunrise_hour': sunrise_hour,
                        'sunset_hour': sunset_hour
                     })
sun_light_features['daylight_hours'] = (
    sun_light_features['sunset_hour'] - sun_light_features['sunrise_hour']
    )
sun_light_features["is_daylight"] = np.where(
    (data.index.hour >= sun_light_features["sunrise_hour"])
    & (data.index.hour < sun_light_features["sunset_hour"]),
    1,
    0,
    )

# 节假日特征
# ==============================================================================
holiday_features = data[['Holiday']].astype(int)
holiday_features['holiday_previous_day'] = holiday_features['Holiday'].shift(24)
holiday_features['holiday_next_day'] = holiday_features['Holiday'].shift(-24)

# 温度的滑动窗口
# ==============================================================================
wf_transformer = WindowFeatures(
    variables   = ["Temperature"],
    window      = ["1D", "7D"],
    functions   = ["mean", "max", "min"],
    freq        = "h",
)
temp_features = wf_transformer.fit_transform(data[['Temperature']])


# 合并所有外生变量
# ==============================================================================
assert all(calendar_features.index == sun_light_features.index)
assert all(calendar_features.index == temp_features.index)
assert all(calendar_features.index == holiday_features.index)
exogenous_features = pd.concat([
                         calendar_features,
                         sun_light_features,
                         temp_features,
                         holiday_features
                     ], axis=1)

# 由于滑动平均的创建，序列开头会有缺失值；holiday_next_day 导致结尾也有缺失值。
exogenous_features = exogenous_features.iloc[7 * 24:, :]
exogenous_features = exogenous_features.iloc[:-24, :]
exogenous_features.head(3)

# 日历和日照特征的周期编码
# ==============================================================================
features_to_encode = [
    "month",
    "week",
    "day_of_week",
    "hour",
    "sunrise_hour",
    "sunset_hour",
]
max_values = {
    "month": 12,
    "week": 52,
    "day_of_week": 6,
    "hour": 24,
    "sunrise_hour": 24,
    "sunset_hour": 24,
}
cyclical_encoder = CyclicalFeatures(
    variables     = features_to_encode,
    max_values    = max_values,
    drop_original = False
)

exogenous_features = cyclical_encoder.fit_transform(exogenous_features)
exogenous_features.head(3)

# 外生变量间的交互作用
# ==============================================================================
transformer_poly = PolynomialFeatures(
                        degree           = 2,
                        interaction_only = True,
                        include_bias     = False
                    ).set_output(transform="pandas")
poly_cols = [
    'month_sin', 
    'month_cos',
    'week_sin',
    'week_cos',
    'day_of_week_sin',
    'day_of_week_cos',
    'hour_sin',
    'hour_cos',
    'sunrise_hour_sin',
    'sunrise_hour_cos',
    'sunset_hour_sin',
    'sunset_hour_cos',
    'daylight_hours',
    'is_daylight',
    'holiday_previous_day',
    'holiday_next_day',
    'Temperature_window_1D_mean',
    'Temperature_window_1D_min',
    'Temperature_window_1D_max',
    'Temperature_window_7D_mean',
    'Temperature_window_7D_min',
    'Temperature_window_7D_max',
    'Temperature',
    'Holiday'
    ]
poly_features = transformer_poly.fit_transform(exogenous_features[poly_cols])
poly_features = poly_features.drop(columns=poly_cols)
poly_features.columns = [f"poly_{col}" for col in poly_features.columns]
poly_features.columns = poly_features.columns.str.replace(" ", "__")
assert all(poly_features.index == exogenous_features.index)
exogenous_features = pd.concat([exogenous_features, poly_features], axis=1)
exogenous_features.head(3)

# 选择要纳入模型的外生变量
# ==============================================================================
exog_features = []
# 选择以 _sin 或 _cos 结尾的列
exog_features.extend(exogenous_features.filter(regex='_sin$|_cos$').columns.tolist())
# 选择以 temp_ 开头的列
exog_features.extend(exogenous_features.filter(regex='^Temperature_.*').columns.tolist())
# 选择以 holiday_ 开头的列
exog_features.extend(exogenous_features.filter(regex='^Holiday_.*').columns.tolist())
# 包含原始特征
exog_features.extend(['Temperature', 'Holiday'])

# 合并目标变量和外生变量到同一 DataFrame
# ==============================================================================
data = data[['Demand']].merge(
           exogenous_features[exog_features],
           left_index  = True,
           right_index = True,
           how         = 'inner'  # 仅使用所有变量均有值的日期
)
data = data.astype('float32')

# 划分数据为 train-val-test
data_train = data.loc[: end_train, :].copy()
data_val   = data.loc[end_train:end_validation, :].copy()
data_test  = data.loc[end_validation:, :].copy()

# 回测模型
# ==============================================================================
metric, predictions = backtesting_forecaster(
                          forecaster = forecaster,
                          y          = data['Demand'],
                          exog       = data[exog_features],
                          cv         = cv,
                          metric     = 'mean_absolute_error'
)
display(metric)
predictions.head()

# 带间隔的回测
# ==============================================================================
forecaster = ForecasterRecursive(
                 regressor       = LGBMRegressor(random_state=15926, verbose=-1),
                 lags            = 24,  # 网格搜索时会被替换
                 window_features = window_features
)

# 用作特征的滞后
lags_grid = [24, (1, 2, 3, 23, 24, 25, 47, 48, 49)]

# 回归器超参数搜索空间
def search_space(trial):
    search_space  = {
        'n_estimators' : trial.suggest_int('n_estimators', 300, 1000, step=100),
        'max_depth'    : trial.suggest_int('max_depth', 3, 10),
        'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.5),
        'reg_alpha'    : trial.suggest_float('reg_alpha', 0, 1),
        'reg_lambda'   : trial.suggest_float('reg_lambda', 0, 1),
        'lags'         : trial.suggest_categorical('lags', lags_grid)
    } 
    return search_space

# 训练和验证折
cv_search = TimeSeriesFold(
                steps              = 24,
                initial_train_size = len(data[:end_train]),
                refit              = False,
            )

results_search, frozen_trial = bayesian_search_forecaster(
                                   forecaster   = forecaster,
                                   y            = data.loc[:end_validation, 'Demand'],
                                   exog         = data.loc[:end_validation, exog_features],
                                   cv           = cv_search,
                                   metric       = 'mean_absolute_error',
                                   search_space = search_space,
                                   n_trials     = 10,  # 增大以获得更全面搜索
                                   return_best  = True
)

# 搜索结果
# ==============================================================================
best_params = results_search.at[0, 'params']
best_params = best_params | {'random_state': 15926, 'verbose': -1}
best_lags = results_search.at[0, 'lags']
results_search.head(3)

# 最优模型
# ==============================================================================
forecaster

# 在测试集上回测最终模型
# ==============================================================================
metric, predictions = backtesting_forecaster(
                          forecaster = forecaster,
                          y          = data['Demand'],
                          exog       = data[exog_features],
                          cv         = cv,
                          metric     = 'mean_absolute_error'
)
display(metric)
predictions.head()

# 创建预测器
# ==============================================================================
regressor = LGBMRegressor(
                n_estimators = 100,
                max_depth    = 4,
                random_state = 15926,
                verbose      = -1
)

forecaster = ForecasterRecursive(
                 regressor       = regressor,
                 lags            = best_lags,
                 window_features = window_features
)

# 递归特征消除与交叉验证
# ==============================================================================
warnings.filterwarnings("ignore", message="X does not have valid feature names.*")
selector = RFECV(
    estimator = regressor,
    step      = 1,
    cv        = 3,
)
lags_select, window_features_select, exog_select = select_features(
    forecaster      = forecaster,
    selector        = selector,
    y               = data_train['Demand'],
    exog            = data_train[exog_features],
    select_only     = None,
    force_inclusion = None,
    subsample       = 0.5,  # 加速过程
    random_state    = 123,
    verbose         = True,
)

# 用选定特征创建预测器
# ==============================================================================
forecaster = ForecasterRecursive(
                regressor       = LGBMRegressor(**best_params),
                lags            = lags_select,
                window_features = window_features
)
# 在验证数据上进行回测以获得样本外残差
# ==============================================================================
metric, predictions = backtesting_forecaster(
                            forecaster = forecaster,
                            y          = data['Demand'],
                            exog       = data[exog_select],
                            cv         = cv,
                            metric     = 'mean_absolute_error'
)                
metric

# 创建并训练预测器
# ==============================================================================
forecaster = ForecasterRecursive(
                 regressor       = LGBMRegressor(**best_params),
                 lags            = lags_select,
                 window_features = window_features,
                 binner_kwargs   = {"n_bins": 5}
             )
forecaster.fit(
    y    = data.loc[:end_train, 'Demand'],
    exog = data.loc[:end_train, exog_select],
    store_in_sample_residuals = True
)

# 预测区间
# ==============================================================================
# 由于模型已使用外生变量进行训练，因此在预测时也必须提供外生变量。
predictions = forecaster.predict_interval(
                  exog     = data.loc[end_train:, exog_select],
                  steps    = 24,
                  interval = [5, 95],
                  method  = 'conformal',
              )
predictions.head()

# 在验证数据上进行回测以获得样本外残差
# ==============================================================================
cv = TimeSeriesFold(
        steps              = 24,
        initial_train_size = len(data.loc[:end_train]),
        refit              = False,
)
_, predictions_val = backtesting_forecaster(
                         forecaster = forecaster,
                         y          = data.loc[:end_validation, 'Demand'],
                         exog       = data.loc[:end_validation, exog_select],
                         cv         = cv,
                         metric     = 'mean_absolute_error'
                     )

# 样本外残差分布
# ============================================================================== 
residuals = data.loc[predictions_val.index, 'Demand'] - predictions_val['pred']
print(pd.Series(np.where(residuals < 0, 'negative', 'positive')).value_counts())
plt.rcParams.update({'font.size': 8})
_ = plot_residuals(residuals=residuals, figsize=(7, 4))

# 在预测器中存储样本外残差
# ============================================================================== 
forecaster.set_out_sample_residuals(
    y_true = data.loc[predictions_val.index, 'Demand'],
    y_pred = predictions_val['pred']
)

# 在测试数据中使用样本外残差进行区间预测的回测
# ============================================================================== 
cv = TimeSeriesFold(
        steps              = 24,
        initial_train_size = len(data.loc[:end_validation]),
        refit              = False,
)
metric, predictions = backtesting_forecaster(
                            forecaster              = forecaster,
                            y                       = data['Demand'],
                            exog                    = data[exog_select],
                            cv                      = cv,
                            metric                  = 'mean_absolute_error',
                            interval                = [5, 95],
                            interval_method         = 'conformal',
                            use_in_sample_residuals = False,  # 样本外残差
                            use_binned_residuals    = True,   # 区间根据预测值分组
                       )
predictions.head(5)

# 绘制预测区间与真实值对比图
# ============================================================================== 
fig = go.Figure([
    go.Scatter(
        name='Prediction', x=predictions.index, y=predictions['pred'], mode='lines',
    ),
    go.Scatter(
        name='Real value', x=data_test.index, y=data_test['Demand'], mode='lines',
    ),
    go.Scatter(
        name='Upper Bound', x=predictions.index, y=predictions['upper_bound'],
        mode='lines', marker=dict(color="#444"), line=dict(width=0), showlegend=False
    ),
    go.Scatter(
        name='Lower Bound', x=predictions.index, y=predictions['lower_bound'],
        marker=dict(color="#444"), line=dict(width=0), mode='lines',
        fillcolor='rgba(68, 68, 68, 0.3)', fill='tonexty', showlegend=False
    )
])
fig.update_layout(
    title="Real value vs predicted in test data",
    xaxis_title="Date time",
    yaxis_title="Demand",
    width=800,
    height=400,
    margin=dict(l=20, r=20, t=35, b=20),
    hovermode="x",
    legend=dict(orientation="h", yanchor="top", y=1.1, xanchor="left", x=0.001)
)
fig.show()

# 预测区间覆盖率（测试数据）
# ============================================================================== 
coverage = calculate_coverage(
              y_true       = data.loc[end_validation:, "Demand"],
              lower_bound  = predictions["lower_bound"],
              upper_bound  = predictions["upper_bound"]
           )
area = (predictions['upper_bound'] - predictions['lower_bound']).sum()
print(f"Total area of the interval: {round(area, 2)}")
print(f"Predicted interval coverage: {round(100 * coverage, 2)} %")

# 创建并训练预测器
# ==============================================================================
forecaster = ForecasterRecursive(
                 regressor       = LGBMRegressor(**best_params),
                 lags            = lags_select,
                 window_features = window_features
             )
forecaster.fit(
    y    = data.loc[:end_validation, 'Demand'],
    exog = data.loc[:end_validation, exog_select]
)

# initial_train_size + 12 位置的末尾
# ==============================================================================
data.iloc[:len(data.loc[:end_validation]) + 12].tail(2)

# 预测器
# ==============================================================================
forecaster = ForecasterRecursive(
                 regressor       = LGBMRegressor(**best_params),
                 lags            = lags_select,
                 window_features = window_features
             )

# 带间隔的回测
# ==============================================================================
cv = TimeSeriesFold(
        steps              = 24,
        initial_train_size = len(data.loc[:end_validation]) + 12,
        refit              = False,
        gap                = 12,
)

metric, predictions = backtesting_forecaster(
                          forecaster = forecaster,
                          y          = data['Demand'],
                          exog       = data[exog_select],
                          cv         = cv,
                          metric     = 'mean_absolute_error'
                      )
display(metric)
predictions.head(5)

import session_info
session_info.show(html=False)