"""
模块功能：该模块用于季节性自回归移动平均 (SARIMA) 预测模型的实施。

函数：
    train_SARIMA_model(endog, exo, exog_pred, steps): 使用 SARIMA 模型对时间序列数据进行预测。
    run(input_data, forecast_target, exog_columns, steps): 执行模型训练和数据预测等步骤，并返回预测结果。
"""
import pandas as pd
import pmdarima as pm
import numpy as np

from typing import List, Union


def train_SARIMA_model(endog: Union[np.ndarray, pd.Series],
                       exog: Union[np.ndarray, pd.DataFrame] = None,
                       exog_pred: Union[np.ndarray, pd.DataFrame] = None,
                       steps: int = 20,
                       information_criterion: str = 'aic') -> np.ndarray:
    """
    使用SARIMA模型对时间序列数据进行预测。

    Args:
        endog (Union[np.ndarray, pd.Series]): 用于训练和预测的数据，可以是numpy数组或者pandas序列。
        exog (Union[np.ndarray, pd.DataFrame], optional): 外生变量，用于提高ARIMA模型的预测精度。
        exog_pred (Union[np.ndarray, pd.DataFrame], optional): 预测阶段使用的外生变量，它的列数必须和训练阶段的外生变量一致。
        steps (int, optional, default=20): 预测期长度。
        information_criterion (str, optional, default='aic'): 用于模型选择的信息准则，可以是 'aic' 或 'bic'。

    Returns:
        np.ndarray: 预测结果。
    """
    model = pm.auto_arima(endog,
                          X=exog,
                          seasonal=True,
                          information_criterion=information_criterion)

    pred = model.predict(n_periods=steps, X=exog_pred)
    return pred


def run(input_data: pd.DataFrame,
        forecast_target: str,
        exog_columns: List[str],
        steps: int = 20) -> pd.DataFrame:
    """
    主运行函数，用以读取数据、训练模型、预测数据。

    Args:
        input_data (pd.DataFrame): 输入的时间序列数据。
        forecast_target (str): 预测目标的列名。
        exog_columns (List[str]): 外生变量的列名列表。
        steps (int, optional, default=20): 预测步长

    Returns:
        pd.DataFrame: 预测结果的DataFrame对象。
    """
    # 创建一个未来日期的索引，用于保存预测数据
    future_index = pd.date_range(start=input_data.index.max() +
                                 pd.Timedelta(days=1),
                                 periods=steps)

    # 创建一个用于保存预测外生变量的空数据帧
    df_exog = pd.DataFrame(index=future_index)

    # 循环每个外生变量，使用ARIMA模型进行训练和预测，然后将预测值保存到df_exog中
    for exog in exog_columns:
        pred = train_SARIMA_model(endog=input_data[exog], steps=steps)
        df_exog[exog] = pred

    # 使用ARIMA模型对目标变量进行训练和预测，注意这里将df_exog作为预测阶段的外生变量传入
    pred = train_SARIMA_model(endog=input_data[forecast_target],
                              exog=input_data[exog_columns],
                              exog_pred=df_exog[exog_columns],
                              steps=steps,
                              information_criterion='bic')

    # 根据预测值创建一个新的数据帧，用于保存预测的目标变量
    forecast_df = pd.DataFrame(pred,
                               index=future_index,
                               columns=[forecast_target])

    return forecast_df