# Path: oscar/agent/learning_structure.py
# class LearningStructure:
# 
#     # must be set by sub classes
#     observation_space = None
#     action_space = None
# 
#     def __init__(self, train_mode=False, shared_memory=None):
#         """
#         Initialize custom agent and set step function according to train mode
#         :param train_mode: indicate if the agent must be constructed in training mode or not
#         """
#         self.train_mode = train_mode
#         self.shared_memory = shared_memory
#         # check shared memory is consistent with train mode
#         if self.train_mode is True and self.shared_memory is None:
#             raise BadSharedMemoryError("When train mode is set to True a valid shared memory must be provided")
#         # check if action/obs space are set by the sub classes
#         if self.train_mode and (self.action_space is None or self.observation_space is None):
#             raise RuntimeError("action_space and observation_space must be set before calling LearningStructure init !")
# 
#         # set which step method must be called
#         if self.train_mode:
#             self.do_step = self._learning_step
#             # set memory with agent action/obs spaces
#             shared_memory.action_space = self.action_space
#             shared_memory.observation_space = self.observation_space
#         else:
#             self.do_step = self._playing_step
#         super().__init__()
# 
#     def step(self, obs, locked_choice=None):
#         """
#         Choose between calling normal step or training step.
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         return self.do_step(obs)
# 
#     def _playing_step(self, obs):
#         """
#         Step function to be used when not in training mode, but when training but must be
#         available.
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         # TODO: error message?
#         raise NotImplementedError()
# 
#     def _learning_step(self, obs):
#         """
#         Step function to be used when in training mode
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         # transform the pysc2 obs into agent specific observation
#         formatted_obs = self._format_observation(obs)
#         # give the obs to the shared memory
#         self.shared_memory.shared_obs = formatted_obs
#         # give the available action to the shared memory
#         self.shared_memory.available_action_mask = self._available_action_mask()
#         # set the semaphore to allow the learning thread to choose an action
#         self.shared_memory.semaphore_obs_ready.release()
#         # wait for the action to be selected
#         self.shared_memory.semaphore_action_set.acquire(blocking=True, timeout=None)
#         # transform the action into an hierarchical compatible action
#         action = self.shared_memory.shared_action
#         return self._transform_action(action_id=action)
# 
#     def _format_observation(self, full_obs):
#         """
#         Format the observation from the pysc2 format into an agent specific format
#         :param full_obs: the observation structure as given by pysc2
#         :return: an numpy array of the obs as wanted by the current agent
#         """
#         raise NotImplementedError()
# 
#     def _transform_action(self, action_id):
#         """
#         Transform the action decided by the agent (tipicaly a number in [0;n]) to a
#         real agent action (dict of action list, callbacks...)
#         :param action_id: number of the action chosen
#         :return: agent play action in the hierarchy format (dict of action list, callbacks
#          and lock)
#         """
#         raise NotImplementedError()
# 
#     def _available_action_mask(self):
#         """
#         A mask to be applied to a vector of action probability to mask to action which are
#         not available at that given time. (Useful for A3C)
#         :return: a mask (vector of the same size as action)
#         """
#         # default behaviour: all actions are available
#         return np.ones(shape=self.action_space.n)
# 
# Path: oscar/agent/custom_agent.py
# class CustomAgent(BaseAgent):
#     """A base agent to write custom scripted agents."""
# 
#     def __init__(self):
#         super().__init__()
#         self._shared = {}
# 
#     def add_shared(self, name, shared):
#         self._shared[name] = shared


from oscar.agent.learning_structure import LearningStructure
from oscar.agent.custom_agent import CustomAgent


class LearningAgent(LearningStructure, CustomAgent):
    """
    An abstract class to build learning agent

    Sub class must implement the following LearningStructure methods:
      - _step
      - _format_observation
      - _transform_action
    (see LearningStructure code for more information on them)
    """

    def __init__(self, train_mode=False, shared_memory=None):
        """
        Constructor of the abstract class LearningAgent
        Sub class must keep the same calling format to work with the hierarchy factory
        self.observation_space and self.action_space must be set before calling this constructor
        :param train_mode: if the agent must train or play (default: False -> play)
        :param shared_memory: the memory used during training to communicate with the environment
                    (useless when playing, indispensable when training)
        """

