from __future__ import annotations

from typing import TYPE_CHECKING

from optuna.samplers._base import _process_constraints_after_trial


if TYPE_CHECKING:
    from collections.abc import Callable
    from collections.abc import Sequence

    from optuna.study import Study
    from optuna.trial import FrozenTrial
    from optuna.trial import TrialState


class NSGAIIAfterTrialStrategy:
    def __init__(
        self, *, constraints_func: Callable[[FrozenTrial], Sequence[float]] | None = None
    ) -> None:
        self._constraints_func = constraints_func

    def __call__(
        self,
        study: Study,
        trial: FrozenTrial,
        state: TrialState,
        values: Sequence[float] | None = None,
    ) -> None:
        """Carry out the after trial process of default NSGA-II.

        This method is called after each trial of the study, examines whether the trial result is
        valid in terms of constraints, and store the results in system_attrs of the study.
        """
        if self._constraints_func is not None:
            _process_constraints_after_trial(self._constraints_func, study, trial, state)