# Run command: python3 main_demo_1d.py

import numpy as np
import sklearn
from sklearn.gaussian_process.kernels import RBF
import matplotlib.pyplot as plt
import time

from function import f_demo_1d
from algorithms import GP_UCB


## An example in 1d ##


# Set inputs
xD = np.linspace(0, 1.5*np.pi, 100)
mu_prior = 0
sigma_prior = 1

# Sigma for the noise epsilon ~ N(0, sigma**2)
sigma_noise = 0.5

# Set the kernel
kernel = 1.0 * RBF(0.3, (1e-2, 1e2))

nplot = 12
nstep = 5
n_iter = nplot*nstep
beta = 50
gp_ucb = GP_UCB(f_demo_1d, kernel=kernel, mu=mu_prior, sigma=sigma_prior, sigma_noise=sigma_noise, beta=beta, coords=(xD,))

# Find the solution xstar
outputs = [(x, gp_ucb.f(x)) for x in gp_ucb.input_space]
xstar, f_xstar = max(outputs, key=lambda item:item[1])

# Run algo
cpu_time = time.time()
print('\n niter: ', n_iter)
for i in range(n_iter):
    if i%nstep == 0:
        print('\n ########## Iteration: ', i, ' ##########')

    # \epsilon ~ N(0, sigma_noise^2)
    epsilon = np.random.rand()*sigma_noise
    
    # x_n = argmax mu + beta*sigma
    gp_ucb.argmax_ucb()
    
    # y_n = f(x_n) + epsilon
    gp_ucb.sample_y(epsilon)
    
    # Bayesian update on mu and sigma
    gp_ucb.bayesian_update()
    
    # Display summary
    #gp_ucb.summary()
    
    # Save data for plot
    gp_ucb.save_plot()

# Plot
gp_ucb.plot(rows=3, step=nstep)

print('\n UCB-algorithm terminated')

plt.show()

print('\n Elapsed time: ', time.time() - cpu_time)
print('\n xstar: ', xstar, '\t f_xstar: ', f_xstar)
print('\n x_n: ', gp_ucb.x_n, '\t f(x_n): ', gp_ucb.f(gp_ucb.x_n))
