# %% Import of packages:

import matplotlib.pyplot as plt
import numpy as np


# %% Definition of objective function:

def fun_calc_objfun(params, S_init, S_obs_array,
                    u_array, N, dt):

    (K, T) = params
    S_pred_k = S_init
    pe_array = np.zeros(N)

    for k in range(0, N):

        # Simulation algorithm (Euler step):
        dS_sim_dt_k = (1/T)*(-S_pred_k + K*u_array[k])
        S_pred_kp1 = S_pred_k + dt*dS_sim_dt_k

        # Updating prediction error (pe):
        pe_array[k] = S_obs_array[k] - S_pred_k

        # Time shift:
        S_pred_k = S_pred_kp1

    sspe = sum(pe_array*pe_array)
    return sspe


# %% Definition of simulation function:

def fun_sim(params, S_init, u_array, N, dt):

    (K, T) = params
    S_sim_k = S_init
    S_sim_array = np.zeros(N)

    for k in range(0, N):
        # Simulation algorithm (Euler step):
        dS_sim_dt_k = (1/T)*(-S_sim_k + K*u_array[k])
        S_sim_kp1 = S_sim_k + dt*dS_sim_dt_k
        S_sim_array[k] = S_sim_k 
        S_sim_k = S_sim_kp1  # Time shift

    return S_sim_array

# %% Time settings:

t_start = 0  # [s]
t_stop = 7
dt = 0.02
N = int(((t_stop - t_start)/dt)) + 1
t_array = np.linspace(t_start, t_stop, N)

# %% Create input signal:

u_array = np.zeros(N)

for k in range(N):
    t_k = k*dt
    if t_start <= t_k < 1:  u_array[k] = 0 
    elif 1 <= t_k < 3:  u_array[k] = 1 
    elif 3 <= t_k < 5:  u_array[k] = -1 
    else: u_array[k] = 0 

# %% Creating simulated observation data:

K_true = 0.15
T_true = 0.3
params = (K_true, T_true)
S_init = 0

S_sim_array = fun_sim(params, S_init, u_array, N, dt)
S_obs_array = S_sim_array

# %% Arrays of candidates of estimated param values:

N_params = 21
K_array = np.linspace(0.05, 0.55, N_params)
T_array = np.linspace(0.1, 1.1, N_params)

# %% For loop implementing estimation with grid optim:

sspe_optim = np.inf  # Initialization of sspe

for K in K_array:

    for T in T_array:
     
        params = (K, T)  # params with candidate K and T

        #Calculating objective function (sspe):
        sspe = fun_calc_objfun(params, S_init,
                               S_obs_array,
                               u_array, N, dt)

        #Improving the previous optim solution:
        if sspe <=  sspe_optim:
            sspe_optim = sspe
            K_optim = K
            T_optim = T
    
# %% Displaying the optimal solution = param estimates:

print('K_true =', f'{K_true:.4f}')
print('T_true =', f'{T_true:.4f}')
print('K_optim =', f'{K_optim:.4f}')
print('T_optim =', f'{T_optim:.4f}')
print('sspe_optim =', f'{sspe_optim:.4f}')

# %% Plotting:

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

plt.subplot(2, 1, 1)
plt.plot(t_array, S_obs_array, 'b', label='S_sim')
plt.legend()
plt.grid()
plt.xlabel('t [s]')
plt.ylabel('[krpm]')


plt.subplot(2, 1, 2)
plt.plot(t_array, u_array, 'g', label='u')
plt.legend()
plt.grid()
plt.xlabel('t [s]')
plt.ylabel('[V]')

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