#%% 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
    f = sum(e_array*e_array)
    return f

#%% Data

x_array = np.arange(1980, 2023)
y_obs_array = np.array([0.26,0.32,0.14,0.31,0.16,0.12,
                        0.18,0.32,0.39,0.27,0.45,0.40,
                        0.22,0.23,0.32,0.45,0.33,0.46,
                        0.61,0.38,0.39,0.54,0.63,0.62,
                        0.53,0.68,0.64,0.66,0.54,0.66,
                        0.72,0.61,0.65,0.68,0.75,0.90,
                        1.02,0.92,0.85,0.98,1.02,0.85,0.89])
    
#%% Initialization

N_resol_param = 10

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

b_lb = -50
b_ub = 0
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]
f_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 results

print(f'a_opt = {a_opt:.4e}')
print(f'b_opt = {b_opt:.4e}')
print(f'f_min = {f_min:.4e}')
print(f'Elapsed time = {t_elapsed:.2e}')

#%% Prediction with adapted model

y_pred = a_opt*x_array + b_opt

#%% Plotting

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

plt.plot(x_array, y_obs_array,'bo', label='y_obs = dT')
plt.plot(x_array, y_pred,'r-', label='y_pred')
plt.legend()
plt.xlim(1980, 2022)
plt.ylim(0, 1.2)
plt.title('Observed and predicted global temp change')
plt.xlabel('x [year]')
plt.ylabel('[deg C]')
plt.grid()

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