from botorch.optim.fit import fit_gpytorch_torch
import gpytorch.settings as gpts
from contextlib import ExitStack
from torch.quasirandom import SobolEngine
from botorch.generation import MaxPosteriorSampling
from botorch.models.gp_regression import *
from botorch.acquisition.analytic import ExpectedImprovement, UpperConfidenceBound
from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
from torch.distributions.normal import Normal
# from matplotlib import pyplot as plt
import time
from EQI import *
from botorch import fit_gpytorch_model
from optim import optimize_acqf
from botorch.acquisition import PosteriorMean
from gpytorch.mlls import ExactMarginalLogLikelihood
import torch
from util import *
from gpytorch.utils.errors import NanError
R_MIN = 50  # num of repeat eval
seeds = [86306, 17607, 51777, 23803, 43045, 19164, 39191, 69259, 86249, 5807, 56237, 42487, 70129, 76726, 93956, 75621, 30496, 49393, 43605, 6680, 2115, 96163, 5056, 12211, 84067, 83843, 33774, 90634, 86481, 94223, 42616, 21395, 5051, 58144, 57238, 74791, 58117, 64132, 97179, 99282, 67791, 69208, 40163, 5851, 68108, 37799, 84578, 73151, 33175, 40519, 21316, 71258, 64937, 80464, 40094, 56345, 69588, 48843, 12421, 47344, 78760, 31895, 88052, 60362, 71642, 11528, 29815, 56509, 99495, 54880, 77859, 46352, 85274, 22173, 85447, 77286, 12991, 96846, 21188, 68496, 90822, 7260, 24035, 82151, 79645, 46267, 76647, 56429, 88091, 24814, 84928, 87323, 72565, 93665, 33478, 81243, 23635, 55844, 78662,
         1605, 25845, 47346, 99311, 74315, 90706, 28630, 34160, 7499, 11940, 31478, 96512, 40343, 85606, 67078, 84144, 38922, 12866, 29183, 91355, 25760, 75534, 81355, 21246, 71085, 23032, 66999, 46176, 85786, 42565, 62064, 76739, 25013, 87786, 28194, 89607, 35603, 71801, 87173, 12266, 17293, 65650, 1511, 83286, 21110, 46579, 38832, 22222, 71993, 86379, 63355, 46352, 95875, 70301, 12766, 61858, 11599, 34013, 6674, 11220, 40804, 90178, 35789, 70081, 58971, 22232, 51168, 35291, 15978, 5777, 85703, 14493, 41285, 26409, 51468, 8435, 7439, 81142, 50216, 74957, 73050, 1397, 14446, 88820, 81722, 81970, 65240, 40474, 94314, 24003, 14215, 99223, 98326, 61501, 69135, 86387, 84661, 117, 17959, 96812, 38316]


def fit_model(train_x, train_obj, train_var):
    homo_model = FixedNoiseGP(train_X= train_x, train_Y=train_obj)
    mll = ExactMarginalLogLikelihood(homo_model.likelihood, homo_model)
    fit_gpytorch_model(mll = mll, optimizer=fit_gpytorch_torch, options={'disp': False})
    predicted_y = homo_model.posterior(train_x)
    
    

    model = HeteroskedasticSingleTaskGP(
        train_x, train_obj, train_var, kernel='gaussian').to(train_x)
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_model(
        mll=mll, optimizer=fit_gpytorch_torch, options={'disp': False})
    return model


def OCBA(mk, all_obj):
    train_obj = torch.tensor([torch.mean(x, dim=-1) for x in all_obj])
    train_var = torch.tensor([torch.var(x, dim=-1) for x in all_obj])
    best_value = train_obj.max().item()
    b = torch.argmax(train_obj).item()
    v = train_obj - best_value
    v = torch.cat((v[:b], v[b+1:]))
    v = 1.0 / torch.square(v)
    u = train_obj - best_value
    u = torch.cat((u[:b], u[b+1:]))

    sigma = torch.cat((train_var[:b], train_var[b+1:]))
    u = torch.square(sigma / u)
    k = mk / (train_var[b] * torch.sqrt(torch.sum(v)) + torch.sum(u))
    res = (u * k).int()

    mb = mk - torch.sum(res)
    res = torch.cat((res[:b], mb.unsqueeze(-1), res[b:]))
    return res


