from gym import Env
from overrides import overrides

from symbols.domain.option import Option
from symbols.domain.option_space import OptionSpace




class Domain(Env):
    """
    An abstract base subclass that uses the same API as Open AI's Gym
    """

    def __init__(self,
                 task_id: int):
        self._task_id = task_id
        self._view = 'problem'

    @property
    def view(self):
        return self._view

    @view.setter
    def view(self, view):
        self._view = view

    @property
    def observation_space(self):
        """
        Return the observation (problem) space
        """
        raise NotImplementedError

    @property
    def agent_space(self):
        """
        Return the observation (problem) space
        """
        raise NotImplementedError

    @property
    def action_space(self):
        raise NotImplementedError

    def _primitive_step(self, action):
        """
        Execute a primitive action in the domain
        :param action: the primitive action to execute
        :return:

            observation (object): agent's observation of the current environment
            reward (float) : amount of reward returned after previous action
            done (boolean): whether the episode has ended, in which case further step() calls will return undefined
            results
            info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
        """
        raise NotImplementedError

    @property
    def current_state(self):
        """
        Get the current state
        :return: the current state of the environment
        """
        raise NotImplementedError

    @property
    def current_observation(self):
        """
        Get the current agent-space observation
        """
        raise NotImplementedError

    def can_execute(self, action):
        """
        Determines whether the given action can be executed at the current state
        :param action: the action to execute
        :return: True, if the action is executable. False otherwise
        """
        raise NotImplementedError

    def init(self, state):
        """
        Initialise the domain to the given state
        :param state: the state to set the domain to
        """
        raise NotImplementedError

    @overrides
    def reset(self, random_starts=False):
        """Resets the state of the environment and returns an initial observation.

        Returns: observation (object): the initial observation of the
            space.
        """
        raise NotImplementedError

    @overrides
    def step(self, option: Option):
        option.initialise()
        state = self.current_state
        reward = 0.0
        done = False
        while not option.done:
            action = option.step(state)
            if action is None:
                break
            state, r, done, _ = self._primitive_step(action)
            reward += r
            if done:
                break
        return state, reward, done, {}

    @property
    def admissible_actions(self):
        return [action for action in self.action_space if self.can_execute(action)]

    def describe(self, option):
        if isinstance(self.action_space, OptionSpace):
            return self.action_space.describe(option)
        return str(option)

    @property
    def task_id(self):
        return self._task_id
