# -*- coding: utf-8 -*-
"""
CODE ADAPTED FROM: https://github.com/sjblim/rmsn_nips_2018

[20180320 Treatment Effects with RNNs] model_rnn
Created on 21/3/2018 4:33 PM

@author: limsi
"""

import numpy as np
import tensorflow as tf
import treatments.RMSN.libs.net_helpers as helpers

_ACTIVATION_MAP = {"sigmoid": tf.nn.sigmoid, "elu": tf.nn.elu, "tanh": tf.nn.tanh, "linear": lambda x: x}


class StateDumpingRNN(tf.contrib.rnn.RNNCell):
    """ This RNNCell dumps out internal states for lstms"""

    def __init__(self, lstm):
        super(StateDumpingRNN, self).__init__()

        # Check that outputs
        self.lstm_cell = lstm

    @property
    def state_size(self):
        return self.lstm_cell.state_size

    @property
    def output_size(self):
        return self.lstm_cell.state_size

    def call(self, inputs, state):
        output, state = self.lstm_cell(inputs, state)

        return state, state


class Seq2SeqDecoderCell(tf.contrib.rnn.RNNCell):
    """ Decoder cell which allows for feedback, and external inputs during training """

    def __init__(self, lstm, W, b, b_training_mode=False):

        super(Seq2SeqDecoderCell, self).__init__()

        self.lstm_cell = lstm
        self.W = W
        self.b = b
        self._output_size = self.W.get_shape().as_list()[-1]

        self.b_training_mode = b_training_mode

    @property
    def state_size(self):

        if self.b_training_mode:  # use actual inputs
            return self.lstm_cell.state_size
        else:
            return self.lstm_cell.state_size + self._output_size

    @property
    def output_size(self):
        return self._output_size

    def call(self, inputs, state):

        # During training time, we assume that the previous input shape is [batch_size, action_vector + output_vector]
        # Output vectors are assumed to be at the end of the input or state vector (depending on train/test mode respectively)

        if self.b_training_mode:

            actual_states = state
            combined_inputs = inputs

        else:
            actual_states, prev_outputs = tf.split(state, [self.lstm_cell.state_size, self._output_size], axis=-1)

            combined_inputs = tf.concat([inputs, prev_outputs], axis=-1)

        # TODO: FIX HACK! THis forces this lstm to be in a different scope
        with tf.variable_scope("seq2seq"):
            output, state = self.lstm_cell(combined_inputs, actual_states)
            output = tf.matmul(output, self.W) + self.b

        if not self.b_training_mode:
            state = tf.concat([state, output], axis=-1)

        return output, state


