import yaml

import tensorboard

import torch

import os

import shutil

import sys

import csv

import argparse

import pickle

from models.resnet_simclr import ResNetSimCLR

from clinical_ts.cpc import CPCModel

import torch.nn.functional as F

from tqdm import tqdm

import numpy as np

import matplotlib.pyplot as plt

from sklearn.decomposition import PCA

from sklearn.manifold import TSNE

from sklearn.metrics import roc_auc_score

from clinical_ts.simclr_dataset_wrapper import SimCLRDataSetWrapper

from clinical_ts.eval_utils_cafa import eval_scores, eval_scores_bootstrap

from clinical_ts.timeseries_utils import aggregate_predictions

import pdb

from copy import deepcopy

from os.path import join, isdir

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def parse_args():

    parser = argparse.ArgumentParser("Finetuning tests")

    parser.add_argument("--model_file")

    parser.add_argument("--method")

    parser.add_argument("--dataset", nargs="+", default="./data/ptb_xl_fs100")

    parser.add_argument("--batch_size", type=int, default=512)

    parser.add_argument("--discriminative_lr", default=False, action="store_true")

    parser.add_argument("--num_workers", type=int, default=8)

    parser.add_argument("--hidden", default=False, action="store_true")

    parser.add_argument("--lr_schedule", default="{}")

    parser.add_argument("--use_pretrained", default=False, action="store_true")

    parser.add_argument("--linear_evaluation",

                        default=False, action="store_true", help="use linear evaluation")

    parser.add_argument("--test_noised", default=False, action="store_true", help="validate also on a noisy dataset")

    parser.add_argument("--noise_level", default=1, type=int, help="level of noise induced to the second validations set")

    parser.add_argument("--folds", default=8, type=int, help="number of folds used in finetuning (between 1-8)")

    parser.add_argument("--tag", default="")

    parser.add_argument("--eval_only", action="store_true", default=False, help="only evaluate mode")

    parser.add_argument("--load_finetuned", action="store_true", default=False)

    parser.add_argument("--test", action="store_true", default=False)

    parser.add_argument("--verbose", action="store_true", default=False)

    parser.add_argument("--cpc", action="store_true", default=False)

    parser.add_argument("--model_location")

    parser.add_argument("--l_epochs", type=int, default=0, help="number of head-only epochs (these are performed first)")

    parser.add_argument("--f_epochs", type=int, default=0, help="number of finetuning epochs (these are perfomed after head-only training")

    parser.add_argument("--normalize", action="store_true", default=False, help="normalize dataset with ptbxl mean and std")

    parser.add_argument("--bn_head", action="store_true", default=False)

    parser.add_argument("--ps_head", type=float, default=0.0)

    parser.add_argument("--conv_encoder", action="store_true", default=False)

    parser.add_argument("--base_model", default="xresnet1d50")

    parser.add_argument("--widen", default=1, type=int, help="use wide xresnet1d50")

    args = parser.parse_args()

    return args

def get_new_state_dict(init_state_dict, lightning_state_dict, method="simclr"):

    # in case of moco model

    from collections import OrderedDict

    # lightning_state_dict = lightning_state_dict["state_dict"]

    new_state_dict = OrderedDict()

    if method != "cpc":

        if method == "moco":

            for key in init_state_dict:

                l_key = "encoder_q." + key

                if l_key in lightning_state_dict.keys():

                    new_state_dict[key] = lightning_state_dict[l_key]

        elif method == "simclr":

            for key in init_state_dict:

                if "features" in key:

                    l_key = key.replace("features", "encoder.features")

                if l_key in lightning_state_dict.keys():

                    new_state_dict[key] = lightning_state_dict[l_key]

        elif method == "swav":

            for key in init_state_dict:

                if "features" in key:

                    l_key = key.replace("features", "model.features")

                if l_key in lightning_state_dict.keys():

                    new_state_dict[key] = lightning_state_dict[l_key]

        elif method == "byol":

            for key in init_state_dict:

                l_key = "online_network.encoder." + key

                if l_key in lightning_state_dict.keys():

                    new_state_dict[key] = lightning_state_dict[l_key]

        else:

            raise("method unknown")

        new_state_dict["l1.weight"] = init_state_dict["l1.weight"]

        new_state_dict["l1.bias"] = init_state_dict["l1.bias"]

        if "l2.weight" in init_state_dict.keys():

            new_state_dict["l2.weight"] = init_state_dict["l2.weight"]

            new_state_dict["l2.bias"] = init_state_dict["l2.bias"]

        assert(len(init_state_dict) == len(new_state_dict))

    else:

        for key in init_state_dict:

            l_key = "model_cpc." + key

            if l_key in lightning_state_dict.keys():

                new_state_dict[key] = lightning_state_dict[l_key]

            if "head" in key:

                new_state_dict[key] = init_state_dict[key]

    return new_state_dict

