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

@@ -14,215 +14,133 @@
 yers
-.convolutional import MaxPooling1D, Convolution1D%0Afrom keras.layers.core import Flatten, Dropout, Dense, Merge%0Afrom keras.layers.normalization import BatchNormalization%0Afrom keras.layers.recurrent import GRU
+ import Input, Dense, GRU, Dropout, BatchNormalization, %5C%0A                         MaxPooling1D, Conv1D, Flatten, Concatenate
 %0Afro
@@ -161,25 +161,20 @@
  import 
-Sequentia
+Mode
 l%0A%0Afrom 
@@ -632,24 +632,25 @@
 f a CNN %22%22%22%0A
+%0A
     NB_FILTE
@@ -713,13 +713,26 @@
 yers
- = %5B%5D
+, inputs = %5B%5D, %5B%5D%0A
 %0A   
@@ -779,60 +779,116 @@
     
-ngram_layer = Sequential()%0A        ngram_layer.add(C
+current_input = Input(shape=(SAMPLE_LENGTH, embedding_size))%0A        inputs.append(current_input)%0A%0A        c
 onvo
@@ -893,16 +893,23 @@
 volution
+ = Conv
 1D(%0A    
@@ -969,90 +969,26 @@
     
-input_dim=embedding_size,%0A            input_length=SAMPLE_LENGTH,%0A            init
+kernel_initializer
 ='le
@@ -1037,25 +1037,40 @@
 ',%0A        )
-)
+(current_input)%0A
 %0A        poo
@@ -1071,22 +1071,20 @@
    pool_
-length
+size
  = SAMPL
@@ -1123,24 +1123,18 @@
     
-ngram_layer.add(
+pooling = 
 MaxP
@@ -1151,27 +1151,35 @@
 ool_
-length=pool_length)
+size=pool_size)(convolution
 )%0A  
@@ -1207,111 +1207,72 @@
 end(
-ngram_layer)%0A%0A    model = Sequential()%0A    model.add(Merge(conv_layers, mode='concat')
+pooling)%0A%0A    merged = Concatenate()(conv_layers
 )%0A
-%0A
     
-model.add(
+dropout = 
 Drop
@@ -1283,92 +1283,169 @@
 0.5)
-)%0A    model.add(Flatten())%0A%0A    model.add(Dense(output_length, activation='sigmoid')
+(merged)%0A    flattened = Flatten()(dropout)%0A    outputs = Dense(output_length, activation='sigmoid')(flattened)%0A%0A    model = Model(inputs=inputs, outputs=outputs
 )%0A%0A 
@@ -1728,44 +1728,73 @@
     
-model = Sequential()%0A%0A    model.add(
+inputs = Input(shape=(SAMPLE_LENGTH, embedding_size))%0A%0A    gru = 
 GRU(
@@ -1839,12 +1839,30 @@
 put_
-dim=
+shape=(SAMPLE_LENGTH, 
 embe
@@ -1871,16 +1871,17 @@
 ing_size
+)
 ,%0A      
@@ -1886,50 +1886,28 @@
     
-input_length=SAMPLE_LENGTH,%0A        init='
+kernel_initializer=%22
 glor
@@ -1908,33 +1908,33 @@
 =%22glorot_uniform
-'
+%22
 ,%0A        inner_
@@ -1931,18 +1931,29 @@
     
-inner_init
+recurrent_initializer
 ='no
@@ -1991,32 +1991,52 @@
 ',%0A    )
+(inputs
 )%0A
+%0A
     
-model.add(
+batch_normalization = 
 BatchNor
@@ -2051,24 +2051,28 @@
 on()
+(gru
 )%0A    
-model.add(
+dropout = 
 Drop
@@ -2083,25 +2083,44 @@
 0.1)
-)%0A%0A    model.add(
+(batch_normalization)%0A    outputs = 
 Dens
@@ -2157,16 +2157,75 @@
 igmoid')
+(dropout)%0A%0A    model = Model(inputs=inputs, outputs=outputs
 )%0A%0A    m

rewrite the NN models to use Keras 2.0 API