# Imports
import math
import gpytorch
import matplotlib.pyplot as plt
import numpy as np
import torch
from botorch import fit_gpytorch_mll
from botorch.acquisition.monte_carlo import qNoisyExpectedImprovement
from botorch.acquisition.objective import ScalarizedPosteriorTransform
from botorch.models.gp_regression import SingleTaskGP
from botorch.models.gpytorch import GPyTorchModel
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.settings import debug as botorch_debug

botorch_debug._set_state(True)

# Five equispaced samples from a simple sine curve over [0,2pi]
# Each sample contains function evaluation and derivative information
#   (I fixed the training data because this consistently causes an error,
#    since different training data sometimes don't cause errors)
def get_fixed_training_data():
    train_X = torch.tensor([[0.        ],
                            [1.57079637],
                            [3.14159274],
                            [4.71238899],
                            [6.28318548]],
                            dtype=torch.float64)
    train_Y_dY = torch.tensor([[ 0.03286297,  0.97644126],
                               [ 0.98860252, -0.06822324],
                               [-0.00680325, -0.90437853],
                               [-0.96307969,  0.01945158],
                               [-0.07107084,  0.98940557]],
                               dtype=torch.float64)
    return train_X, train_Y_dY[:,0].unsqueeze(1), train_Y_dY[:,1].unsqueeze(1), train_Y_dY


# Obtain the value of our acquisition function at x
def get_acqf(initialize_model, get_transformed_training_data, scale_tensor):
    # Obtain training data
    train_X, train_Y = get_transformed_training_data()

    # Get and fit our model
    mll, model = initialize_model(train_X, train_Y)
    fit_gpytorch_mll(mll)

    print("Lengthscale: ", model.covar_module.base_kernel.lengthscale.item())
    print("Noise:       ", model.likelihood.noise.item())
    
    # Transforms the model's outputs to a single value, in this case, a weighted sum
    #   (we want this because we're optimizing over the function values, but
    #    using derivative information with something like KG allows the posterior to
    #    benefit from fantasy points that have derivative information as well)
    scal_transf = ScalarizedPosteriorTransform(weights=scale_tensor)
    
    # Define qNEI acquisition function
    # I got the sampler from here:
    #   https://botorch.org/api/_modules/botorch/acquisition/monte_carlo.html#qNoisyExpectedImprovement
    sampler = SobolQMCNormalSampler(1024) 
    qNEI = qNoisyExpectedImprovement(model,\
                train_X,\
                sampler,\
                posterior_transform=scal_transf)
    
    return qNEI

# Evaluate our acquisition function over the list x
def eval_acqf(x, initialize_model, get_transformed_training_data, scale_tensor):
    qNEI = get_acqf(initialize_model, get_transformed_training_data, scale_tensor)
    print("\n======================")
    print("= Obtained qNEI acqf =")
    print("======================\n")
    return qNEI(x)

# Plot the acquisition function over many x
def plot_acqf(filename, initialize_model, get_transformed_training_data, scale_tensor):
    acqf_X = np.linspace(0,2*math.pi,200)
    acqf_X = torch.tensor(acqf_X, dtype=torch.float64).unsqueeze(-1).unsqueeze(-1)
    acqf_Y = eval_acqf(acqf_X, initialize_model, get_transformed_training_data, scale_tensor)

    acqf_X = acqf_X.detach().numpy().squeeze()
    acqf_Y = acqf_Y.detach().numpy().squeeze()

    plt.plot(acqf_X, acqf_Y, label="acqf", color='green', linestyle="-")

    plt.legend(loc='upper right')
    plt.savefig(f"{filename}.pdf")
    plt.close()

    print(acqf_X)
    print(acqf_Y)

# Plot underlying GP
def plot_gp(filename, initialize_model, get_transformed_training_data):
    # Obtain training data
    train_X, train_Y = get_transformed_training_data()

    # Points we're plotting at
    n = 200
    test_x = torch.linspace(0,2*math.pi,200)

    # Get and fit our model
    mll, model = initialize_model(train_X, train_Y)
    fit_gpytorch_mll(mll)

    train_Y = train_Y.squeeze().unsqueeze(-1)
    print(train_X.shape)
    print(train_Y.shape)

    # Compute the posterior
    posterior = model.posterior(test_x)
    mean = posterior.mean.detach().numpy()
    variance1 = posterior.variance.detach().numpy()
    stddev1 = np.sqrt(variance1)

    # Compute likelihood
    with torch.no_grad(), gpytorch.settings.fast_pred_var():
        variance2 = model.likelihood(posterior.mvn).variance.detach().reshape(n,-1).numpy()
        stddev2 = np.sqrt(variance2)*2

    # Make plot
    plt.plot(test_x, mean[:,0], label="gp mean", color='blue', linestyle="-")
    plt.plot(train_X[:,0], train_Y[:,0], 'k*', label="test points")
    print(mean.shape)
    print(stddev1.shape)
    print(stddev2.shape)
    plt.fill_between(test_x, (mean[:,0] - 2*stddev1[:,0]).squeeze(),\
        (mean[:,0] + 2*stddev1[:,0]).squeeze(), color='blue', label="gp 2*stddev", alpha=0.4, linewidth=0)
    plt.fill_between(test_x, (mean[:,0] - 2*stddev2[:,0]).squeeze(),\
        (mean[:,0] + 2*stddev2[:,0]).squeeze(), color='orange', label="gp 2*stddev likelihood", alpha=0.4, linewidth=0)
    plt.title(filename)

    plt.legend(loc='upper right')
    plt.savefig(f"{filename}.pdf")
    plt.close()