from collections import namedtuple
import decimal
from emoji import emojize
import math
import numpy as np
import numpy.random as nrand
from recordclass import recordclass



# Transition is a `namedtuple` representing a sample of observations obtained by
# the agent while exploring and interacting with its environment. It holds
# following four attributes:
# * `state` - The starting state the action was taking from
# * `action` - The action taken by the agent based on some policy
# * `next_state` - The state reached by taking `action` in `state`
# * `reward` - The reward observed while taking the `action` in the `state`
# and reaching the `next_state`.
Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward', 'terminal'))

# The A3C Transition tuple is an extended version of the `Transition` tuple
# declared above, containing the following additional fields:
# `policy`, `value`, `hidden_state`
A3CTransition = namedtuple('A3CTransition',
                           ('state', 'action', 'next_state', 'reward',
                            'terminal', 'value', 'policy', 'hidden_state'))

# The `StepData` recordclass is used to keep track of the current episode, the
# state-action value and the reward of a current step. TODO: implement clear method
StepData = recordclass('StepData', (
                                    'episode',
                                    'q_value',
                                    'reward',
                                    'loss',
                                    'cx',
                                    'cv',
                                    'py',
                                    'px',
                                    'pav'))


def default_fallback_message(var):
    '''
    This function spins up a message that no value was passed into
    a certain function or method, leading to a fallback to the default
    value set by the developer.

    :return: An emojized message, that a fallback has been activated

    Usage:
    ```
    msg = default_fallback_message('my_var_name')
    logging.info(msg)
    ```
    '''
    return emojize(('\n :sparkles: `{}` not provided - '
                    'Spinning up the default for you :raised_hands:')
                   .format(var), use_aliases=True)


def info_message(str):
    '''
    This function spins up an info message prefixed with a megaphone emoji

    Usage:
    ```
    msg = info_message('Doing this and that')
    logging.info(msg)
    ```
    '''
    return emojize('\n :mega: {}'.format(msg), use_aliases=True)


def train_start_message():
    '''
    Train message returns an emojized message that training has started.

    :return: An emojized string, indicating the training has started.
    ```
    msg = train_start_message()
    print(msg)
    ```
    '''
    return emojize(('\n :rocket:  Starting training process. Prepare '
                   'for take off  :rocket: \n'))


def error_message(err_msg):
    '''
    Error message returns an emojized message that something bad has
    happened.
    :param err_msg: The string describing the error message
    :return: An emojized version of the error message
    ```
    msg = train_start_message()
    print(msg)
    ```
    '''
    return emojize(('\n :boom:  {} \n').format(err_msg), use_aliases=True)


def missing_kwarg(var, expected_type, fn_name, cls_name):
    '''
    This function generates a generic message that some mandatory kwarg
    is missing. Takes the variable name as an input argument. Pass this
    message into `error_message` in order to emojize it. The emojization and
    generation of this message is kept seperate for consistency reasons.

    :param var: The variable which caused the error_message
    :param expected_type: The type the variable should be
    :param cls_name: The class that the error occured in

    :return: A unified missing keyword argument error message
    '''
    return (('No value provided for `{}` in {} while executing `{}`! '
            ':cold_sweat:\n => Please pass in a valid arg of type `{}`\n')
            .format(var, cls_name, fn_name, expected_type))


def type_error(var, provided_type, expected_type, fn_name, cls_name):
    '''
    This function generates a generic message that some mandatory kwarg is of
    wrong type. Pass this message into `error_message` in order to emojize it.
    The emojization and generation of this message is kept seperate for
    consistency reasons.

    Args:
    :param var: The variable which caused the error_message
    :param provided_type: The type provided
    :param expected_type: The expected type
    :param cls_name: The class that the error occured in
    :param fn_name: The function/method the error occured in

    :return: A unified invalid type error message
    '''
    return (('The argument `{}` received an invalid type in {} while '
             'exeuting `{}` :cold_sweat: :\n => Provided: {}, Expected: {}')
            .format(var, cls_name, fn_name, provided_type, expected_type))


class FrameBuffer(object):
    '''
    The ``FrameBuffer`` class was opensourced by the openai team as part of their
    baselines blog post. The class addresses one key issue:
    When training an agent, in many cases, a replay memory is leveraged to perform
    batch updates on previously seen transitions. The original paper by Mnih et al.
    describes a replay memory, which stores the past 1 million seen transitions.
    Passing in the actual numpy arrays is practically infeasable as the python process
    quickly takes up gigabytes of memory, leading to termination.
    Passing in a ``FrameBuffer`` instead leads to a pure reference, ensuring
    that the frames are only stored once.

    Original implementation:
    `https://github.com/openai/baselines/blob/0778e9f10f4c51417070f24ab12594466443e3a1/baselines/common/atari_wrappers_deprecated.py#L152-L167`_
    '''
    def __init__(self, frames):
        self._frames = frames

    def __array__(self, dtype=None):
        out = np.concatenate(self._frames, axis=2)
        # Transpose to torch notation BCHW
        out = out.transpose((2, 0, 1))
        if dtype is not None:
            out = out.astype(dtype)
        return out
