"""An explorer that explores by solving tasks with bilevel planning."""

from typing import List, Set

from gym.spaces import Box

from predicators.explorers.bilevel_planning_explorer import \
    BilevelPlanningExplorer
from predicators.explorers.random_options_explorer import RandomOptionsExplorer
from predicators.option_model import _OptionModelBase
from predicators.planning import PlanningFailure, PlanningTimeout
from predicators.structs import NSRT, ExplorationStrategy, \
    ParameterizedOption, Predicate, Task, Type


class ExploitBilevelPlanningExplorer(BilevelPlanningExplorer):
    """ExploitBilevelPlanningExplorer implementation."""

    def __init__(self, predicates: Set[Predicate],
                 options: Set[ParameterizedOption], types: Set[Type],
                 action_space: Box, train_tasks: List[Task], nsrts: Set[NSRT],
                 option_model: _OptionModelBase) -> None:
        super().__init__(predicates, options, types, action_space, train_tasks,
                         nsrts, option_model)
        # Falls back to random options.
        self._fallback_explorer = RandomOptionsExplorer(
            predicates, options, types, action_space, train_tasks)

    @classmethod
    def get_name(cls) -> str:
        return "exploit_planning"

    def get_exploration_strategy(self, train_task_idx: int,
                                 timeout: int) -> ExplorationStrategy:
        task = self._train_tasks[train_task_idx]
        try:
            return self._solve(task, timeout)
        except (PlanningFailure, PlanningTimeout):
            return self._fallback_explorer.get_exploration_strategy(
                train_task_idx, timeout)
