# 导入库
import pandas as pd # pandas库
import numpy as np # numpy库
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf # acf和pacf展示库
from statsmodels.tsa.stattools import adfuller # adf检验库
from statsmodels.stats.diagnostic import acorr_ljungbox # 随机性检验库
from statsmodels.tsa.arima_model import ARMA # ARMA库
import matplotlib.pyplot as plt # matplotlib图形展示库
import prettytable # 导入表格库

# 多次用到的表格
def pre_table(table_name, table_rows):
    '''
    :param table_name: 表格名称，字符串列表
    :param table_rows: 表格内容，嵌套列表
    :return: 展示表格对象
    '''
    table = prettytable.PrettyTable() # 创建表格实例
    table.field_names = table_name # 定义表格列名
    for i in table_rows: # 循环读多条数据
    table.add_row(i) # 增加数据
    return table

# 数据平稳处理
def get_best_log(ts, max_log=5, rule1=True, rule2=True):
    '''
    :param ts: 时间序列数据，Series类型
    :param max_log: 最大log处理的次数，int型
    :param rule1: rule1规则布尔值，布尔型
    :param rule2: rule2规则布尔值，布尔型
    :return: 达到平稳处理的最佳次数值和处理后的时间序列
    '''
    if rule1 and rule2: # 如果两个规则同时满足
        return 0, ts # 直接返回0和原始时间序列数据
    else: # 只要有一个规则不满足
        for i in range(1, max_log): # 循环做log处理
            ts = np.log(ts) # log处理
            lbvalue, pvalue1 = acorr_ljungbox(ts, lags=1) # 白噪声检验结果
            adf, pvalue2, usedlag, nobs, critical_values, icbest = adfuller(ts) # ADF检验
            rule_1 = (adf < critical_values['1%'] and adf < critical_values['5%'] and adf < critical_values['10%'] and pvalue1 < 0.01) # 稳定性检验
            rule_2 = (pvalue2 < 0.05) # 白噪声检验
            rule_3 = (i < 5)
            if rule_1 and rule_2 and rule_3: # 如果同时满足条件
                print ('The best log n is: {0}'.format(i)) # 打印输出最佳次数
                return i, ts # 返回最佳次数和处理后的时间序列

# 还原经过平稳处理的数据
def recover_log(ts, log_n):
    '''
    :param ts: 经过log方法平稳处理的时间序列，Series类型
    :param log_n: log方法处理的次数，int型
    :return: 还原后的时间序列
    '''
    for i in range(1, log_n + 1): # 循环多次
        ts = np.exp(ts) # log方法还原
        return ts # 返回时间序列

# 平稳性检验
def adf_val(ts, ts_title, acf_title, pacf_title):
    '''
    :param ts: 时间序列数据，Series类型
    :param ts_title: 时间序列图的标题名称，字符串
    :param acf_title: acf图的标题名称，字符串
    :param pacf_title: pacf图的标题名称，字符串
    :return: adf值、adf的p值、三种状态的检验值
    '''
    plt.figure()
    plt.plot(ts) # 时间序列图
    plt.title(ts_title) # 时间序列标题
    plt.show()
    plot_acf(ts, lags=20, title=acf_title).show() # 自相关检测
    plot_pacf(ts, lags=20, title=pacf_title).show() # 偏相关检测
    adf, pvalue, usedlag, nobs, critical_values, icbest = adfuller(ts) # 平稳性检验
    table_name = ['adf', 'pvalue', 'usedlag', 'nobs', 'critical_values', 'icbest'] # 表格列名列表
    table_rows = [[adf, pvalue, usedlag, nobs, critical_values, icbest]] # 表格行数据，嵌套列表
    adf_table = pre_table(table_name, table_rows) # 获得平稳性展示表格对象
    print ('stochastic score') # 打印标题
    print (adf_table) # 打印展示表格
    return adf, pvalue, critical_values, # 返回adf值、adf的p值、三种状态的检验值

