from tensorflow.keras.optimizers import Adam

from tensorflow.keras.callbacks import (ModelCheckpoint, TensorBoard, ReduceLROnPlateau,

                                        CSVLogger, EarlyStopping)

from model import get_model

import argparse

from datasets import ECGSequence

if __name__ == "__main__":

    # Get data and train

    parser = argparse.ArgumentParser(description='Train neural network.')

    parser.add_argument('path_to_hdf5', type=str,

                        help='path to hdf5 file containing tracings')

    parser.add_argument('path_to_csv', type=str,

                        help='path to csv file containing annotations')

    parser.add_argument('--val_split', type=float, default=0.02,

                        help='number between 0 and 1 determining how much of'

                             ' the data is to be used for validation. The remaining '

                             'is used for validation. Default: 0.02')

    parser.add_argument('--dataset_name', type=str, default='tracings',

                        help='name of the hdf5 dataset containing tracings')

    args = parser.parse_args()

    # Optimization settings

    loss = 'binary_crossentropy'

    lr = 0.001

    batch_size = 64

    opt = Adam(lr)

    callbacks = [ReduceLROnPlateau(monitor='val_loss',

                                   factor=0.1,

                                   patience=7,

                                   min_lr=lr / 100),

                 EarlyStopping(patience=9,  # Patience should be larger than the one in ReduceLROnPlateau

                               min_delta=0.00001)]

    train_seq, valid_seq = ECGSequence.get_train_and_val(

        args.path_to_hdf5, args.dataset_name, args.path_to_csv, batch_size, args.val_split)

    # If you are continuing an interrupted section, uncomment line bellow:

    #   model = keras.models.load_model(PATH_TO_PREV_MODEL, compile=False)

    model = get_model(train_seq.n_classes)

    model.compile(loss=loss, optimizer=opt)

    # Create log

    callbacks += [TensorBoard(log_dir='./logs', write_graph=False),

                  CSVLogger('training.log', append=False)]  # Change append to true if continuing training

    # Save the BEST and LAST model

    callbacks += [ModelCheckpoint('./backup_model_last.hdf5'),

                  ModelCheckpoint('./backup_model_best.hdf5', save_best_only=True)]

    # Train neural network

    history = model.fit(train_seq,

                        epochs=70,

                        initial_epoch=0,  # If you are continuing a interrupted section change here

                        callbacks=callbacks,

                        validation_data=valid_seq,

                        verbose=1)

    # Save final result

    model.save("./final_model.hdf5")