from .Probability import ProbabilityMeasure
from abc import ABC, abstractmethod
from modules.utils.Log import Logger


class Agent(ABC):
    def __init__(self, parameters: dict):
        self._policy = None
        self._epsilon = None
        self.set_trustregion(parameters["epsilon"])
        self._name = None
        self._is_training = False

    @abstractmethod
    def policy(self, copy=True):
        pass

    @abstractmethod
    def update_policy(self, advantage, rho: ProbabilityMeasure, logger: Logger) -> None:
        pass

    @abstractmethod
    def state_space(self, copy=True):
        pass

    @abstractmethod
    def action_space(self, copy=True):
        pass

    def take_action(self, state):
        return self.policy(copy=False).sample(state)

    def is_training(self, set=None):
        if set is not None:
            self._is_training = set
        return self._is_training

    def epsilon(self) -> float:
        return self._epsilon

    def name(self) -> str:
        return self._name

    def set_trustregion(self, epsilon: float) -> None:
        assert epsilon >= 0, "the radius of the trust region should be non-negative"
        self._epsilon = epsilon

    def save(self):
        return {"state_space": self.state_space(),
                "action_space": self.action_space(),
                "policy": self.policy().save(),
                "name": self.name(),
                "parameters": {"epsilon": self.epsilon()}}
