import tensorflow as tf

def model_BiLSTM_multi(x, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length=-1):
    cell_fw = []
    cell_bw = []
    for i in range(layer_num):
        lstm_fw_cell = tf.contrib.rnn.BasicLSTMCell(lstmUnitNum, forget_bias=forget_bias, state_is_tuple=True,
                                                    activation=None, reuse=None)
        lstm_bw_cell = tf.contrib.rnn.BasicLSTMCell(lstmUnitNum, forget_bias=forget_bias, state_is_tuple=True,
                                                    activation=None, reuse=None)
        if attn_length != -1:
            lstm_fw_cell = tf.contrib.rnn.AttentionCellWrapper(lstm_fw_cell, attn_length, state_is_tuple=True)
            lstm_bw_cell = tf.contrib.rnn.AttentionCellWrapper(lstm_bw_cell, attn_length, state_is_tuple=True)
        lstm_fw_cell = tf.contrib.rnn.DropoutWrapper(cell=lstm_fw_cell, input_keep_prob=input_keep_prob, output_keep_prob=output_keep_prob)
        lstm_bw_cell = tf.contrib.rnn.DropoutWrapper(cell=lstm_bw_cell, input_keep_prob=input_keep_prob, output_keep_prob=output_keep_prob)
        cell_fw.append(lstm_fw_cell)
        cell_bw.append(lstm_bw_cell)

    Mul_cell_fw = tf.contrib.rnn.MultiRNNCell(cell_fw, state_is_tuple=True)
    Mul_cell_bw = tf.contrib.rnn.MultiRNNCell(cell_bw, state_is_tuple=True)

    # The input 'x' is [batch_size, n_steps, (n_input)n_dimensions]
    hiddens, state = tf.nn.bidirectional_dynamic_rnn(Mul_cell_fw, Mul_cell_bw, inputs=x, dtype=tf.float32)
    # fw:hiddens[0], bw:hiddens[1]
    hiddens = tf.concat(hiddens, axis=2)
    # The hiddens is [batch_size, n_steps, lstmUnitNum*2]
    # hiddens = tf.transpose(hiddens, [1, 0, 2])
    return hiddens

def model_HBiLSTM(input, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length=-1):
    with tf.variable_scope("layer1"):
        output_1 = model_BiLSTM_multi(input, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length)
    with tf.variable_scope("layer2"):
        output_2 = model_BiLSTM_multi(output_1, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length)
    highway_output = tf.concat([output_1, output_2], axis=1)
    highway_output = tf.transpose(highway_output, [1, 0, 2])
    return highway_output

def model_HBiLSTM_input(input, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length=-1):
    with tf.variable_scope("layer1"):
        output_1 = model_BiLSTM_multi(input, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length)
    output_2 = tf.concat([output_1, input], axis=1)
    with tf.variable_scope("layer2"):
        highway_output = model_BiLSTM_multi(output_2, lstmUnitNum, layer_num, forget_bias, input_keep_prob, output_keep_prob, attn_length)
    highway_output = tf.transpose(highway_output, [1, 0, 2])
    return highway_output
