'Define model architecture.'

def build_model(self, hps):

    if hps.is_training:
        self.global_step = tf.Variable(0, name='global_step', trainable=False)
    if (hps.dec_model == 'lstm'):
        cell_fn = rnn.LSTMCell
    elif (hps.dec_model == 'layer_norm'):
        cell_fn = rnn.LayerNormLSTMCell
    elif (hps.dec_model == 'hyper'):
        cell_fn = rnn.HyperLSTMCell
    else:
        assert False, 'please    choose    a    respectable    cell'
    if (hps.enc_model == 'lstm'):
        enc_cell_fn = rnn.LSTMCell
    elif (hps.enc_model == 'layer_norm'):
        enc_cell_fn = rnn.LayerNormLSTMCell
    elif (hps.enc_model == 'hyper'):
        enc_cell_fn = rnn.HyperLSTMCell
    else:
        assert False, 'please    choose    a    respectable    cell'
    use_recurrent_dropout = self.hps.use_recurrent_dropout
    use_input_dropout = self.hps.use_input_dropout
    use_output_dropout = self.hps.use_output_dropout
    if (hps.dec_model == 'hyper'):
        cell = cell_fn(hps.dec_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
    else:
        cell = cell_fn(hps.dec_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
    if hps.conditional:
        if (hps.enc_model == 'hyper'):
            self.enc_cell_fw = enc_cell_fn(hps.enc_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
            self.enc_cell_bw = enc_cell_fn(hps.enc_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
        else:
            self.enc_cell_fw = enc_cell_fn(hps.enc_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
            self.enc_cell_bw = enc_cell_fn(hps.enc_rnn_size, use_recurrent_dropout=use_recurrent_dropout, dropout_keep_prob=self.hps.recurrent_dropout_prob)
    tf.logging.info('Input    dropout    mode    =    %s.', use_input_dropout)
    tf.logging.info('Output    dropout    mode    =    %s.', use_output_dropout)
    tf.logging.info('Recurrent    dropout    mode    =    %s.', use_recurrent_dropout)
    if use_input_dropout:
        tf.logging.info('Dropout    to    input    w/    keep_prob    =    %4.4f.', self.hps.input_dropout_prob)
        cell = tf.contrib.rnn.DropoutWrapper(cell, input_keep_prob=self.hps.input_dropout_prob)
    if use_output_dropout:
        tf.logging.info('Dropout    to    output    w/    keep_prob    =    %4.4f.', self.hps.output_dropout_prob)
        cell = tf.contrib.rnn.DropoutWrapper(cell, output_keep_prob=self.hps.output_dropout_prob)
    self.cell = cell
    self.sequence_lengths = tf.placeholder(dtype=tf.int32, shape=[self.hps.batch_size])
    self.input_data = tf.placeholder(dtype=tf.float32, shape=[self.hps.batch_size, (self.hps.max_seq_len + 1), 5])
    self.output_x = self.input_data[:, 1:(self.hps.max_seq_len + 1), :]
    self.input_x = self.input_data[:, :self.hps.max_seq_len, :]
    if hps.conditional:
        (self.mean, self.presig) = self.encoder(self.output_x, self.sequence_lengths)
        self.sigma = tf.exp((self.presig / 2.0))
        eps = tf.random_normal((self.hps.batch_size, self.hps.z_size), 0.0, 1.0, dtype=tf.float32)
        self.batch_z = (self.mean + tf.multiply(self.sigma, eps))
        self.kl_cost = ((-0.5) * tf.reduce_mean((((1 + self.presig) - tf.square(self.mean)) - tf.exp(self.presig))))
        self.kl_cost = tf.maximum(self.kl_cost, self.hps.kl_tolerance)
        pre_tile_y = tf.reshape(self.batch_z, [self.hps.batch_size, 1, self.hps.z_size])
        overlay_x = tf.tile(pre_tile_y, [1, self.hps.max_seq_len, 1])
        actual_input_x = tf.concat([self.input_x, overlay_x], 2)
        self.initial_state = tf.nn.tanh(rnn.super_linear(self.batch_z, cell.state_size, init_w='gaussian', weight_start=0.001, input_size=self.hps.z_size))
    else:
        self.batch_z = tf.zeros((self.hps.batch_size, self.hps.z_size), dtype=tf.float32)
        self.kl_cost = tf.zeros([], dtype=tf.float32)
        actual_input_x = self.input_x
        self.initial_state = cell.zero_state(batch_size=hps.batch_size, dtype=tf.float32)
    self.num_mixture = hps.num_mixture
    n_out = (3 + (self.num_mixture * 6))
    with tf.variable_scope('RNN'):
        output_w = tf.get_variable('output_w', [self.hps.dec_rnn_size, n_out])
        output_b = tf.get_variable('output_b', [n_out])
    (output, last_state) = tf.nn.dynamic_rnn(cell, actual_input_x, initial_state=self.initial_state, time_major=False, swap_memory=True, dtype=tf.float32, scope='RNN')
    output = tf.reshape(output, [(-1), hps.dec_rnn_size])
    output = tf.nn.xw_plus_b(output, output_w, output_b)
    self.final_state = last_state
    def tf_2d_normal(x1, x2, mu1, mu2, s1, s2, rho):
        'Returns    result    of    eq    #    24    of    http://arxiv.org/abs/1308.0850.'
        norm1 = tf.subtract(x1, mu1)
        norm2 = tf.subtract(x2, mu2)
        s1s2 = tf.multiply(s1, s2)
        z = ((tf.square(tf.div(norm1, s1)) + tf.square(tf.div(norm2, s2))) - (2 * tf.div(tf.multiply(rho, tf.multiply(norm1, norm2)), s1s2)))
        neg_rho = (1 - tf.square(rho))
        result = tf.exp(tf.div((- z), (2 * neg_rho)))
        denom = ((2 * np.pi) * tf.multiply(s1s2, tf.sqrt(neg_rho)))
        result = tf.div(result, denom)
        return result
    def get_lossfunc(z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr, z_pen_logits, x1_data, x2_data, pen_data):
        'Returns    a    loss    fn    based    on    eq    #26    of    http://arxiv.org/abs/1308.0850.'
        result0 = tf_2d_normal(x1_data, x2_data, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr)
        epsilon = 1e-06
        result1 = tf.multiply(result0, z_pi)
        result1 = tf.reduce_sum(result1, 1, keep_dims=True)
        result1 = (- tf.log((result1 + epsilon)))
        fs = (1.0 - pen_data[:, 2])
        fs = tf.reshape(fs, [(-1), 1])
        result1 = tf.multiply(result1, fs)
        result2 = tf.nn.softmax_cross_entropy_with_logits(labels=pen_data, logits=z_pen_logits)
        result2 = tf.reshape(result2, [(-1), 1])
        if (not self.hps.is_training):
            result2 = tf.multiply(result2, fs)
        result = (result1 + result2)
        return result
    def get_mixture_coef(output):
        'Returns    the    tf    slices    containing    mdn    dist    params.'
        z = output
        z_pen_logits = z[:, 0:3]
        (z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr) = tf.split(z[:, 3:], 6, 1)
        z_pi = tf.nn.softmax(z_pi)
        z_pen = tf.nn.softmax(z_pen_logits)
        z_sigma1 = tf.exp(z_sigma1)
        z_sigma2 = tf.exp(z_sigma2)
        z_corr = tf.tanh(z_corr)
        r = [z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr, z_pen, z_pen_logits]
        return r
    out = get_mixture_coef(output)
    [o_pi, o_mu1, o_mu2, o_sigma1, o_sigma2, o_corr, o_pen, o_pen_logits] = out
    self.pi = o_pi
    self.mu1 = o_mu1
    self.mu2 = o_mu2
    self.sigma1 = o_sigma1
    self.sigma2 = o_sigma2
    self.corr = o_corr
    self.pen_logits = o_pen_logits
    self.pen = o_pen
    target = tf.reshape(self.output_x, [(-1), 5])
    [x1_data, x2_data, eos_data, eoc_data, cont_data] = tf.split(target, 5, 1)
    pen_data = tf.concat([eos_data, eoc_data, cont_data], 1)
    lossfunc = get_lossfunc(o_pi, o_mu1, o_mu2, o_sigma1, o_sigma2, o_corr, o_pen_logits, x1_data, x2_data, pen_data)
    self.r_cost = tf.reduce_mean(lossfunc)
    if self.hps.is_training:
        self.lr = tf.Variable(self.hps.learning_rate, trainable=False)
        optimizer = tf.train.AdamOptimizer(self.lr)
        self.kl_weight = tf.Variable(self.hps.kl_weight_start, trainable=False)
        self.cost = (self.r_cost + (self.kl_cost * self.kl_weight))
        gvs = optimizer.compute_gradients(self.cost)
        g = self.hps.grad_clip
        capped_gvs = [(tf.clip_by_value(grad, (- g), g), var) for (grad, var) in gvs]
        self.train_op = optimizer.apply_gradients(capped_gvs, global_step=self.global_step, name='train_step')
