from keras import optimizers, losses, activations, models

from keras.layers import Dense, Input, Dropout, Convolution1D, MaxPool1D, GlobalMaxPool1D, GlobalAveragePooling1D, \

    concatenate, SpatialDropout1D, TimeDistributed, Bidirectional, LSTM

from keras_contrib.layers import CRF

from utils import WINDOW_SIZE

def get_model():

    nclass = 5

    inp = Input(shape=(3000, 1))

    img_1 = Convolution1D(16, kernel_size=5, activation=activations.relu, padding="valid")(inp)

    img_1 = Convolution1D(16, kernel_size=5, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(256, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(256, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = GlobalMaxPool1D()(img_1)

    img_1 = Dropout(rate=0.01)(img_1)

    dense_1 = Dropout(rate=0.01)(Dense(64, activation=activations.relu, name="dense_1")(img_1))

    dense_1 = Dropout(rate=0.05)(Dense(64, activation=activations.relu, name="dense_2")(dense_1))

    dense_1 = Dense(nclass, activation=activations.softmax, name="dense_3")(dense_1)

    model = models.Model(inputs=inp, outputs=dense_1)

    opt = optimizers.Adam(0.001)

    model.compile(optimizer=opt, loss=losses.sparse_categorical_crossentropy, metrics=['acc'])

    model.summary()

    return model

def get_base_model():

    inp = Input(shape=(3000, 1))

    img_1 = Convolution1D(16, kernel_size=5, activation=activations.relu, padding="valid")(inp)

    img_1 = Convolution1D(16, kernel_size=5, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(32, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = MaxPool1D(pool_size=2)(img_1)

    img_1 = SpatialDropout1D(rate=0.01)(img_1)

    img_1 = Convolution1D(256, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = Convolution1D(256, kernel_size=3, activation=activations.relu, padding="valid")(img_1)

    img_1 = GlobalMaxPool1D()(img_1)

    img_1 = Dropout(rate=0.01)(img_1)

    dense_1 = Dropout(0.01)(Dense(64, activation=activations.relu, name="dense_1")(img_1))

    base_model = models.Model(inputs=inp, outputs=dense_1)

    opt = optimizers.Adam(0.001)

    base_model.compile(optimizer=opt, loss=losses.sparse_categorical_crossentropy, metrics=['acc'])

    #model.summary()

    return base_model

def get_model_cnn():

    nclass = 5

    seq_input = Input(shape=(None, 3000, 1))

    base_model = get_base_model()

    # for layer in base_model.layers:

    #     layer.trainable = False

    encoded_sequence = TimeDistributed(base_model)(seq_input)

    encoded_sequence = SpatialDropout1D(rate=0.01)(Convolution1D(128,

                                                               kernel_size=3,

                                                               activation="relu",

                                                               padding="same")(encoded_sequence))

    encoded_sequence = Dropout(rate=0.05)(Convolution1D(128,

                                                               kernel_size=3,

                                                               activation="relu",

                                                               padding="same")(encoded_sequence))

    #out = TimeDistributed(Dense(nclass, activation="softmax"))(encoded_sequence)

    out = Convolution1D(nclass, kernel_size=3, activation="softmax", padding="same")(encoded_sequence)

    model = models.Model(seq_input, out)

    model.compile(optimizers.Adam(0.001), losses.sparse_categorical_crossentropy, metrics=['acc'])

    model.summary()

    return model

def get_model_lstm():

    nclass = 5

    seq_input = Input(shape=(None, 3000, 1))

    base_model = get_base_model()

    for layer in base_model.layers:

        layer.trainable = False

    encoded_sequence = TimeDistributed(base_model)(seq_input)

    encoded_sequence = Bidirectional(LSTM(100, return_sequences=True))(encoded_sequence)

    encoded_sequence = Dropout(rate=0.5)(encoded_sequence)

    encoded_sequence = Bidirectional(LSTM(100, return_sequences=True))(encoded_sequence)

    #out = TimeDistributed(Dense(nclass, activation="softmax"))(encoded_sequence)

    out = Convolution1D(nclass, kernel_size=1, activation="softmax", padding="same")(encoded_sequence)

    model = models.Model(seq_input, out)

    model.compile(optimizers.Adam(0.001), losses.sparse_categorical_crossentropy, metrics=['acc'])

    model.summary()

    return model

def get_model_cnn_crf(lr=0.001):

    nclass = 5

    seq_input = Input(shape=(None, 3000, 1))

    base_model = get_base_model()

    # for layer in base_model.layers:

    #     layer.trainable = False

    encoded_sequence = TimeDistributed(base_model)(seq_input)

    encoded_sequence = SpatialDropout1D(rate=0.01)(Convolution1D(128,

                                                               kernel_size=3,

                                                               activation="relu",

                                                               padding="same")(encoded_sequence))

    encoded_sequence = Dropout(rate=0.05)(Convolution1D(128,

                                                               kernel_size=3,

                                                               activation="linear",

                                                               padding="same")(encoded_sequence))

    #out = TimeDistributed(Dense(nclass, activation="softmax"))(encoded_sequence)

    # out = Convolution1D(nclass, kernel_size=3, activation="linear", padding="same")(encoded_sequence)

    crf = CRF(nclass, sparse_target=True)

    out = crf(encoded_sequence)

    model = models.Model(seq_input, out)

    model.compile(optimizers.Adam(lr), crf.loss_function, metrics=[crf.accuracy])

    model.summary()

    return model