def adjust(model, num_classes, hidden=False):

    in_features = model.l1.in_features

    last_layer = torch.nn.modules.linear.Linear(

        in_features, num_classes).to(device)

    if hidden:

        model.l1 = torch.nn.modules.linear.Linear(

            in_features, in_features).to(device)

        model.l2 = last_layer

    else:

        model.l1 = last_layer

    def def_forward(self):

        def new_forward(x):

            h = self.features(x)

            h = h.squeeze()

            x = self.l1(h)

            if hidden:

                x = F.relu(x)

                x = self.l2(x)

            return x

        return new_forward

    model.forward = def_forward(model)

def configure_optimizer(model, batch_size, head_only=False, discriminative_lr=False, base_model="xresnet1d", optimizer="adam", discriminative_lr_factor=1):

    loss_fn = F.binary_cross_entropy_with_logits

    if base_model == "xresnet1d":

        wd = 1e-1

        if head_only:

            lr = (8e-3*(batch_size/256))

            optimizer = torch.optim.AdamW(

                model.l1.parameters(), lr=lr, weight_decay=wd)

        else:

            lr = 0.01

            if not discriminative_lr:

                optimizer = torch.optim.AdamW(

                    model.parameters(), lr=lr, weight_decay=wd)

            else:

                param_dict = dict(model.named_parameters())

                keys = param_dict.keys()

                weight_layer_nrs = set()

                for key in keys:

                    if "features" in key:

                        # parameter names have the form features.x

                        weight_layer_nrs.add(key[9])

                weight_layer_nrs = sorted(weight_layer_nrs, reverse=True)

                features_groups = []

                while len(weight_layer_nrs) > 0:

                    if len(weight_layer_nrs) > 1:

                        features_groups.append(list(filter(

                            lambda x: "features." + weight_layer_nrs[0] in x or "features." + weight_layer_nrs[1] in x,  keys)))

                        del weight_layer_nrs[:2]

                    else:

                        features_groups.append(

                            list(filter(lambda x: "features." + weight_layer_nrs[0] in x,  keys)))

                        del weight_layer_nrs[0]

                # filter linear layers

                linears = list(filter(lambda x: "l" in x, keys))

                groups = [linears] + features_groups

                optimizer_param_list = []

                tmp_lr = lr

                for layers in groups:

                    layer_params = [param_dict[param_name]

                                    for param_name in layers]

                    optimizer_param_list.append(

                        {"params": layer_params, "lr": tmp_lr})

                    tmp_lr /= 4

                optimizer = torch.optim.AdamW(

                    optimizer_param_list, lr=lr, weight_decay=wd)

        print("lr", lr)

        print("wd", wd)

        print("batch size", batch_size)

    elif base_model == "cpc":

        if(optimizer == "sgd"):

            opt = torch.optim.SGD

        elif(optimizer == "adam"):

            opt = torch.optim.AdamW

        else:

            raise NotImplementedError("Unknown Optimizer.")

        lr = 1e-4

        wd = 1e-3

        if(head_only):

            lr = 1e-3

            print("Linear eval: model head", model.head)

            optimizer = opt(model.head.parameters(), lr, weight_decay=wd)

        elif(discriminative_lr_factor != 1.):  # discrimative lrs

            optimizer = opt([{"params": model.encoder.parameters(), "lr": lr*discriminative_lr_factor*discriminative_lr_factor}, {

                            "params": model.rnn.parameters(), "lr": lr*discriminative_lr_factor}, {"params": model.head.parameters(), "lr": lr}], lr, weight_decay=wd)

            print("Finetuning: model head", model.head)

            print("discriminative lr: ", discriminative_lr_factor)

        else:

            lr = 1e-3

            print("normal supervised training")

            optimizer = opt(model.parameters(), lr, weight_decay=wd)

    else:

        raise("model unknown")

    return loss_fn, optimizer

