from scipy.signal import lfilter
import scipy.stats as stats
import matplotlib
from IPython.core.display import display, HTML, Image

# Choose the filter
c1, c2 = (1.5,-0.90)
# Plot
c_range = np.linspace(-2,2,101)
fig = plt.figure('coeff', figsize=(9,5))
ax = fig.add_subplot(111, aspect='equal', title='AR(2) coefficient space',
                     xlabel='coeff $\phi_1$', ylabel='coeff $\phi_2$')
# Plot the coeff point:
ax.plot(c1, c2, 'r*', ms = 10)

ax.plot(c_range, -c_range**2/4, 'r--', label='$\Delta=0$', color='#FF0000')
ax.plot(c_range, -np.ones_like(c_range), 'b-',  label='stability')
ax.plot(c_range, 1-np.abs(c_range), 'b-',  label='stability')
ax.fill_between(c_range, -np.ones_like(c_range), 1-np.abs(c_range), color='#0077FF', alpha=0.05)
ax.legend(loc='upper right')
ax.set_xlim(-2.5,2.5); ax.set_ylim(-1.5,1.5)
ax.grid(True)

def Var_AR2(c1,c2):
    '''Variance of an AR(2) process
    X_k = c1 X_{k-1} + c2 X_{k-2}  + Z_k
    with Z_k of unit variance
    '''
    return 1/(1-c2**2 - c1**2 * (1+c2)/(1-c2))

def std_AR2(c1,c2):
    '''standard deviation of an AR(2) process'''
    return np.sqrt(Var_AR2(c1,c2))

# Apply the computation to the current coefficients
Vx = Var_AR2(c1,c2)

print('Variance for c1,c2 = (%.2f,%.2f) : Vx = %.2f' % (c1,c2, Vx))
display(HTML('<strong>Standard deviation</strong> sqrt(Vx) = <strong>%.2f</strong> * sigma_eps' % std_AR2(c1,c2)))