# 白噪声（随机性）检验
def acorr_val(ts):
    '''
    :param ts: 时间序列数据，Series类型
    :return: 白噪声检验的P值和展示数据表格对象
    '''
    lbvalue, pvalue = acorr_ljungbox(ts, lags=1) # 白噪声检验
    table_name = ['lbvalue', 'pvalue'] # 表格列名列表
    table_rows = [[lbvalue, pvalue]] # 表格行数据，嵌套列表
    acorr_ljungbox_table = pre_table(table_name, table_rows)
    print ('stationarity score') # 打印标题
    print (acorr_ljungbox_table) # 打印展示表格
    return pvalue # 返回白噪声检验的P值和展示数据表格对象

# arma最优模型训练
def arma_fit(ts):
    '''
    :param ts: 时间序列数据，Series类型
    :return: 最优状态下的p值、q值、arma模型对象、pdq数据框和展示参数表格对象
    '''
    max_count = int(len(ts) / 10) # 最大循环次数最大定义为记录数的10%
    bic = float('inf') # 初始值为正无穷
    tmp_score = [] # 临时p、q、aic、bic和hqic的值的列表
    for tmp_p in range(max_count + 1): # p循环max_count+1次
        for tmp_q in range(max_count + 1): # q循环max_count+1次
            model = ARMA(ts, order=(tmp_p, tmp_q)) # 创建ARMA模型对象
            try:
                results_ARMA = model.fit(disp=-1, method='css') # ARMA模型训练
            except:
                continue # 遇到报错继续
            finally:
                tmp_aic = results_ARMA.aic # 模型的获得aic
                tmp_bic = results_ARMA.bic # 模型的获得bic
                tmp_hqic = results_ARMA.hqic # 模型的获得hqic
                tmp_score.append([tmp_p, tmp_q, tmp_aic, tmp_bic, tmp_hqic]) # 追加每个模型的训练参数和结果
                if tmp_bic < bic: # 如果模型bic小于最小值，那么获得最优模型ARMA的下列参数：
                    p = tmp_p # 最优模型ARMA的p值
                    q = tmp_q # 最优模型ARMA的q值
                    model_arma = results_ARMA # 最优模型ARMA的模型对象
                    aic = tmp_bic # 最优模型ARMA的aic
                    bic = tmp_bic # 最优模型ARMA的bic
                    hqic = tmp_bic # 最优模型ARMA的hqic
                pdq_metrix = np.array(tmp_score) # 将嵌套列表转换为矩阵
                pdq_pd = pd.DataFrame(pdq_metrix, columns=['p', 'q', 'aic', 'bic', 'hqic']) # 基于矩阵创建数据框
                table_name = ['p', 'q', 'aic', 'bic', 'hqic'] # 表格列名列表
                table_rows = [[p, q, aic, bic, hqic]] # 表格行数据，嵌套列表
                parameter_table = pre_table(table_name, table_rows) # 获得最佳ARMA模型结果展示表格对象
                print ('each p/q traning record') # 打印标题
                print (pdq_pd) # 打印输出每次ARMA拟合结果，包含p、d、q以及对应的AIC、BIC、HQIC
                print ('best p and q') # 打印标题
                print (parameter_table) # 输出最佳ARMA模型结果展示表格对象
                return model_arma # 最优状态下的arma模型对象

