import copy
from torch.nn.modules import loss
# from model.hash_model import DCMHT as DCMHT
import os
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.utils.data as data
import scipy.io as scio
import numpy as np

from .base import TrainBase
from torch.nn import functional as F
from utils import get_args, calc_neighbor, cosine_similarity, euclidean_similarity, find_indices
from utils.calc_utils import cal_map, cal_pr
from dataset.dataloader import dataloader
import clip
import re
from transformers import BertForMaskedLM
from model.bert_tokenizer import BertTokenizer
from transformers import AutoTokenizer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

filter_words = ['a', 'about', 'above', 'across', 'after', 'afterwards', 'again', 'against', 'ain', 'all', 'almost',
                'alone', 'along', 'already', 'also', 'although', 'am', 'among', 'amongst', 'an', 'and', 'another',
                'any', 'anyhow', 'anyone', 'anything', 'anyway', 'anywhere', 'are', 'aren', "aren't", 'around', 'as',
                'at', 'back', 'been', 'before', 'beforehand', 'behind', 'being', 'below', 'beside', 'besides',
                'between', 'beyond', 'both', 'but', 'by', 'can', 'cannot', 'could', 'couldn', "couldn't", 'd', 'didn',
                "didn't", 'doesn', "doesn't", 'don', "don't", 'down', 'due', 'during', 'either', 'else', 'elsewhere',
                'empty', 'enough', 'even', 'ever', 'everyone', 'everything', 'everywhere', 'except', 'first', 'for',
                'former', 'formerly', 'from', 'hadn', "hadn't", 'hasn', "hasn't", 'haven', "haven't", 'he', 'hence',
                'her', 'here', 'hereafter', 'hereby', 'herein', 'hereupon', 'hers', 'herself', 'him', 'himself', 'his',
                'how', 'however', 'hundred', 'i', 'if', 'in', 'indeed', 'into', 'is', 'isn', "isn't", 'it', "it's",
                'its', 'itself', 'just', 'latter', 'latterly', 'least', 'll', 'may', 'me', 'meanwhile', 'mightn',
                "mightn't", 'mine', 'more', 'moreover', 'most', 'mostly', 'must', 'mustn', "mustn't", 'my', 'myself',
                'namely', 'needn', "needn't", 'neither', 'never', 'nevertheless', 'next', 'no', 'nobody', 'none',
                'noone', 'nor', 'not', 'nothing', 'now', 'nowhere', 'o', 'of', 'off', 'on', 'once', 'one', 'only',
                'onto', 'or', 'other', 'others', 'otherwise', 'our', 'ours', 'ourselves', 'out', 'over', 'per',
                'please', 's', 'same', 'shan', "shan't", 'she', "she's", "should've", 'shouldn', "shouldn't", 'somehow',
                'something', 'sometime', 'somewhere', 'such', 't', 'than', 'that', "that'll", 'the', 'their', 'theirs',
                'them', 'themselves', 'then', 'thence', 'there', 'thereafter', 'thereby', 'therefore', 'therein',
                'thereupon', 'these', 'they', 'this', 'those', 'through', 'throughout', 'thru', 'thus', 'to', 'too',
                'toward', 'towards', 'under', 'unless', 'until', 'up', 'upon', 'used', 've', 'was', 'wasn', "wasn't",
                'we', 'were', 'weren', "weren't", 'what', 'whatever', 'when', 'whence', 'whenever', 'where',
                'whereafter', 'whereas', 'whereby', 'wherein', 'whereupon', 'wherever', 'whether', 'which', 'while',
                'whither', 'who', 'whoever', 'whole', 'whom', 'whose', 'why', 'with', 'within', 'without', 'won',
                "won't", 'would', 'wouldn', "wouldn't", 'y', 'yet', 'you', "you'd", "you'll", "you're", "you've",
                'your', 'yours', 'yourself', 'yourselves', '.', '-', 'a the', '/', '?', 'some', '"', ',', 'b', '&', '!',
                '@', '%', '^', '*', '(', ')', "-", '-', '+', '=', '<', '>', '|', ':', ";", '～', '·']
filter_words = set(filter_words)


def text_filter(text):
    text = re.findall(r"<|startoftext|>(.+)<|endoftext|>", text)
    text = text[2]
    text = re.sub(r'</w>', ' ', text)
    return text


