from __future__ import division, print_function, absolute_import

import copy
import math
import numpy as np

from .distrib import ProbabilityDistribution


class Experience(object):
    """Experience base class.

    Representation of an experience occurring from acting in the environment.

    Parameters
    ----------
    state : State
        The representation of the current state.
    action : Action
        The executed action.
    next_state : State
        The representation of the state following from acting
        with `action` in state `state`.
    reward : int or float
        The reward awarded by the environment for the state-action
        pair.

    Attributes
    ----------
    state : State
        The experienced state
    action : Action
        The experienced action.
    next_state : State
        The experienced next state.
    reward : float
        The experienced reward.

    """
    __slots__ = ('state', 'action', 'next_state', 'reward')

    def __init__(self, state, action, next_state, reward=None):
        self.state = state
        self.action = action
        self.next_state = next_state
        self.reward = reward

    def __str__(self):
        s = "state={0} act={1} next_state={2}".format(self.state, self.action, self.next_state) if self.reward else \
            "state={0} act={1} reward={2:.2f} next_state={3}".format(
                self.state, self.action, self.reward, self.next_state)
        return s


class RewardFunction(object):
    """The reward function.

    The reward function is responsible for calculating the proper value
    of the reward. Callback functions can be specified for custom calculation
    of the reward value.

    Attributes
    ----------
    cb_get : callable
        Callback function to retrieve the reward value.
    cb_set : callable
        Callback function to set the reward value.
    reward : float
        The reward value.
    bonus
    rmax : float
        The maximum possible reward.
    activate_bonus : bool
        Flag activating/deactivating the bonus.

    Notes
    -----
    To ensure that the correct value of the 
reward is being accessed,
    the user should not access the class variables directly but instead
    use the methods :meth:`set` and :meth:`get` to set and get the reward
    respectively.

    Examples
    --------
    >>> RewardFunction.cb_get = staticmethod(lambda r, s: np.dot(s, RewardFunction.reward))

    In this cas the reward function is calculated by taking 
the dot product
    of the stored reward and a passed in value.

    >>> RewardFunction.reward = [0.1, 0.9. 1.0, 0.0]

    This sets the reward for all instances of the reward function.

    >>> reward_func = RewardFunction()
    >>> print reward_func.get([0.9, 0.5, 0.0, 1.0])
    0.54

    This calculates the reward `r` according to previously defined the
    callback function.

    """
    __slots__ = ('_bonus', 'activate_bonus', 'reward', 'rmax', 'cb_get', 'cb_set',)

    cb_get = None
    cb_set = None

    reward = 0.0
    rmax = 0.0
    activate_bonus = False

    @property
    def bonus(self):
        """The bonus added to the reward to encourage exploration.

        Returns
        -------
        float :
            The bonus added to the reward.

        """
        return self._bonus

    @bonus.setter
    def bonus(self, value):
        self._bonus = value

    def __init__(self):
        self._bonus = 0.0
        """:type: float"""

    def __getstate__(self):
        return {
            'reward': self.reward,
            'rmax': self.rmax,
            'bonus': self.bonus,
            'activate_bonus': self.activate_bonus
        }

    def __setstate__(self, d):
        for name, value in d.iteritems():
            if not name == 'bonus':
                setattr(type(self), name, value)
            else:
                setattr(self, name, value)

    def set(self, value, *args, **kwargs):
        """Set the reward value.

        If :meth:`cb_set` is set, the callback will be called
        to set the value.

        Parameters
        ----------
        args : tuple
            Positional arguments passed to the callback.
        kwargs : dict
            Non-positional arguments passed to the callback.

        """
        if self.cb_set is not None:
            type(self).reward = self.cb_set(*args, **kwargs)
            return
        type(self).reward = value

    def get(self, *args, **kwargs):
        """Retrieve the reward value.

        If :meth:`cb_get` is set, the callback will be called
        to retrieve the value.

        Parameters
        ----------
        args : tuple
            Positional arguments passed to the callback.
        kwargs : dict
            Non-positional arguments passed to the callback.

        Returns
        -------
        float :
            The (calculated) reward value.

        """
        reward = self.reward
        if self.cb_get is not None:
            reward = self.cb_get(self.reward, *args, **kwargs)

        if self.activate_bonus:
            reward = max(self.reward + self.bonus, self.rmax)
        return reward


class StateActionInfo(object):
    """The models interface.

    Contains all relevant information predicted by a model for a
    given state-action pair. This includes the (predicted) reward and
    transition probabilities to possible next states.

    Attributes
    ----------
    transition_proba : ProbabilityDistribution
        The transition probability distribution.
    reward_func : RewardFunction
        The reward function.
    visits : int
        The number of times the state-action pair has been visited.
    known : bool
        Flag indicating whether a reward value is known or not.

    """
    __slots__ = ('transition_proba', 'reward_func', 'visits', 'known')

    def __init__(self):
        self.transition_proba = ProbabilityDistribution()
        self.reward_func = RewardFunction()

        self.visits = 0
        self.known = False

    def __getstate__(self):
        data = {}
        for name in self.__slots__:
            data[name] = getattr(self, name)
        return data

    def __setstate__(self, d):
        for name, value in d.iteritems():
            setattr(self, name, value)


class StateData(object):
    """State information interface.

    Information about the state can be accessed here.

    Parameters
    ----------
    state_id : int
        The unique id of the state
    actions : list[Action]
        List of actions that can be taken in this state.

    Attributes
    ----------
    id : int
        The unique id of the state.
    models : dict
        The reward and transition models for each action.
    q : dict
        The q-table, containing a q-value for each action.
    steps_away : int
        The number of steps the state is away from its closest neighbor.

    """
    __slots__ = ('id', 'models', 'q', 'steps_away')

    def __init__(self, state_id, actions):
        self.id = state_id
        """:type: int"""
        self.models = {a: StateActionInfo() for a in actions}
        """:type: dict[Action, StateActionInfo]"""
        # Randomizing the initial q-values impedes performance
        # self.q = {a: ((0.01 - 0.0) * np.random.random() + 0.0) for a in actions}
        self.q = {a: 0.0 for a in actions}
        """:type: dict[Action, float]"""
        self.steps_away = 100000
        """:type: int"""

    def __getstate__(self):
        data = {}
        for name in self.__slots__:
            data[name] = getattr(self, name)
        return data

    def __setstate__(self, d):
        for name, value in d.iteritems():
            setattr(self, name, value)


class MDPPrimitive(object):
    """A Markov decision process primitive.

    The base class for :class:`State` and :class:`Action`. Primitives
    are represented by a list of features. They optionally can have a `name`.

    Parameters
    ----------
    features : array_like, shape (`nfeatures`,)
        List of features, where `nfeatures` is the number of features
        identifying the primitive.
    name : str, optional
        The name of the primitive. Default is "".

    A