"""
modeling
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time
import shutil
import numpy as np
import tensorflow as tf
from tensorflow.contrib import rnn
from tensorflow.contrib import layers
from utils_cnn import Norm


class C_MODEL(object):
    """ 
    """

    def __init__(self, config, graph):
        """ TO build up the graph
        Inputs
        ------
        config : 
            * batch_size : mini batch size
            * log_dir : path to save training summary
            * learning_rate : adam's learning rate
            * hidden_size : number of hidden units in LSTM
            * rnn_layers : number of stacked LSTM 
            * seq_length : length of LSTM
            * num_features : dimensions of input feature
            * latent_dims : dimensions of latent feature
            * penalty_lambda = gradient penalty's weight, ref from  paper of 'improved-wgan'
        graph : 
            tensorflow default graph
        """
        self.normer = Norm()
        # hyper-parameters
        self.batch_size = config.batch_size
        self.log_dir = config.log_dir
        self.learning_rate = config.learning_rate
        self.hidden_size = config.hidden_size
        self.rnn_layers = config.rnn_layers
        self.seq_length = config.seq_length
        self.num_features = config.num_features
        self.latent_dims = config.latent_dims
        self.penalty_lambda = config.penalty_lambda
        self.if_log_histogram = config.if_log_histogram
        # steps
        self.__global_steps = tf.train.get_or_create_global_step(graph=graph)
        self.__D_steps = 0
        # data
        self.__G_samples = tf.placeholder(dtype=tf.float32, shape=[
            self.batch_size, self.seq_length, self.normer.PLAYERS, self.normer.COLS, self.normer.ROWS], name='G_samples')
        self.__X = tf.placeholder(dtype=tf.float32, shape=[
            self.batch_size, self.seq_length, self.normer.PLAYERS, self.normer.COLS, self.normer.ROWS], name='real_data')
        # adversarial learning : wgan
        self.__build_wgan()

        # summary
        self.__summary_D_op = tf.summary.merge(tf.get_collection('D'))
        self.__summary_D_valid_op = tf.summary.merge(
            tf.get_collection('D_valid'))
        self.D_summary_writer = tf.summary.FileWriter(
            self.log_dir + 'D')
        self.D_valid_summary_writer = tf.summary.FileWriter(
            self.log_dir + 'D_valid')

    def __build_wgan(self):
        with tf.name_scope('WGAN'):
            D_real = self.inference(self.__X, seq_len=None)
            __D_fake = self.inference(
                self.__G_samples, seq_len=None, reuse=True)
            # loss function
            self.__D_loss, F_real, F_fake, grad_pen = self.__D_loss_fn(
                self.__X, self.__G_samples, __D_fake, D_real, self.penalty_lambda)
            theta_D = self.__get_var_list()
            with tf.name_scope('D_optimizer') as scope:
                D_optimizer = tf.train.AdamOptimizer(
                    learning_rate=self.learning_rate, beta1=0.5, beta2=0.9)
                D_grads = tf.gradients(self.__D_loss, theta_D)
                D_grads = list(zip(D_grads, theta_D))
                self.__D_train_op = D_optimizer.apply_gradients(
                    grads_and_vars=D_grads, global_step=self.__global_steps)
            # logging
            for grad, var in D_grads:
                self.__summarize(var.name, grad, collections='D',
                                 postfix='gradient')
            tf.summary.scalar('D_loss', self.__D_loss,
                              collections=['D', 'D_valid'])
            tf.summary.scalar('F_real', F_real, collections=['D'])
            tf.summary.scalar('F_fake', F_fake, collections=['D'])
            tf.summary.scalar('grad_pen', grad_pen, collections=['D'])

    def __summarize(self, name, value, collections, postfix=''):
        """ Helper to create summaries for activations.
        Creates a summary that provides a histogram of activations.
        Creates a summary that measures the sparsity of activations.
        Args
        ----
        name : string
        value : Tensor
        collections : list of string
        postfix : string
        Returns
        -------
            nothing
        """
        if self.if_log_histogram:
            tensor_name = name + '/' + postfix
            tf.summary.histogram(tensor_name,
                                 value, collections=collections)
            # tf.summary.scalar(tensor_name + '/sparsity',
            #                   tf.nn.zero_fraction(x), collections=collections)

    def __get_var_list(self):
        """ to get both Generator's and Discriminator's trainable variables
        and add trainable variables into histogram
        """
        trainable_V = tf.trainable_variables()
        theta_D = []
        for _, v in enumerate(trainable_V):
            if v.name.startswith('D'):
                theta_D.append(v)
                self.__summarize(v.op.name, v, collections='D',
                                 postfix='Trainable')
        return theta_D

    def __leaky_relu(self, features, alpha=0.7):
        return tf.maximum(features, alpha * features)

    def __lstm_cell(self):
        return rnn.LSTMCell(self.hidden_size, use_peepholes=True, initializer=None,
                            forget_bias=1.0, state_is_tuple=True,
                            # activation=self.__leaky_relu, cell_clip=2,
                            activation=tf.nn.tanh, reuse=tf.get_variable_scope().reuse)

    def inference(self, inputs, seq_len=None, reuse=False):
        """
        Inputs
        ------
        inputs : float, shape=[batch_size, seq_length=100, PLAYERS=11, COLS=98, ROWS=46]
            real(from data) or fake(from G)
        seq_len : 
            temparily not used

        Return
        ------
        decision : bool
            real(from data) or fake(from G)
        """
        with tf.variable_scope('D', reuse=reuse) as scope:
            # unstack, axis=1 -> [batch, time, feature]
            print(inputs)
            inputs = tf.transpose(inputs, perm=[0, 1, 3, 4, 2])
            print(inputs)
            inputs = tf.unstack(inputs, num=self.seq_length, axis=1)
            blstm_input = []
            output_list = []
            for time_step in range(self.seq_length):
                with tf.variable_scope('conv') as scope:
                    if time_step > 0:
                        tf.get_variable_scope().reuse_variables()
                    filters_list = [32, 64, 128, 256]
                    next_input = inputs[time_step]
                    for i in range(len(filters_list)):
                        with tf.variable_scope('conv' + str(i)) as scope:
                            conv = layers.conv2d(
                                inputs=next_input,
                                num_outputs=filters_list[i],
                                kernel_size=[5, 5],
                                stride=2,
                                padding='SAME',
                                activation_fn=tf.nn.relu,
                                weights_initializer=layers.xavier_initializer(
                                    uniform=False),
                                weights_regularizer=None,
                                biases_initializer=tf.zeros_initializer(),
                                reuse=scope.reuse,
                                scope=scope
                            )
                            next_input = conv
                    with tf.variable_scope('fc') as scope:
                        flat_input = layers.flatten(next_input)
                        fc = layers.fully_connected(
                            inputs=flat_input,
                            num_outputs=self.hidden_size,
                            activation_fn=tf.nn.relu,
                            weights_initializer=layers.xavier_initializer(
                                uniform=False),
                            biases_initializer=tf.zeros_initializer(),
                            reuse=scope.reuse,
                            scope=scope
                        )
                        blstm_input.append(fc)
            with tf.variable_scope('stack_blstm') as scope:
                stack_blstm, _, _ = rnn.stack_bidirectional_rnn(
                    cells_fw=[self.__lstm_cell()
                              for _ in range(self.rnn_layers)],
                    cells_bw=[self.__lstm_cell()
                              for _ in range(self.rnn_layers)],
                    inputs=blstm_input,
                    dtype=tf.float32,
                    sequence_length=seq_len
                )
            with tf.variable_scope('output') as scope:
                for i, out_blstm in enumerate(stack_blstm):
                    if i > 0:
                        tf.get_variable_scope().reuse_variables()
                    with tf.variable_scope('fc') as scope:
                        fc = layers.fully_connected(
                            inputs=out_blstm,
                            num_outputs=1,
                            activation_fn=self.__leaky_relu,
                            weights_initializer=layers.xavier_initializer(
                                uniform=False),
                            biases_initializer=tf.zeros_initializer(),
                            reuse=scope.reuse,
                            scope=scope
                        )
                        output_list.append(fc)
            # stack, axis=1 -> [batch, time, feature]
            decisions = tf.stack(output_list, axis=1)
            print('decisions', decisions)
            decision = tf.reduce_mean(decisions, axis=1)
            print('decision', decision)
            return decision

    def __D_loss_fn(self, __X, __G_sample, D_fake, D_real, penalty_lambda):
        """ D loss
        """
        with tf.name_scope('D_loss') as scope:
            # grad_pen, base on paper (Improved WGAN)
            epsilon = tf.random_uniform(
                [self.batch_size, 1, 1, 1, 1], minval=0.0, maxval=1.0)
            __X_inter = epsilon * __X + (1.0 - epsilon) * __G_sample
            grad = tf.gradients(
                self.inference(__X_inter, seq_len=None, reuse=True), [__X_inter])[0]
            print(grad)
            sum_ = tf.reduce_sum(tf.square(grad), axis=[1, 2, 3, 4])
            print(sum_)
            grad_norm = tf.sqrt(sum_)
            grad_pen = penalty_lambda * tf.reduce_mean(
                tf.square(grad_norm - 1.0))
            f_fake = tf.reduce_mean(D_fake)
            f_real = tf.reduce_mean(D_real)

            loss = f_fake - f_real + grad_pen
        return loss, f_real, f_fake, grad_pen

    def step(self, sess, G_samples, real_data):
        """ train one batch on D
        """
        self.__D_steps += 1
        feed_dict = {self.__G_samples: G_samples,
                     self.__X: real_data}
        loss, global_steps, _ = sess.run(
            [self.__D_loss, self.__global_steps, self.__D_train_op], feed_dict=feed_dict)
        if not self.if_log_histogram or self.__D_steps % 500 == 0:  # % 500 to save space
            summary = sess.run(self.__summary_D_op, feed_dict=feed_dict)
            # log
            self.D_summary_writer.add_summary(
                summary, global_step=global_steps)
        return loss, global_steps

    def D_log_valid_loss(self, sess, G_samples, real_data):
        """ one batch valid loss
        """
        feed_dict = {self.__G_samples: G_samples,
                     self.__X: real_data}
        loss, global_steps = sess.run(
            [self.__D_loss, self.__global_steps], feed_dict=feed_dict)
        if not self.if_log_histogram or self.__D_steps % 500 == 0:  # % 500 to save space
            summary = sess.run(self.__summary_D_valid_op, feed_dict=feed_dict)
            # log
            self.D_valid_summary_writer.add_summary(
                summary, global_step=global_steps)
        return loss

[(39, 'utils_cnn.Norm', 'Norm', False, 'from utils_cnn import Norm\n'), (52, 'tensorflow.train.get_or_create_global_step', 'tf.train.get_or_create_global_step', True, 'import tensorflow as tf\n'), (55, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (57, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (66, 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', True, 'import tensorflow as tf\n'), (68, 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', True, 'import tensorflow as tf\n'), (122, 'tensorflow.trainable_variables', 'tf.trainable_variables', True, 'import tensorflow as tf\n'), (132, 'tensorflow.maximum', 'tf.maximum', True, 'import tensorflow as tf\n'), (63, 'tensorflow.get_collection', 'tf.get_collection', True, 'import tensorflow as tf\n'), (65, 'tensorflow.get_collection', 'tf.get_collection', True, 'import tensorflow as tf\n'), (72, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (91, 'tensorflow.summary.scalar', 'tf.summary.scalar', True, 'import tensorflow as tf\n'), (93, 'tensorflow.summary.scalar', 'tf.summary.scalar', True, 'import tensorflow as tf\n'), (94, 'tensorflow.summary.scalar', 'tf.summary.scalar', True, 'import tensorflow as tf\n'), (95, 'tensorflow.summary.scalar', 'tf.summary.scalar', True, 'import tensorflow as tf\n'), (113, 'tensorflow.summary.histogram', 'tf.summary.histogram', True, 'import tensorflow as tf\n'), (154, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (157, 'tensorflow.transpose', 'tf.transpose', True, 'import tensorflow as tf\n'), (159, 'tensorflow.unstack', 'tf.unstack', True, 'import tensorflow as tf\n'), (225, 'tensorflow.stack', 'tf.stack', True, 'import tensorflow as tf\n'), (227, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (234, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (236, 'tensorflow.random_uniform', 'tf.random_uniform', True, 'import tensorflow as tf\n'), (244, 'tensorflow.sqrt', 'tf.sqrt', True, 'import tensorflow as tf\n'), (247, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (248, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (80, 'tensorflow.name_scope', 'tf.name_scope', True, 'import tensorflow as tf\n'), (81, 'tensorflow.train.AdamOptimizer', 'tf.train.AdamOptimizer', True, 'import tensorflow as tf\n'), (83, 'tensorflow.gradients', 'tf.gradients', True, 'import tensorflow as tf\n'), (198, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (208, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (242, 'tensorflow.square', 'tf.square', True, 'import tensorflow as tf\n'), (138, 'tensorflow.get_variable_scope', 'tf.get_variable_scope', True, 'import tensorflow as tf\n'), (163, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (246, 'tensorflow.square', 'tf.square', True, 'import tensorflow as tf\n'), (185, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (186, 'tensorflow.contrib.layers.flatten', 'layers.flatten', False, 'from tensorflow.contrib import layers\n'), (212, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (165, 'tensorflow.get_variable_scope', 'tf.get_variable_scope', True, 'import tensorflow as tf\n'), (191, 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', False, 'from tensorflow.contrib import layers\n'), (193, 'tensorflow.zeros_initializer', 'tf.zeros_initializer', True, 'import tensorflow as tf\n'), (211, 'tensorflow.get_variable_scope', 'tf.get_variable_scope', True, 'import tensorflow as tf\n'), (217, 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', False, 'from tensorflow.contrib import layers\n'), (219, 'tensorflow.zeros_initializer', 'tf.zeros_initializer', True, 'import tensorflow as tf\n'), (177, 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', False, 'from tensorflow.contrib import layers\n'), (180, 'tensorflow.zeros_initializer', 'tf.zeros_initializer', True, 'import tensorflow as tf\n')]