# %% Import of packages:

import numpy as np
import matplotlib.pyplot as plt

# %% Definition of objective function:

def fun_obj(params):
    a = params[0]
    b = params[1]
    kpi_pred_array = a*t_array + b
    e_array = kpi_obs_array - kpi_pred_array
    sspe = sum(e_array*e_array)
    return sspe

# %% Data:

t_array = np.arange(2000, 2009+1)
kpi_obs_array = np.array(
    [75.5, 77.7, 78.7, 80.7, 81.0,
     82.3, 84.2, 84.8, 88.0, 89.9])
    
# %% Initialization:

N_resol_param = 1500

a_lb = 1
a_ub = 2
a_array = np.linspace(a_lb, a_ub, N_resol_param)

b_lb = -4000
b_ub = -2000
b_array = np.linspace(b_lb, b_ub, N_resol_param)

sspe_min = np.inf
a_opt = 0
b_opt = 0

# %% Grid optimization:

for a in a_array:

    for b in b_array:

        # Calc of objective function:
        params = np.array([a, b])
        sspe = fun_obj(params)  

        # Improvement of solution:
        if (sspe < sspe_min):
            sspe_min = sspe
            a_opt = a
            b_opt = b

# %% Presentation of result:

print(f'a_opt = {a_opt:.3e}')
print(f'b_opt = {b_opt:.3e}')
print(f'sspe_min = {sspe_min:.3e}')

# %% Calculation of predicted KPI:

kpi_pred = a_opt*t_array + b_opt

# %% Plotting:

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

plt.plot(t_array, kpi_obs_array,'bo-', label='kpi_real')
plt.plot(t_array, kpi_pred,'ro-', label='kpi_estim')
plt.legend()
plt.grid()
plt.xlabel('t [year]')

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