from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import tensorflow as tf
import tensorflow.contrib.rnn as rnn

from .rnn_cell_util import (
    GeneratorCellBuilder, GeneratorRNNCellBuilder, mean_field_cell,
    OutputRangeWrapper
)
from functools import partial
from itertools import product
from tensorflow.python.ops import variable_scope


class RNNCellGenerator(object):
    """Train a model to emit action sequences"""
    def __init__(self, m, seq_len, cell_builder, name='gen', reuse=False,
        **kwargs):
        self.m          = m
        self.seq_len    = seq_len
        self.cell_build = cell_builder
        self.name       = name
        self.action_seq = None
        self.logits     = None
        with tf.variable_scope(name, reuse=reuse):
            self.batch_size = tf.placeholder(tf.int32, [])
            self._build_generator(**kwargs)

    def _build_generator(self, feed_actions=True, init_type='zeros'):
        """Build the RNN TF sequence generator

        @feed_actions: Feed one-hot actions taken in previous step, rather than
            logits (from which action was sampled)
        @init_type: How to initialize the sequence generation:
            * train: Train a variable to init with
            * zeros: Send in an all-zeros vector
        """
        # Build the cell
        cell, state = self.cell_build.build_cell_and_init_state(
            self.batch_size, feed_actions
        )
        # If train input, train a variable input. Otherwise, all zeros
        if init_type.lower().startswith('train'):
            input_var = tf.Variable(tf.zeros((1, self.m), dtype=tf.float32))
            feed      = tf.tile(input_var, (self.batch_size, 1))
        elif init_type.lower().startswith('zero'):
            feed = tf.zeros((self.batch_size, self.m), dtype=tf.float32)
        else:
            raise ValueError(
                "RNN cell generator init_type %s not recognized" % init_type)

        # Placeholders to recover policy for updates
        self.rerun         = tf.placeholder_with_default(False, [])
        self.input_actions = tf.placeholder_with_default(
            tf.zeros((1, self.seq_len), dtype=tf.int32), (None, self.seq_len)
        )
        self.coo_actions   = tf.placeholder(tf.int32, (None, 3))
        # Run loopy feed forward
        actions_arr = tf.TensorArray(tf.int32, self.seq_len)
        logits_arr  = tf.TensorArray(tf.float32, self.seq_len)
        for t in range(self.seq_len):
            if t > 0:
                variable_scope.get_variable_scope().reuse_variables()
            # Compute logits for next action using RNN cell
            logits, state = cell(feed, state)
            # Samplers to draw actions
            def sample():
                return tf.to_int32(tf.multinomial(logits, 1))
            def rerun_sample():
                return self.input_actions[:, t]
            # If rerunning to apply policy gradients, draw is the input
            draw = tf.reshape(tf.cond(self.rerun, rerun_sample, sample), (-1,))
            # Write to arrays
            logits_arr  = logits_arr.write(t, logits)
            actions_arr = actions_arr.write(t, draw)
            # Update feed- either with the action taken (default), or with
            # the logits output at the previous timestep
            if feed_actions:
                feed = tf.one_hot(draw, self.m)
            else:
                feed = logits
        # Reshape logits to [batch_size, seq_len, n_actions]
        self.logits = tf.transpose(logits_arr.stack(), (1, 0, 2))
        # Reshape action_seq to [batch_size, seq_len]
        self.action_seq = tf.transpose(actions_arr.stack())

    def _get_generated_probabilities(self):
        """Returns a [batch_size, seq_len] Tensor with probabilities for each
           action that was drawn
        """
        input_batch_size = tf.shape(self.input_actions)[0]
        dists            = tf.nn.softmax(self.logits)
        r_dists          = tf.gather_nd(dists, self.coo_actions)
        return tf.reshape(r_dists, (input_batch_size, self.seq_len))

    def _build_discounts_matrix(self, T, gamma):
        """Build lower-triangular matrix of discounts.
        For example for T = 3: D = [[1,       0,     0]
                                   [gamma,   1,     0]
                                   [gamma^2, gamma, 1]]
        Then with R, our N x T incremental rewards matrix, the discounted sum is
            R * D
        """
        power_ltri  = tf.cumsum(
            tf.sequence_mask(tf.range(T)+1, T, dtype=tf.float32), exclusive=True
        )
        gamma_ltri  = tf.pow(gamma, power_ltri) 
        gamma_ltri *= tf.sequence_mask(tf.range(T)+1, T, dtype=tf.float32)
        return gamma_ltri

    def get_policy_loss_op(self, incremental_rewards, gamma):
        """Input is a [batch_size, seq_len] Tensor where each entry represents
           the incremental reward for an action on a data point
        """
        T = tf.shape(incremental_rewards)[1]
        # Form matrix of discounts to apply
        gamma_ltri = self._build_discounts_matrix(T, gamma)
        # Compute future discounted rewards as [batch_size x seq_len] matrix
        future_rewards = tf.matmul(incremental_rewards, gamma_ltri)
        # Compute baseline and advantage
        baseline   = tf.reduce_mean(future_rewards, axis=0)
        advantages = future_rewards - baseline
        # Apply advantage to policy
        policy = self._get_generated_probabilities()
        return tf.reduce_sum(tf.log(policy) * tf.stop_gradient(advantages))

    def get_action_sequence(self, session, batch_size):
        """Sample action sequences"""
        return session.run(self.action_seq, {self.batch_size: batch_size})

    def get_feed(self, actions, **kwargs):
        """Get the feed_dict for the training step.
        @action_seqs: The sequence of actions taken to generate the transformed
            data in this training step.

        Note that we feed `action_seqs` back in and set rerun=True to indicate
        that the exact same sequence of actions should be used in all other
        operations in this step!
        """
        coord   = product(range(actions.shape[0]), range(actions.shape[1]))
        feed    = {
            self.batch_size   : actions.shape[0],
            self.input_actions: actions,
            self.coo_actions:   [[i, j, actions[i, j]] for i, j in coord],
            self.rerun:         True,
        }
        kwargs.update(feed)
        return kwargs


