import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from statsmodels.tsa.arima_model import ARIMA
from statsmodels.graphics.tsaplots import plot_acf
from statsmodels.graphics.tsaplots import plot_pacf
from statsmodels.tsa.seasonal import seasonal_decompose

# 读取SP500数据
sp500_data = pd.read_csv('SP500.csv')

# 转换日期格式
sp500_data['Date'] = pd.to_datetime(sp500_data['Date'], format='%Y-%m-%d')

# 将日期设置为索引
sp500_data.set_index('Date', inplace=True)

# 将数据按照月份进行重采样
sp500_month = sp500_data.resample('M').last()

# 对数据进行季节性分解
decomposition = seasonal_decompose(sp500_month['Adj Close'], freq=12)

# 绘制季节性分解图
fig, axes = plt.subplots(4, 1, figsize=(10, 8))
axes[0].plot(sp500_month['Adj Close'], label='Original')
axes[0].legend(loc='best')
axes[1].plot(decomposition.trend, label='Trend')
axes[1].legend(loc='best')
axes[2].plot(decomposition.seasonal,label='Seasonality')
axes[2].legend(loc='best')
axes[3].plot(decomposition.resid, label='Residuals')
axes[3].legend(loc='best')
plt.tight_layout()

# 绘制ACF和PACF图
fig, axes = plt.subplots(2, 1, figsize=(10, 6))
plot_acf(sp500_month['Adj Close'], lags=30, ax=axes[0])
plot_pacf(sp500_month['Adj Close'], lags=30, ax=axes[1])
plt.tight_layout()

# 拟合ARIMA模型
model = ARIMA(sp500_month['Adj Close'], order=(2, 1, 2))
results = model.fit()

# 预测未来12个月的数据
forecast = results.forecast(12)

# 绘制预测结果
plt.figure(figsize=(10, 6))
plt.plot(sp500_month['Adj Close'], label='Original')
plt.plot(forecast, label='Forecast')
plt.legend(loc='best')
plt.tight_layout()
plt.show()