from sklearn.model_selection import train_test_split
from keras.layers import Dense, Dropout, Flatten, Conv1D, Input, MaxPooling1D
from keras.models import Model
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras import backend as K
from matplotlib import pyplot
from keras.models import load_model
history_array = []

for set_size in range(2000, 8000, 1000):
    print("TRAINING FOR SET SIZE: ", set_size)

    y=le.fit_transform(all_label)
    y=np_utils.to_categorical(y, num_classes=len(labels))[:set_size]
    all_wave_i = np.array(all_wave).reshape(-1,set_size,1)

    x_tr, x_val, y_tr, y_val = train_test_split(np.array(all_wave_i), np.array(y), stratify=y, test_size = 0.2, random_state=777,shuffle=True)


    K.clear_session()

    inputs = Input(shape=(set_size,1))

    #First Conv1D layer
    conv = Conv1D(8,13, padding='valid', activation='relu', strides=1)(inputs)
    conv = MaxPooling1D(3)(conv)
    conv = Dropout(0.3)(conv)

    #Second Conv1D layer
    conv = Conv1D(16, 11, padding='valid', activation='relu', strides=1)(conv)
    conv = MaxPooling1D(3)(conv)
    conv = Dropout(0.3)(conv)

    #Third Conv1D layer
    conv = Conv1D(32, 9, padding='valid', activation='relu', strides=1)(conv)
    conv = MaxPooling1D(3)(conv)
    conv = Dropout(0.3)(conv)

    #Fourth Conv1D layer
    conv = Conv1D(64, 7, padding='valid', activation='relu', strides=1)(conv)
    conv = MaxPooling1D(3)(conv)
    conv = Dropout(0.3)(conv)

    #Flatten layer
    conv = Flatten()(conv)

    #Dense Layer 1
    conv = Dense(256, activation='relu')(conv)
    conv = Dropout(0.3)(conv)

    #Dense Layer 2
    conv = Dense(128, activation='relu')(conv)
    conv = Dropout(0.3)(conv)

    outputs = Dense(len(labels), activation='softmax')(conv)

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

    model.compile(loss='categorical_crossentropy',optimizer='adam',metrics=['accuracy'])

    es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=10, min_delta=0.0001)
    model_path = '/Users/fryzu/Documents/systemyGlosowe/lab1/src/klasyfikacja/best_model_' + str(set_size) + '.hdf5'
    mc = ModelCheckpoint(model_path, monitor='val_acc', verbose=1, save_best_only=True, mode='max')

    history=model.fit(x_tr, y_tr ,epochs=3, callbacks=[es,mc], batch_size=32, validation_data=(x_val,y_val))

    model.save(model_path)

    pyplot.plot(history.history['loss'], label='train')
    pyplot.plot(history.history['val_loss'], label='test')
    pyplot.legend()
    plt.ylabel('Loss [%]', fontsize=12)
    pyplot.show()

    model=load_model(model_path)

    history_array.append(history)