import keras
from keras.engine import Input, Model
from keras.layers import Embedding, Conv1D, GlobalMaxPooling1D, Dense, Dropout, Activation, Reshape


def get_shared_cnn(vocabSize, embeddingSize, input_length, filters, kernel_size,
                   hidden_dims, drop_out):
    x = y = Input(shape=(input_length,))
    y = Embedding(input_dim=vocabSize, output_dim=embeddingSize)(y)
    y = Conv1D(filters, kernel_size, activation='relu')(y)
    y = GlobalMaxPooling1D()(y)
    y = Dense(hidden_dims)(y)
    y = Dropout(drop_out)(y)
    y = Activation('relu')(y)
    return Model(x, y)


def get_full_model(vocabSize, embeddingSize, maxLen, domainFeatures, flowFeatures,
                   filters, h1, h2, dropout, dense):
    pass


def get_top_cnn(cnn, numFeatures, maxLen, windowSize, domainFeatures, filters, kernel_size, cnnHiddenDims, cnnDropout):
    inputList = []
    encodedList = []
    # TODO: ???
    for i in range(windowSize):
        inputList.append(Input(shape=(maxLen,)))
        encodedList.append(cnn(inputList[-1]))  # add shared domain model
        inputList.append(Input(shape=(numFeatures,)))
    # TODO: ???
    merge_layer_input = []
    for i in range(windowSize):
        merge_layer_input.append(encodedList[i])
        merge_layer_input.append(inputList[(2 * i) + 1])
    # We can then concatenate the two vectors:
    merged_vector = keras.layers.concatenate(merge_layer_input, axis=-1)
    reshape = Reshape((windowSize, domainFeatures + numFeatures))(merged_vector)
    # add second cnn
    cnn = Conv1D(filters,
                 kernel_size,
                 activation='relu',
                 input_shape=(windowSize, domainFeatures + numFeatures))(reshape)
    # we use max pooling:
    maxPool = GlobalMaxPooling1D()(cnn)
    cnnDropout = Dropout(cnnDropout)(maxPool)
    cnnDense = Dense(cnnHiddenDims, activation='relu')(cnnDropout)
    cnnOutput1 = Dense(2, activation='softmax')(cnnDense)
    cnnOutput2 = Dense(2, activation='softmax')(cnnDense)

    # We define a trainable model linking the
    # tweet inputs to the predictions
    return Model(inputs=inputList, outputs=(cnnOutput1, cnnOutput2))