import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

def model(input_shape):

    x_input = layers.Input(shape = input_shape)

    x = layers.Conv1D(196, kernel_size=15, strides=4)(x_input)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.Dropout(0.8)(x)

    x = layers.GRU(units = 128, return_sequences = True)(x)
    x = layers.Dropout(0.8)(x)
    x = layers.BatchNormalization()(x)
    
    x = layers.GRU(units = 128, return_sequences = True)(x)
    x = layers.Dropout(0.8)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.8)(x)
    
    x = layers.TimeDistributed(layers.Dense(1, activation = "sigmoid"))(x)
    model = keras.Model(inputs = x_input, outputs = x)
    
    return model  


#model = model(input_shape = (Tx, n_freq))
model = model(input_shape = (5511, 101))
model.summary()

opt = keras.optimizers.Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, decay=0.01)
model.compile(loss='binary_crossentropy', optimizer=opt, metrics=["accuracy"])

X = np.load("./Data/XY_train/X.npy")
Y = np.load("./Data/XY_train/Y.npy")

X_dev = np.load("./Data/XY_dev/X_dev.npy")
Y_dev = np.load("./Data/XY_dev/Y_dev.npy")

model.fit(X, Y, batch_size = 64, epochs=20000)

# save model
model.save('my_model.h5')  # creates a HDF5 file 'my_model.h5'
# del model  # deletes the existing model

# returns a compiled model
# identical to the previous one
# model = load_model('my_model.h5')