import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import control as ct
from math import sqrt, exp

# First order system
a = 1
c = 1
sys = ct.tf(c, [1, a])

# Create the time vector that we want to use
Tf = 5
T = np.linspace(0, Tf, 1000)
dt = T[1] - T[0]

# Create the basis for a white noise signal
# Note: use sqrt(Q/dt) for desired covariance
Q = np.array([[0.1]])
# V = np.random.normal(0, sqrt(Q[0,0]/dt), T.shape)
V = ct.white_noise(T, Q)

plt.plot(T, V[0])
plt.xlabel('Time [s]')
plt.ylabel('$V$');

# Calculate the sample properties and make sure they match
print("mean(V) [0.0] = ", np.mean(V))
print("cov(V) * dt [%0.3g] = " % Q, np.round(np.cov(V), decimals=3) * dt)

# Response of the first order system
# Scale white noise by sqrt(dt) to account for impulse
T, Y = ct.forced_response(sys, T, V)
plt.plot(T, Y)
plt.xlabel('Time [s]')
plt.ylabel('$Y$');

# Compare static properties to what we expect analytically
def r(tau):
    return c**2 * Q / (2 * a) * exp(-a * abs(tau))
    
print("* mean(Y) [%0.3g] = %0.3g" % (0, np.mean(Y).item()))
print("* cov(Y) [%0.3g] = %0.3g" % (r(0).item(), np.cov(Y).item()))

# Correlation function for the input
# Scale by dt to take time step into account
# r_V = sp.signal.correlate(V, V) * dt / Tf
# tau = sp.signal.correlation_lags(len(V), len(V)) * dt
tau, r_V = ct.correlation(T, V)

plt.plot(tau, r_V, 'r-')
plt.xlabel(r'$\tau$')
plt.ylabel(r'$r_V(\tau)$');

# Correlation function for the output
# r_Y = sp.signal.correlate(Y, Y) * dt / Tf
# tau = sp.signal.correlation_lags(len(Y), len(Y)) * dt
tau, r_Y = ct.correlation(T, Y)
plt.plot(tau, r_Y)
plt.xlabel(r'$\tau$')
plt.ylabel(r'$r_Y(\tau)$')

# Compare to the analytical answer
plt.plot(tau, [r(t)[0, 0] for t in tau], 'k--');

# As a crude approximation, compute the average correlation
r_avg = np.zeros_like(r_Y)
for i in range(100):
    V = ct.white_noise(T, Q)
    _, Y = ct.forced_response(sys, T, V)
    tau, r_Y = ct.correlation(T, Y)
    r_avg = r_avg + r_Y
r_avg = r_avg / i
plt.plot(tau, r_avg)
plt.xlabel(r'$\tau$')
plt.ylabel(r'$r_Y(\tau)$')

# Compare to the analytical answer
plt.plot(tau, [r(t)[0, 0] for t in tau], 'k--');

sigma_z = 1
T = 1
filter = ct.ss([[0, 1], [-1/T**2, -2/T]], [[0], [1]], [[1/T**2, sqrt(3)/T]], 0)

timepts = np.linspace(0, 10, 1000)
V = ct.white_noise(timepts, sigma_z**2)
resp = ct.input_output_response(filter, timepts, V)

plt.plot(resp.time, resp.outputs);

# Compute the correlation function
tau, R = ct.correlation(resp.time, resp.outputs)

# Analytical expression for the correlation function (see Friedland)
def dryden_corrfcn(tau, sigma_z=1, T=1):
    return sigma_z**2 * np.exp(-np.abs(tau)/T) * (1- np.abs(tau)/(2*T))

# Plot the correlation function
fig, axs = plt.subplots(1, 2)
axs[0].plot(tau, R)
axs[0].plot(tau, dryden_corrfcn(tau))
axs[0].set_xlabel(r"$\tau$")
axs[0].set_ylabel(r"$r(\tau)$")
axs[0].set_title("Correlation function")

# Compute the power spectral density
dt = timepts[1] - timepts[0]
S = sp.fft.rfft(R) * dt * 2            # rfft returns omega >= 0 => muliple mag by 2
omega = sp.fft.rfftfreq(R.size, dt)

# Analytical expression for the correlation function (see Friedland)
def dryden_psd(omega, sigma_z=1., T=1.):
    return sigma_z**2 * T * (1 + 3 * (omega * T)**2) / (1 + (omega * T)**2)**2

# Plot the power spectral density
axs[1].loglog(omega[1:], np.abs(S[1:]))
axs[1].loglog(omega[1:], dryden_psd(omega[1:]))
axs[1].set_xlabel(r"$\omega$ [rad/sec]")
axs[1].set_ylabel(r"$S(\omega)$")
axs[1].set_title("Power spectral density")

plt.tight_layout()