"""
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

[((1065, 1071), 'utils_cnn.Norm', 'Norm', ([], {}), '()\n', (1069, 1071), False, 'from utils_cnn import Norm\n'), ((1613, 1660), 'tensorflow.train.get_or_create_global_step', 'tf.train.get_or_create_global_step', ([], {'graph': 'graph'}), '(graph=graph)\n', (1647, 1660), True, 'import tensorflow as tf\n'), ((1730, 1883), 'tensorflow.placeholder', 'tf.placeholder', ([], {'dtype': 'tf.float32', 'shape': '[self.batch_size, self.seq_length, self.normer.PLAYERS, self.normer.COLS,\n    self.normer.ROWS]', 'name': '"""G_samples"""'}), "(dtype=tf.float32, shape=[self.batch_size, self.seq_length,\n    self.normer.PLAYERS, self.normer.COLS, self.normer.ROWS], name='G_samples')\n", (1744, 1883), True, 'import tensorflow as tf\n'), ((1912, 2065), 'tensorflow.placeholder', 'tf.placeholder', ([], {'dtype': 'tf.float32', 'shape': '[self.batch_size, self.seq_length, self.normer.PLAYERS, self.normer.COLS,\n    self.normer.ROWS]', 'name': '"""real_data"""'}), "(dtype=tf.float32, shape=[self.batch_size, self.seq_length,\n    self.normer.PLAYERS, self.normer.COLS, self.normer.ROWS], name='real_data')\n", (1926, 2065), True, 'import tensorflow as tf\n'), ((2359, 2400), 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', (["(self.log_dir + 'D')"], {}), "(self.log_dir + 'D')\n", (2380, 2400), True, 'import tensorflow as tf\n'), ((2452, 2499), 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', (["(self.log_dir + 'D_valid')"], {}), "(self.log_dir + 'D_valid')\n", (2473, 2499), True, 'import tensorflow as tf\n'), ((4888, 4912), 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), '()\n', (4910, 4912), True, 'import tensorflow as tf\n'), ((5257, 5295), 'tensorflow.maximum', 'tf.maximum', (['features', '(alpha * features)'], {}), '(features, alpha * features)\n', (5267, 5295), True, 'import tensorflow as tf\n'), ((2207, 2229), 'tensorflow.get_collection', 'tf.get_collection', (['"""D"""'], {}), "('D')\n", (2224, 2229), True, 'import tensorflow as tf\n'), ((2297, 2325), 'tensorflow.get_collection', 'tf.get_collection', (['"""D_valid"""'], {}), "('D_valid')\n", (2314, 2325), True, 'import tensorflow as tf\n'), ((2555, 2576), 'tensorflow.name_scope', 'tf.name_scope', (['"""WGAN"""'], {}), "('WGAN')\n", (2568, 2576), True, 'import tensorflow as tf\n'), ((3600, 3672), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""D_loss"""', 'self.__D_loss'], {'collections': "['D', 'D_valid']"}), "('D_loss', self.__D_loss, collections=['D', 'D_valid'])\n", (3617, 3672), True, 'import tensorflow as tf\n'), ((3715, 3769), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""F_real"""', 'F_real'], {'collections': "['D']"}), "('F_real', F_real, collections=['D'])\n", (3732, 3769), True, 'import tensorflow as tf\n'), ((3782, 3836), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""F_fake"""', 'F_fake'], {'collections': "['D']"}), "('F_fake', F_fake, collections=['D'])\n", (3799, 3836), True, 'import tensorflow as tf\n'), ((3849, 3907), 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""grad_pen"""', 'grad_pen'], {'collections': "['D']"}), "('grad_pen', grad_pen, collections=['D'])\n", (3866, 3907), True, 'import tensorflow as tf\n'), ((4457, 4522), 'tensorflow.summary.histogram', 'tf.summary.histogram', (['tensor_name', 'value'], {'collections': 'collections'}), '(tensor_name, value, collections=collections)\n', (4477, 4522), True, 'import tensorflow as tf\n'), ((6045, 6080), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""D"""'], {'reuse': 'reuse'}), "('D', reuse=reuse)\n", (6062, 6080), True, 'import tensorflow as tf\n'), ((6194, 6236), 'tensorflow.transpose', 'tf.transpose', (['inputs'], {'perm': '[0, 1, 3, 4, 2]'}), '(inputs, perm=[0, 1, 3, 4, 2])\n', (6206, 6236), True, 'import tensorflow as tf\n'), ((6284, 6331), 'tensorflow.unstack', 'tf.unstack', (['inputs'], {'num': 'self.seq_length', 'axis': '(1)'}), '(inputs, num=self.seq_length, axis=1)\n', (6294, 6331), True, 'import tensorflow as tf\n'), ((9720, 9749), 'tensorflow.stack', 'tf.stack', (['output_list'], {'axis': '(1)'}), '(output_list, axis=1)\n', (9728, 9749), True, 'import tensorflow as tf\n'), ((9815, 9848), 'tensorflow.reduce_mean', 'tf.reduce_mean', (['decisions'], {'axis': '(1)'}), '(decisions, axis=1)\n', (9829, 9848), True, 'import tensorflow as tf\n'), ((10038, 10061), 'tensorflow.name_scope', 'tf.name_scope', (['"""D_loss"""'], {}), "('D_loss')\n", (10051, 10061), True, 'import tensorflow as tf\n'), ((10148, 10220), 'tensorflow.random_uniform', 'tf.random_uniform', (['[self.batch_size, 1, 1, 1, 1]'], {'minval': '(0.0)', 'maxval': '(1.0)'}), '([self.batch_size, 1, 1, 1, 1], minval=0.0, maxval=1.0)\n', (10165, 10220), True, 'import tensorflow as tf\n'), ((10566, 10579), 'tensorflow.sqrt', 'tf.sqrt', (['sum_'], {}), '(sum_)\n', (10573, 10579), True, 'import tensorflow as tf\n'), ((10701, 10723), 'tensorflow.reduce_mean', 'tf.reduce_mean', (['D_fake'], {}), '(D_fake)\n', (10715, 10723), True, 'import tensorflow as tf\n'), ((10745, 10767), 'tensorflow.reduce_mean', 'tf.reduce_mean', (['D_real'], {}), '(D_real)\n', (10759, 10767), True, 'import tensorflow as tf\n'), ((2981, 3009), 'tensorflow.name_scope', 'tf.name_scope', (['"""D_optimizer"""'], {}), "('D_optimizer')\n", (2994, 3009), True, 'import tensorflow as tf\n'), ((3050, 3128), 'tensorflow.train.AdamOptimizer', 'tf.train.AdamOptimizer', ([], {'learning_rate': 'self.learning_rate', 'beta1': '(0.5)', 'beta2': '(0.9)'}), '(learning_rate=self.learning_rate, beta1=0.5, beta2=0.9)\n', (3072, 3128), True, 'import tensorflow as tf\n'), ((3176, 3212), 'tensorflow.gradients', 'tf.gradients', (['self.__D_loss', 'theta_D'], {}), '(self.__D_loss, theta_D)\n', (3188, 3212), True, 'import tensorflow as tf\n'), ((8341, 8373), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""stack_blstm"""'], {}), "('stack_blstm')\n", (8358, 8373), True, 'import tensorflow as tf\n'), ((8832, 8859), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""output"""'], {}), "('output')\n", (8849, 8859), True, 'import tensorflow as tf\n'), ((10482, 10497), 'tensorflow.square', 'tf.square', (['grad'], {}), '(grad)\n', (10491, 10497), True, 'import tensorflow as tf\n'), ((5604, 5627), 'tensorflow.get_variable_scope', 'tf.get_variable_scope', ([], {}), '()\n', (5625, 5627), True, 'import tensorflow as tf\n'), ((6464, 6489), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""conv"""'], {}), "('conv')\n", (6481, 6489), True, 'import tensorflow as tf\n'), ((10652, 10678), 'tensorflow.square', 'tf.square', (['(grad_norm - 1.0)'], {}), '(grad_norm - 1.0)\n', (10661, 10678), True, 'import tensorflow as tf\n'), ((7660, 7683), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""fc"""'], {}), "('fc')\n", (7677, 7683), True, 'import tensorflow as tf\n'), ((7731, 7757), 'tensorflow.contrib.layers.flatten', 'layers.flatten', (['next_input'], {}), '(next_input)\n', (7745, 7757), False, 'from tensorflow.contrib import layers\n'), ((9051, 9074), 'tensorflow.variable_scope', 'tf.variable_scope', (['"""fc"""'], {}), "('fc')\n", (9068, 9074), True, 'import tensorflow as tf\n'), ((6562, 6585), 'tensorflow.get_variable_scope', 'tf.get_variable_scope', ([], {}), '()\n', (6583, 6585), True, 'import tensorflow as tf\n'), ((8018, 8058), 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', ([], {'uniform': '(False)'}), '(uniform=False)\n', (8043, 8058), False, 'from tensorflow.contrib import layers\n'), ((8140, 8162), 'tensorflow.zeros_initializer', 'tf.zeros_initializer', ([], {}), '()\n', (8160, 8162), True, 'import tensorflow as tf\n'), ((8984, 9007), 'tensorflow.get_variable_scope', 'tf.get_variable_scope', ([], {}), '()\n', (9005, 9007), True, 'import tensorflow as tf\n'), ((9336, 9376), 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', ([], {'uniform': '(False)'}), '(uniform=False)\n', (9361, 9376), False, 'from tensorflow.contrib import layers\n'), ((9458, 9480), 'tensorflow.zeros_initializer', 'tf.zeros_initializer', ([], {}), '()\n', (9478, 9480), True, 'import tensorflow as tf\n'), ((7252, 7292), 'tensorflow.contrib.layers.xavier_initializer', 'layers.xavier_initializer', ([], {'uniform': '(False)'}), '(uniform=False)\n', (7277, 7292), False, 'from tensorflow.contrib import layers\n'), ((7440, 7462), 'tensorflow.zeros_initializer', 'tf.zeros_initializer', ([], {}), '()\n', (7460, 7462), True, 'import tensorflow as tf\n')]