def load_model(linear_evaluation, num_classes, use_pretrained, discriminative_lr=False, hidden=False, conv_encoder=False, bn_head=False, ps_head=0.5, location="./checkpoints/moco_baselinewonder200.ckpt", method="simclr", base_model="xresnet1d50", out_dim=16, widen=1):

    discriminative_lr_factor = 1

    if use_pretrained:

        print("load model from " + location)

        discriminative_lr_factor = 0.1

        if base_model == "cpc":

            lightning_state_dict = torch.load(location, map_location=device)

            # num_head = np.sum([1 if 'proj' in f else 0 for f in lightning_state_dict.keys()])

            if linear_evaluation:

                lin_ftrs_head = []

                bn_head = False

                ps_head = 0.0

            else:

                if hidden:

                    lin_ftrs_head = [512]

                else:

                    lin_ftrs_head = []

            if conv_encoder:

                strides = [2, 2, 2, 2]

                kss = [10, 4, 4, 4]

            else:

                strides = [1]*4

                kss = [1]*4

            model = CPCModel(input_channels=12, strides=strides, kss=kss, features=[512]*4, n_hidden=512, n_layers=2, mlp=False, lstm=True, bias_proj=False,

                             num_classes=num_classes, skip_encoder=False, bn_encoder=True, lin_ftrs_head=lin_ftrs_head, ps_head=ps_head, bn_head=bn_head).to(device)

            if "state_dict" in lightning_state_dict.keys():

                print("load pretrained model")

                model_state_dict = get_new_state_dict(

                    model.state_dict(), lightning_state_dict["state_dict"], method="cpc")

            else:

                print("load already finetuned model")

                model_state_dict = lightning_state_dict

            model.load_state_dict(model_state_dict)

        else:

            model = ResNetSimCLR(base_model, out_dim, hidden=hidden, widen=widen).to(device)

            model_state_dict = torch.load(location, map_location=device)

            if "state_dict" in model_state_dict.keys():

                model_state_dict = model_state_dict["state_dict"]

            if "l1.weight" in model_state_dict.keys():  # load already fine-tuned model

                model_classes = model_state_dict["l1.weight"].shape[0]

                if model_classes != num_classes:

                    raise Exception("Loaded model has different output dim ({}) than needed ({})".format(

                        model_classes, num_classes))

                adjust(model, num_classes, hidden=hidden)

                if not hidden and "l2.weight" in model_state_dict:

                    del model_state_dict["l2.weight"]

                    del model_state_dict["l2.bias"]

                model.load_state_dict(model_state_dict)

            else:  # load pretrained model

                base_dict = model.state_dict()

                model_state_dict = get_new_state_dict(

                    base_dict, model_state_dict, method=method)

                model.load_state_dict(model_state_dict)

                adjust(model, num_classes, hidden=hidden)

    else:

        if "xresnet1d" in base_model:

            model = ResNetSimCLR(base_model, out_dim, hidden=hidden, widen=widen).to(device)

            adjust(model, num_classes, hidden=hidden)

        elif base_model == "cpc":

            if linear_evaluation:

                lin_ftrs_head = []

                bn_head = False

                ps_head = 0.0

            else:

                if hidden:

                    lin_ftrs_head = [512]

                else:

                    lin_ftrs_head = []

            if conv_encoder:

                strides = [2, 2, 2, 2]

                kss = [10, 4, 4, 4]

            else:

                strides = [1]*4

                kss = [1]*4

            model = CPCModel(input_channels=12, strides=strides, kss=kss, features=[512]*4, n_hidden=512, n_layers=2, mlp=False, lstm=True, bias_proj=False,

                             num_classes=num_classes, skip_encoder=False, bn_encoder=True, lin_ftrs_head=lin_ftrs_head, ps_head=ps_head, bn_head=bn_head).to(device)

        else:

            raise Exception("model unknown")

    return model

def evaluate(model, dataloader, idmap, lbl_itos, cpc=False):

    preds, targs = eval_model(model, dataloader, cpc=cpc)

    scores = eval_scores(targs, preds, classes=lbl_itos, parallel=True)

    preds_agg, targs_agg = aggregate_predictions(preds, targs, idmap)

    scores_agg = eval_scores(targs_agg, preds_agg,

                             classes=lbl_itos, parallel=True)

    macro = scores["label_AUC"]["macro"]

    macro_agg = scores_agg["label_AUC"]["macro"]

    return preds, macro, macro_agg

