import argparse
import os
from pathlib import Path

import numpy as np
import tensorflow as tf

from sklearn.model_selection import StratifiedKFold

from finetuning.utils import ecg_feature_extractor, train_test_split
from transplant.evaluation import auc, f1, multi_f1, CustomCheckpoint
from transplant.utils import (create_predictions_frame, load_pkl, is_multiclass)


def _create_dataset_from_data(data):
    """
    input:  data = {'x': x,
                    'y': labels.to_numpy(),
                    'record_ids': labels.index.to_numpy(),
                    'classes': labels.columns.to_numpy()}
    return: data and label
    """
    return tf.data.Dataset.from_tensor_slices((data['x'], data['y']))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--job-dir', type=Path, required=True, help='Job output directory.')
    parser.add_argument('--train', type=Path, required=True, help='Path to the train file.')
    parser.add_argument('--val', type=Path, help='Path to the validation file.\n'
                                                 'Overrides --val-size.')
    parser.add_argument('--test', type=Path, help='Path to the test file.')
    parser.add_argument('--weights-file', type=Path, help='Path to pretrained weights or a checkpoint of the model.')
    parser.add_argument('--val-size', type=float, default=None,
                        help='Size of the validation set or proportion of the train set.')
    parser.add_argument('--subset', type=float, default=None, help='Size of a subset of the train set '
                                                                   'or proportion of the train set.')
    parser.add_argument('--batch-size', type=int, default=32, help='Batch size.')
    parser.add_argument('--val-metric', default='loss',
                        help='Validation metric used to find the best model at each epoch. Supported metrics are:'
                             '`loss`, `acc`, `f1`, `auc`.')
    parser.add_argument('--channel', type=int, default=None, help='Use only the selected channel. '
                                                                  'By default use all available channels.')
    parser.add_argument('--epochs', type=int, default=1, help='Number of epochs.')
    parser.add_argument('--seed', type=int, default=None, help='Random state.')
    parser.add_argument('--verbose', action='store_true', help='Show debug messages.')
    args, _ = parser.parse_known_args()

    if args.val_metric not in ['loss', 'acc', 'f1', 'auc']:
        raise ValueError('Unknown metric: {}'.format(args.val_metric))

    os.makedirs(name=str(args.job_dir), exist_ok=True)
    print('Creating working directory in {}'.format(args.job_dir))

    seed = args.seed or np.random.randint(2 ** 16)
    print('Setting random state {}'.format(seed))
    np.random.seed(seed)

    # Không sử dụng val file riêng biệt mà chia val từ train set
    if not args.val and args.val_size:
        if args.val_size >= 1:  # Lấy theo số lượng của patients, nếu <= 1 thì đó là theo tỷ lệ của train set
            args.val_size = int(args.val_size)

    # Tiếp tục chia train set sau khi đã chia thành train và val set
    # Lớn hơn hoặc bằng 1 thì sẽ là lấy theo số lượng patient, nếu <= 1 thì đó là theo tỷ lệ của train set
    if args.subset and args.subset >= 1:
        args.subset = int(args.subset)

    print('Loading train data from {} ...'.format(args.train))
    train = load_pkl(file=str(args.train))

    if args.val:    # Loading val set từ val file
        print('Loading validation data from {} ...'.format(args.val))
        val = load_pkl(file=str(args.val))
    elif args.val_size:     # Chia tỷ lệ val set từ train set mà không phải là load từ file
        original_train_size = len(train['x'])   # Kích thước của toàn bộ dữ liệu dataset
        train, val = train_test_split(train, test_size=args.val_size, stratify=train['y'])  # Chia thành train và val set
        new_train_size = len(train['x'])    # Trả về kích thước của train set mới
        new_val_size = len(val['x'])    # trả về kích thước của val set mới
        print('Split data into train {:.2%} and validation {:.2%}'.format(
            new_train_size / original_train_size, new_val_size / original_train_size))
    else:   # Không sử dụng val set
        val = None

    if args.test:   # Sử dụng test set file riêng biệt
        print('Loading test data from {} ...'.format(args.test))
        test = load_pkl(str(args.test))
    else:   # Không sử dụng test set
        test = None

    if args.subset:     # Tiếp tục chia train set sau khi đã chia train dataset ban đầu thành new train set và val set
        original_train_size = len(train['x'])   # Trả về kích thước của train set
        train, _ = train_test_split(train, train_size=args.subset, stratify=train['y'])     # Trả về new train set
        new_train_size = len(train['x'])    # Trả về kích thước của new train set
        print('Using only {:.2%} of train data'.format(new_train_size / original_train_size))

    if args.channel is not None:
        train['x'] = train['x'][:, :, args.channel:args.channel + 1]
        if val:
            val['x'] = val['x'][:, :, args.channel:args.channel + 1]
        if test:
            test['x'] = test['x'][:, :, args.channel:args.channel + 1]

    print('Train data shape:', train['x'].shape)

    train_data = _create_dataset_from_data(train).shuffle(len(train['x'])).batch(args.batch_size)
    val_data = _create_dataset_from_data(val).batch(args.batch_size) if val else None
    test_data = _create_dataset_from_data(test).batch(args.batch_size) if test else None

    strategy = tf.distribute.MirroredStrategy()

    with strategy.scope():
        print('Building model ...')
        num_classes = len(train['classes'])

        if is_multiclass(train['y']):
            activation = 'sigmoid'
            loss = tf.keras.losses.BinaryCrossentropy()
            accuracy = tf.keras.metrics.BinaryAccuracy(name='acc')
        else:
            activation = 'softmax'
            loss = tf.keras.losses.CategoricalCrossentropy()
            accuracy = tf.keras.metrics.CategoricalAccuracy(name='acc')

        # not include fc layer
        model = ecg_feature_extractor(arch=args.arch)
        model.add(tf.keras.layers.Dense(units=num_classes, activation=activation))

        # initialize the weights of the model
        inputs = tf.keras.layers.Input(shape=train['x'].shape[1:], dtype=train['x'].dtype)
        model(inputs)   # complete model

        print('# model parameters: {:,d}'.format(model.count_params()))

        if args.weights_file:   # Sử dụng trọng số đã được pre-trained
            # initialize weights (excluding the optimizer state) to load the pretrained resnet
            # the optimizer state is randomly initialized in the `model.compile` function
            print('Loading weights from file {} ...'.format(args.weights_file))
            model.load_weights(str(args.weights_file))

        model.compile(optimizer=tf.keras.optimizers.Adam(),
                      loss=loss,
                      metrics=[accuracy])

        callbacks = []

        logger = tf.keras.callbacks.CSVLogger(filename=str(args.job_dir / 'history.csv'))
        callbacks.append(logger)

        if args.val_metric in ['loss', 'acc']:
            monitor = ('val_' + args.val_metric) if val else args.val_metric
            checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath=str(args.job_dir / 'best_model.weights'),
                                                            monitor=monitor,
                                                            save_best_only=True,
                                                            save_weights_only=True,
                                                            mode='auto',
                                                            verbose=1)
        elif args.val_metric == 'f1':
            if is_multiclass(train['y']):
                score_fn = multi_f1
            else:
                score_fn = f1

            checkpoint = CustomCheckpoint(filepath=str(args.job_dir / 'best_model.weights'),
                                          data=(val_data, val['y']) if val else (train_data, train['y']),
                                          score_fn=score_fn,
                                          save_best_only=True,
                                          verbose=1)

        elif args.val_metric == 'auc':
            checkpoint = CustomCheckpoint(filepath=str(args.job_dir / 'best_model.weights'),
                                          data=(val_data, val['y']) if val else (train_data, train['y']),
                                          score_fn=auc,
                                          save_best_only=True,
                                          verbose=1)
        else:
            raise ValueError('Unknown metric: {}'.format(args.val_metric))

        callbacks.append(checkpoint)

        if val:
            # new adding
            rl_stopping = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=7,
                                                               verbose=1, min_lr=1e-7)
            callbacks.append(rl_stopping)

            early_stopping = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=30, verbose=1)
            callbacks.append(early_stopping)

        model.fit(train_data, epochs=args.epochs, verbose=2, validation_data=val_data, callbacks=callbacks)

        # load best model for inference
        print('Loading the best weights from file {} ...'.format(str(args.job_dir / 'best_model.weights')))
        model.load_weights(filepath=str(args.job_dir / 'best_model.weights'))

        print('Predicting training data ...')
        train_y_prob = model.predict(x=train['x'], batch_size=args.batch_size)
        train_predictions = create_predictions_frame(y_prob=train_y_prob,
                                                     y_true=train['y'],
                                                     class_names=train['classes'],
                                                     record_ids=train['record_ids'])
        train_predictions.to_csv(path_or_buf=args.job_dir / 'train_predictions.csv', index=False)

        if val:
            print('Predicting validation data ...')
            val_y_prob = model.predict(x=val['x'], batch_size=args.batch_size)
            val_predictions = create_predictions_frame(y_prob=val_y_prob, y_true=val['y'],
                                                       class_names=train['classes'],
                                                       record_ids=val['record_ids'])
            val_predictions.to_csv(path_or_buf=args.job_dir / 'val_predictions.csv', index=False)

        if test:
            print('Predicting test data ...')
            test_y_prob = model.predict(x=test['x'], batch_size=args.batch_size)
            test_predictions = create_predictions_frame(y_prob=test_y_prob, y_true=test['y'],
                                                        class_names=train['classes'],
                                                        record_ids=test['record_ids'])
            test_predictions.to_csv(path_or_buf=args.job_dir / 'test_predictions.csv', index=False)

