import numpy as np
from abc import abstractmethod
from modules.utils.Log import Logger
from modules.agents.Probability import ProbabilityMeasure
from modules.agents.Policy import Policy, PolicyFiniteActionsFiniteStates
from modules.agents.Agent import Agent


class DiscreteAgent(Agent):
    def __init__(self, policy: Policy or None, parameters: dict) -> None:
        super().__init__(parameters)
        self._n_state = parameters["n_state"]
        self._n_action = parameters["n_action"]
        self._state_space = np.arange(self._n_state)
        self._action_space = np.arange(self._n_action)
        self._name = None

        if policy is None:
            self._policy = PolicyFiniteActionsFiniteStates(self._state_space, self._action_space, None)
        else:
            self._policy = policy

    def policy(self, copy=True):
        if copy:
            return self._policy.copy()
        else:
            return self._policy

    def state_space(self, copy=True):
        if copy:
            return self._state_space.copy()
        else:
            return self._state_space

    def action_space(self, copy=True):
        if copy:
            return self._action_space.copy()
        else:
            return self._action_space

    def n_state(self):
        return self._n_state

    def n_action(self):
        return self._n_action

    def set_policy(self, pi):
        for s in self.state_space(copy=False):
            self.policy(copy=False).set_policy(s, pi)

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