class GRUGenerator(RNNCellGenerator):
    def __init__(self, m, seq_len, name='gen', reuse=False, n_stack=1,
        logit_range=4.0, **kwargs):
        # Get GRU cell builder
        range_wrapper = partial(OutputRangeWrapper, output_range=logit_range)
        cb = GeneratorRNNCellBuilder(
            rnn.GRUCell, m=m, n_stack=n_stack, wrappers=[range_wrapper]
        )
        # Super constructor
        super(GRUGenerator, self).__init__(
            m, seq_len, name=name, cell_builder=cb, reuse=reuse, **kwargs
        )


class LSTMGenerator(RNNCellGenerator):
    def __init__(self, m, seq_len, name='gen', reuse=False, n_stack=1,
        logit_range=4.0, **kwargs):
        # Get LSTM cell builder
        def norm(x):
            return 0.5 * (x + 1.)
        range_wrapper = partial(
            OutputRangeWrapper, output_range=logit_range, norm_op=norm
        )
        cb = GeneratorRNNCellBuilder(
            rnn.BasicLSTMCell, m=m, n_stack=n_stack, wrappers=[range_wrapper]
        )
        # Super constructor
        super(LSTMGenerator, self).__init__(
            m, seq_len, name=name, cell_builder=cb, reuse=reuse, **kwargs
        )


class MeanFieldGenerator(RNNCellGenerator):
    def __init__(self, m, seq_len, name='gen', reuse=False, **kwargs):
        # Get mean field cell builder
        cb = GeneratorCellBuilder(mean_field_cell)
        # Super constructor
        super(MeanFieldGenerator, self).__init__(
            m, seq_len, name=name, cell_builder=cb, reuse=reuse, **kwargs
        )
