# %% Import of packages:

import matplotlib.pyplot as plt
import numpy as np

# %% Filter function:

def fun_ma_filter(ym_k, ma_array, Nf):
    
    # Updating meas array with new measurement:
    ma_array = np.roll(ma_array,1)
    ma_array[0] = ym_k
    yf_k = np.sum(ma_array)/Nf
    
    return (yf_k, ma_array)

# %% Simulation time settings:

dt = 0.01  # [s]
t_start = 0  # [s]
t_stop = 20  # [s]
N_sim = int((t_stop - t_start)/dt) + 1  # Num time-steps

# %% Preallocation of arrays for plotting:

t_array = np.zeros(N_sim)
yp_array = np.zeros(N_sim)
yn_array = np.zeros(N_sim)
ym_array = np.zeros(N_sim)
yf_array = np.zeros(N_sim)

# %% Params of signals:

P = 10  # [s] Period of sine
a_yn = 0.2  # Ampl of uniformly distributed random noise
samples_yn = 1

# %% Filter param:

Nf = 25  # Filter length
ma_array = np.zeros(Nf)  # Array of measurements

# %% Simulation loop:

for k in range(0, N_sim):

    t_k = k*dt  # Time

    # Signals:
    yp_k = np.sin(2*np.pi*(1/P)*t_k)
    yn_k = np.random.uniform(-a_yn, a_yn, samples_yn)[0]
    ym_k = yp_k + yn_k

    # MA filter:
    (yf_k, ma_array) = fun_ma_filter(ym_k, ma_array, Nf)
    
    # Arrays for plotting:
    t_array[k] = t_k
    yp_array[k] = yp_k
    yn_array[k] = yn_k
    ym_array[k] = ym_k
    yf_array[k] = yf_k
    
# %% Plotting:

plt.close('all')
plt.figure(1, figsize=(12, 9))

plt.plot(t_array, yp_array, 'g', label='yp')
plt.plot(t_array, ym_array, 'b', label='ym')
plt.plot(t_array, yf_array, 'r', label='yf')
plt.legend()
plt.xlabel('t_k [s]')
plt.grid()

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