"""
Date: 2017-11-23, 2018-02-19, 2023-07-05
Christopher Ting
"""

from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt
import statsmodels.api as sm

n = 5000 
theta, lambda1, alpha = 0.1, 0.6, 0.5
su2 = 3

np.random.seed(20230705)

u = np.random.normal(0, 1, n)
mu, su = np.mean(u), np.std(u, ddof=1)
u -=  mu
u /- su
u *= np.sqrt(su2)

Y = np.zeros(n, dtype=float)
Y[0] = theta + u[0]
for t in range(1,n):
   Y[t] = theta + lambda1 * Y[t-1] + alpha * u[t-1] + u[t]

fig = plt.figure(figsize=(8,12))

ax1 = fig.add_subplot(211)
sm.graphics.tsa.plot_acf(Y, lags=10, zero=False, ax=ax1, color='red')                             
ax1.set_ylim([-0.1, 0.8])
plt.grid(which='major', color='green', alpha=0.4)


ax2 = fig.add_subplot(212)
fig = sm.graphics.tsa.plot_pacf(Y, lags=10, zero=False, ax=ax2, 
                                color = 'darkgreen', method='ywm')
ax2.set_ylim([-0.4, 0.8])
plt.grid(which='major',color='green', alpha=0.4)
plt.show()