# %% Import of packages:

import numpy as np
import matplotlib.pyplot as plt
import time

# %% Definition of objective function:

def fun_obj(params):
    a = params[0]
    b = params[1]
    y_pred_array = a*x_array + b
    e_array = y_obs_array - y_pred_array
    sspe = sum(e_array*e_array)
    return sspe

# %% Data:

x_array = np.arange(2000, 2019)
y_obs_array = np.array([54.8, 56.1, 55.0, 55.7, 56.2, 
                        55.4, 55.3, 57.0, 55.6, 53.2, 
                        55.5, 54.6, 54.1, 54.0, 54.1, 
                        54.4, 53.6, 52.7, 52.9])
    
# %% Initialization:

N_resol_param = 1000

a_lb = -0.3
a_ub = 0
a_array = np.linspace(a_lb, a_ub, N_resol_param)

b_lb = 200
b_ub = 500
b_array = np.linspace(b_lb, b_ub, N_resol_param)

sspe_min = np.inf
a_opt = 0
b_opt = 0

# %% Starting timer to get execution time:
tic = time.time()

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

# %% Stopping timer:
toc = time.time()
t_elapsed = toc-tic

# %% Presentation of result:

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

# %% Plotting:

plt.close('all')

y_pred = a_opt*x_array + b_opt

plt.figure(num=1, figsize=(24/2.54, 18/2.54))  #Inches

plt.plot(x_array, y_obs_array,'ro', label='y_obs')
plt.plot(x_array, y_pred,'b-', label='y_pred')
plt.legend()
plt.xlim(2000, 2018)
plt.title('Greenhouse gas emissions in Norway')
plt.xlabel('x [year]')
plt.ylabel('[mill tons CO2-equiv]')
plt.grid()

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