# Initialize

train_dataname = 'data_processed.npy'
data_processed = np.load(train_dataname)

# Data process
n, o, p = np.shape(data_processed)
data = np.empty((p, n, o, 1))

for i in range(p):
    data[i, :, :, 0] = data_processed[:, :, i]

# Input shape
img_input = Input(shape=(n, o, 1))  # input shape:[251,16,1]

# Encoder model for CAE
# ---------------------------------------------------------------------------
# intial depth increases by 2x each layer
depth      = 8
strides    = 1
activation = 'relu'
# ReLU initial weights
kernel_initializer='glorot_uniform'
latent_dim = 14

# Crop to dimensions of [120,40,1] to allow for reconstruction
# e = Cropping2D(cropping = ((0, 6), (0, 1)))(img_input)
# e = Conv2D(depth*2**0, (5,5), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same') (e)
e = Conv2D(depth*2**0, (5,5), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same') (img_input)
e = Conv2D(depth*2**1, (5,5), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same') (e)
e = Conv2D(depth*2**2, (3,3), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same') (e)
e = Conv2D(depth*2**3, (3,3), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same')(e)

shape_before_flattening = K.int_shape(e)
x = Flatten()(e)
# Embedded latent space: 14 dimensions (features)
encoded = Dense(latent_dim, activation=activation, name='encoded')(x)

# Decoder model for CAE (Reverse Operations)
# ---------------------------------------------------------------------------
d = Dense(np.prod(shape_before_flattening[1:]), activation='relu')(encoded)
d = Reshape(shape_before_flattening[1:])(d)
d = Conv2DTranspose(depth*2**2, (3,3), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same')(d)
d = Conv2DTranspose(depth*2**1, (5,5), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same')(d)
d = Conv2DTranspose(depth*2**0, (5,5), strides=strides, activation=activation, kernel_initializer=kernel_initializer, padding='same')(d)


decoded = Conv2DTranspose(1, (5,5), strides=strides, activation='linear', kernel_initializer=kernel_initializer, padding='same')(d)

# Define Autoencoder Input: Sprectrograms & Output: Reconstructions
autoencoder = Model(inputs=img_input, outputs=decoded, name='autoencoder')

# Define Encoder Input: Sprectrograms & Output: Compressed image representations (embedded latent space)
encoder = Model(inputs=img_input, outputs=encoded, name='encoder')

# Model architecture
autoencoder.summary()
date_name = '2023'
architecture_fname = 'CAE_Model_{}.png'.format(date_name)
plot_model(autoencoder, to_file=architecture_fname, show_shapes=True)

# Model training parameter
LR = 0.001  # Learning rate
n_epochs = 50  # Number of epochs
batch_sz = 1024 # Batch size

# create log file to record training & validation loss
logger_fname = 'HistoryLog_LearningCurve.csv'.format(date_name)
csv_logger = CSVLogger(logger_fname)

checkpoint = ModelCheckpoint("model.h5", monitor="val_loss", save_best_only=True)

# Early stopping halts training after validation loss stops decreasing for 10 consectutive epochs
early_stop = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1,
                           mode='min', restore_best_weights=True)

optim = tf.keras.optimizers.Adam(lr=LR)             # Adaptive learning rate optimization algorithm (Adam)
loss = 'mse'                               # Mean Squared Error Loss function

# Compile Encoder & Autoencoder(initialize random filter weights)
encoder.compile(loss=loss, optimizer=optim)
autoencoder.compile(loss=loss,
                  optimizer=optim,
                  metrics=[metrics.mae])

# Model training
tic = time.time()
autoencoder.fit(data, data, batch_size=batch_sz, epochs=n_epochs, callbacks=[csv_logger, early_stop, checkpoint])
toc = time.time()
print('Elapsed Time : {0:4.1f} minutes'.format((toc-tic)/60))