N = 100
# Dirac Impulse:
imp = np.zeros(N)
imp[0] = 1
# Impulse response:
imp_res = lfilter(b=[1], a=[1,-c1, -c2], x = imp)
# Plot
plt.figure('coeff', figsize=(10,5))
plt.plot(np.arange(N), imp_res, '-o')
plt.title('AR(2) Impulse Response\n$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1,c2))
plt.xlabel('instant k')
plt.xlim(0,30)
plt.grid(True)

from numpy.polynomial import Polynomial
p_ar = Polynomial([-c2, -c1, 1])
roots = p_ar.roots()
if roots.dtype == np.float:
    print('Real roots: %.2f and %.2f' % tuple(roots) )
else:
    rho = np.abs(roots[0])
    theta = np.abs(np.angle(roots[0]))
    print('Complex roots: %.2f exp(+/- %.2f)' % (rho, theta))
    print('Pseudo-period: %.1f' %(2*np.pi/theta))
    print('Frequency: %.2f (normalized /pi)' % (theta/np.pi))

delta = c1**2 + 4*c2
print('delta = %.2f' % delta)
if delta > 0:
    display(HTML('Delta is <strong> positive </strong>'))
elif delta < 0:
    display(HTML('Delta is <strong> negative </strong>'))
else:
    display(HTML('Delta is <strong> zero </strong>'))

# Check Stability:
if delta > 0:
    if np.abs(c1)<= 2 and c2 <= 1 - np.abs(c1):
        display(HTML('system is <strong> stable </strong>'))
    else:
        display(HTML('system is <strong> unstable </strong>'))

# Check Stability:
if delta < 0:
    if np.abs(c2)<= 1:
        display(HTML('system is <strong> stable </strong>'))
    else:
        display(HTML('system is <strong> unstable </strong>'))

# Information about oscillations:
if delta < 0:
    # Damping factor `r`
    r = np.sqrt(-c2)
    display(HTML('<strong>damping factor</strong> : r = %.3f' % r))
    display(HTML('-> damping <strong>time</strong> (at 10 %%) : T_10 = %.1f' % (np.log(0.1)/np.log(r))))
    # complex angle theta
    theta = np.arccos(c1/(2*r))
    display(HTML('<strong>oscillation frequency</strong> : theta = %.2f (rad) = %.2f (normalized /pi)' % (theta, theta/np.pi)))
    display(HTML('-> oscillation <strong>period</strong> : T = %.1f' % (2*np.pi/theta)))
    

c1_max = 2.5
c1_range = np.linspace(-c1_max,c1_max,1001).reshape(1,-1)
c2_max = 1.5
c2_range = np.linspace(-c2_max,c2_max,801).reshape(-1,1)
# Compute delta on a grid:
delta_grid = c1_range**2 + 4*c2_range
# Compute the magnitude:
magnitude = (delta_grid<0) *  np.sqrt(np.abs(c2_range)) + \
            (delta_grid>=0)* (np.abs(c1_range) + np.sqrt(np.abs(delta_grid)))/2
### Plot:
fig = plt.figure('delta_img', figsize=(15,5))
ax = fig.add_subplot(111, title='Damping factor', aspect='equal',
                     xlabel='coeff $\phi_1$', ylabel='coeff $\phi_2$')
axim = ax.imshow(magnitude, origin='lower', aspect='equal', interpolation='nearest',
                 extent=(-c1_max,c1_max, -c2_max,c2_max), vmax=2)
ax.plot(c_range, -np.ones_like(c_range), 'b-',  label='stability')
ax.plot(c_range, 1-np.abs(c_range), 'b-',  label='stability')
ax.set_xlim(-c1_max, c1_max)
fig.colorbar(axim);

# Compute the cos(theta):
cos_theta_grid = np.where(delta_grid<0, c1_range/(2*np.sqrt(-c2_range)) , np.nan)
theta_grid = np.arccos(cos_theta_grid)
freq_grid = theta_grid/np.pi
### Plot:
fig = plt.figure('delta_img', figsize=(15,5))
ax = fig.add_subplot(111, title='normalized frequency map', aspect='equal',
                     xlabel='coeff $\phi_1$', ylabel='coeff $\phi_2$')
axim = ax.imshow(freq_grid, origin='lower', aspect='equal', interpolation='nearest',
                 extent=(-c1_max,c1_max, -c2_max,c2_max))
ax.plot(c_range, -np.ones_like(c_range), 'b-',  label='stability')
ax.plot(c_range, 1-np.abs(c_range), 'b-',  label='stability')
ax.set_xlim(-c1_max, c1_max)
fig.colorbar(axim);

N0 = 0 # Burn-in instants
N = 2000
# White noise input
noise = np.random.randn(N0+N)
X_traj = lfilter(b=[1], a=[1,-c1, -c2], x = noise)[N0:]
# Normalize with the theoretical standard deviation
X_traj = X_traj/np.sqrt(Vx)
# Plot
fig = plt.figure('stochastic trajectory', figsize=(10,8))
fig.add_subplot(211, title='AR(2) Stochastic Trajectory\n$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1,c2),
                     xlabel='')
plt.plot(np.arange(N), X_traj, '-')
fig.add_subplot(212, title='Stochastic Trajectory : zoom',
                     xlabel='instant k')
plt.plot(np.arange(N), X_traj, '-o')
plt.xlim(N-100,N)
plt.grid(True)

# Histogram plot
fig = plt.figure('histogram', figsize=(15,5))
fig.add_subplot(122, title='Histogram\n$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1,c2),
                     xlabel='$x_k$')
plt.hist(h, bins=50, normed=True)
print('Observed standard deviation, with %d instants : %.2f (1.0 expected)' % (N, h.std()))

# Appendix : Get this result with sympy
from sympy import symbols, simplify, factor, expand, I 
from sympy import exp, cos, sin
c1sym, c2sym, wsym, x = symbols('phi1 phi2 omega C', real=True)
fsym = 1-c1sym*exp(I*wsym) - c2sym*exp(2*I*wsym)
f2sym = fsym *fsym.conjugate() # f2 = |f|^2
f2sym  = simplify(expand(f2sym, complex=True,trig=True))
f2sym  = f2sym.subs({cos(wsym):x, sin(wsym)**2:1-x**2})
1/factor(f2sym, x)

def spectrum_ar2(c1,c2,w, normed=True):
    '''spectrum of an AR(2) process with unit variance input
    http://en.wikipedia.org/wiki/Autoregressive_model#AR.282.29
    '''
    cosw = np.cos(w)
    f = 1/(1 + 2*c2 + c1**2 + c2**2 + 2*c1*(c2-1)*cosw - 4*c2*cosw**2)
    if normed:
        return f/Var_AR2(c1, c2)
    else:
        return f

# Peak frequency
def peak_ar2(c1,c2):
    '''peak frequency of an AR2 spectrum
    (when there is indeed a peak... )'''
    return np.arccos(c1*(c2-1)/(4*c2))
omega_max = peak_ar2(c1,c2)
display(HTML('<strong>Peak frequency</strong> : %.2f (normalized /pi)' % (omega_max/np.pi)))

N_omega = 4000
omega_list = np.linspace(0, np.pi, N_omega)
# Spectral Density

c1a, c2a = (1.5,-0.90)
poly_a = Polynomial([1, -c1a, -c2a])
f_ar2a = 1/np.abs(poly_a(np.exp(1j * omega_list)))**2 /Var_AR2(c1a, c2a)
omega_max_a = peak_ar2(c1a,c2a)
# another set of parameters
c1b, c2b = (1.8,-0.90)
poly_b = Polynomial([1, -c1b, -c2b])
f_ar2b = 1/np.abs(poly_b(np.exp(1j * omega_list)))**2 /Var_AR2(c1b, c2b)
omega_max_b = peak_ar2(c1b,c2b)
# another set of parameters
c1c, c2c = (1.4,-0.80)
poly_c = Polynomial([1, -c1c, -c2c])
f_ar2c = 1/np.abs(poly_c(np.exp(1j * omega_list)))**2 /Var_AR2(c1c, c2c)
omega_max_c = peak_ar2(c1c,c2c)

# Plot
fig = plt.figure('spectrum', figsize=(10,5))
fig.add_subplot(111, title='Spectral density of an AR(2) process (normalized)',
                     xlabel='normalized frequency $\omega/\pi$' )
plt.plot(omega_list/np.pi, f_ar2a, 'b-', label='$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1a,c2a))
plt.plot(omega_max_a/np.pi, spectrum_ar2(c1a,c2a,omega_max_a), 'b*')

plt.plot(omega_list/np.pi, f_ar2b, 'r-', label='$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1b,c2b))
plt.plot(omega_max_b/np.pi, spectrum_ar2(c1b,c2b,omega_max_b), 'r*')

plt.plot(omega_list/np.pi, f_ar2c, 'g-', label='$(\phi_1, \phi_2) = (%.2f, %.2f)$' % (c1c,c2c))
plt.plot(omega_max_c/np.pi, spectrum_ar2(c1c,c2c,omega_max_c), 'g*')

plt.legend();

# Check that the intergral is one
print f_ar2a.sum()/N_omega 
print f_ar2b.sum()/N_omega
print f_ar2c.sum()/N_omega

Hs = 4. # (m)
Tp = 10. # (s)
fs = 10. # (Hz) sampling freq
A = 5/16 * Hs**2 /Tp**4
B = 5/4 / Tp**4
f_list = omega_list/(2*pi)*fs
f_Bret = A/f_list**5 * np.exp(-B/f_list**4)
f_Bret[0] = 0
print('frequency resolution : %.3f Hz' % (f_list[1]-f_list[0]))

# Convert 1/Tp in normalized radians:
omega_p_Bret = 1/Tp/fs*2*np.pi
# Manual correction
omega_p_Bret *= 1.00
print omega_p_Bret
c2_Bret = -0.971
c1_Bret = -4*c2_Bret * np.cos(omega_p_Bret) / (1-c2_Bret)
print(c1_Bret)

# Plot
fig = plt.figure('Wave spectrum', figsize=(10,5))
fig.add_subplot(111, title='Bretschneider Spectral density (Hs= %.1f, Tp=%.1f)' % (Hs, Tp),
                     xlabel='frequency (Hz)' )
plt.plot(f_list, f_Bret * fs/2, 'b-', label='Bretschneider')
plt.plot(f_list, spectrum_ar2(c1_Bret,c2_Bret,omega_list), 'r--',
         label='AR(2) with (%.3f, %.3f)$' % (c1_Bret,c2_Bret))
plt.xlim(0,0.5)
plt.legend(loc='upper right')

# Going back to the autocorrelation

def f2acf(f):
    '''spectrum to acf, using ifft on a symetrized spectrum
    TODO : check that it is the right formula !!!
    '''
    N = len(f)
    # Create the symetric :
    f_symetric = np.concatenate((f, f[::-1]))
    acf = np.fft.ifft(f_symetric)
    # Take the real part
    acf = np.real(acf)
    # return only one half
    return acf[:N]

# Time vector
t = np.arange(N_omega, dtype=float)/fs

fig = plt.figure('Wave acf', figsize=(10,5))
fig.add_subplot(111, title='Theoretical Wave autocorrelation (Hs= %.1f, Tp=%.1f)' % (Hs, Tp),
                xlabel='time (s)')
plt.plot(t, f2acf(spectrum_ar2(c1_Bret,c2_Bret,omega_list) ), 'b-', label='Bretschneider')
plt.plot(t, f2acf(f_Bret*5), 'r--',
         label='AR(2) with (%.3f, %.3f)$' % (c1_Bret,c2_Bret))
plt.xlim(0,50)
plt.ylim(-1,1)
plt.legend(loc='upper right');

Image(filename='interactive_AR2_response.png')

# Numerical computation of the prediction error
pred_std = np.sqrt((imp_res**2).cumsum())
# add zero for h=0
pred_std = np.concatenate(([0], pred_std))
h_range = np.arange(len(pred_std))
# Normalization
pred_std /= std_AR2(c1,c2)

### Plot the prediction error
fig = plt.figure('pred error', figsize=(10,5))
fig.add_subplot(111, title='standard deviation of the prediction error',
                xlabel='prediction step h', ylabel='normalized std.')
plt.plot(h_range, pred_std, '-d')
plt.hlines(1, 0, h_range[-1], colors='b', linestyles='dashed')
plt.xlim(0,20)
plt.ylim(0,1.05);

# Compute the prediction
H = 30 # Prediction maximum horizon
X_pred = np.zeros(H+2)
# Initialize the recursion using the last two observations:
X_pred[0:2] = X_traj[[-2,-1]]
for h in range(2,H+2):
    X_pred[h] = c1*X_pred[h-1] + c2*X_pred[h-2]
# trash the first elements (that is : keep only the last observation)
X_pred = X_pred[1:]

### Plot the prediction ###
fig = plt.figure('AR prediction', figsize=(12,7))
ax = fig.add_subplot(111, title='Stochastic trajectory : Prediction of %d instants' % H,
                     xlabel='prediction horizon h = k-N')
plt.plot(np.arange(N)-N+1, X_traj, 'b-o', label='observations')
plt.plot(np.arange(H+1), X_pred, 'b-o', label='predictions', mfc='#BBBBFF')

# Add some confidence intervals
level1 = 0.95 # 95 % confidence level
color1 = '#FFFF88'
a1 = stats.norm.interval(level1)[1] # = 1.96
plt.fill_between(np.arange(H+1), X_pred-a1*pred_std[:H+1], X_pred + a1*pred_std[:H+1],
                 color=color1, lw=0)
level2 = 0.70 # a tighter interval
color2 = '#FFDD66'
a2 = stats.norm.interval(level2)[1] 
plt.fill_between(np.arange(H+1), X_pred-a2*pred_std[:H+1], X_pred + a2*pred_std[:H+1],
                 color=color2, lw=0)
# Legend
conf_rect1 = matplotlib.patches.Rectangle((0, 0), 1, 1, fc=color1)
conf_rect2 = matplotlib.patches.Rectangle((0, 0), 1, 1, fc=color2)
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles + [conf_rect1, conf_rect2],
          labels + ["%.0f %% conf. interval" % (level1*100), "%.0f %% conf. interval" % (level2*100)],
          loc='upper left')

plt.xlim(-30,H)
plt.ylim(-3.5,3.5)
plt.grid(True)

N_scenar = 100 # Number of sceniaros to generate
X_scenar_list = []
for i in range(N_scenar):
    noise_scenar = np.concatenate((noise, np.random.randn(H)))
    X_scenar = lfilter(b=[1], a=[1,-c1, -c2], x = noise_scenar)[-H-1:]
    X_scenar /= std_AR2(c1,c2)
    X_scenar_list.append(X_scenar)

### Plot the scenarios ###
fig = plt.figure('AR scenarios', figsize=(12,7))
ax = fig.add_subplot(111, title='Stochastic trajectories : %d scenarios along %d instants ' % (N_scenar, H),
                     xlabel='prediction horizon h = k-N')
plt.plot(np.arange(N)-N+1, X_traj, 'b-o', label='observations')
# 95 % confidence interval
plt.fill_between(np.arange(H+1), X_pred-a1*pred_std[:H+1], X_pred + a1*pred_std[:H+1],
                 color='#FFFFAA')
for i,X_scenar in enumerate(X_scenar_list):
    if i==0:
        zorder=2
        lineformat='^-'
        alpha=1
    else:
        zorder=1
        lineformat = '-'
        alpha=0.3
    plt.plot(np.arange(H+1), X_scenar, lineformat, label='scenario %d/%d' % (i+1,N_scenar), 
             color=(0,i/N_scenar,1), alpha=alpha, zorder = zorder)

# Legend
conf_rect1 = matplotlib.patches.Rectangle((0, 0), 1, 1, fc='#FFFFAA', lw=1)
handles, labels = ax.get_legend_handles_labels()
plt.legend(handles[0:3]+[conf_rect1], 
           labels[0:3] + ['%.0f %% conf. interval' % (level1*100)],
           loc='upper left')
plt.xlim(-30,H)
plt.ylim(-3.5,3.5)
plt.grid(True)