"""
    Policies for contextual bandit problems.

"""

import numpy as np
import tensorflow as tf
import contextlib
from termcolor import cprint

from rnlps.policies.base import Trial, Policy
from rnlps.policies.base import BaseOracle, BaseFixed, BaseRandom
from rnlps.policies.base import BaseThompsonRecurrentNetwork, BaseThompsonSinFeedforwardNetwork



@contextlib.contextmanager
def _printoptions(*args, **kwargs):
    original = np.get_printoptions()
    np.set_printoptions(*args, **kwargs)
    try:
        yield
    finally:
        np.set_printoptions(**original)

class Oracle(BaseOracle):
    pass


class Fixed(BaseFixed):
    pass


class Random(BaseRandom):
    pass


class ThompsonRecurrentNetwork(BaseThompsonRecurrentNetwork):
    """ Recurrent neural-linear: Thompson sampling based policy by using
        Bayesian linear regression on the representation(context) generated by
        the penultimate layer of the recurrent architecture. """

    def __init__(self, bandit, n_units, learning_rate, regularise_lambda, epochs, train_every,
                 std_targets, std_weights, verbose, seed):

        self.context_dims = bandit.context_dims

        BaseThompsonRecurrentNetwork.__init__(self, bandit, n_units,
        learning_rate, regularise_lambda, epochs, train_every, std_targets,
        std_weights, verbose, seed)

    def _setup_input(self):
        """ Returns the input and input size for reward prediction by the RNN. """

        self._arms = tf.placeholder(tf.int32, shape=None, name='arms')
        self._contexts = tf.placeholder(tf.float32, shape=[None, self.context_dims], name='contexts')

        rewards = tf.reshape(self._rewards[:-1], (-1, 1))
        # past rewards seen as features for prediction, first value is 0, last one can't be used as a feature for the next step

        arms_oh = tf.one_hot(self._arms, depth=self.bandit.n_arms,
                                dtype=tf.float32, name='arms_oh')

        contexts = self._contexts
        # Don't have a reward for the last context, ignore it while passing feed dict

        inputs = tf.concat([rewards, arms_oh, contexts], axis=1, name='inputs')
        isize = self.bandit.n_arms + 1 + self.context_dims

        return inputs, isize

    def _get_pred_from_sampled_model(self, trial, w, state):
        """ Obtains the predicted reward for every action under the sampled
        model(w). """

        pred = np.zeros(trial.n_arms)
        for arm in range(trial.n_arms):
            feed = {self._istate: state, self._arms: [arm],
                    self._rewards: [trial.rewards[-1], 0], self._contexts: [trial.contexts[-1]]}

            arm_features = self.session.run(self._h_output, feed)[0]
            pred[arm] = arm_features.dot(w)

        return pred

    def _get_feed_for_rnn(self, trial):
        """ Returns the input feed for the action selection by the policy. """

        feed = {self._arms: trial.arms, self._rewards: [0.] + trial.rewards, self._contexts: trial.contexts[:-1]}
        # Include contexts while dealing with contextual bandits

        return feed



class ThompsonSinFeedforwardNetwork(BaseThompsonSinFeedforwardNetwork):
    """ Neural-linear: Thompson sampling with Bayesian linear regression on
        the representation generated by a FFNN fed with handcrafted context and
         sinusoidal input units."""

    def __init__(self, bandit, order, periods, periods_dim, n_units,
                 learning_rate, regularise_lambda, epochs, train_every,
                 std_targets, std_weights, verbose, seed):

        self.context_dims = bandit.context_dims

        BaseThompsonSinFeedforwardNetwork.__init__(self, bandit, order, periods, periods_dim, n_units,
                     learning_rate, regularise_lambda, epochs, train_every,
                     std_targets, std_weights, verbose, seed)

        self._setup(seed)

    def _setup_input_size(self):
        """ Returns the number of units in the input layer, before and after
        including the sinusoidal units. """

        isize = len(self.periods)*2
        isize += self.bandit.n_arms + self.context_dims + (self.bandit.n_arms + self.context_dims+ 1) * self.order

        #add more units for the periodic transfromation
        isize_pre = isize + 1
        isize_post = isize + self.periods_dim

        return isize_pre, isize_post

    def _example(self, arms, rewards, contexts, t):
        """ Creates one sample of the handcrafted context-target pair. Used to
        generate the training data for training the NN representation model. """

        if t < self.order:
            raise Exception('Incomplete observation.')

        arms_oh = np.eye(self.bandit.n_arms)

        observation = []

        # Can be used for hard-coded periods
        for p in self.periods:
            angle = (2*np.pi*t)/p
            observation += [np.cos(angle), np.sin(angle)]

        observation.extend(arms_oh[arms[t]])

        observation.extend(contexts[t]) #Add the current context

        for i in range(self.order):
            observation.extend(arms_oh[arms[t - i - 1]])
            observation.append(rewards[t - i - 1])

            #Add the previous contexts
            observation.extend(contexts[t - i - 1])

        observation.append(t) #add the raw time-step as the new input
        target = rewards[t]

        return observation, target



contextual_policies = {'Oracle': Oracle,
            'Fixed': Fixed,
            'Random': Random,
            'ThompsonRecurrentNetwork': ThompsonRecurrentNetwork,
            'ThompsonSinFeedforwardNetwork': ThompsonSinFeedforwardNetwork}
