# %% 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_samples, dt):

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

    for k in range(0, N_samples):

        # Simulation algorithm (Euler step):
        dS_sim_dt_k = (1/T)*(-S_pred_k +K*(u_array[k]+L))
        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, L) = 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]+L))
        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

# %% Loading experimental data:

data = np.loadtxt("data_dc_motor.txt", float)
dt = 0.02  # Sampling interval
t_obs_array = data[:, 0]
u_array = data[:, 1]
S_tacho_array = data[:, 2]
N_samples = len(t_obs_array)

# Converting from V to krpm:
Kt = 5  # [V/krpm]
S_obs_array = S_tacho_array/Kt

S_init = S_obs_array[0]

# Generating time array starting with 0:
t_array = t_obs_array - t_obs_array[0]

# %% Arrays of candidates of estimated param values:

N_params = 21
K_array = np.linspace(0.05, 0.5, N_params)
T_array = np.linspace(0.1, 1., N_params)
L_array = np.linspace(-0.25, 0.25, 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:
     
        for L in L_array:

            params = (K, T, L)  # Candidates K, T, L
    
            #Calculating objective function (sspe):
            sspe = fun_calc_objfun(params, S_init,
                                   S_obs_array,
                                   u_array,
                                   N_samples, dt)
    
            #Improving the previous optim solution:
            if sspe <=  sspe_optim:
                sspe_optim = sspe
                K_optim = K
                T_optim = T
                L_optim = L
    
# %% Displaying the optimal solution = param estimates:

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

# %% Simulation with estimated params:

params = (K_optim, T_optim, L_optim)
S_sim_array = fun_sim(params, S_init, u_array,
                      N_samples, dt)

# %% Plotting:

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

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

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

plt.show()
plt.savefig('plot_grid_K_T_L_dcmotor_real.pdf')