def set_train_eval(model, cpc, linear_evaluation):

    if linear_evaluation:

        if cpc:

            model.encoder.eval()

        else:

            model.features.eval()

    else:

        model.train()

def train_model(model, train_loader, valid_loader, test_loader, epochs, loss_fn, optimizer, head_only=True, linear_evaluation=False, percentage=1, lr_schedule=None, save_model_at=None, val_idmap=None, test_idmap=None, lbl_itos=None, cpc=False):

    if head_only:

        if linear_evaluation:

            print("linear evaluation for {} epochs".format(epochs))

        else:

            print("head-only for {} epochs".format(epochs))

    else:

        print("fine tuning for {} epochs".format(epochs))

    if head_only:

        for key, param in model.named_parameters():

            if "l1." not in key and "head." not in key:

                param.requires_grad = False

        print("copying state dict before training for sanity check after training")

    else:

        for param in model.parameters():

            param.requires_grad = True

    if cpc:

        data_type = model.encoder[0][0].weight.type()

    else:

        data_type = model.features[0][0].weight.type()

    set_train_eval(model, cpc, linear_evaluation)

    state_dict_pre = deepcopy(model.state_dict())

    print("epoch", "batch", "loss\n========================")

    loss_per_epoch = []

    macro_agg_per_epoch = []

    max_batches = len(train_loader)

    break_point = int(percentage*max_batches)

    best_macro = 0

    best_macro_agg = 0

    best_epoch = 0

    best_preds = None

    test_macro = 0

    test_macro_agg = 0

    for epoch in tqdm(range(epochs)):

        if type(lr_schedule) == dict:

            if epoch in lr_schedule.keys():

                for param_group in optimizer.param_groups:

                    param_group['lr'] /= lr_schedule[epoch]

        total_loss_one_epoch = 0

        for batch_idx, samples in enumerate(train_loader):

            if batch_idx == break_point:

                print("break at batch nr.", batch_idx)

                break

            data = samples[0].to(device).type(data_type)

            labels = samples[1].to(device).type(data_type)

            optimizer.zero_grad()

            preds = model(data)

            loss = loss_fn(preds, labels)

            loss.backward()

            optimizer.step()

            total_loss_one_epoch += loss.item()

            if(batch_idx % 100 == 0):

                print(epoch, batch_idx, loss.item())

        loss_per_epoch.append(total_loss_one_epoch)

        preds, macro, macro_agg = evaluate(

            model, valid_loader, val_idmap, lbl_itos, cpc=cpc)

        macro_agg_per_epoch.append(macro_agg)

        print("loss:", total_loss_one_epoch)

        print("aggregated macro:", macro_agg)

        if macro_agg > best_macro_agg:

            torch.save(model.state_dict(), save_model_at)

            best_macro_agg = macro_agg

            best_macro = macro

            best_epoch = epoch

            best_preds = preds

            _, test_macro, test_macro_agg = evaluate(

                model, test_loader, test_idmap, lbl_itos, cpc=cpc)

        set_train_eval(model, cpc, linear_evaluation)

    if epochs > 0:

        sanity_check(model, state_dict_pre, linear_evaluation, head_only)

    return loss_per_epoch, macro_agg_per_epoch, best_macro, best_macro_agg, test_macro, test_macro_agg, best_epoch, best_preds

def sanity_check(model, state_dict_pre, linear_evaluation, head_only):

    """

    Linear classifier should not change any weights other than the linear layer.

    This sanity check asserts nothing wrong happens (e.g., BN stats updated).

    """

    print("=> loading state dict for sanity check")

    state_dict = model.state_dict()

    if linear_evaluation:

        for k in list(state_dict.keys()):

            # only ignore fc layer

            if 'fc.' in k or 'head.' in k or 'l1.' in k:

                continue

            equals = (state_dict[k].cpu() == state_dict_pre[k].cpu()).all()

            if (linear_evaluation != equals):

                raise Exception(

                    '=> failed sanity check in {}'.format("linear_evaluation"))

    elif head_only:

        for k in list(state_dict.keys()):

            # only ignore fc layer

            if 'fc.' in k or 'head.' in k:

                continue

            equals = (state_dict[k].cpu() == state_dict_pre[k].cpu()).all()

            if (equals and "running_mean" in k):

                raise Exception(

                    '=> failed sanity check in {}'.format("head-only"))

    # else:

    #     for k in list(state_dict.keys()):

    #         equals=(state_dict[k].cpu() == state_dict_pre[k].cpu()).all()

    #         if equals:

    #             pdb.set_trace()

    #             raise Exception('=> failed sanity check in {}'.format("fine_tuning"))

    print("=> sanity check passed.")

