from typing import List

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor


def random_forest_model(train_data: pd.DataFrame,
                        forecast_target: str,
                        exog_columns: List[str],
                        future_data: pd.DataFrame,
                        steps: int = 20) -> pd.DataFrame:
    """
    使用随机森林模型根据给定的特征和目标数据进行训练，并预测未来数据。

    Args:
        train_data (pd.DataFrame): 训练数据集。
        forecast_target (str): 训练数据集中的目标列的列名。
        exog_columns (List[str): 训练数据集用于预测的特征列名的列表。
        future_data (pd.DataFrame): 存储未来预测所用的外生变量的数据集。
        steps (int, optional, default=20): 要进行预测的天数。

    Returns:
        pd.DataFrame: 存储预测结果的数据表。
    """
    # 制作输入特征和目标变量
    X = train_data[exog_columns].values
    y = train_data[forecast_target].values
    X_test = future_data[exog_columns].values

    model = RandomForestRegressor(n_estimators=1200,
                                  max_depth=8,
                                  min_samples_split=2,
                                  random_state=0)

    model.fit(X, y)

    pred = model.predict(X_test[-steps:])

    forecast_df = pd.DataFrame(
        pred,
        index=pd.date_range(start=train_data.index.max() +
                            pd.Timedelta(days=1),
                            periods=steps),
        columns=[forecast_target])

    return forecast_df


def forecast_future(data: np.ndarray, steps: int = 20) -> List:
    """
    使用随机森林预测未来的数据。

    Args:
        data (np.ndarray): 已知的用于预测的数据。
        steps (int, optional, default=20): 要进行预测的天数。

    Returns:
        List: 存放预测结果的列表。
    """
    # 制作输入特征和目标变量
    X = data[:-1].reshape(-1, 1)
    y = data[1:]
    X_test = [y[-1]]

    # 创建和训练随机森林模型
    model = RandomForestRegressor(n_estimators=1200,
                                  max_depth=8,
                                  min_samples_split=2,
                                  random_state=0)

    model.fit(X, y)

    # 创建一个列表保存预测结果
    pred = []

    # 迭代预测下一个数据点
    for _ in range(steps):
        y_pred = model.predict(np.array([X_test[-1]]).reshape(-1, 1))
        pred.append(y_pred)

        # 将预测的数据点添加到下一轮的输入
        X_test.append(y_pred)
    return pred


def RF_run(input_data: pd.DataFrame,
        forecast_target: str,
        exog_columns: List[str],
        steps: int = 20) -> pd.DataFrame:
    """
    执行数据读取、预处理、模型训练、预测并绘图等一系列步骤的主函数。

    Args:
        input_data (pd.DataFrame): 存储原始数据的DataFrame。
        forecast_target (str): 需要被预测的目标列名。
        exog_columns (List[str]): 特征列名的列表。
        steps (int, optional, default=20): 需要进行预测的天数。

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

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

    for exog in exog_columns:
        pred = forecast_future(input_data[exog].values, steps=steps)
        df_exog[exog] = pred

    df_processed = random_forest_model(input_data, forecast_target,
                                       exog_columns, df_exog, steps)

    return df_processed