import tensorflow as tf

from deepr.layers import base


@base.layer(n_in=2, n_out=3)
def LSTM(tensors, num_units: int, bidirectional: bool = False, **kwargs):
    """LSTM layer."""
    words, nwords = tensors
    t = tf.transpose(words, perm=[1, 0, 2])
    lstm_cell_fw = tf.contrib.rnn.LSTMBlockFusedCell(num_units=num_units, **kwargs)
    outputs_fw, (hidden_fw, output_fw) = lstm_cell_fw(t, dtype=tf.float32, sequence_length=nwords)

    if bidirectional:
        lstm_cell_bw = tf.contrib.rnn.LSTMBlockFusedCell(num_units=num_units, **kwargs)
        lstm_cell_bw = tf.contrib.rnn.TimeReversedFusedRNN(lstm_cell_bw)