def eval_model(model, valid_loader, cpc=False):

    if cpc:

        data_type = model.encoder[0][0].weight.type()

    else:

        data_type = model.features[0][0].weight.type()

    model.eval()

    preds = []

    targs = []

    with torch.no_grad():

        for batch_idx, samples in tqdm(enumerate(valid_loader)):

            data = samples[0].to(device).type(data_type)

            preds_tmp = torch.sigmoid(model(data))

            targs.append(samples[1])

            preds.append(preds_tmp.cpu())

        preds = torch.cat(preds).numpy()

        targs = torch.cat(targs).numpy()

    return preds, targs

def get_dataset(batch_size, num_workers, target_folder, apply_noise=False, percentage=1.0, folds=8, t_params=None, test=False, normalize=False):

    if apply_noise:

        transformations = ["BaselineWander",

                           "PowerlineNoise", "EMNoise", "BaselineShift"]

        if normalize:

            transformations.append("Normalize")

        dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,

                                       mode="linear_evaluation", transformations=transformations, percentage=percentage, folds=folds, t_params=t_params, test=test, ptb_xl_label="label_all")

    else:

        if normalize:

            # always use PTB-XL stats

            transformations = ["Normalize"]

            dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,

            mode="linear_evaluation", percentage=percentage, folds=folds, test=test, transformations=transformations, ptb_xl_label="label_all")

        else:

            dataset = SimCLRDataSetWrapper(batch_size,num_workers,None,"(12, 250)",None,target_folder,[target_folder],None,None,

                                           mode="linear_evaluation", percentage=percentage, folds=folds, test=test, ptb_xl_label="label_all")

    train_loader, valid_loader = dataset.get_data_loaders()

    return dataset, train_loader, valid_loader

