import tensorflow as tf

def model_CNN(input_r, numDimensions, windowSize, filterNum, keepratio):
    w_conv = tf.Variable(tf.random_normal([windowSize, numDimensions, 1, filterNum], stddev=0.1))
    b_conv = tf.Variable(tf.random_normal([filterNum], stddev=0.1))
    conv = tf.nn.conv2d(input_r, w_conv, strides=[1, 1, 1, 1], padding='VALID')
    conv = tf.nn.relu(tf.nn.bias_add(conv, b_conv))
    conv = tf.reshape(conv, shape=[conv.shape[0].value, -1, filterNum])
    # pool = tf.nn.max_pool(conv, ksize=[1, input_r.shape[1].value-windowSize+1, 1, 1], strides=[1, 1, 1, 1], padding='VALID')
    conv = tf.nn.dropout(conv, keepratio)
    return conv

def model_BiLSTM(x, lstmUnitNum, forget_bias, input_keep_prob, output_keep_prob, attn_length=-1):
    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)

    # The input 'x' is [batch_size, n_steps, (n_input)n_dimensions]
    hiddens, state = tf.nn.bidirectional_dynamic_rnn(cell_fw=lstm_fw_cell, cell_bw=lstm_bw_cell, 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_CRNN(input_r, numDimensions, filterNum, lstmUnitNum, keepratio_cnn, input_keep_prob, output_keep_prob, forget_bias_rnn):
    p3 = model_CNN(input_r, numDimensions, 3, filterNum, keepratio_cnn)
    p5 = model_CNN(input_r, numDimensions, 5, filterNum, keepratio_cnn)
    output_cnn = tf.concat([p3, p5], axis=1)
    output_rnn = model_BiLSTM(output_cnn, lstmUnitNum, forget_bias_rnn, input_keep_prob, output_keep_prob)
    return output_rnn
