# %% Import of packages:

import matplotlib.pyplot as plt
import numpy as np

# %% Filter function:

def fun_timeconstant_filter(ym_k, yf_km1, a):
    
    yf_k = (1 - a)*yf_km1 + a*ym_k
    
    return yf_k

# %% 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:

Tf = 0.5  # [s] Filter time constant
a = dt/(Tf + dt)

# %% Initialization:

yf_km1 = yf_init = 0

# %% 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

    # Recursive MA filter:
    yf_k = fun_timeconstant_filter(ym_k, yf_km1, a)
    
    # 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
    
    # Time shift:
    yf_km1 = 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_timeconst_filter.pdf')
plt.show()