import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_pacf
from statsmodels.tsa.stattools import adfuller
from statsmodels.tsa.arima_model import ARIMA
from arch import arch_model
import seaborn as sns
import warnings

warnings.filterwarnings("ignore")
sns.set()  # 设置默认绘图样式


def Pre_processing(dataset):
    """
    数据预处理
    :param dataset: 数据集
    :return:
    """
    # 把数据转化为时间序列数据
    series = dataset.copy()
    series = pd.Series(series)
    # 计算 日收益率
    returns = 100 * series.pct_change(1).dropna()

    # 绘制 returns 图  根据图 判断是否具有波动性聚集
    plt.figure(figsize=(10, 4))
    plt.plot(returns)
    plt.title('Returns', fontsize=16)
    plt.show()

    # 绘制 returns的平方的 PACF 图来判断阶数
    plot_pacf(returns ** 2)
    plt.show()

    return returns


def ARCH_Model(returns, train_len, p):
    """
    ARCH 模型
    returns:日益率
    train_len：训练集长度
    :param p:
    :return:
    """
    rolling_predictions = []
    test_len = len(returns) - train_len

    for i in range(test_len):
        train = returns[:-(test_len - i)]
        model = arch_model(train, mean="Constant", p=p, vol='ARCH')
        model_fit = model.fit(disp='off')
        pred = model_fit.forecast(horizon=1)
        rolling_predictions.append(np.sqrt(pred.variance.values[-1, :][0]))

    return rolling_predictions


def GARCH_Model(returns, train_len, p, q):
    rolling_predictions = []
    test_len = len(returns) - train_len

    for i in range(test_len):
        train = returns[:-(test_len - i)]
        model = arch_model(train, mean="Constant", vol='GARCH', p=p, q=q)
        model_fit = model.fit(disp='off')
        pred = model_fit.forecast(horizon=1)
        rolling_predictions.append(np.sqrt(pred.variance.values[-1, :][0]))

    return rolling_predictions


def show(test, rolling_predictions,model_name):
    # rolling_predictions = pd.Series(rolling_predictions)
    plt.figure(figsize=(10, 4))
    plt.plot(test, label='True Returns')
    plt.plot(rolling_predictions, label=f'{model_name} Predicted Volatility')
    plt.legend(loc=0)
    plt.show()


if __name__ == '__main__':
    raw_csv_data = pd.read_csv("data/Index2018.csv")
    train_len = int(len(raw_csv_data) * 0.8)
    print(train_len)
    returns = Pre_processing(raw_csv_data['ftse'])
    test = returns.tolist()[train_len:]
    # ARCH模型
    # ARCH_Pre = ARCH_Model(returns, train_len, 5)
    # show(test, ARCH_Pre)

    # GARCH模型

    # model = arch_model(returns, mean="Constant", vol='GARCH', p=5, q=1).fit(disp='off')
    # print(model.summary())

    GARCH_Pre = GARCH_Model(returns, train_len, 5, 1)
    show(test, GARCH_Pre,'GARCH')