class GoalFunctionStatus(object):
    SUCCEEDED = 0  # attack succeeded
    SEARCHING = 1  # In process of searching for a success
    FAILED = 2  # attack failed


class GoalFunctionResult(object):
    goal_score = 1

    def __init__(self, text, score=0, similarity=None):
        self.status = GoalFunctionStatus.SEARCHING
        self.text = text
        self.score = score
        self.similarity = similarity

    @property
    def score(self):
        return self.__score

    @score.setter
    def score(self, value):
        self.__score = value
        if value >= self.goal_score:
            self.status = GoalFunctionStatus.SUCCEEDED

    def __eq__(self, __o):
        return self.text == __o.text

    def __hash__(self):
        return hash(self.text)


class Trainer(TrainBase):

    def __init__(self,
                 rank=0):
        args = get_args()
        super(Trainer, self).__init__(args, rank)
        self.logger.info("dataset len: {}".format(len(self.train_loader.dataset)))
        image_mean, image_var = self.generate_mapping()
        self.image_mean = image_mean
        self.image_var = image_var
        self.device = rank
        self.clip_tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32", model_max_length=77,
                                                            truncation=True)
        self.bert_tokenizer = BertTokenizer.from_pretrained(self.args.text_encoder, do_lower_case=True)
        self.ref_net = BertForMaskedLM.from_pretrained(self.args.text_encoder).to(device)
        self.attack_thred = self.args.attack_thred
        # self.run()

    def _init_model(self):
        self.logger.info("init model.")
        model_clip, preprocess = clip.load(self.args.victim, device=device)
        self.model = model_clip
        self.model.eval()
        self.model.float()

    def _init_dataset(self):
        self.logger.info("init dataset.")
        self.logger.info(f"Using {self.args.dataset} dataset.")
        self.args.index_file = os.path.join("./dataset", self.args.dataset, self.args.index_file)
        self.args.caption_file = os.path.join("./dataset", self.args.dataset, self.args.caption_file)
        self.args.label_file = os.path.join("./dataset", self.args.dataset, self.args.label_file)
        train_data, query_data, retrieval_data = dataloader(captionFile=self.args.caption_file,
                                                            indexFile=self.args.index_file,
                                                            labelFile=self.args.label_file,
                                                            maxWords=self.args.max_words,
                                                            imageResolution=self.args.resolution,
                                                            query_num=self.args.query_num,
                                                            train_num=self.args.train_num,
                                                            seed=self.args.seed)
        self.train_labels = train_data.get_all_label()
        self.query_labels = query_data.get_all_label()
        self.retrieval_labels = retrieval_data.get_all_label()
        self.args.retrieval_num = len(self.retrieval_labels)
        self.logger.info(f"query shape: {self.query_labels.shape}")
        self.logger.info(f"retrieval shape: {self.retrieval_labels.shape}")
        self.train_loader = DataLoader(
            dataset=train_data,
            batch_size=self.args.batch_size,
            num_workers=self.args.num_workers,
            pin_memory=True,
            shuffle=True
        )
        self.query_loader = DataLoader(
            dataset=query_data,
            batch_size=self.args.batch_size,
            num_workers=self.args.num_workers,
            pin_memory=True,
            shuffle=True
        )
        self.retrieval_loader = DataLoader(
            dataset=retrieval_data,
            batch_size=self.args.batch_size,
            num_workers=self.args.num_workers,
            pin_memory=True,
            shuffle=True
        )
        self.train_data = train_data

    def _tokenize(self, text):
        words = text.split(' ')

        sub_words = []
        keys = []
        index = 0
        for word in words:
            sub = self.bert_tokenizer.tokenize(word)
            sub_words += sub
            keys.append([index, index + len(sub)])
            index += len(sub)

        return words, sub_words, keys

    def get_important_scores(self, text, origin_embeds, batch_size, max_length):
        # device = origin_embeds.device

        masked_words = self._get_masked(text)
        masked_texts = [' '.join(words) for words in masked_words]  # list of text of masked words

        masked_embeds = []
        for i in range(0, len(masked_texts), batch_size):
            masked_text_input = self.bert_tokenizer(masked_texts[i:i + batch_size], padding='max_length',
                                                    truncation=True, max_length=max_length, return_tensors='pt').to(
                device)
            masked_embed = self.ref_net(masked_text_input.text_inputs, attention_mask=masked_text_input.attention_mask)
            masked_embeds.append(masked_embed)
        masked_embeds = torch.cat(masked_embeds, dim=0)

        criterion = torch.nn.KLDivLoss(reduction='none')

        import_scores = criterion(masked_embeds.log_softmax(dim=-1),
                                  origin_embeds.softmax(dim=-1).repeat(len(masked_texts), 1))

        return import_scores.sum(dim=-1)

    def _get_masked(self, text):
        words = text.split(' ')
        len_text = len(words)
        masked_words = []
        for i in range(len_text):
            masked_words.append(words[0:i] + ['[UNK]'] + words[i + 1:])
        # list of words
        return masked_words

    def get_transformations(self, text, idx, substitutes):
        words = text.split(' ')

        trans_text = []
        for sub in substitutes:
            words[idx] = sub
            trans_text.append(' '.join(words))
        return trans_text

    def get_word_predictions(self, text):
        _, _, keys = self._tokenize(text)

        inputs = self.bert_tokenizer.encode_plus(text, add_special_tokens=True, max_length=self.args.max_words,
                                                 truncation=True, return_tensors="pt")
        input_ids = inputs["input_ids"].to(self.device)
        attention_mask = inputs['attention_mask']
        with torch.no_grad():
            word_predictions = self.ref_net(input_ids)['logits'].squeeze(0)  # (seq_len, vocab_size)
        # print(self.ref_net(input_ids)['logits'].shape)
        word_pred_scores_all, word_predictions = torch.topk(word_predictions, self.args.max_candidate, -1)

        word_predictions = word_predictions[1:-1, :]  # remove [CLS] and [SEP]
        word_pred_scores_all = word_pred_scores_all[1:-1, :]

        return keys, word_predictions, word_pred_scores_all, attention_mask

    def get_bpe_substitutes(self, substitutes):
        # substitutes L, k
        substitutes = substitutes[0:12, 0:4]  # maximum BPE candidates

        # find all possible candidates
        all_substitutes = []
        for i in range(substitutes.size(0)):
            if len(all_substitutes) == 0:
                lev_i = substitutes[i]
                all_substitutes = [[int(c)] for c in lev_i]
            else:
                lev_i = []
                for all_sub in all_substitutes:
                    for j in substitutes[i]:
                        lev_i.append(all_sub + [int(j)])
                all_substitutes = lev_i

        # all substitutes: list of list of token-id (all candidates)
        cross_entropy_loss = nn.CrossEntropyLoss(reduction='none')

        word_list = []
        all_substitutes = torch.tensor(all_substitutes)  # [ N, L ]
        all_substitutes = all_substitutes[:24].to(self.device)

        N, L = all_substitutes.size()
        word_predictions = self.ref_net(all_substitutes)[0]  # N L vocab-size
        ppl = cross_entropy_loss(word_predictions.view(N * L, -1), all_substitutes.view(-1))  # [ N*L ]
        ppl = torch.exp(torch.mean(ppl.view(N, L), dim=-1))  # N

        _, word_list = torch.sort(ppl)
        word_list = [all_substitutes[i] for i in word_list]
        final_words = []
        for word in word_list:
            tokens = [self.bert_tokenizer.convert_ids_to_tokens(int(i)) for i in word]
            text = ' '.join([t.strip() for t in tokens])
            final_words.append(text)
        return final_words

    def get_substitutes(self, substitutes, substitutes_score, threshold=3.0):
        ret = []
        num_sub, _ = substitutes.size()
        if num_sub == 0:
            ret = []
        elif num_sub == 1:
            for id, score in zip(substitutes[0], substitutes_score[0]):
                if threshold != 0 and score < threshold:
                    break
                ret.append(self.bert_tokenizer.convert_ids_to_tokens(int(id)))
        elif self.args.enable_bpe:
            ret = self.get_bpe_substitutes(substitutes)
        return ret

    def filter_substitutes(self, substitues):

        ret = []
        for word in substitues:
            if word.lower() in filter_words:
                continue
            if '##' in word:
                continue

            ret.append(word)
        return ret

    def get_goal_results(self, trans_texts, negetive_code, negetive_mean, negative_var, positive_code, positive_mean,
                         positive_var, beta=10, temperature=0.05):
        # print(trans_texts)
        trans_feature = clip.tokenize(trans_texts,context_length=77,truncate=True).to(device)
        anchor = self.model.encode_text(trans_feature)
        batch_size=anchor.shape[0]
        loss1 = F.triplet_margin_with_distance_loss(anchor, positive_code.repeat(batch_size, 1), negetive_code.repeat(batch_size, 1),
                                                    distance_function=nn.CosineSimilarity(), reduction='none')
        sim=F.cosine_similarity(anchor,positive_code.unsqueeze(0), dim=1, eps=1e-8).unsqueeze(1)
        negative_dist = (anchor - negetive_mean) ** 2 / (negative_var+ 1e-5)
        positive_dist = (anchor - positive_mean) ** 2 / (positive_var+ 1e-5)
        negatives = torch.exp(negative_dist / temperature)
        positives = torch.exp(positive_dist / temperature)
        loss = torch.log(positives / (positives + negatives)).mean(dim=1, keepdim=True) + beta * loss1
        results = []
        # print(loss.shape)
        for i in range(len(trans_texts)):
            if loss[i].shape[0] >1 or sim[i] >self.args.sim_threshold:
                continue
            results.append(GoalFunctionResult(trans_texts[i], score=loss[i], similarity=sim[i]))
        return results

    def generate_mapping(self):
        image_train = []
        label_train = []
        for image, text, label, index in self.train_loader:
            # raw_text=[self.clip_tokenizer.decode(token) for token in text]
            image = image.to(device, non_blocking=True)
            # print(self.model.vocab_size)
            temp_image = self.model.encode_image(image)
            image_train.append(temp_image.cpu().detach().numpy())
            label_train.append(label.detach().numpy())
        image_train = np.concatenate(image_train, axis=0)
        label_train = np.concatenate(label_train, axis=0)
        label_unipue = np.unique(label_train, axis=0)
        image_centroids = np.stack(
            [image_train[find_indices(label_train, label_unipue[i])].mean(axis=0) for i in range(len(label_unipue))],
            axis=0)
        image_var = np.stack(
            [image_train[find_indices(label_train, label_unipue[i])].var(axis=0) for i in range(len(label_unipue))],
            axis=0)

        image_representation = {}
        image_var_representation = {}
        for i, centroid in enumerate(label_unipue):
            image_representation[str(centroid.astype(int))] = image_centroids[i]
            image_var_representation[str(centroid.astype(int))] = image_var[i]
        return image_representation, image_var_representation

    def target_adv(self, raw_text, negetive_code, negetive_mean, negative_var, positive_code, positive_mean,
                   positive_var, beta=10, temperature=0.05):
        # print(raw_text)
        keys, word_predictions, word_pred_scores_all, mask = self.get_word_predictions(raw_text)

        #clean state
        # clean_embeds=self.ref_net(bert_inputs.input_ids, attention_mask=bert_inputs.attention_mask)
        cur_result = GoalFunctionResult(raw_text)
        mask_idx = np.where(mask.cpu().numpy() == 1)[0]

        for idx in mask_idx:
            predictions = word_predictions[keys[idx][0]: keys[idx][1]]
            predictions_socre = word_pred_scores_all[keys[idx][0]: keys[idx][1]]
            substitutes = self.get_substitutes(predictions, predictions_socre)
            substitutes = self.filter_substitutes(substitutes)
            trans_texts = self.get_transformations(raw_text, idx, substitutes)
            if len(trans_texts) == 0:
                continue
            # loss function
            results = self.get_goal_results(trans_texts, negetive_code, negetive_mean, negative_var, positive_code,
                                            positive_mean, positive_var, beta, temperature)
            results = sorted(results, key=lambda x: x.score, reverse=True)

            if len(results) > 0 and results[0].score > cur_result.score:
                cur_result = results[0]
            else:
                continue

            if cur_result.status == GoalFunctionStatus.SUCCEEDED:
                max_similarity = cur_result.similarity
                if max_similarity is None:
                    # similarity is not calculated
                    continue

                for result in results[1:]:
                    if result.status != GoalFunctionStatus.SUCCEEDED:
                        break
                    if result.similarity > max_similarity:
                        max_similarity = result.similarity
                        cur_result = result
                return cur_result
        if cur_result.status == GoalFunctionStatus.SEARCHING:
            cur_result.status = GoalFunctionStatus.FAILED
        return cur_result

    def train_epoch(self):
        # self.change_state(mode="valid")
        save_dir = os.path.join(self.args.save_dir, "adv_PR_cruve")
        all_loss = 0
        times = 0
        adv_codes = []
        adv_label = []
        for image, text, label, index in self.train_loader:
            self.global_step += 1
            times += 1
            print(times)
            image.float()

            image = image.to(self.rank, non_blocking=True)
            text = text.to(self.rank, non_blocking=True)
            negetive_mean = np.stack([self.image_mean[str(i.astype(int))] for i in label.detach().cpu().numpy()])
            negative_var = np.stack([self.image_var[str(i.astype(int))] for i in label.detach().cpu().numpy()])
            negetive_mean = torch.from_numpy(negetive_mean).to(self.rank, non_blocking=True)
            negative_var = torch.from_numpy(negative_var).to(self.rank, non_blocking=True)
            negetive_code = self.model.encode_image(image)

            #targeted sample
            np.random.seed(times)
            select_index = np.random.choice(len(self.train_data), size=self.args.batch_size)
            target_dataset = data.Subset(self.train_data, select_index)
            target_subset = torch.utils.data.DataLoader(target_dataset, batch_size=self.args.batch_size)
            target_image, _, target_label, _ = next(iter(target_subset))
            target_image = target_image.to(self.rank, non_blocking=True)
            positive_mean = np.stack([self.image_mean[str(i.astype(int))] for i in target_label.detach().cpu().numpy()])
            positive_var = np.stack([self.image_var[str(i.astype(int))] for i in target_label.detach().cpu().numpy()])
            positive_mean = torch.from_numpy(positive_mean).to(self.rank, non_blocking=True)
            positive_var = torch.from_numpy(positive_var).to(self.rank, non_blocking=True)
            positive_code = self.model.encode_image(target_image)
            # print(self.clip_tokenizer.my_encode('This day is good!'))
            raw_text = [self.clip_tokenizer.convert_ids_to_tokens(token.cpu()) for token in text]
            raw_text = [text_filter(self.clip_tokenizer.convert_tokens_to_string(txt)) for txt in raw_text]
            final_texts=[]
            for i in range(self.args.batch_size):
                adv_txt=self.target_adv( raw_text[i], negetive_code[i], negetive_mean[i], negative_var[i],
                                            positive_code[i], positive_mean[i], positive_var[i])
                final_texts.append(adv_txt.text)
            # final_adverse = self.target_adv( raw_text, negetive_code, negetive_mean, negative_var,
            #                                 positive_code, positive_mean, positive_var)
            final_text = clip.tokenize(final_texts,context_length=77,truncate=True).to(self.rank, non_blocking=True)
            adv_code = self.model.encode_text(final_text)
            adv_codes.append(adv_code.cpu().detach().numpy())
            adv_label.append(target_label.numpy())
        adv_img = np.concatenate(adv_codes)
        adv_labels = np.concatenate(adv_label)

        _, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num)

        retrieval_txt = retrieval_txt.cpu().detach().numpy()
        retrieval_labels = self.retrieval_labels.numpy()

        mAP_t = cal_map(adv_img, adv_labels, retrieval_txt, retrieval_labels)
        self.logger.info(f">>>>>> MAP_t: {mAP_t}")
        result_dict = {
            'adv_img': adv_img,
            'r_txt': retrieval_txt,
            'adv_l': adv_labels,
            'r_l': retrieval_labels
            # 'q_l':query_labels
            # 'pr': pr,
            # 'pr_t': pr_t
        }
        scio.savemat(os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-adv-" + self.args.dataset + ".mat"),
                     result_dict)
        self.logger.info(">>>>>> save all data!")

    def train(self):
        self.logger.info("Start train.")

        for epoch in range(self.args.epochs):
            self.train_epoch(epoch)
            self.valid(epoch)
            self.save_model(epoch)

        self.logger.info(
            f">>>>>>> FINISHED >>>>>> Best epoch, I-T: {self.best_epoch_i}, mAP: {self.max_mapi2t}, T-I: {self.best_epoch_t}, mAP: {self.max_mapt2i}")

    def make_hash_code(self, code: list) -> torch.Tensor:

        code = torch.stack(code)
        # print(code.shape)
        code = code.permute(1, 0, 2)
        hash_code = torch.argmax(code, dim=-1)
        hash_code[torch.where(hash_code == 0)] = -1
        hash_code = hash_code.float()

        return hash_code

    def get_code(self, data_loader, length: int):

        img_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank)
        text_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank)

        for image, text, label, index in tqdm(data_loader):
            image = image.to(self.device, non_blocking=True)
            text = text.to(self.device, non_blocking=True)
            index = index.numpy()
            with torch.no_grad():
                image_feature = self.model.encode_image(image)
                text_features = self.model.encode_text(text)
            img_buffer[index, :] = image_feature.detach()
            text_buffer[index, :] = text_features.detach()

        return img_buffer, text_buffer  # img_buffer.to(self.rank), text_buffer.to(self.rank)

    def valid_attack(self, adv_images, texts, adv_labels):
        save_dir = os.path.join(self.args.save_dir, "adv_PR_cruve")
        os.makedirs(save_dir, exist_ok=True)

    def test(self, mode_name="i2t"):
        self.logger.info("Valid Clean.")
        save_dir = os.path.join(self.args.save_dir, "PR_cruve")
        os.makedirs(save_dir, exist_ok=True)
        query_img, query_txt = self.get_code(self.query_loader, self.args.query_num)
        retrieval_img, retrieval_txt = self.get_code(self.retrieval_loader, self.args.retrieval_num)

        query_img = query_img.cpu().detach().numpy()
        query_txt = query_txt.cpu().detach().numpy()
        retrieval_img = retrieval_img.cpu().detach().numpy()
        retrieval_txt = retrieval_txt.cpu().detach().numpy()
        query_labels = self.query_labels.numpy()
        retrieval_labels = self.retrieval_labels.numpy()
        mAPi2t = cal_map(query_img, query_labels, retrieval_txt, retrieval_labels)
        mAPt2i = cal_map(query_txt, query_labels, retrieval_img, retrieval_labels)
        # pr_i2t=cal_pr(retrieval_txt,query_img,query_labels,retrieval_labels)
        # pr_t2i=cal_pr(retrieval_img,query_txt,query_labels,retrieval_labels)
        self.max_mapt2i = max(self.max_mapt2i, mAPi2t)
        self.logger.info(f">>>>>> MAP(i->t): {mAPi2t}, MAP(t->i): {mAPt2i}")
        result_dict = {
            'q_img': query_img,
            'q_txt': query_txt,
            'r_img': retrieval_img,
            'r_txt': retrieval_txt,
            'q_l': query_labels,
            'r_l': retrieval_labels
        }
        scio.savemat(os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-ours-" + self.args.dataset + ".mat"),
                     result_dict)
        self.logger.info(">>>>>> save all data!")



    def save_mat(self, query_img, query_txt, retrieval_img, retrieval_txt, mode_name="i2t"):

        save_dir = os.path.join(self.args.save_dir, "PR_cruve")
        os.makedirs(save_dir, exist_ok=True)

        query_img = query_img.cpu().detach().numpy()
        query_txt = query_txt.cpu().detach().numpy()
        retrieval_img = retrieval_img.cpu().detach().numpy()
        retrieval_txt = retrieval_txt.cpu().detach().numpy()
        query_labels = self.query_labels.numpy()
        retrieval_labels = self.retrieval_labels.numpy()

        result_dict = {
            'q_img': query_img,
            'q_txt': query_txt,
            'r_img': retrieval_img,
            'r_txt': retrieval_txt,
            'q_l': query_labels,
            'r_l': retrieval_labels
        }
        scio.savemat(
            os.path.join(save_dir, str(self.args.victim).replace("/", "_") + "-ours-" + self.args.dataset + "-" + mode_name + ".mat"),
            result_dict)
        self.logger.info(f">>>>>> save best {mode_name} data!")
