from bayes_opt import BayesianOptimization
from bayes_opt.util import load_logs
from argparse import ArgumentParser
from pprint import pprint


def parse_args():
    p = ArgumentParser()
    p.add_argument("logs", nargs="+", help="json logs to load")
    return p.parse_args()


def get_optimizer(logs):
    bounds = {
        "H0": (100, 600),
        "H1": (100, 600),
        "batch_size": (1, 32),
        "S": (1, 64),
        "min_lr": (0.0001, 0.1),
        "max_lr": (0.001, 0.5),
        "epochs_per_cycle": (1, 40),
        "n_parents": (1, 8),
        "n_children": (1, 7),
        "n_generations": (1, 3),
        "crossover": (0.0, 1.0),
        "analytical_pi_updates": (0.0, 1.0),
        "analytical_sigma_updates": (0.0, 1.0),
    }
    def evaluate(H0, H1, batch_size, S, min_lr, max_lr, epochs_per_cycle,
                 n_parents, n_children, n_generations, crossover,
                 analytical_pi_updates, analytical_sigma_updates):
        return 0.
    optimizer = BayesianOptimization(f=evaluate, pbounds=bounds)
    load_logs(optimizer, logs=logs)
    return optimizer


if __name__ == "__main__":
    args = parse_args()
    opt = get_optimizer(args.logs)
    pprint(opt.max)
