import numpy as np
import matplotlib.pyplot as plt
import optuna

def objective(trial):
    params = np.empty(5)
    for i in range(5):
        params[i] = trial.suggest_int(f"x{i}", 1, 100, log=True)
    return np.sum((np.log(params) - np.log(10)) ** 2)

N_STUDIES = 100
N_TRIALS = 100

y = [[] for _ in  range(N_STUDIES)]
for seed in range(N_STUDIES):
    sampler = optuna.samplers.TPESampler(seed=seed)
    study = optuna.create_study(sampler=sampler)
    study.optimize(objective, n_trials=N_TRIALS)
    y[seed].append(study.trials[0].value)
    for i in range(1, N_TRIALS):
        y[seed].append(min(y[seed][-1], study.trials[i].value))

y = np.asarray(y)
np.save("master.npy", y)
# np.save("pr.npy", y)

y = np.load("master.npy")
mean = np.mean(y, axis=0)
std = np.std(y, axis=0)
plt.fill_between(range(N_TRIALS), mean + std, mean - std, alpha=0.1)
plt.plot(range(N_TRIALS), mean, label="master")

y = np.load("pr.npy")
mean = np.mean(y, axis=0)
std = np.std(y, axis=0)
plt.fill_between(range(N_TRIALS), mean + std, mean - std, alpha=0.1)
plt.plot(range(N_TRIALS), mean, label="fixed")

plt.xlabel("trial number")
plt.ylabel("loss")
plt.legend()
plt.savefig("logint.png")