import torch

from rlberry.envs import Wrapper

import numpy as np
from rlberry.utils.factory import load

import rlberry

logger = rlberry.logger


class UncertaintyEstimatorWrapper(Wrapper):
    """
    Adds exploration bonuses to the info output of env.step(), according to an
    instance of UncertaintyEstimator.

    Example
    -------

    ```
    observation, reward, terminated, truncated, info  = env.step(action)
    bonus = info['exploration_bonus']
    ```

    Parameters
    ----------
    uncertainty_estimator_fn : function(observation_space,
                                        action_space, **kwargs)
        Function that gives an instance of UncertaintyEstimator,
        used to compute bonus.
    uncertainty_estimator_kwargs:
        kwargs for uncertainty_estimator_fn
    bonus_scale_factor : double
            Scale factor for the bonus.
    """

    def __init__(
        self,
        env,
        uncertainty_estimator_fn,
        uncertainty_estimator_kwargs=None,
        bonus_scale_factor=1.0,
        bonus_max=np.inf,
    ):
        Wrapper.__init__(self, env)

        self.bonus_scale_factor = bonus_scale_factor
        self.bonus_max = bonus_max
        uncertainty_estimator_kwargs = uncertainty_estimator_kwargs or {}

        uncertainty_estimator_fn = (
            load(uncertainty_estimator_fn)
            if isinstance(uncertainty_estimator_fn, str)
            else uncertainty_estimator_fn
        )
        self.uncertainty_estimator = uncertainty_estimator_fn(
            env.observation_space, env.action_space, **uncertainty_estimator_kwargs
        )
        self.previous_obs = None

    def reset(self, seed=None, options=None):
        self.previous_obs, info = self.env.reset(seed=seed, options=options)
        return self.previous_obs, info

    def _update_and_get_bonus(self, state, action, next_state, reward):
        if self.previous_obs is None:
            return 0.0
        #
        self.uncertainty_estimator.update(state, action, next_state, reward)
        return self.bonus(state, action)

    def step(self, action):
        observation, reward, terminated, truncated, info = self.env.step(action)
        done = terminated or truncated

        # update uncertainty and compute bonus
        bonus = self._update_and_get_bonus(
            self.previous_obs, action, observation, reward
        )
        #
        self.previous_obs = observation

        # add bonus to info
        if info is None:
            info = {}
        else:
            if "exploration_bonus" in info:
                logger.error(
                    "UncertaintyEstimatorWrapper Error: info has"
                    + "  already a key named exploration_bonus!"
                )

        info["exploration_bonus"] = bonus

        return observation, reward, terminated, truncated, info

    def sample(self, state, action):
        logger.warning(
            "[UncertaintyEstimatorWrapper]: sample()"
            + " method does not consider nor update bonuses."
        )
        return self.env.sample(state, action)

    def bonus(self, state, action=None):
        bonus = self.bonus_scale_factor * self.uncertainty_estimator.measure(
            state, action
        )
        return np.clip(bonus, 0, self.bonus_max)

    def bonus_batch(self, states, actions=None):
        bonus = self.bonus_scale_factor * self.uncertainty_estimator.measure_batch(
            states, actions
        )
        return (
            np.clip(bonus, 0, self.bonus_max)
            if isinstance(bonus, np.ndarray)
            else torch.clamp(bonus, 0, self.bonus_max)
        )