# Path: oscar/agent/learning_structure.py
# class LearningStructure:
# 
#     # must be set by sub classes
#     observation_space = None
#     action_space = None
# 
#     def __init__(self, train_mode=False, shared_memory=None):
#         """
#         Initialize custom agent and set step function according to train mode
#         :param train_mode: indicate if the agent must be constructed in training mode or not
#         """
#         self.train_mode = train_mode
#         self.shared_memory = shared_memory
#         # check shared memory is consistent with train mode
#         if self.train_mode is True and self.shared_memory is None:
#             raise BadSharedMemoryError("When train mode is set to True a valid shared memory must be provided")
#         # check if action/obs space are set by the sub classes
#         if self.train_mode and (self.action_space is None or self.observation_space is None):
#             raise RuntimeError("action_space and observation_space must be set before calling LearningStructure init !")
# 
#         # set which step method must be called
#         if self.train_mode:
#             self.do_step = self._learning_step
#             # set memory with agent action/obs spaces
#             shared_memory.action_space = self.action_space
#             shared_memory.observation_space = self.observation_space
#         else:
#             self.do_step = self._playing_step
#         super().__init__()
# 
#     def step(self, obs, locked_choice=None):
#         """
#         Choose between calling normal step or training step.
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         return self.do_step(obs)
# 
#     def _playing_step(self, obs):
#         """
#         Step function to be used when not in training mode, but when training but must be
#         available.
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         # TODO: error message?
#         raise NotImplementedError()
# 
#     def _learning_step(self, obs):
#         """
#         Step function to be used when in training mode
#         :param obs: observation used to choose the action to do
#         :return: a dict of the agent's choice for this step (action list, callbacks)
#         """
#         # transform the pysc2 obs into agent specific observation
#         formatted_obs = self._format_observation(obs)
#         # give the obs to the shared memory
#         self.shared_memory.shared_obs = formatted_obs
#         # give the available action to the shared memory
#         self.shared_memory.available_action_mask = self._available_action_mask()
#         # set the semaphore to allow the learning thread to choose an action
#         self.shared_memory.semaphore_obs_ready.release()
#         # wait for the action to be selected
#         self.shared_memory.semaphore_action_set.acquire(blocking=True, timeout=None)
#         # transform the action into an hierarchical compatible action
#         action = self.shared_memory.shared_action
#         return self._transform_action(action_id=action)
# 
#     def _format_observation(self, full_obs):
#         """
#         Format the observation from the pysc2 format into an agent specific format
#         :param full_obs: the observation structure as given by pysc2
#         :return: an numpy array of the obs as wanted by the current agent
#         """
#         raise NotImplementedError()
# 
#     def _transform_action(self, action_id):
#         """
#         Transform the action decided by the agent (tipicaly a number in [0;n]) to a
#         real agent action (dict of action list, callbacks...)
#         :param action_id: number of the action chosen
#         :return: agent play action in the hierarchy format (dict of action list, callbacks
#          and lock)
#         """
#         raise NotImplementedError()
# 
#     def _available_action_mask(self):
#         """
#         A mask to be applied to a vector of action probability to mask to action which are
#         not available at that given time. (Useful for A3C)
#         :return: a mask (vector of the same size as action)
#         """
#         # default behaviour: all actions are available
#         return np.ones(shape=self.action_space.n)
# 
# Path: oscar/agent/custom_agent.py
# class CustomAgent(BaseAgent):
#     """A base agent to write custom scripted agents."""
# 
#     def __init__(self):
#         super().__init__()
#         self._shared = {}
# 
#     def add_shared(self, name, shared):
#         self._shared[name] = shared


# Path: oscar/agent/learning_agent.py
from oscar.agent.learning_structure import LearningStructure
from oscar.agent.custom_agent import CustomAgent



class LearningAgent(LearningStructure, CustomAgent):
    """
    An abstract class to build learning agent

    Sub class must implement the following LearningStructure methods:
      - _step
      - _format_observation
      - _transform_action
    (see LearningStructure code for more information on them)
    """

    def __init__(self, train_mode=False, shared_memory=None):
        """
        Constructor of the abstract class LearningAgent
        Sub class must keep the same calling format to work with the hierarchy factory
        self.observation_space and self.action_space must be set before calling this constructor
        :param train_mode: if the agent must train or play (default: False -> play)
        :param shared_memory: the memory used during training to communicate with the environment
                    (useless when playing, indispensable when training)
        """

        LearningStructure.__init__(self, train_mode, shared_memory)