tensorflow.keras.metrics.BinaryAccuracy tensorflow.keras.losses.CategoricalCrossentropy numpy.random.seed tensorflow.distribute.MirroredStrategy tensorflow.keras.layers.Dense tensorflow.data.Dataset.from_tensor_slices tensorflow.keras.callbacks.ReduceLROnPlateau tensorflow.keras.losses.BinaryCrossentropy tensorflow.keras.optimizers.Adam numpy.random.randint tensorflow.keras.metrics.CategoricalAccuracy tensorflow.keras.callbacks.EarlyStopping tensorflow.keras.layers.Input
[(23, 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (["(data['x'], data['y'])"], {}), True, 'import tensorflow as tf\n'), (27, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (57, 'numpy.random.seed', 'np.random.seed', (['seed'], {}), True, 'import numpy as np\n'), (110, 'tensorflow.distribute.MirroredStrategy', 'tf.distribute.MirroredStrategy', ([], {}), True, 'import tensorflow as tf\n'), (55, 'numpy.random.randint', 'np.random.randint', (['(2 ** 16)'], {}), True, 'import numpy as np\n'), (93, 'finetuning.utils.train_test_split', 'train_test_split', (['train'], {'train_size': 'args.subset', 'stratify': "train['y']"}), False, 'from finetuning.utils import ecg_feature_extractor, train_test_split\n'), (116, 'transplant.utils.is_multiclass', 'is_multiclass', (["train['y']"], {}), False, 'from transplant.utils import create_predictions_frame, load_pkl, is_multiclass\n'), (126, 'finetuning.utils.ecg_feature_extractor', 'ecg_feature_extractor', ([], {'arch': 'args.arch'}), False, 'from finetuning.utils import ecg_feature_extractor, train_test_split\n'), (130, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', ([], {'shape': "train['x'].shape[1:]", 'dtype': "train['x'].dtype"}), True, 'import tensorflow as tf\n'), (198, 'transplant.utils.create_predictions_frame', 'create_predictions_frame', ([], {'y_prob': 'train_y_prob', 'y_true': "train['y']", 'class_names': "train['classes']", 'record_ids': "train['record_ids']"}), False, 'from transplant.utils import create_predictions_frame, load_pkl, is_multiclass\n'), (77, 'finetuning.utils.train_test_split', 'train_test_split', (['train'], {'test_size': 'args.val_size', 'stratify': "train['y']"}), False, 'from finetuning.utils import ecg_feature_extractor, train_test_split\n'), (118, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (119, 'tensorflow.keras.metrics.BinaryAccuracy', 'tf.keras.metrics.BinaryAccuracy', ([], {'name': '"""acc"""'}), True, 'import tensorflow as tf\n'), (122, 'tensorflow.keras.losses.CategoricalCrossentropy', 'tf.keras.losses.CategoricalCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (123, 'tensorflow.keras.metrics.CategoricalAccuracy', 'tf.keras.metrics.CategoricalAccuracy', ([], {'name': '"""acc"""'}), True, 'import tensorflow as tf\n'), (127, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'num_classes', 'activation': 'activation'}), True, 'import tensorflow as tf\n'), (183, 'tensorflow.keras.callbacks.ReduceLROnPlateau', 'tf.keras.callbacks.ReduceLROnPlateau', ([], {'monitor': '"""val_loss"""', 'factor': '(0.5)', 'patience': '(7)', 'verbose': '(1)', 'min_lr': '(1e-07)'}), True, 'import tensorflow as tf\n'), (187, 'tensorflow.keras.callbacks.EarlyStopping', 'tf.keras.callbacks.EarlyStopping', ([], {'monitor': '"""val_loss"""', 'patience': '(30)', 'verbose': '(1)'}), True, 'import tensorflow as tf\n'), (207, 'transplant.utils.create_predictions_frame', 'create_predictions_frame', ([], {'y_prob': 'val_y_prob', 'y_true': "val['y']", 'class_names': "train['classes']", 'record_ids': "val['record_ids']"}), False, 'from transplant.utils import create_predictions_frame, load_pkl, is_multiclass\n'), (215, 'transplant.utils.create_predictions_frame', 'create_predictions_frame', ([], {'y_prob': 'test_y_prob', 'y_true': "test['y']", 'class_names': "train['classes']", 'record_ids': "test['record_ids']"}), False, 'from transplant.utils import create_predictions_frame, load_pkl, is_multiclass\n'), (141, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {}), True, 'import tensorflow as tf\n'), (159, 'transplant.utils.is_multiclass', 'is_multiclass', (["train['y']"], {}), False, 'from transplant.utils import create_predictions_frame, load_pkl, is_multiclass\n')]