# %% Import of packages:

import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
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 = 100

a_ub = 0
a_lb = -0.3
a_step = (a_ub - a_lb)/(N_resol_param - 1)

b_ub = 500
b_lb = 300
b_step = (b_ub - b_lb)/(N_resol_param - 1)

params_ranges = (slice(a_lb, a_ub, a_step),
                slice(b_lb, b_ub, b_step))

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

# %% Solving the optim problem with optimize.brute():
# finish_setting = None
finish_setting = optimize.fmin
result_optim = optimize.brute(fun_obj, params_ranges,
                              full_output=True,
                              finish=finish_setting)

params_optim = result_optim[0]
sspe_min = result_optim[1]

# %% Optimal parameter values:
a_opt = params_optim[0]
b_opt = params_optim[1]

# %% 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='Klimagassutslipp i Norge pr. år',
           figsize=(24/2.54, 18/2.54)) #  Inches

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

plt.legend(labels=('y_obs', 'y_pred'), loc='upper left')

plt.show()

# plt.savefig('prog_optim_grid_greenhouse_gas.pdf')