def eTSSO(config, fun):
    DIM = config.dim
    FUNCTION = config.function
    NOISE_SE = config.noise_level
    BATCH_SIZE = config.batch_size  # eval in parallel
    device = config.cuda
    NUM_RESTARTS = config.num_restarts  # num_restarts of ac function optimizer
    RAW_SAMPLES = 512  # num of samples for initiallization

    N_TRIALS = config.num_trials  # round of trials
    N_BATCH = config.num_batch  # round of optim for a trial
    MC_SAMPLES = config.num_mc_samples
    KERNEL = config.kernel
    R_MIN = 20
    bounds = torch.tensor([(0, 1) for _ in range(DIM)],
                          device=device, dtype=torch.double).T
    cnt = 0
    for trial in range(1, 3 * N_TRIALS + 1):
        try:
            if cnt >= N_TRIALS:
                break
            print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="\n")
            set_seed(seeds[trial])
            inital_batch, train_x, train_obj, best_value, train_var, all_obj = generate_initial_data(
                n=DIM, fun=fun, device=device, all_value=True)
            model = fit_model(config, train_x, train_obj, train_var)
            all_obj = [all_obj[i] for i in range(all_obj.shape[0])]
            cur = 0
            la_m = R_MIN
            num_batch = N_BATCH - inital_batch

            print(
                f"best_value =  {best_value:>4.10f}, "
                f"time = 0.00, "
                f"num_eval={cur},"
                f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                # f"new_x = {new_x}"
                # f"new_obj = {mean_obj}"
            )
            while cur < num_batch:
                t0 = time.time()
                x_max, fx_max = optimize_acqf(acq_function=PosteriorMean(
                    model), bounds=bounds, q=1, num_restarts=NUM_RESTARTS, raw_samples=2048)
                etsso = ExpectedImprovement(model, best_f=fx_max)
                candidates, _ = optimize_acqf(
                    acq_function=etsso,
                    bounds=bounds,
                    q=BATCH_SIZE,
                    num_restarts=NUM_RESTARTS,
                    raw_samples=RAW_SAMPLES,  # used for intialization heuristic
                    options={"batch_limit": 5, "maxiter": 200},
                )
                # observe new values
                r_min = R_MIN
                new_x = candidates.detach()
                new_obj = fun(new_x)
                cur += 1
                for x in range(r_min-1):
                    cur += 1
                    new_obj = torch.cat((new_obj, fun(new_x)), dim=-1)
                    if cur >= num_batch:
                        break
                var_obj = torch.var(new_obj, dim=-1)
                posterior = model.posterior(new_x)
                var = posterior.variance
                cur_m = la_m * (1 + (var_obj.item()) /
                                (var_obj.item() + var))
                cur_m = cur_m.int().item()
                train_x = torch.cat([train_x, new_x])
                all_obj.append(new_obj)
                if cur_m - r_min > 0:
                    mk = cur_m - r_min
                    num_it = OCBA(mk, all_obj)
                    for i, num in enumerate(num_it):
                        for x in range(num):
                            cur += 1
                            all_obj[i] = torch.cat(
                                (all_obj[i],  fun(new_x)), dim=-1)

                        if cur >= num_batch:
                            break

                train_obj = torch.tensor(
                    [torch.mean(x, dim=-1) for x in all_obj]).unsqueeze(-1)
                train_var = torch.tensor(
                    [torch.var(x, dim=-1) for x in all_obj]).unsqueeze(-1)
                # update progress
                best_value = fun.exact(train_x[train_obj.argmax().item()])
                model = fit_model(config, train_x, train_obj, train_var)
                t1 = time.time()
                print(
                    f"best_value =  {best_value:>4.10f}, "
                    f"time = {t1-t0:>4.2f}, "
                    f"num_eval={cur},"
                    f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                    # f"new_x = {new_x}"
                    # f"new_obj = {mean_obj}"
                )
                la_m = cur_m
            cnt += 1
        except NanError:
            print('error')
            pass


