"""
    Defines the interaction between policy and bandit.
    Also includes the base policies for the different types of problems -
    non-contextual, contextual and linear bandits.

"""

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


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

class Trial:

    def __init__(self, n_arms):
        self.n_arms = n_arms

        self.returns = np.zeros(n_arms, dtype=np.float)
        self.pulls = np.zeros(n_arms, dtype=np.int)

        self.arms = []
        self.rewards = []
        self.contexts = []
        self.regrets = []

        self.length = 0

    def update_contexts(self, context):
        self.contexts.append(context)

    def append(self, arm, reward, context, regret):
        if (arm >= self.n_arms) or (arm < 0):
            raise Exception('Invalid arm.')

        self.arms.append(arm)
        self.rewards.append(reward)
        self.update_contexts(context)
        self.regrets.append(regret)

        self.length += 1

        self.returns[arm] += reward
        self.pulls[arm] += 1

    def average_rewards(self):
        out = np.full(self.n_arms, float('inf'))
        where = (self.pulls > 0)

        return np.divide(self.returns, self.pulls, out=out, where=where)

    def cumulative_rewards(self):
        return np.cumsum(self.rewards)

    def cumulative_regret(self):
        return np.cumsum(self.regrets)


class Policy:

    def __init__(self, bandit):
        self.bandit = bandit
        self.contextual_bandit = 0

        if hasattr(bandit, 'reset'):
            self.contextual_bandit = 1

    def select(self, trial):
        raise NotImplementedError()

    def interact(self, trial_length):
        trial = Trial(self.bandit.n_arms)

        if self.contextual_bandit:
            reset_context = self.bandit.reset()
            trial.update_contexts(reset_context)

        for i in range(trial_length):
            arm = self.select(trial)
            reward, context, regret = self.bandit.pull(arm)

            trial.append(arm, reward, context, regret)

        return trial


class BaseOracle(Policy):
    """Selects the best available action at every time step. """

    def __init__(self, bandit):
        Policy.__init__(self, bandit)

    def select(self, trial):
        return int(self.bandit.best_arms()[0])

    def __repr__(self):
        return 'Oracle()'


class BaseFixed(Policy):
    """Selects a fixed chosen action at every time step. """

    def __init__(self, bandit, arm):
        Policy.__init__(self, bandit)
        self.arm = arm

    def select(self, trial):
        return self.arm

    def __repr__(self):
        return 'Fixed(arm={0})'.format(self.arm)


class BaseRandom(Policy):
    """Selects a random action at every time step. """

    def __init__(self, bandit, seed):
        Policy.__init__(self, bandit)
        self.random_state = np.random.RandomState(seed)

    def select(self, trial):
        return self.random_state.randint(trial.n_arms)

    def __repr__(self):
        return 'Random()'

