import time
import optuna
from tqdm import tqdm

N_TRIAL = 1000
N_BATCH = 10

for multivariate in (False, True):
    for constant_liar in (False, True):
        start = time.time()
        sampler = optuna.samplers.TPESampler(
            seed=42,
            multivariate=multivariate,
            constant_liar=constant_liar,
        )
        study = optuna.create_study(sampler=sampler, storage="sqlite:///tmp.db")

        for i in range(0, N_TRIAL, N_BATCH):
            trials = []
            for j in range(N_BATCH):
                trials.append(study.ask())
            X = [trial.suggest_float("x", -10, 10) for trial in trials]
            Y = [trial.suggest_float("y", -10, 10) for trial in trials]
            for j in range(N_BATCH):
                study.tell(trials[j], X[j] ** 2 + Y[j] ** 2)
        print(f"{multivariate=} {constant_liar=} {time.time()-start}")

from multiprocessing import Pool
import optuna
import optunahub
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
from scipy.stats import mannwhitneyu

optuna.logging.set_verbosity(optuna.logging.WARNING)

N_STUDY = 100
N_TRIAL = 1000
N_BATCH = 50
function_id = 3

bbob = optunahub.load_module("benchmarks/bbob")


def run_study(args):
    n_dim, constant_liar, seed = args
    objective = bbob.Problem(function_id=function_id, dimension=n_dim, instance_id=1)
    assert objective.directions == [optuna.study.StudyDirection.MINIMIZE]

    study_name = f"{n_dim}_{constant_liar}_{seed:02d}"
    sampler = optuna.samplers.TPESampler(
        seed=seed,
        multivariate=True,
        constant_liar=constant_liar,
    )
    study = optuna.create_study(study_name=study_name, sampler=sampler)

    for i in range(0, N_TRIAL, N_BATCH):
        trials = [study.ask() for _ in range(N_BATCH)]
        for d in range(n_dim):
            for trial in trials:
                trial.suggest_float(f"x{d}", -5, 5)
        values = [objective(trial) for trial in trials]
        for trial, value in zip(trials, values):
            study.tell(trial, value)

    storage = optuna.storages.JournalStorage(
        optuna.storages.journal.JournalFileBackend(f"./benchmark_{function_id:02d}/{study_name}.log")
    )
    optuna.copy_study(
        from_study_name=study_name,
        from_storage=study._storage,
        to_storage=storage,
    )


def main():
    for n_dim in [2, 3, 5, 10, 20, 40]:
        args = []
        for constant_liar in (True, False):
            for seed in range(N_STUDY):
                args.append((n_dim, constant_liar, seed))
        with Pool(processes=10) as pool:
            with tqdm(total=len(args)) as t:
                for _ in pool.imap_unordered(run_study, args):
                    t.update(1)
        all_best_values = [[], []]
        for constant_liar in (True, False):
            for seed in range(N_STUDY):
                study_name = f"{n_dim}_{constant_liar}_{seed:02d}"
                storage = optuna.storages.JournalStorage(
                    optuna.storages.journal.JournalFileBackend(f"./benchmark_{function_id:02d}/{study_name}.log")
                )
                study = optuna.load_study(storage=storage, study_name=study_name)
                best_values = []
                best_value = float("inf")
                for trial in study.trials:
                    if best_value > trial.value:
                        best_value = trial.value
                    best_values.append(best_value)
                all_best_values[int(constant_liar)].append(best_values)
        all_best_values = np.asarray(all_best_values)
        p_values = []
        for i in range(N_TRIAL):
            p_value = mannwhitneyu(
                all_best_values[1, :, i],
                all_best_values[0, :, i],
                alternative='less',
            ).pvalue
            p_values.append(p_value)
        plt.plot(range(N_TRIAL), p_values, label=f"D={n_dim}")
    plt.ylim(0, 1)
    plt.legend(loc=1)
    plt.savefig(f"benchmark_{function_id:02}.png")


if __name__ == "__main__":
    main()