class RnnModel:
    def __init__(self, params):

        # Generic params
        self.net_name = params["net_name"]
        self.experiment_name = params["experiment_name"]

        # Data params
        self.training_data = params["training_dataset"]
        self.validation_data = params["validation_dataset"]
        self.test_data = params["test_dataset"]
        self.input_size = params["input_size"]
        self.output_size = params["output_size"]

        # Network params
        self.softmax_size = params["softmax_size"]
        self.dropout_rate = params["dropout_rate"]
        self.hidden_layer_size = params["hidden_layer_size"]
        self.memory_activation_type = params["hidden_activation"]
        self.output_activation_type = params["output_activation"]
        self.b_use_seq2seq_feedback = params["use_seq2seq_feedback"]
        self.b_use_seq2seq_training_mode = params["use_seq2seq_training_mode"]

        # Memory Adapter params
        self.b_use_memory_adapter = False if "use_memory_adapter" not in params else params["use_memory_adapter"]
        self.memory_adapter_size = 0 if "memory_adapter_size" not in params else params["memory_adapter_size"]
        self.encoder_state_size = None

        # TODO: FIX THIS HACK FOR LOADING
        # Change scope for seq2seq network - so weights can be loaded later...
        variable_scope_name = "seq2seq" if "seq2seq" in self.net_name else "network"

        with tf.variable_scope(variable_scope_name):
            self.rnn_cell = tf.contrib.rnn.BasicLSTMCell(
                self.hidden_layer_size,
                activation=_ACTIVATION_MAP[self.memory_activation_type],
                state_is_tuple=False,
                name=variable_scope_name if variable_scope_name != "network" else None,
            )
            self.output_activation = _ACTIVATION_MAP[self.output_activation_type]
            self.output_w = tf.get_variable(
                "Output_W",
                [self.hidden_layer_size, self.output_size],
                dtype=tf.float32,
                initializer=tf.contrib.layers.xavier_initializer(),
            )
            self.output_b = tf.get_variable(
                "Output_b", [self.output_size], dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer()
            )

        # Training params
        self.performance_metric = params["performance_metric"]
        self.epochs = params["num_epochs"]
        self.minibatch_size = params["minibatch_size"]
        self.learning_rate = params["learning_rate"]
        self.max_global_norm = params["max_norm"]
        self.backprop_length = params["backprop_length"]
        self.global_step = tf.get_variable("global_step_tfrnn", initializer=0, dtype=np.int32, trainable=False)

        # Test params
        self.num_prediction_samples = 500

        # Saving params
        self.model_folder = params["model_folder"]
        relevant_name_parts = [
            self.experiment_name,
            # self.net_name,
            self.dropout_rate,
            self.hidden_layer_size,
            self.epochs,
            self.minibatch_size,
            self.learning_rate,
            self.max_global_norm,
            self.backprop_length,
        ]

        # Check
        if not (self.memory_activation_type == "elu" and self.output_activation_type == "linear"):
            relevant_name_parts += [self.memory_activation_type, self.output_activation_type]

        if self.memory_adapter_size > 0:
            relevant_name_parts += [self.memory_adapter_size]

        self.serialisation_name = "_".join([str(s) for s in relevant_name_parts])

    def _apply_memory_adapter(self, encoder_states):

        b_single_layer = self.memory_adapter_size == 0  # since externally checked that memory adapter should be applied

        if self.encoder_state_size is None:

            encoder_size = encoder_states.get_shape().as_list()[-1]
            self.encoder_state_size = encoder_size

            if b_single_layer:
                self.memory_adapter_layer = {
                    "W1": tf.get_variable(
                        "Adapter_Layer1_W",
                        [self.encoder_state_size, self.hidden_layer_size * 2],
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                    "b1": tf.get_variable(
                        "Adapter_Layer1_b",
                        [self.hidden_layer_size * 2],
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                }
            else:
                self.memory_adapter_layer = {
                    "W1": tf.get_variable(
                        "Adapter_Layer1_W",
                        [self.encoder_state_size, self.memory_adapter_size],
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                    "b1": tf.get_variable(
                        "Adapter_Layer1_b",
                        [self.memory_adapter_size],
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                    "W2": tf.get_variable(
                        "Adapter_Layer2_W",
                        [self.memory_adapter_size, self.hidden_layer_size * 2],
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                    "b2": tf.get_variable(
                        "Adapter_Layer2_b",
                        [self.hidden_layer_size * 2],
                        # LSTM memory is double concated
                        dtype=tf.float32,
                        initializer=tf.contrib.layers.xavier_initializer(),
                    ),
                }

        # Use elu and linear to avoid placing any restrictions on the range of internal activations
        memory_activation_fxn = _ACTIVATION_MAP[self.memory_activation_type]
        decoder_states = memory_activation_fxn(
            tf.matmul(encoder_states, self.memory_adapter_layer["W1"]) + self.memory_adapter_layer["b1"]
        )
        if not b_single_layer:
            decoder_states = memory_activation_fxn(
                tf.matmul(decoder_states, self.memory_adapter_layer["W2"]) + self.memory_adapter_layer["b2"]
            )

        return decoder_states

    def get_prediction_graph(
        self,
        use_validation_set,
        with_dropout=True,
        placeholder_time_steps=None,
        b_use_state_initialisation=False,
        b_dump_all_states=False,
    ):

        if placeholder_time_steps:
            data_chunk = {}
            data_chunk["inputs"] = tf.placeholder(tf.float32, [None, placeholder_time_steps, self.input_size])
            data_chunk["sequence_lengths"] = tf.placeholder(tf.float32, [None])  # Length
        else:
            if use_validation_set is None:
                dataset = self.training_data.batch(self.minibatch_size)
            elif use_validation_set:
                dataset = self.validation_data.batch(self.minibatch_size)
            else:
                dataset = self.test_data.batch(self.minibatch_size)

            iterator = tf.data.Iterator.from_structure(dataset.output_types, dataset.output_shapes)

            initializer = iterator.make_initializer(dataset)
            data_chunk = iterator.get_next()

        if b_use_state_initialisation:
            if "initial_states" not in data_chunk:
                raise ValueError("State initialisations not present!")

            initial_states = tf.cast(data_chunk["initial_states"], tf.float32)
        else:
            initial_states = None

        output = self._build_prediction_graph(
            data_chunk, with_dropout=with_dropout, initial_states=initial_states, b_dump_all_states=b_dump_all_states
        )

        if placeholder_time_steps:
            output["input_holder"] = data_chunk["inputs"]
            output["sequence_length_holder"] = data_chunk["sequence_lengths"]
        else:
            output["initializer"] = initializer

        return output

    def _build_prediction_graph(self, data_chunk, with_dropout=True, initial_states=None, b_dump_all_states=False):

        # output_minibatch = tf.cast(data_chunk['outputs'], tf.float32)
        # active_entries = tf.cast(data_chunk['active_entries'], tf.float32)
        input_minibatch = tf.cast(data_chunk["inputs"], tf.float32)
        sequence_lengths = tf.cast(data_chunk["sequence_lengths"], tf.int32)

        time_steps = input_minibatch.get_shape().as_list()[1]

        # Setup graph now
        outputs = []
        states_list = []
        if with_dropout:
            num_samples = self.num_prediction_samples
            keep_probs = 1 - self.dropout_rate
        else:
            num_samples = 1
            keep_probs = 1.0

        lstm_additional_size = (
            self.output_size if not self.b_use_seq2seq_training_mode and self.b_use_seq2seq_feedback else 0
        )
        cell = tf.nn.rnn_cell.DropoutWrapper(
            self.rnn_cell,
            input_keep_prob=keep_probs,
            output_keep_prob=keep_probs,
            state_keep_prob=keep_probs,
            variational_recurrent=True,
            input_size=input_minibatch.shape[2] + lstm_additional_size,
            dtype=tf.float32,
        )

        # Extension for feedback loops in seq2seq architecture
        if self.b_use_seq2seq_feedback:
            cell = Seq2SeqDecoderCell(cell, self.output_w, self.output_b, b_training_mode=False)

        # Extension for memory adapter
        if self.b_use_memory_adapter:

            if initial_states is None:
                raise ValueError("Memory adapter requires initial states!")

            initial_states = self._apply_memory_adapter(initial_states)

        for i in range(num_samples):

            val, states = tf.nn.dynamic_rnn(
                cell,
                input_minibatch,
                initial_state=initial_states,  # None for default
                dtype=tf.float32,
                sequence_length=sequence_lengths,
            )

            if b_dump_all_states:
                state_dumping_cell = StateDumpingRNN(cell)
                all_states, dumped_states = tf.nn.dynamic_rnn(
                    state_dumping_cell,
                    input_minibatch,
                    initial_state=initial_states,  # None for default
                    dtype=tf.float32,
                    sequence_length=sequence_lengths,
                )
            else:
                all_states = states  # just dump one state - used to speed up training while enforcing function params

            # Linear output layer
            flattened_val = tf.reshape(val, [-1, self.hidden_layer_size])

            if self.b_use_seq2seq_feedback:
                logits = flattened_val
            else:
                logits = tf.matmul(flattened_val, self.output_w) + self.output_b

            if self.softmax_size != 0:

                logits = tf.reshape(logits, [-1, time_steps, self.output_size])

                core_outputs, softmax_outputs = tf.split(
                    logits, [self.output_size - self.softmax_size, self.softmax_size], axis=2
                )

                output = tf.concat(
                    [self.output_activation(core_outputs), tf.nn.softmax(softmax_outputs, axis=2)], axis=2
                )

            else:

                output = self.output_activation(logits)
                output = tf.reshape(output, [-1, time_steps, self.output_size])

            outputs.append(tf.expand_dims(output, 0))
            states_list.append(tf.expand_dims(all_states, 0))

        # Dumping output
        samples = tf.concat(outputs, axis=0)
        mean_estimate = tf.reduce_mean(samples, axis=0)
        upper_bound = tf.contrib.distributions.percentile(samples, q=95.0, axis=0)
        lower_bound = tf.contrib.distributions.percentile(samples, q=5.0, axis=0)

        # Averages across all samples - no difference for single sample
        ave_state = tf.reduce_mean(tf.concat(states_list, axis=0), axis=0)

        return {"mean": mean_estimate, "upper_bound": upper_bound, "lower_bound": lower_bound, "ave_states": ave_state}

    def get_training_graph(self, use_truncated_bptt=True, b_stub_front=True, b_use_state_initialisation=True):

        training_dataset = self.training_data.shuffle(buffer_size=10000).batch(self.minibatch_size).repeat(self.epochs)

        iterator = training_dataset.make_one_shot_iterator()
        data_chunk = iterator.get_next()

        input_minibatch = tf.cast(data_chunk["inputs"], tf.float32)
        output_minibatch = tf.cast(data_chunk["outputs"], tf.float32)
        active_entries = tf.cast(data_chunk["active_entries"], tf.float32)
        sequence_lengths = tf.cast(data_chunk["sequence_lengths"], tf.int32)

        if b_use_state_initialisation:
            if "initial_states" not in data_chunk:
                raise ValueError("State initialisations not present!")

            initial_states = tf.cast(data_chunk["initial_states"], tf.float32)

            # Extension for memory adapter
            if self.b_use_memory_adapter:
                if initial_states is None:
                    raise ValueError("Memory adapter requires initial states!")
                initial_states = self._apply_memory_adapter(initial_states)

        else:
            initial_states = None

        if "propensity_weights" in data_chunk:
            weights = tf.cast(data_chunk["propensity_weights"], tf.float32)
        else:
            weights = 1

        keep_probs = 1 - self.dropout_rate

        # Setup graph now
        lstm_additional_size = (
            self.output_size if not self.b_use_seq2seq_training_mode and self.b_use_seq2seq_feedback else 0
        )
        cell = tf.nn.rnn_cell.DropoutWrapper(
            self.rnn_cell,
            input_keep_prob=keep_probs,
            output_keep_prob=keep_probs,
            state_keep_prob=keep_probs,
            variational_recurrent=True,
            input_size=input_minibatch.shape[2] + lstm_additional_size,
            dtype=tf.float32,
        )

        if self.b_use_seq2seq_feedback:
            cell = Seq2SeqDecoderCell(
                cell, self.output_w, self.output_b, b_training_mode=self.b_use_seq2seq_training_mode
            )

        # Stack up the dynamic RNNs for T-BPTT.

        # Splitting it up
        total_timesteps = input_minibatch.get_shape().as_list()[1]
        num_slices = int(total_timesteps / self.backprop_length)
        chunk_sizes = [self.backprop_length for i in range(num_slices)]
        odd_size = total_timesteps - self.backprop_length * num_slices

        # get all the chunks
        if odd_size > 0:
            if b_stub_front:
                chunk_sizes = [odd_size] + chunk_sizes
            else:
                chunk_sizes = chunk_sizes + [odd_size]

        # Implement TF style Truncated-backprop through time
        outputs = []
        start = 0
        states = initial_states

        for chunk_size in chunk_sizes:

            input_chunk = tf.slice(input_minibatch, [0, start, 0], [-1, chunk_size, self.input_size])

            if states is not None and use_truncated_bptt:
                val, states = tf.nn.dynamic_rnn(
                    cell, input_chunk, sequence_length=sequence_lengths, dtype=tf.float32, initial_state=states
                )
            else:
                val, states = tf.nn.dynamic_rnn(cell, input_chunk, sequence_length=sequence_lengths, dtype=tf.float32)

            # Linear output layer
            flattened_val = tf.reshape(val, [-1, self.hidden_layer_size])
            if self.b_use_seq2seq_feedback:
                logits = flattened_val
            else:
                logits = tf.matmul(flattened_val, self.output_w) + self.output_b

            if self.softmax_size != 0:

                logits = tf.reshape(logits, [-1, chunk_size, self.output_size])

                core_outputs, softmax_outputs = tf.split(
                    logits, [self.output_size - self.softmax_size, self.softmax_size], axis=2
                )

                output = tf.concat(
                    [self.output_activation(core_outputs), tf.nn.softmax(softmax_outputs, axis=2)], axis=2
                )

            else:
                output = self.output_activation(logits)
                output = tf.reshape(output, [-1, chunk_size, self.output_size])

            outputs.append(output)

            # break links between states for truncated bptt
            states = tf.identity(states)

            # Starting point
            start += chunk_size

        # Dumping output
        predictions = tf.concat(outputs, axis=1)

        # Split out the softmax components
        if self.softmax_size > 0:
            original_vs_softmax_size = [self.output_size - self.softmax_size, self.softmax_size]
            predictions, softmax_predictions = tf.split(predictions, original_vs_softmax_size, axis=2)
            output_minibatch, softmax_output_minibatch = tf.split(output_minibatch, original_vs_softmax_size, axis=2)
            active_entries, softmax_active = tf.split(active_entries, original_vs_softmax_size, axis=2)

        # Compute loss function
        if self.performance_metric == "mse":
            loss = tf.reduce_sum(tf.square(predictions - output_minibatch) * active_entries * weights) / tf.reduce_sum(
                active_entries
            )  # cos some zero entires

        elif self.performance_metric == "xentropy":

            loss = tf.reduce_sum(
                (
                    output_minibatch * -tf.log(predictions + 1e-8)
                    + (1 - output_minibatch) * -tf.log(1 - predictions + 1e-8)
                )
                * active_entries
                * weights
            ) / tf.reduce_sum(active_entries)
        else:
            raise ValueError("Unknown performance metric {}".format(self.performance_metric))

        if self.softmax_size > 0:
            loss += tf.reduce_sum(
                softmax_output_minibatch * -tf.log(softmax_predictions + 1e-8) * softmax_active * weights
            ) / tf.reduce_sum(softmax_active)

        optimiser = helpers.get_optimization_graph(
            loss, learning_rate=self.learning_rate, max_global_norm=self.max_global_norm, global_step=self.global_step
        )
        # Parcel outputs
        handles = {"loss": loss, "optimiser": optimiser}

        return handles
