total_budget = 0
  for stage in _n_configs_in_stage.keys():
      for b, c in zip(_budgets_in_stage[stage], _n_configs_in_stage[stage]):
          total_budget += b * c

from __future__ import annotations

import numpy as np
from ConfigSpace import Configuration, ConfigurationSpace, Float
from matplotlib import pyplot as plt

from smac import MultiFidelityFacade, RunHistory, Scenario
from smac.intensifier.hyperband_utils import get_n_trials_for_hyperband_multifidelity

__copyright__ = "Copyright 2021, AutoML.org Freiburg-Hannover"
__license__ = "3-clause BSD"


class QuadraticFunction:
    max_budget = 500

    @property
    def configspace(self) -> ConfigurationSpace:
        cs = ConfigurationSpace(seed=0)
        x = Float("x", (-5, 5), default=-5)
        cs.add([x])

        return cs

    def train(self, config: Configuration, seed: int = 0, budget: float | None = None) -> float:
        """Returns the y value of a quadratic function with a minimum we know to be at x=0."""
        x = config["x"]

        if budget is None:
            multiplier = 1
        else:
            multiplier = 1 + budget / self.max_budget

        return x**2 * multiplier


if __name__ == "__main__":
    model = QuadraticFunction()

    total_budget = 10000
    min_budget = 10  # minimum budget per trial
    max_budget = 500  # maximum budget per trial
    eta = 3  # standard HB parameter influencing the number of stages

    # Let's calculate how many trials we need to exhaust the total optimization budget (in terms of
    # fidelity units)
    n_trials = get_n_trials_for_hyperband_multifidelity(
        total_budget=total_budget,  # this is the total optimization budget we specify in terms of fidelity units
        min_budget=min_budget,  # This influences the Hyperband rounds, minimum budget per trial
        max_budget=max_budget,  # This influences the Hyperband rounds, maximum budget per trial
        eta=eta,  # This influences the Hyperband rounds
        print_summary=True,
    )

    # Scenario object specifying the optimization "environment"
    scenario = Scenario(
        model.configspace, deterministic=True, n_trials=n_trials, min_budget=min_budget, max_budget=max_budget
    )

    # Now we use SMAC to find the best hyperparameters
    smac = MultiFidelityFacade(
        scenario,
        model.train,  # We pass the target function here
        overwrite=True,  # Overrides any previous results that are found that are inconsistent with the meta-data
        intensifier=MultiFidelityFacade.get_intensifier(scenario=scenario, eta=eta),
    )

    incumbent = smac.optimize()

    total_budget_used = sum([trial.budget for trial in smac.runhistory])

    print(f"Total budget defined: {total_budget}")
    print(f"Total budget used: {total_budget_used}")