from collections import defaultdict

import optuna as ot


def objective(t):
    x = t.suggest_float("x", 0, 1)

    for step in range(100):
        t.report(step=step, value=x * step)
        if t.should_prune():
            raise ot.TrialPruned

    return x


if __name__ == "__main__":
    p = ot.pruners.HyperbandPruner()
    study = ot.create_study(
        study_name="test_hyperband_3083",
        pruner=p,  # You actually don't need to prune, for this demonstration.
        load_if_exists=True,
    )
    study.optimize(objective, n_trials=100)

    # Check bracket splits. Using private APIs, only use for testing.
    brackets_to_trial_numbers = defaultdict(list)
    for t in study.trials:
        brackets_to_trial_numbers[study.pruner._get_bracket_id(study, t)].append(t.number)
    print({k: len(v) for k, v in brackets_to_trial_numbers.items()})
    # {0: 81, 1: 19}