"""Generate constrained search space. Alert if constraints are not satisfiable."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/02_search_space_helpers.ipynb.

# %% auto 0
__all__ = ['Trial', 'ConstrainedSearchSpace']

# %% ../../nbs/utils/02_search_space_helpers.ipynb 3
import numpy as np
from typing import Generator, List, Tuple, Protocol
from dataclasses import dataclass
from collections import OrderedDict

# %% ../../nbs/utils/02_search_space_helpers.ipynb 4
class Trial(Protocol):
    "Protocol for a trial object"
    @staticmethod
    def suggest_float(name: str, *args, **kwargs) -> float:
        ...

# %% ../../nbs/utils/02_search_space_helpers.ipynb 5
@dataclass
class ConstrainedSearchSpace:
    """
    A class that generates a search space with constraints
    """
    bounds: dict[str, tuple[float, float]] # bounds of each channel
    constraint: tuple[float, float] # constraint of the sum of all channels

    def __post_init__(self):
        self.bounds = OrderedDict(sorted(self.bounds.items(), key=lambda x: x[1]))
    
    def __call__(
        self, 
        trial: Trial # trial object
        ) -> dict[str, float]: # selected budget
        "Sample from constrained search space"
        selected_budget = {}
        bounds_values = list(self.bounds.values())
        bounds_items = list(self.bounds.items())
        for n, (name, bound) in enumerate(bounds_items[:-1]):
            curr_total = sum(selected_budget.values())
            new_min_bound = self.constraint[0]-(curr_total+sum(b[1] for b in bounds_values[n+1:]))
            new_max_bound = self.constraint[1]-(curr_total+sum(b[0] for b in bounds_values[n+1:]))
            updated_bounds = (
                    max(bound[0], new_min_bound), 
                    min(bound[1], new_max_bound)
            )
            selection = trial.suggest_float(name, *updated_bounds)
            selected_budget[name] = selection
        last = bounds_values[-1]
        last_name = bounds_items[-1][0]
        choice = trial.suggest_float(
            last_name, 
            max(last[0], self.constraint[0]-sum(selected_budget.values())),
            min(last[1], self.constraint[1]-sum(selected_budget.values()))
        )
        selected_budget[last_name] = choice
        return selected_budget
