import numpy as np


from rlberry.agents import AgentWithSimplePolicy
from rlberry_scool.agents.dynprog.utils import backward_induction, value_iteration
from gymnasium.spaces import Discrete

import rlberry

logger = rlberry.logger


class MBQVIAgent(AgentWithSimplePolicy):
    """
    Model-Basel Q-Value iteration (MBQVI).

    Builds an empirical MDP and runs value iteration on it.
    Corresponds to the "indirect" algorithm studied by Kearns and Singh (1999).

    Parameters
    -----------
    env : Model
        generative model with finite state-action space
    n_samples : int
        number of samples *per state-action pair* used to estimate
        the empirical MDP.
    gamma : double
        discount factor in [0, 1]
    horizon : int
        horizon, if the problem is finite-horizon. if None, the discounted
        problem is solved. default = None
    epsilon : double
        precision of value iteration, only used in discounted problems
        (when horizon is None).
    **kwargs : Keyword Arguments
         Arguments to be passed to `AgentWithSimplePolicy.__init__(self, env, **kwargs)` (:class:`~rlberry.agents.AgentWithSimplePolicy`).

    References
    ----------
    Kearns, Michael J., and Satinder P. Singh.
    "Finite-sample convergence rates for Q-learning and indirect algorithms."
    Advances in neural information processing systems. 1999.
    """

    name = "MBQVI"

    def __init__(
        self, env, n_samples=10, gamma=0.99, horizon=None, epsilon=1e-6, **kwargs
    ):
        AgentWithSimplePolicy.__init__(self, env, **kwargs)

        # initialize base class
        assert self.env.is_generative(), "MBQVI requires a generative model."
        assert isinstance(
            self.env.observation_space, Discrete
        ), "MBQVI requires a finite state space."
        assert isinstance(
            self.env.action_space, Discrete
        ), "MBQVI requires a finite action space."

        #
        self.n_samples = n_samples
        self.gamma = gamma
        self.horizon = horizon
        self.epsilon = epsilon

        # empirical MDP, created in fit()
        self.R_hat = None
        self.P_hat = None

        # value functions
        self.V = None
        self.Q = None

    def _update(self, state, action, next_state, reward):
        """Update model statistics."""
        self.N_sa[state, action] += 1
        self.N_sas[state, action, next_state] += 1
        self.S_sa[state, action] += reward

    def fit(self, budget=None, **kwargs):
        """
        Build empirical MDP and run value iteration.

        Parameters
        ----------
        budget: None
            Not used. Only defined for compatibility purpose with rlberry.
            Changing `budget` value has no effect.
        **kwargs : Keyword Arguments
            Extra arguments. Not used for this agent.
        """
        del kwargs
        S = self.env.observation_space.n
        A = self.env.action_space.n
        self.N_sa = np.zeros((S, A))
        self.N_sas = np.zeros((S, A, S))
        self.S_sa = np.zeros((S, A))

        # collect data
        total_samples = S * A * self.n_samples
        count = 0
        logger.debug(
            f"[{self.name}] collecting {self.n_samples} samples per (s,a)"
            f", total = {total_samples} samples."
        )
        for ss in range(S):
            for aa in range(A):
                for _ in range(self.n_samples):
                    next_state, reward, _, _, _ = self.env.sample(ss, aa)
                    self._update(ss, aa, next_state, reward)

                    count += 1
                    if count % 10000 == 0:
                        completed = 100 * count / total_samples
                        logger.debug(
                            "[{}] ... {}/{} ({:0.0f}%)".format(
                                self.name, count, total_samples, completed
                            )
                        )

        # build model and run VI
        logger.debug(f"{self.name} building model and running backward induction...")

        N_sa = np.maximum(self.N_sa, 1)
        self.R_hat = self.S_sa / N_sa
        self.P_hat = np.zeros((S, A, S))
        for ss in range(S):
            self.P_hat[:, :, ss] = self.N_sas[:, :, ss] / N_sa

        info = {}
        info["n_samples"] = self.n_samples
        info["total_samples"] = total_samples
        if self.horizon is None:
            assert self.gamma < 1.0, "The discounted setting requires gamma < 1.0"
            self.Q, self.V, n_it = value_iteration(
                self.R_hat, self.P_hat, self.gamma, self.epsilon
            )
            info["n_iterations"] = n_it
            info["precision"] = self.epsilon
        else:
            self.Q, self.V = backward_induction(
                self.R_hat, self.P_hat, self.horizon, self.gamma
            )
            info["n_iterations"] = self.horizon
            info["precision"] = 0.0
        return info

    def policy(self, observation):
        state = observation
        assert self.env.observation_space.contains(state)
        if self.horizon is None:
            return self.Q[state, :].argmax()
        else:
            return self.Q[0, state, :].argmax()
