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


def get_shared_cnn(vocab_size, embedding_size, input_length, filters, kernel_size,
                   hidden_dims, drop_out):
    x = y = Input(shape=(input_length,))
    y = Embedding(input_dim=vocab_size, output_dim=embedding_size)(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):
    ipt_domains = Input(shape=(windowSize, maxLen), name="ipt_domains")
    encoded = TimeDistributed(cnn)(ipt_domains)
    ipt_flows = Input(shape=(windowSize, numFeatures), name="ipt_flows")
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    # add second cnn
    y = Conv1D(filters,
               kernel_size,
               activation='relu',
               input_shape=(windowSize, domainFeatures + numFeatures))(merged)
    # TODO: why global pooling? -> 3D to 2D
    # we use max pooling:
    y = GlobalMaxPooling1D()(y)
    y = Dropout(cnnDropout)(y)
    y = Dense(cnnHiddenDims, activation='relu')(y)
    y1 = Dense(2, activation='softmax', name="client")(y)
    y2 = Dense(2, activation='softmax', name="server")(y)

    return Model(inputs=[ipt_domains, ipt_flows], outputs=(y1, y2))