import pandas as pd
from mil_data_generator import *
from mil_models_pytorch import*
from mil_trainer_torch import *
import argparse

np.random.seed(42)
random.seed(42)
torch.manual_seed(42)


def main(args):

    metrics = []
    for i_iteration in np.arange(0, args.iterations):
        id = str(i_iteration) + '_'
        df = pd.read_excel(args.dir_data_frame)

        # Set data generators
        dataset_train = MILDataset(args.dir_images, df[df['Partition'] == 'train'], args.classes,
                                   bag_id='slide_name', input_shape=args.input_shape,
                                   data_augmentation=args.data_augmentation, images_on_ram=args.images_on_ram,
                                   pMIL=args.pMIL, proportions=args.proportions)
        data_generator_train = MILDataGenerator(dataset_train, batch_size=1, shuffle=True, max_instances=512)

        dataset_val = MILDataset(args.dir_images, df[df['Partition'] == 'val'], args.classes,
                                 bag_id='slide_name', input_shape=args.input_shape,
                                 data_augmentation=args.data_augmentation, images_on_ram=args.images_on_ram,
                                 pMIL=args.pMIL, proportions=args.proportions)
        data_generator_val = MILDataGenerator(dataset_val, batch_size=1, shuffle=False, max_instances=512)

        dataset_test = MILDataset(args.dir_images, df[df['Partition'] == 'test'], args.classes,
                                  bag_id='slide_name', input_shape=args.input_shape,
                                  data_augmentation=args.data_augmentation, images_on_ram=args.images_on_ram,
                                  pMIL=args.pMIL, proportions=args.proportions,
                                  dataframe_instances=pd.read_excel(args.dir_data_frame_test))
        data_generator_test = MILDataGenerator(dataset_test, batch_size=1, shuffle=False, max_instances=512)

        # Set network architecture
        network = MILArchitecture(args.classes, mode=args.mode, aggregation=args.aggregation,
                                  backbone='vgg19', include_background=args.include_background)

        # Perform training
        trainer = MILTrainer(args.dir_results + args.experiment_name + '/', network,
                             lr=args.lr, pMIL=args.pMIL, margin=args.margin,
                             alpha_ic=args.alpha_ic, alpha_pc=args.alpha_pc, t_ic=args.t_ic,
                             t_pc=args.t_pc, alpha_ce=args.alpha_ce, id=id,
                             early_stopping=args.early_stopping, scheduler=args.scheduler,
                             virtual_batch_size=args.virtual_batch_size,
                             criterion=args.criterion,
                             alpha_H=args.alpha_H)
        trainer.train(train_generator=data_generator_train, val_generator=data_generator_val,
                      test_generator=data_generator_test, epochs=args.epochs)

        metrics.append([list(trainer.metrics.values())[1:]])

    # Get overall metrics
    metrics = np.squeeze(np.array(metrics))

    mu = np.mean(metrics, axis=0)
    std = np.std(metrics, axis=0)

    info = "AUCtest={:.4f}({:.4f}) ; AUCval={:.4f}({:.4f})  ; acc={:.4f}({:.4f}) ; f1-score={:.4f}({:.4f}) ; k2={:.4f}({:.4f})".format(
          mu[0], std[0], mu[1], std[1], mu[2], std[2], mu[3], std[3], mu[4], std[4])
    if args.alpha_pc > 0:
        info += " ; constrain_cumpliment={:.4f}({:.4f}) ; constrain_proportion={:.4f}({:.4f})".format(
          mu[5], std[5], mu[6], std[6])

    f = open(args.dir_results + args.experiment_name + '/' + 'method_metrics.txt', 'w')
    f.write(info)
    f.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dir_images", default='../data/SICAP_MIL/patches/', type=str)
    parser.add_argument("--dir_data_frame", default='../data/SICAP_MIL/dataframes/gt_global_slides.xlsx', type=str)
    parser.add_argument("--dir_data_frame_test", default='../data/SICAP_MIL/dataframes/gt_test_patches.xlsx', type=str)
    parser.add_argument("--dir_results", default='../data/results/', type=str)
    parser.add_argument("--criterion", default='z', type=str)
    parser.add_argument("--experiment_name", default="test_test_test", type=str)
    parser.add_argument("--classes", default=['G3', 'G4', 'G5'], type=list)
    parser.add_argument("--proportions", default=['Primary', 'Secondary'], type=list)
    parser.add_argument("--input_shape", default=(3, 224, 224), type=list)
    parser.add_argument("--images_on_ram", default=False, type=bool)
    parser.add_argument("--epochs", default=100, type=int)
    parser.add_argument("--aggregation", default="max", type=str)
    parser.add_argument("--mode", default="instance", type=str)
    parser.add_argument("--include_background", default=True, type=bool)
    parser.add_argument("--lr", default=1*1e-2, type=float)
    parser.add_argument("--pMIL", default=False, type=bool)

    parser.add_argument("--alpha_ce", default=1., type=float)
    parser.add_argument("--margin", default=0., type=float)
    parser.add_argument("--alpha_ic", default=1, type=float)
    parser.add_argument("--alpha_pc", default=1, type=float)
    parser.add_argument("--alpha_H", default=0, type=float)
    parser.add_argument("--t_ic", default=15, type=float)
    parser.add_argument("--t_pc", default=5, type=float)
    parser.add_argument("--data_augmentation", default=True, type=bool)
    parser.add_argument("--iterations", default=3, type=int)

    parser.add_argument("--early_stopping", default=True, type=bool)
    parser.add_argument("--scheduler", default=True, type=bool)

    parser.add_argument("--virtual_batch_size", default=1, type=int)

    args = parser.parse_args()
    main(args)