if __name__ == "__main__":

    args = parse_args()

    dataset, train_loader, _ = get_dataset(

        args.batch_size, args.num_workers, args.dataset, folds=args.folds, test=args.test, normalize=args.normalize)

    _, _, valid_loader = get_dataset(

        args.batch_size, args.num_workers, args.dataset, folds=args.folds, test=False, normalize=args.normalize)

    val_idmap = dataset.val_ds_idmap

    dataset, _, test_loader = get_dataset(

        args.batch_size, args.num_workers, args.dataset, test=True, normalize=args.normalize)

    test_idmap = dataset.val_ds_idmap

    lbl_itos = dataset.lbl_itos

    tag = "f=" + str(args.folds) + "_" + args.tag

    tag = tag if args.use_pretrained else "ran_" + tag

    tag = "eval_" + tag if args.eval_only else tag

    model_tag = "finetuned" if args.load_finetuned else "ckpt"

    if args.test_noised:

        t_params_by_level = {

            1: {"bw_cmax": 0.05, "em_cmax": 0.25, "pl_cmax": 0.1, "bs_cmax": 0.5},

            2: {"bw_cmax": 0.1, "em_cmax": 0.5, "pl_cmax": 0.2, "bs_cmax": 1},

            3: {"bw_cmax": 0.1, "em_cmax": 1, "pl_cmax": 0.2, "bs_cmax": 2},

            4: {"bw_cmax": 0.2, "em_cmax": 1, "pl_cmax": 0.4, "bs_cmax": 2},

            5: {"bw_cmax": 0.2, "em_cmax": 1.5, "pl_cmax": 0.4, "bs_cmax": 2.5},

            6: {"bw_cmax": 0.3, "em_cmax": 2, "pl_cmax": 0.5, "bs_cmax": 3},

        }

        if args.noise_level not in t_params_by_level.keys():

            raise("noise level does not exist")

        t_params = t_params_by_level[args.noise_level]

        dataset, _, noise_valid_loader = get_dataset(

            args.batch_size, args.num_workers, args.dataset, apply_noise=True, t_params=t_params, test=args.test)

    else:

        noise_valid_loader = None

    losses, macros, predss, result_macros, result_macros_agg, test_macros, test_macros_agg, noised_macros, noised_macros_agg = [

    ], [], [], [], [], [], [], [], []

    ckpt_epoch_lin=0

    ckpt_epoch_fin=0

    if args.f_epochs == 0:

        save_model_at = os.path.join(os.path.dirname(

            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "lin_finetuned")

        filename = os.path.join(os.path.dirname(

            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "res_lin.pkl")

    else:

        save_model_at = os.path.join(os.path.dirname(

            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "fin_finetuned")

        filename = os.path.join(os.path.dirname(

            args.model_file), "n=" + str(args.noise_level) + "_"+tag + "res_fin.pkl")

    model = load_model(

        args.linear_evaluation, 71, args.use_pretrained or args.load_finetuned, hidden=args.hidden,

        location=args.model_file, discriminative_lr=args.discriminative_lr, method=args.method)

    loss_fn, optimizer = configure_optimizer(

        model, args.batch_size, head_only=True, discriminative_lr=args.discriminative_lr, discriminative_lr_factor=0.1 if args.use_pretrained and args.discriminative_lr else 1)

    if not args.eval_only:

        print("train model...")

        if not isdir(save_model_at):

            os.mkdir(save_model_at)

        l1, m1, bm, bm_agg, tm, tm_agg, ckpt_epoch_lin, preds = train_model(model, train_loader, valid_loader, test_loader, args.l_epochs, loss_fn,

                                                                            optimizer, head_only=True, linear_evaluation=args.linear_evaluation, lr_schedule=args.lr_schedule, save_model_at=join(save_model_at, "finetuned.pt"),

                                                                            val_idmap=val_idmap, test_idmap=test_idmap, lbl_itos=lbl_itos, cpc=(args.method == "cpc"))

        if bm != 0:

            print("best macro after head-only training:", bm_agg)

        l2 = []

        m2 = []

        if args.f_epochs != 0:

            if args.l_epochs != 0:

                model = load_model(

                    False, 71, True, hidden=args.hidden,

                    location=join(save_model_at, "finetuned.pt"), discriminative_lr=args.discriminative_lr, method=args.method)

            loss_fn, optimizer = configure_optimizer(

                model, args.batch_size, head_only=False, discriminative_lr=args.discriminative_lr, discriminative_lr_factor=0.1 if args.use_pretrained and args.discriminative_lr else 1)

            l2, m2, bm, bm_agg, tm, tm_agg, ckpt_epoch_fin, preds = train_model(model, train_loader, valid_loader, test_loader, args.f_epochs, loss_fn,

                                                                                optimizer, head_only=False, linear_evaluation=False, lr_schedule=args.lr_schedule, save_model_at=join(save_model_at, "finetuned.pt"),

                                                                                val_idmap=val_idmap, test_idmap=test_idmap, lbl_itos=lbl_itos, cpc=(args.method == "cpc"))

        losses.append(l1+l2)

        macros.append(m1+m2)

        test_macros.append(tm)

        test_macros_agg.append(tm_agg)

        result_macros.append(bm)

        result_macros_agg.append(bm_agg)

    else:

        preds, eval_macro, eval_macro_agg = evaluate(

            model, test_loader, test_idmap, lbl_itos, cpc=(args.method == "cpc"))

        result_macros.append(eval_macro)

        result_macros_agg.append(eval_macro_agg)

        if args.verbose:

            print("macro:", eval_macro)

    predss.append(preds)

    if noise_valid_loader is not None:

        _, noise_macro, noise_macro_agg = evaluate(

            model, noise_valid_loader, val_idmap, lbl_itos)

        noised_macros.append(noise_macro)

        noised_macros_agg.append(noise_macro_agg)

    res = {"filename": filename, "epochs": args.l_epochs+args.f_epochs, "model_location": args.model_location,

           "losses": losses, "macros": macros, "predss": predss, "result_macros": result_macros, "result_macros_agg": result_macros_agg,

           "test_macros": test_macros, "test_macros_agg": test_macros_agg, "noised_macros": noised_macros, "noised_macros_agg": noised_macros_agg, "ckpt_epoch_lin": ckpt_epoch_lin, "ckpt_epoch_fin": ckpt_epoch_fin,

           "discriminative_lr": args.discriminative_lr, "hidden": args.hidden, "lr_schedule": args.lr_schedule,

           "use_pretrained": args.use_pretrained, "linear_evaluation": args.linear_evaluation, "loaded_finetuned": args.load_finetuned,

           "eval_only": args.eval_only, "noise_level": args.noise_level, "test_noised": args.test_noised, "normalized": args.normalize}

    pickle.dump(res, open(filename, "wb"))

    print("dumped results to", filename)

    print(res)

    print("Done!")