"""
模块功能：
本模块核心功能是实现向量自回归（VAR）预测模型的训练和运行。

函数:
    - convert_timestamp_index(data, to_period): 转换时间序列数据的时间索引至DatetimeIndex或PeriodIndex。
    - train_VAR_model(data, max_lags): 利用输入的时间序列数据训练 VAR 模型。
    - run(input_data, steps): 执行时间索引转换、模型训练和数据预测等步骤，并返回预测结果。
"""
import pandas as pd
import statsmodels.api as sm

from typing import List


def convert_timestamp_index(data: pd.DataFrame,
                            to_period: bool) -> pd.DataFrame:
    """
    根据to_period参数，选择将数据的时间索引转换为DatetimeIndex或PeriodIndex。

    Args:
        data (pd.DataFrame): 输入的数据。
        to_period (bool): 如果为True，则将DatetimeIndex转换为PeriodIndex；
                          如果为False，则将PeriodIndex转换为DatetimeIndex。

    Returns:
        pd.DataFrame: 索引被转换后的数据。
    """
    if to_period:
        data.index = pd.DatetimeIndex(data.index).to_period('D')
    else:
        data.index = data.index.to_timestamp()

    return data


def train_VAR_model(data: pd.DataFrame, max_lags: int = 30):
    """
    利用输入的时间序列数据训练VAR模型，通过比较BIC值确定最优滞后阶数。

    Args:
        data (pd.DataFrame): 用于模型训练的时间序列数据。
        max_lags (int, default=30): 最大滞后阶数，默认为 30。

    Returns:
        VARResultsWrapper: 训练得到的VAR模型。
    """
    model = sm.tsa.VAR(data)
    criteria = []
    lags = range(1, max_lags + 1)

    # 通过比较每个滞后阶数模型的BIC值，选择最优滞后阶数
    for lag in lags:
        result = model.fit(maxlags=lag)
        criteria.append(result.bic)

    # 使用最优滞后阶数再次训练模型
    best_lag = lags[criteria.index(min(criteria))]
    results = model.fit(maxlags=best_lag)

    return results


def run(input_data: pd.DataFrame,
        forecast_target: str,
        _: List[str],
        steps: int = 20) -> pd.DataFrame:
    """
    运行函数，执行一系列步骤，包括索引转换、训练模型、数据预测。

    Args:
        input_data (pd.DataFrame): 输入的DataFrame数据。
        forecast_target (str): 需要被预测的目标变量的列名。
        _ (List[str]): 占位参数，用于保持和其他模型函数的接口一致性。
        steps (int, default=20): 预测步数。

    Returns:
        pd.DataFrame: 预测结果的DataFrame对象。
    """
    # 将DataFrame对象的时间索引转换为PeriodIndex
    input_data = convert_timestamp_index(input_data, to_period=True)

    # 训练 VAR 模型
    model = train_VAR_model(input_data, max_lags=60)

    # 将DataFrame对象的时间索引转回原样
    input_data = convert_timestamp_index(input_data, to_period=False)

    # 利用VAR模型进行预测
    pred = model.forecast(input_data.values[-model.k_ar:], steps=steps)
    forecast_df = pd.DataFrame(
        pred,
        index=pd.date_range(start=input_data.index.max() +
                            pd.Timedelta(days=1),
                            periods=steps),
        columns=input_data.columns)

    return forecast_df[forecast_target]