"""
Autocovariance of colored noise generated as LP-filtered white noise
Finn Aakre Haugen, TechTeach. finn@techteach.no
2025 06 09
"""

#%% Imports

import numpy as np
import matplotlib.pyplot as plt

#%% Defining EWMA fitler (time constant filter)

def fun_filter(x, xf_prev, tf, ts, k):
    
    Nf = (tf/ts) + 1
    a = 1/Nf

    if k == 0:
        xf = x  # Initially, filter out = filter in
    else:
        xf = (1 - a)*xf_prev + a*x
    
    return xf


#%% Simulating filtered random signal

ts = 0.01  # [s]
t_start = 0  # [s]
t_stop = 5  # [s]
N_data = int((t_stop - t_start)/ts) + 1  # Num time-steps

tf = 0.1;


#%% Input signal generation
x_min = -1
x_max = 1
C = 0

t_array = np.zeros(N_data)
x_array = np.zeros(N_data)
xf_array = np.zeros(N_data)

xf_prev = 0
k = 0
rng = np.random.default_rng(2)  # Seed of random generator

#%% Filtering the input signal

for k in range(0, N_data):
    
    t = k*ts
    
    x = C + rng.uniform(x_min, x_max, 1)[0]
    xf = fun_filter(x, xf_prev, tf, ts, k)
    
    t_array[k] = t
    x_array[k] = x
    xf_array[k] = xf
    
    xf_prev = xf

#%% Raw autocorrelation

Rxx_raw = np.correlate(xf_array, xf_array, mode='full')

#%% Normalize autocorrelation

Rxx = Rxx_raw/Rxx_raw[N_data-1]  # Normalize by zero-lag peak

#%% Extract autocorrelation for a range of lags

N_lags = 50
L_array = np.arange(0, N_lags)
length_Rxx = len(Rxx)
index_middle_Rxx = int((length_Rxx-1)/2)
Rxx_plot = Rxx[index_middle_Rxx:index_middle_Rxx+N_lags]

#%% Plotting

plt.figure(figsize=(8, 6))

plt.subplot(2,1,1)
plt.plot(t_array, x_array, 'b', label='x')
plt.plot(t_array, xf_array, 'r', label='xf')
plt.legend()

plt.subplot(2,1,2)
plt.stem(L_array, Rxx_plot, 'g', label=('R_xx(L)'))
plt.legend()
plt.xlabel("Lag")
plt.grid()

# plt.savefig('autocorr.pdf')
plt.show()