import statsmodels.api as sm
import numpy as np
import matplotlib.pyplot as plt
import japanize_matplotlib

def create_ARdata(phis=[0.1], N=500, init=1, c=1, sigma=0.3):
    """AR過程のデータを作成する"""
    print(f"==AR({len(phis)})過程の長さ{N}のデータを作成==")
    data = np.zeros(N)
    data[0] = init + np.random.normal(0, sigma)

    for t in range(2, N):
        res = c + np.random.normal(0, sigma)
        for j, phi_j in enumerate(phis):
            res += phi_j * data[t - j - 1]
        data[t] = res
    return data

plt.figure(figsize=(12, 6))
phis = [0.1]
ar1_1 = create_ARdata(phis=phis)
plt.plot(ar1_1)
plt.title(f"係数が{phis[0]}である場合", fontsize=15)
plt.show()

plt.figure(figsize=(12, 6))
phis = [1]
ar1_2 = create_ARdata(phis=phis)
plt.plot(ar1_2)
plt.title(f"係数が{phis[0]}である場合", fontsize=15)
plt.show()

plt.figure(figsize=(12, 6))
phis = [1.04]
ar1_2 = create_ARdata(phis=phis)
plt.plot(ar1_2)
plt.title(f"係数が{phis[0]}である場合", fontsize=15)
plt.show()

plt.figure(figsize=(12, 6))
phis = [0.1, 0.3]
ar2_1 = create_ARdata(phis=phis, N=100)
plt.plot(ar2_1)
plt.title(f"係数が{phis}である場合", fontsize=15)
plt.show()

plt.figure(figsize=(12, 6))
phis = [0.1, -1.11]
ar2_1 = create_ARdata(phis=phis)
plt.plot(ar2_1)
plt.title(f"係数が{phis}である場合", fontsize=15)
plt.show()

from statsmodels.tsa.ar_model import AutoReg

res = AutoReg(ar1_1, lags=1).fit()

out = "AIC: {0:0.3f}, HQIC: {1:0.3f}, BIC: {2:0.3f}"
print(out.format(res.aic, res.hqic, res.bic))

print(res.params)
print(res.sigma2)
res.summary()