def repeat(config, fun, ac):
    DIM = config.dim
    FUNCTION = config.function
    NOISE_SE = config.noise_level
    BATCH_SIZE = config.batch_size  # eval in parallel
    device = config.cuda
    NUM_RESTARTS = config.num_restarts  # num_restarts of ac function optimizer
    RAW_SAMPLES = 512  # num of samples for initiallization
    R_MIN = 50
    N_TRIALS = config.num_trials  # round of trials
    N_BATCH = config.num_batch  # round of optim for a trial
    MC_SAMPLES = config.num_mc_samples
    bounds = torch.tensor([(0, 1) for _ in range(DIM)],
                          device=device, dtype=torch.double).T
    cnt = 0

    for trial in range(1, 3 * N_TRIALS + 1):
        try:
            if cnt >= N_TRIALS:
                break
            print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="\n")
            best_observed = []
            set_seed(seeds[trial])
            # call helper functions to generate initial training data and initialize model
            inital_batch, train_x, train_obj, best_value, train_var = generate_initial_data(
                n=DIM, fun=fun, device=device)
            model = fit_model(config, train_x, train_obj, train_var)
            print(
                f"best_value =  {best_value:>4.10f}, "
                f"time = 0.00,"
                f"num_eval= 0, "
                f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                # f"new_x = {new_x}"
                # f"new_obj = {mean_obj}"
            )

            best_observed.append(best_value)
            num_batch = N_BATCH - inital_batch
            for iteration in range(1, num_batch // R_MIN + 1):
                torch.set_printoptions(precision=10)
                t0 = time.time()
                ac_fun = ac(model, best_f=best_observed, train_x=train_x)
                candidates, _ = optimize_acqf(
                    acq_function=ac_fun,
                    bounds=bounds,
                    q=BATCH_SIZE,
                    num_restarts=NUM_RESTARTS,
                    raw_samples=RAW_SAMPLES,  # used for intialization heuristic
                    options={"batch_limit": 5, "maxiter": 200},
                )
                # observe new values
                new_x = candidates.detach()
                xs = new_x.expand(R_MIN, -1)
                new_obj = fun(xs)
                var_obj = torch.var(new_obj, dim=-1).unsqueeze(-1)
                mean_obj = torch.mean(new_obj, dim=-1).unsqueeze(-1)
                # update training points
                train_x = torch.cat([train_x, new_x])
                # print(train_obj, mean_obj.unsqueeze(-1))
                train_obj = torch.cat([train_obj, mean_obj.unsqueeze(-1)])
                train_var = torch.cat([train_var, var_obj.unsqueeze(-1)])

                # update progress
                best_value = fun.exact(train_x[train_obj.argmax().item()])
                best_observed.append(best_value)
                model = fit_model(config, train_x, train_obj, train_var)
                t1 = time.time()
                print(
                    f"best_value =  {best_value:>4.10f}, "
                    f"time = {t1-t0:>4.2f}, "
                    f"num_eval= {R_MIN * iteration},"
                    f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                    # f"new_x = {new_x}"
                    # f"new_obj = {mean_obj} "
                )
                # best_observed_all.append(best_observed)
            cnt += 1
        except NanError as e:
            print('error')
            print(e)
            pass


def UCB(config, fun):
    Nor = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
    beta = Nor.icdf(torch.tensor([0.9]))
    return repeat(config=config, fun=fun, ac=lambda model, best_f, train_x:  UpperConfidenceBound(model=model,  beta=beta))


def EQI(config, fun):
    gamma = 0.8  # for repeat
    DIM = config.dim
    FUNCTION = config.function
    NOISE_SE = config.noise_level
    BATCH_SIZE = config.batch_size  # eval in parallel
    device = config.cuda
    NUM_RESTARTS = config.num_restarts  # num_restarts of ac function optimizer
    RAW_SAMPLES = 512  # num of samples for initiallization

    N_TRIALS = config.num_trials  # round of trials
    N_BATCH = config.num_batch  # round of optim for a trial
    MC_SAMPLES = config.num_mc_samples
    R_MIN = 20
    bounds = torch.tensor([(0, 1) for _ in range(DIM)],
                          device=device, dtype=torch.double).T
    cnt = 0
    for trial in range(1, 3 * N_TRIALS + 1):
        try:
            if cnt >= N_TRIALS:
                break
            print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="\n")
            best_observed = []

            # call helper functions to generate initial training data and initialize model
            set_seed(seeds[trial])
            inital_batch, train_x, train_obj, best_value, train_var = generate_initial_data(
                n=DIM, fun=fun, device=device)

            model = fit_model(config, train_x, train_obj, train_var)
            print(
                f"best_value =  {best_value:>4.10f}, "
                f"time = 0.00, "
                f"num_eval= 0,"
                f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                # f"new_x = {new_x}"
                # f"new_obj = {mean_obj}"
            )
            best_observed.append(best_value)

            # run N_BATCH rounds of BayesOpt after the initial random batch
            cur = 0
            la_m = 10
            num_batch = N_BATCH - inital_batch
            print(inital_batch, num_batch, R_MIN)
            print(num_batch // R_MIN + 1)
            iteration = 1
            while iteration <= num_batch // R_MIN + 1:
                iteration += 1
                # torch.set_printoptions(precision=10)
                # print(train_obj, train_var, train_x, sep='\n')
                t0 = time.time()
                # fit the models

                ac_fun = ExpectedQuantileImprovement(
                    model=model, train=train_x, budget=iteration)
                candidates, value = optimize_acqf(
                    acq_function=ac_fun,
                    bounds=bounds,
                    q=BATCH_SIZE,
                    num_restarts=NUM_RESTARTS,
                    raw_samples=RAW_SAMPLES,  # used for intialization heuristic
                    options={"batch_limit": 5, "maxiter": 200},
                )
                # observe new values
                new_x = candidates.detach()
                xs = new_x.expand(R_MIN, -1)
                new_obj = fun(xs)
                var_obj = torch.var(new_obj, dim=-1).unsqueeze(-1)
                mean_obj = torch.mean(new_obj, dim=-1).unsqueeze(-1)
                # update training points
                train_x = torch.cat([train_x, new_x])
                train_obj = torch.cat([train_obj, mean_obj.unsqueeze(-1)])
                train_var = torch.cat([train_var, var_obj.unsqueeze(-1)])

                # update progress
                best_value = fun.exact(train_x[train_obj.argmax().item()])
                best_observed.append(best_value)
                model = fit_model(config, train_x, train_obj, train_var)
                t1 = time.time()
                print(
                    f"best_value =  {best_value:>4.10f}, "
                    f"time = {t1-t0:>4.2f}, "
                    f"num_eval= {iteration * R_MIN},"
                    f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                    # f"new_x = {new_x}"
                    # f"new_obj = {mean_obj} "
                )
                ac_fun = ExpectedQuantileImprovement(
                    model=model, train=train_x, budget=num_batch // R_MIN + 1 - iteration)
                # print(ac_fun(train_x[-1].unsqueeze(0)), gamma * value)
                num_repeat = 0
                while ac_fun(train_x[-1].unsqueeze(0)) > gamma * value and num_repeat <= 10:
                    num_repeat += 1
                    iteration += 1
                    # print(ac_fun(train_x[-1].unsqueeze(0)), gamma * value)
                    print('repeat')
                    t0 = time.time()
                    xs = new_x.expand(R_MIN, -1)
                    new_obj = torch.cat((new_obj, fun(xs)), dim=-1)
                    var_obj = torch.var(new_obj, dim=-1)
                    mean_obj = torch.mean(new_obj, dim=-1)
                    train_obj[-1] = mean_obj.unsqueeze(-1)
                    train_var[-1] = var_obj.unsqueeze(-1)
                    model = fit_model(config, train_x, train_obj, train_var)
                    t1 = time.time()
                    print(
                        f"best_value =  {best_value:>4.10f}, "
                        f"time = {t1-t0:>4.2f}, "
                        f"num_eval= {iteration * R_MIN},"
                        f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                        # f"new_x = {new_x}"
                        # f"new_obj = {mean_obj} "
                    )
                    ac_fun = ExpectedQuantileImprovement(
                        model=model, train=train_x)
                    if iteration >= num_batch // R_MIN + 1:
                        break
            cnt += 1
        except NanError:
            print('error')
            pass


def MES(config, fun):
    DIM = config.dim
    device = config.cuda
    bounds = torch.tensor([(0, 1) for _ in range(DIM)],
                          device=device, dtype=torch.double).T
    candi = torch.rand(2048, DIM, device=device, dtype=torch.double) * \
        (bounds[1] - bounds[0]) + bounds[0]
    return repeat(config=config, fun=fun, ac=lambda model, best_f, train_x: qMaxValueEntropy(model=model, candidate_set=candi))


def KG(config, fun):
    return repeat(config=config, fun=fun, ac=lambda model, best_f, train_x: qKnowledgeGradient(model=model, num_fantasies=32))


def thompson(
    X,
    Y,
    batch_size,
    model, device,
    sampler="cholesky",  # "cholesky", "ciq", "rff",
):
    n_candidates = 5000
    assert torch.all(torch.isfinite(Y))
    X = X - X.min()
    X = X / X.max()
    sobol = SobolEngine(X.shape[-1], scramble=True)
    X_cand = sobol.draw(n_candidates).to(dtype=torch.double, device=device)
    # Thompson sample
    with ExitStack() as es:
        if sampler == "cholesky":
            es.enter_context(gpts.max_cholesky_size(float("inf")))
        elif sampler == "ciq":
            es.enter_context(gpts.fast_computations(
                covar_root_decomposition=True))
            es.enter_context(gpts.max_cholesky_size(0))
            es.enter_context(gpts.ciq_samples(True))
            # Controls accuracy and runtime
            es.enter_context(gpts.minres_tolerance(2e-3))
            es.enter_context(gpts.num_contour_quadrature(15))
        elif sampler == "lanczos":
            es.enter_context(gpts.fast_computations(
                covar_root_decomposition=True))
            es.enter_context(gpts.max_cholesky_size(0))
            es.enter_context(gpts.ciq_samples(False))
        elif sampler == "rff":
            es.enter_context(gpts.fast_computations(
                covar_root_decomposition=True))

        thompson_sampling = MaxPosteriorSampling(
            model=model, replacement=False)
        X_next = thompson_sampling(
            X_cand, num_samples=batch_size, observation_noise=True)

    return X_next


def TS(config, fun):

    DIM = config.dim
    FUNCTION = config.function
    NOISE_SE = config.noise_level
    BATCH_SIZE = config.batch_size  # eval in parallel
    device = config.cuda
    NUM_RESTARTS = config.num_restarts  # num_restarts of ac function optimizer
    RAW_SAMPLES = 512  # num of samples for initiallization
    R_MIN = 50
    N_TRIALS = config.num_trials  # round of trials
    N_BATCH = config.num_batch  # round of optim for a trial
    MC_SAMPLES = config.num_mc_samples
    cnt = 0
    for trial in range(1, 3 * N_TRIALS + 1):
        try:
            if cnt >= N_TRIALS:
                break
            print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="\n")
            best_observed = []
            # call helper functions to generate initial training data and initialize model
            set_seed(seeds[trial])

            inital_batch, train_x, train_obj, best_value, train_var = generate_initial_data(
                n=DIM, fun=fun, device=device)
            model = fit_model(config, train_x, train_obj, train_var)
            print(
                f"best_value =  {best_value:>4.10f}, "
                f"time = 0, "
                f"num_eval= 0,"
                f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                # f"new_x = {new_x}"
                # f"new_obj = {mean_obj} "
            )

            best_observed.append(best_value)

            # run N_BATCH rounds of BayesOpt after the initial random batch
            cur = 0
            la_m = 10
            num_batch = N_BATCH - inital_batch
            for iteration in range(1, num_batch//R_MIN + 1):
                t0 = time.time()
                candidates = thompson(
                    X=train_x, Y=train_obj, model=model, batch_size=BATCH_SIZE, device=device)
                new_x = candidates.detach()

                new_obj = fun(new_x).unsqueeze(-1)

                # for x in range(R_MIN-1):
                #     new_obj = torch.cat(
                #         (new_obj, fun(new_x).unsqueeze(-1)), dim=-1)
                # var_obj = torch.var(new_obj, dim=-1)
                # mean_obj = torch.mean(new_obj, dim=-1)
                xs = new_x.expand(R_MIN, -1)
                new_obj = fun(xs)
                var_obj = torch.var(new_obj, dim=-1).unsqueeze(-1)
                mean_obj = torch.mean(new_obj, dim=-1).unsqueeze(-1)
                # update training points
                train_x = torch.cat([train_x, new_x])
                train_obj = torch.cat([train_obj, mean_obj.unsqueeze(-1)])
                train_var = torch.cat([train_var, var_obj.unsqueeze(-1)])
                # update progress
                best_value = fun.exact(train_x[train_obj.argmax().item()])
                best_observed.append(best_value)
                model = fit_model(config, train_x, train_obj, train_var)
                t1 = time.time()
                print(
                    f"best_value =  {best_value:>4.10f}, "
                    f"time = {t1-t0:>4.2f}, "
                    f"num_eval= {iteration * R_MIN},"
                    f"length_scale= {model.covar_module.base_kernel.raw_lengthscale.item()}, "
                )
            cnt += 1
        except NanError:
            print('error')
            pass
