import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.statespace.sarimax import SARIMAX

rng = np.random.default_rng(32)
dates = pd.date_range("2018-01-01", periods=6 * 12, freq="M")
trend = 1.2 * np.arange(len(dates))
seasonal = 15 * np.sin(2 * np.pi * np.arange(len(dates)) / 12)
noise = rng.normal(0, 3, len(dates))
series = pd.Series(200 + trend + seasonal + noise, index=dates)

train = series.iloc[:-12]
test_index = series.index[-12:]

model = SARIMAX(train, order=(1, 1, 1), seasonal_order=(1, 1, 1, 12))
result = model.fit(disp=False)
forecast = result.get_forecast(steps=12)
pred_mean = forecast.predicted_mean
pred_ci = forecast.conf_int()

fig, ax = plt.subplots(figsize=(7.5, 4))
ax.plot(series.index, series, color="#cbd5f5", label="実測値（全体）")
ax.plot(train.index, train, color="#2563eb", linewidth=1.2, label="学習区間")
ax.plot(test_index, pred_mean, color="#f97316", linewidth=1.6, label="SARIMAX 予測")
ax.fill_between(
    test_index,
    pred_ci.iloc[:, 0],
    pred_ci.iloc[:, 1],
    color="#f97316",
    alpha=0.2,
    label="95% 信頼区間",
)
ax.set_title("SARIMAX モデルの予測")
ax.set_xlabel("年月")
ax.set_ylabel("値")
ax.legend()
ax.grid(alpha=0.3)

fig.tight_layout()
fig.savefig("static/images/timeseries/sarimax.svg")