# 模型训练和效果评估
def train_test(model_arma, ts, log_n, rule1=True, rule2=True):
    '''
    :param model_arma: 最优ARMA模型对象
    :param ts: 时间序列数据，Series类型
    :param log_n: 平稳性处理的log的次数，int型
    :param rule1: rule1规则布尔值，布尔型
    :param rule2: rule2规则布尔值，布尔型
    :return: 还原后的时间序列
    '''
    train_predict = model_arma.predict() # 得到训练集的预测时间序列
    if not (rule1 and rule2): # 如果两个条件有任意一个不满足
    train_predict = recover_log(train_predict, log_n) # 恢复平稳性处理前的真实时间序列值
    ts = recover_log(ts, log_n) # 时间序列还原处理
    ts_data_new = ts[train_predict.index] # 将原始时间序列数据的长度与预测的周期对齐
    RMSE = np.sqrt(np.sum((train_predict - ts_data_new) ** 2) / ts_data_new.size) # 求RMSE
    # 对比训练集的预测和真实数据
    plt.figure() # 创建画布
    train_predict.plot(label='predicted data', style='--') # 以虚线展示预测数据
    ts_data_new.plot(label='raw data') # 以实线展示原始数据
    plt.legend(loc='best') # 设置图例位置
    plt.title('raw data and predicted data with RMSE of %.2f' % RMSE) # 设置标题
    plt.show() # 展示图像
    return ts # 返回还原后的时间序列

# 预测未来指定时间项的数据
def predict_data(model_arma, ts, log_n, start, end, rule1=True, rule2=True):
    '''
    :param model_arma: 最优ARMA模型对象
    :param ts: 时间序列数据，Series类型
    :param log_n: 平稳性处理的log的次数，int型
    :param start: 要预测数据的开始时间索引
    :param end: 要预测数据的结束时间索引
    :param rule1: rule1规则布尔值，布尔型
    :param rule2: rule2规则布尔值，布尔型
    :return: 无
    '''
    predict_ts = model_arma.predict(start=start, end=end) # 预测未来指定时间项的数据
    print ('-----------predict data----------') # 打印标题
    if not (rule1 and rule2): # 如果两个条件有任意一个不满足
        predict_ts = recover_log(predict_ts, log_n) # 还原数据
    print (predict_ts) # 展示预测数据
    # 展示预测趋势
    plt.figure() # 创建画布
    ts.plot(label='raw time series') # 设置推向标签
    predict_ts.plot(label='predicted data', style='--') # 以虚线展示预测数据
    plt.legend(loc='best') # 设置图例位置
    plt.title('predicted time series') # 设置标题
    plt.show() # 展示图像

# 读取数据
date_parse = lambda dates: pd.datetime.strptime(dates, '%m-%d-%Y') # 创建解析列的功能对象
df = pd.read_table('time_series.txt', delimiter='\t', index_col='date', date_parser=date_parse)
ts_data = df['number'].astype('float32') # 将列转换为float32类型
print ('data summary') # 打印标题
print (ts_data.describe()) # 打印输出时间序列数据概况

  # 原始数据检验 
  adf, pvalue1, critical_values = adf_val(ts_data, 'raw time series', 'raw acf', 'raw pacf') # 稳定性检验 
  pvalue2 = acorr_val(ts_data) # 白噪声检验

# 创建用于区分是否进行平稳性处理的规则 
rule1 = (adf < critical_values['1%'] and adf < critical_values['5%'] and adf < critical_values['10%'] and pvalue1 < 0.01) # 稳定性检验 
rule2 = (pvalue2[0,] < 0.05) # 白噪声检验

# 对时间序列做稳定性处理 
log_n, ts_data = get_best_log(ts_data, max_log=5, rule1=rule1, rule2=rule2)

adf, pvalue1, critical_values = adf_val(ts_data, 'final time series', 'final acf', 'final pacf') # 稳定性检验 
pvalue2 = acorr_val(ts_data) # 白噪声检验

# 训练最佳ARMA模型并输出相关参数和对象 
model_arma = arma_fit(ts_data)

ts_data = train_test(model_arma, ts_data, log_n, rule1=rule1, rule2=rule2)

start = '1991-07-28' # 设置预测开始的时间索引 
end = '1991-08-02' # 设置预测结束的时间索引 
predict_data(model_arma, ts_data, log_n, start, end, rule1=rule1, rule2=rule2) # 预测时间序列数据