from keras.layers.convolutional import MaxPooling1D, Convolution1D
from keras.layers.core import Flatten, Dropout, Dense, Merge
from keras.layers.normalization import BatchNormalization
from keras.layers.recurrent import GRU
from keras.models import Sequential

from magpie.config import SAMPLE_LENGTH


def get_nn_model(nn_model, embedding, output_length):
    if nn_model == 'cnn':
        return cnn(embedding_size=embedding, output_length=output_length)
    elif nn_model == 'rnn':
        return rnn(embedding_size=embedding, output_length=output_length)
    else:
        raise ValueError("Unknown NN type: {}".format(nn_model))


def cnn(embedding_size, output_length):
    """ Create and return a keras model of a CNN """
    NB_FILTER = 256
    NGRAM_LENGTHS = [1, 2, 3, 4, 5]

    conv_layers = []
    for ngram_length in NGRAM_LENGTHS:
        ngram_layer = Sequential()
        ngram_layer.add(Convolution1D(
            NB_FILTER,
            ngram_length,
            input_dim=embedding_size,
            input_length=SAMPLE_LENGTH,
            init='lecun_uniform',
            activation='tanh',
        ))
        pool_length = SAMPLE_LENGTH - ngram_length + 1
        ngram_layer.add(MaxPooling1D(pool_length=pool_length))
        conv_layers.append(ngram_layer)

    model = Sequential()
    model.add(Merge(conv_layers, mode='concat'))

    model.add(Dropout(0.5))
    model.add(Flatten())

    model.add(Dense(output_length, activation='sigmoid'))

    model.compile(
        loss='binary_crossentropy',
        optimizer='adam',
        metrics=['top_k_categorical_accuracy'],
    )

    return model


def rnn(embedding_size, output_length):
    """ Create and return a keras model of a RNN """
    HIDDEN_LAYER_SIZE = 256

    model = Sequential()

    model.add(GRU(
        HIDDEN_LAYER_SIZE,
        input_dim=embedding_size,
        input_length=SAMPLE_LENGTH,
        init='glorot_uniform',
        inner_init='normal',
        activation='relu',
    ))
    model.add(BatchNormalization())
    model.add(Dropout(0.1))

    model.add(Dense(output_length, activation='sigmoid'))

    model.compile(
        loss='binary_crossentropy',
        optimizer='adam',
        metrics=['top_k_categorical_accuracy'],
    )

    return model

from keras.layers import Input, Dense, GRU, Dropout, BatchNormalization, \
                         MaxPooling1D, Conv1D, Flatten, Concatenate
from keras.models import Model

from magpie.config import SAMPLE_LENGTH


def get_nn_model(nn_model, embedding, output_length):
    if nn_model == 'cnn':
        return cnn(embedding_size=embedding, output_length=output_length)
    elif nn_model == 'rnn':
        return rnn(embedding_size=embedding, output_length=output_length)
    else:
        raise ValueError("Unknown NN type: {}".format(nn_model))


def cnn(embedding_size, output_length):
    """ Create and return a keras model of a CNN """

    NB_FILTER = 256
    NGRAM_LENGTHS = [1, 2, 3, 4, 5]

    conv_layers, inputs = [], []

    for ngram_length in NGRAM_LENGTHS:
        current_input = Input(shape=(SAMPLE_LENGTH, embedding_size))
        inputs.append(current_input)

        convolution = Conv1D(
            NB_FILTER,
            ngram_length,
            kernel_initializer='lecun_uniform',
            activation='tanh',
        )(current_input)

        pool_size = SAMPLE_LENGTH - ngram_length + 1
        pooling = MaxPooling1D(pool_size=pool_size)(convolution)
        conv_layers.append(pooling)

    merged = Concatenate()(conv_layers)
    dropout = Dropout(0.5)(merged)
    flattened = Flatten()(dropout)
    outputs = Dense(output_length, activation='sigmoid')(flattened)

    model = Model(inputs=inputs, outputs=outputs)

    model.compile(
        loss='binary_crossentropy',
        optimizer='adam',
        metrics=['top_k_categorical_accuracy'],
    )

    return model


def rnn(embedding_size, output_length):
    """ Create and return a keras model of a RNN """
    HIDDEN_LAYER_SIZE = 256

    inputs = Input(shape=(SAMPLE_LENGTH, embedding_size))

    gru = GRU(
        HIDDEN_LAYER_SIZE,
        input_shape=(SAMPLE_LENGTH, embedding_size),
        kernel_initializer="glorot_uniform",
        recurrent_initializer='normal',
        activation='relu',
    )(inputs)

    batch_normalization = BatchNormalization()(gru)
    dropout = Dropout(0.1)(batch_normalization)
    outputs = Dense(output_length, activation='sigmoid')(dropout)

    model = Model(inputs=inputs, outputs=outputs)

    model.compile(
        loss='binary_crossentropy',
        optimizer='adam',
        metrics=['top_k_categorical_accuracy'],
    )

    return model

rewrite the NN models to use Keras 2.0 API