class BaseThompsonRecurrentNetwork(Policy):
    """ 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):
        Policy.__init__(self, bandit)

        self.n_units = n_units
        if len(self.n_units) < 2:
            raise Exception('Invalid number of layers.')

        self.learning_rate = learning_rate
        self.regularise_lambda = regularise_lambda
        self.epochs = epochs
        self.train_every = train_every

        self.verbose = verbose

        self.random_state = np.random.RandomState(seed=seed)

        self.var_targets = std_targets**2.
        self.var_weights = std_weights**2.
        self.one_over_lambda = self.var_weights/self.var_targets

        self._setup(seed)

    def _setup(self, seed):
        """ Creates the recurrent representation architecture. """

        tf.reset_default_graph()
        tf.set_random_seed(seed)

        self._rewards = tf.placeholder(tf.float32, shape=None, name='rewards')

        inputs, isize = self._setup_input()

        W = tf.Variable(tf.truncated_normal(shape=(isize, self.n_units[0])),
                                            name='W0')
        b = tf.Variable(tf.zeros(shape=(self.n_units[0])), name='b0')


        rnn_inputs = tf.matmul(inputs, W) + b
        rnn_inputs = tf.expand_dims(rnn_inputs, axis=0)

        cell = tf.contrib.rnn.LSTMCell(num_units=self.n_units[1])

        self._istate = cell.zero_state(1, dtype=tf.float32)

        rnn_outputs, \
            self._final_state = tf.nn.dynamic_rnn(cell, rnn_inputs,
                                                  initial_state=self._istate)

        self._h_output = tf.reshape(rnn_outputs, (-1, self.n_units[1]))
        for i in range(2, len(self.n_units)):
            W = tf.Variable(tf.truncated_normal(shape=(self.n_units[i - 1],
                                                       self.n_units[i])),
                            name='W{0}'.format(i))
            b = tf.Variable(tf.zeros(shape=(self.n_units[i])),
                            name='b{0}'.format(i))

            self._h_output = tf.tanh(tf.matmul(self._h_output, W) + b)


        W = tf.Variable(tf.truncated_normal(shape=(self.n_units[-1], 1)),
                        name='Wout')


        # Note: No bias in the output layer
        self._pred = tf.matmul(self._h_output, W)
        self._pred = tf.reshape(self._pred, (-1,))

        self._reg_loss = sum(tf.nn.l2_loss(tf_var)
                for tf_var in tf.trainable_variables()
                if not ("b" in tf_var.name))

        self._loss = tf.reduce_mean((self._pred - self._rewards[1:])**2) + self.regularise_lambda * self._reg_loss
        # Predictions compared with reward from 1 step ahead

        optimizer = tf.train.AdamOptimizer(self.learning_rate)
        self._train = optimizer.minimize(self._loss)

        config = tf.ConfigProto(device_count={'GPU': 0})
        self.session = tf.Session(config=config)
        self.session.run(tf.global_variables_initializer())


    def select(self, trial):
        """ Selects which arm to play in the current round. """

        if trial.length < trial.n_arms:
            # the case when insufficient number of arms have been played.
            return np.argmin(trial.pulls)

        else:
            return self._select_from_policy(trial)


    def _select_from_policy(self, trial):
        """ Selects which arm to play in the current round - according to the
        policy. """

        feed = self._get_feed_for_rnn(trial)


        if trial.length % self.train_every == 0:
            # Train the RNN weights when this condition is true
            for _ in range(self.epochs):
                self.session.run(self._train, feed)

        loss, state, observations = self.session.run([self._loss,
                                                      self._final_state,
                                                      self._h_output], feed)

        targets = np.array(trial.rewards)

        # Update the posterior for Bayesian linear regression

        mean, cov = self._update_bayesian_lr_posterior(observations, targets)

        # Sample from posterior for Bayesian linear regression
        w = self.random_state.multivariate_normal(mean, cov)

        # Exploration with posterior sampling
        # Pull the arm with the highest prediction under the sampled model

        pred = self._get_pred_from_sampled_model(trial, w, state)

        if self.verbose:
            msg = 'Pull: {0:4d}. Loss: {1:.4f}. Prediction: {2}. Policy: {3}.'

            p = np.zeros(trial.n_arms)
            p[np.argmax(pred)] = 1.0

            with _printoptions(precision=4, suppress=True):
                if np.argmax(p) in self.bandit.best_arms():
                    cprint(msg.format(trial.length + 1, loss, pred, p), 'green')
                else:
                    print(msg.format(trial.length + 1, loss, pred, p))

        return np.argmax(pred)

    def _update_bayesian_lr_posterior(self, observations, targets):
        """ Updates the posterior for Bayesian linear regression. """

        M = self.one_over_lambda * observations.T.dot(observations)
        M = np.linalg.inv(M + np.eye(M.shape[0]))

        mean = (self.one_over_lambda * M).dot(observations.T.dot(targets))
        cov = self.var_weights*M

        return mean, cov

    def _setup_input(self):
        raise NotImplementedError()

    def _get_feed_for_rnn(self, trial):
        raise NotImplementedError()

    def _get_pred_from_sampled_model(self, trial, w, state):
        raise NotImplementedError()

    def __repr__(self):
        r = 'ThompsonRecurrentNetwork(n_units={0}, learning_rate={1}, '
        r += 'regL2={2}, epochs={3}, train_every={4}, std_targets={5}, std_weights={6})'
        return r.format(self.n_units, self.learning_rate,
                        self.regularise_lambda, self.epochs,
                        self.train_every, np.sqrt(self.var_targets),
                        np.sqrt(self.var_weights))


class BaseThompsonSinFeedforwardNetwork(Policy):
    """ Neural-linear: Posterior 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):

        Policy.__init__(self, bandit)

        self.order = order
        self.periods = np.array(periods) # Can be used to provide hard-coded periods

        self.periods_dim = periods_dim # Number of sinudoidal units
        self.n_units = n_units

        self.learning_rate = learning_rate
        self.regularise_lambda = regularise_lambda
        self.epochs = epochs
        self.train_every = train_every

        self.var_targets = std_targets**2.
        self.var_weights = std_weights**2.
        self.one_over_lambda = self.var_weights/self.var_targets

        self.verbose = verbose

        self.random_state = np.random.RandomState(seed=seed)

        self._setup(seed)

    def _setup(self, seed):
        """ Creates the NN architecture used in neural-linear. """

        tf.reset_default_graph()
        tf.set_random_seed(seed)

        isize_pre, isize_post = self._setup_input_size()

        self._observations = tf.placeholder(tf.float32, shape=[None, isize_pre],
                                            name='observations')
        self._targets = tf.placeholder(tf.float32, shape=None, name='targets')

        # Re-create _observations to include the periodic transformations

        self._observations_orig = self._observations[:, :-1]
        self._raw_t = self._observations[:, -1]

        W = tf.Variable(tf.truncated_normal(shape=(1, self.periods_dim)),
                        name='W_periodic')
        b = tf.Variable(tf.zeros(shape=(self.periods_dim)),
                        name='b_periodic')

        # Sinusoidal activation
        self._raw_t = tf.reshape(self._raw_t, [-1, 1])
        t_op = tf.sin(tf.matmul(self._raw_t, W) + b)

        # Observations now include the periodic functions of the time step.
        self._observations_new = tf.concat([self._observations_orig, t_op], axis = -1)

        n_units = [isize_post] + self.n_units
        self._h_output = self._observations_new

        for i in range(1, len(n_units)):
            W = tf.Variable(tf.truncated_normal(shape=(n_units[i - 1],
                                                       n_units[i])),
                            name='W{0}'.format(i))
            b = tf.Variable(tf.zeros(shape=(n_units[i])),
                            name='b{0}'.format(i))

            if i == 1:
                # No activation in the first layer
                self._h_output = tf.matmul(self._h_output, W) + b
            else:
                self._h_output = tf.tanh(tf.matmul(self._h_output, W) + b)


        # Note: No bias in the output layer
        W = tf.Variable(tf.truncated_normal(shape=(n_units[-1], 1)),
                        name='Wout')

        self._pred = tf.matmul(self._h_output, W)
        self._pred = tf.reshape(self._pred, (-1,))

        # L2 regularization on weights
        self._reg_loss = sum(tf.nn.l2_loss(tf_var)
                for tf_var in tf.trainable_variables()
                if not ("b" in tf_var.name))

        self._loss = tf.reduce_mean((self._pred - self._targets)**2) + \
                    self.regularise_lambda * self._reg_loss

        optimizer = tf.train.AdamOptimizer(self.learning_rate)
        self._train = optimizer.minimize(self._loss)

        config = tf.ConfigProto(device_count={'GPU': 0})
        self.session = tf.Session(config=config)
        self.session.run(tf.global_variables_initializer())

    def select(self, trial):
        """ Selects which arm to play in the current round. """

        if trial.length < max(self.order + 1, trial.n_arms):
            # the case when insufficient number of arms have been played.
            return np.argmin(trial.pulls)

        else:
            return self._select_from_policy(trial)

    def _select_from_policy(self, trial):
        """ Selects which arm to play in the current round - according to the
        policy. """

        observations, targets = self._get_observation_target_pair(trial)

        feed = {self._observations: observations, self._targets: targets}

        if trial.length % self.train_every == 0:
            # Train the NN weights when this condition is true
            for _ in range(self.epochs):
                self.session.run(self._train, feed)

        loss, observations = self.session.run([self._loss, self._h_output],
                                              feed)


        # Observations now contain the output activation of the penultimate layer of the NN,
        # on which we perform Bayesian Linear Regression

        # Update the posterior for Bayesian linear regression

        mean, cov = self._update_bayesian_lr_posterior(observations, targets)

        # Sample from posterior for Bayesian linear regression

        w = self.random_state.multivariate_normal(mean, cov)

        # Exploration with posterior sampling
        # Pull the arm with the highest prediction under the sampled model

        pred = self._get_pred_from_sampled_model(trial, w)

        if self.verbose:
            msg = 'Pull: {0:4d}. Loss: {1:.4f}. Prediction: {2}. Policy: {3}.'

            p = np.zeros(trial.n_arms)
            p[np.argmax(pred)] = 1.0

            with _printoptions(precision=4, suppress=True):
                if np.argmax(p) in self.bandit.best_arms():
                    cprint(msg.format(trial.length + 1, loss, pred, p), 'green')
                else:
                    print(msg.format(trial.length + 1, loss, pred, p))

        return np.argmax(pred)

    def _get_observation_target_pair(self, trial):
        """ Creates the dataset for training the NN representation model. """
        observations = []
        targets = []

        for t in range(self.order, trial.length):
            observation, target = self._example(trial.arms, trial.rewards, trial.contexts[:-1], t)
            # remove the last context for which you don't have a pull.

            observations.append(observation)
            targets.append(target)

        return observations, targets

    def _get_pred_from_sampled_model(self, trial, w):
        """ 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):
            observation, _ = self._example(trial.arms + [arm],
                                           trial.rewards + [0.],
                                           trial.contexts,
                                           trial.length)
            feed = {self._observations: [observation]}
            arm_features = self.session.run(self._h_output, feed)[0]
            pred[arm] = arm_features.dot(w)

        return pred

    def _update_bayesian_lr_posterior(self, observations, targets):
        """ Updates the posterior for Bayesian linear regression weights. """

        M = self.one_over_lambda * observations.T.dot(observations)
        M = np.linalg.inv(M + np.eye(M.shape[0]))

        mean = (self.one_over_lambda * M).dot(observations.T.dot(targets))
        cov = self.var_weights*M

        return mean, cov

    def _example(self, arms, rewards, contexts, t):
        raise NotImplementedError()

    def _setup_input_size(self):
        raise NotImplementedError()

    def __repr__(self):
        r = 'ThompsonSinFeedforwardNetwork(order={0}, periods={1}, periods_dim={2},'
        r += ' n_units={3}, learning_rate={4}, regL2={5}, epochs={6}, train_every={7},'
        r += ' std_targets={8}, std_weights={9})'
        return r.format(self.order, self.periods, self.periods_dim, self.n_units,
                        self.learning_rate, self.regularise_lambda,
                        self.epochs, self.train_every,
                        np.sqrt(self.var_targets), np.sqrt(self.var_weights))
