import warnings
warnings.filterwarnings("ignore")

import imp
from random import sample
from dassl.engine import TRAINER_REGISTRY, TrainerX
import os.path as osp
import os
import time
import datetime
import numpy as np
from tqdm import tqdm
import json

import pandas as pd

from sklearn.metrics import confusion_matrix
import seaborn as sns
import matplotlib 
matplotlib.use('Agg')
import matplotlib.pylab as plt
from sklearn.mixture import GaussianMixture

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.cuda.amp import GradScaler, autocast

from dassl.engine import TRAINER_REGISTRY, TrainerX
from dassl.metrics import compute_accuracy
from dassl.utils import load_pretrained_weights, load_checkpoint
from dassl.utils import (
    MetricMeter, AverageMeter, tolist_if_not, count_num_param, load_checkpoint,
    save_checkpoint, mkdir_if_missing, resume_from_checkpoint,
    load_pretrained_weights
)
from dassl.optim import build_optimizer, build_lr_scheduler
from dassl.data import DataManager

from clip import clip
from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer

from datasets.data_manager import MYDataManager, IndexDataset
from evaluation.evaluator import MYClassification
from .hhzsclip import ZeroshotCLIP
from .utils import (select_top_k_similarity_per_class, caculate_noise_rate, save_outputs,
select_top_k_similarity, select_top_by_value, caculate_noise_rate_analyze, gen_noise_label_asym, gen_noise_label_sym)

_tokenizer = _Tokenizer()
from trainers.loss import GeneralizedCrossEntropy

from description import *

# KL Divergence
def kl_divergence(p, q):
    return (p * ((p+1e-10) / (q+1e-10)).log()).sum(dim=1)

## Jensen-Shannon Divergence 
class Jensen_Shannon(nn.Module):
    def __init__(self):
        super(Jensen_Shannon,self).__init__()
        pass
    def forward(self, p,q):
        m = (p+q)/2
        return 0.5*kl_divergence(p, m) + 0.5*kl_divergence(q, m)


def rand_argmax(tens):
    max_inds, = torch.where(tens == tens.max())
    return np.random.choice(max_inds)

## Calculate JSD
def Calculate_JSD(outputs, targets, num_classes):  
    JS_dist = Jensen_Shannon()
    outputs = torch.softmax(outputs,dim=1)
    JSD = JS_dist(outputs,  F.one_hot(targets, num_classes = num_classes))
    return JSD

def MA(matrix, alpha):
    ret = np.zeros_like(matrix)
    for epoch in range(ret.shape[1]):
        if epoch == 0:
            ret[:,epoch] = matrix[:,epoch]
        else:
            ret[:,epoch] = ret[:,epoch-1] * alpha + matrix[:,epoch] * (1-alpha)
    
    return ret

def calc_entropy(input_tensor):
    lsm = nn.LogSoftmax()
    log_probs = lsm(input_tensor)
    probs = torch.exp(log_probs)
    p_log_p = log_probs * probs
    entropy = -p_log_p.mean(dim=1)
    return entropy


CUSTOM_TEMPLATES = {
    "OxfordPets": "a photo of a {}, a type of pet.",
    "OxfordFlowers": "a photo of a {}, a type of flower.",
    "FGVCAircraft": "a photo of a {}, a type of aircraft.",
    "DescribableTextures": "{} texture.",
    "EuroSAT": "a centered satellite photo of {}.",
    "StanfordCars": "a photo of a {}.",
    "Food101": "a photo of {}, a type of food.",
    "SUN397": "a photo of a {}.",
    "Caltech101": "a photo of a {}.",
    "UCF101": "a photo of a person doing {}.",
    "ImageNet": "a photo of a {}.",
    "ImageNetSketch": "a photo of a {}.",
    "ImageNetV2": "a photo of a {}.",
    "ImageNetA": "a photo of a {}.",
    "ImageNetR": "a photo of a {}.",
    # semi-supervised templates
    "MYOxfordPets": "a photo of a {}, a type of pet.",
    "MYOxfordFlowers": "a photo of a {}, a type of flower.",
    "MYFGVCAircraft": "a photo of a {}, a type of aircraft.",
    "MYDescribableTextures": "{} texture.",
    "MYEuroSAT": "a centered satellite photo of {}.",
    "MYStanfordCars": "a photo of a {}.",
    "MYFood101": "a photo of {}, a type of food.",
    "MYSUN397": "a photo of a {}.",
    "MYCaltech101": "a photo of a {}.",
    "MYUCF101": "a photo of a person doing {}.",
    "MYImageNet": "a photo of a {}.",
}


def load_clip_to_cpu(cfg):
    backbone_name = cfg.MODEL.BACKBONE.NAME
    url = clip._MODELS[backbone_name]
    model_path = clip._download(url)
    try:
        # loading JIT archive
        model = torch.jit.load(model_path, map_location="cpu").eval()
        state_dict = None

    except RuntimeError:
        state_dict = torch.load(model_path, map_location="cpu")

    model = clip.build_model(state_dict or model.state_dict())

    return model


class TextEncoder(nn.Module):
    def __init__(self, clip_model):
        super().__init__()
        self.transformer = clip_model.transformer
        self.positional_embedding = clip_model.positional_embedding
        self.ln_final = clip_model.ln_final
        self.text_projection = clip_model.text_projection
        self.dtype = clip_model.dtype

    def forward(self, prompts, tokenized_prompts):
        x = prompts + self.positional_embedding.type(self.dtype)
        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD
        x = self.ln_final(x).type(self.dtype)

        # x.shape = [batch_size, n_ctx, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection

        return x


################## Synonyms ##############


class PromptLearner(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        n_cls = len(classnames)
        n_ctx = cfg.TRAINER.MYTrainer.N_CTX
        ctx_init = cfg.TRAINER.MYTrainer.CTX_INIT
        dtype = clip_model.dtype
        ctx_dim = clip_model.ln_final.weight.shape[0]
        clip_imsize = clip_model.visual.input_resolution
        cfg_imsize = cfg.INPUT.SIZE[0]
        assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})"

        if ctx_init:
            # use given words to initialize context vectors
            ctx_init = ctx_init.replace("_", " ")
            n_ctx = len(ctx_init.split(" "))
            prompt = clip.tokenize(ctx_init)
            with torch.no_grad():
                embedding = clip_model.token_embedding(prompt).type(dtype)
            ctx_vectors = embedding[0, 1 : 1 + n_ctx, :]
            prompt_prefix = ctx_init

        else:
            # random initialization
            if cfg.TRAINER.MYTrainer.CSC:
                print("Initializing class-specific contexts")
                ctx_vectors = torch.empty(n_cls, n_ctx, ctx_dim, dtype=dtype)
            else:
                print("Initializing a generic context")
                ctx_vectors = torch.empty(n_ctx, ctx_dim, dtype=dtype)
            nn.init.normal_(ctx_vectors, std=0.02)
            # nn.init.zeros_(ctx_vectors)
            prompt_prefix = " ".join(["X"] * n_ctx)

        print(f'Initial context: "{prompt_prefix}"')
        print(f"Number of context words (tokens): {n_ctx}")

        self.ctx = nn.Parameter(ctx_vectors)  # to be optimized

        classnames = [name.replace("_", " ") for name in classnames]
        name_lens = [len(_tokenizer.encode(name)) for name in classnames]
        prompts = [prompt_prefix + " " + name + "." for name in classnames]

        tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts])
        with torch.no_grad():
            embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)

        # These token vectors will be saved when in save_model(),
        # but they should be ignored in load_model() as we want to use
        # those computed using the current class names
        self.register_buffer("token_prefix", embedding[:, :1, :])  # SOS
        self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :])  # CLS, EOS

        self.n_cls = n_cls
        self.n_ctx = n_ctx
        self.tokenized_prompts = tokenized_prompts  # torch.Tensor
        self.name_lens = name_lens
        self.class_token_position = cfg.TRAINER.MYTrainer.CLASS_TOKEN_POSITION

    def forward(self):
        ctx = self.ctx
        if ctx.dim() == 2:
            ctx = ctx.unsqueeze(0).expand(self.n_cls, -1, -1)

        prefix = self.token_prefix
        suffix = self.token_suffix

        if self.class_token_position == "end":
            prompts = torch.cat(
                [
                    prefix,  # (n_cls, 1, dim)
                    ctx,     # (n_cls, n_ctx, dim)
                    suffix,  # (n_cls, *, dim)
                ],
                dim=1,
            )

        elif self.class_token_position == "middle":
            half_n_ctx = self.n_ctx // 2
            prompts = []
            for i in range(self.n_cls):
                name_len = self.name_lens[i]
                prefix_i = prefix[i : i + 1, :, :]
                class_i = suffix[i : i + 1, :name_len, :]
                suffix_i = suffix[i : i + 1, name_len:, :]
                ctx_i_half1 = ctx[i : i + 1, :half_n_ctx, :]
                ctx_i_half2 = ctx[i : i + 1, half_n_ctx:, :]
                prompt = torch.cat(
                    [
                        prefix_i,     # (1, 1, dim)
                        ctx_i_half1,  # (1, n_ctx//2, dim)
                        class_i,      # (1, name_len, dim)
                        ctx_i_half2,  # (1, n_ctx//2, dim)
                        suffix_i,     # (1, *, dim)
                    ],
                    dim=1,
                )
                prompts.append(prompt)
            prompts = torch.cat(prompts, dim=0)

        elif self.class_token_position == "front":
            prompts = []
            for i in range(self.n_cls):
                name_len = self.name_lens[i]
                prefix_i = prefix[i : i + 1, :, :]
                class_i = suffix[i : i + 1, :name_len, :]
                suffix_i = suffix[i : i + 1, name_len:, :]
                ctx_i = ctx[i : i + 1, :, :]
                prompt = torch.cat(
                    [
                        prefix_i,  # (1, 1, dim)
                        class_i,   # (1, name_len, dim)
                        ctx_i,     # (1, n_ctx, dim)
                        suffix_i,  # (1, *, dim)
                    ],
                    dim=1,
                )
                prompts.append(prompt)
            prompts = torch.cat(prompts, dim=0)

        else:
            raise ValueError

        return prompts, self.tokenized_prompts


class CustomCLIP(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
        self.prompt_learner = PromptLearner(cfg, classnames, clip_model)
        # self.tokenized_prompts = self.prompt_learner.tokenized_prompts
        self.image_encoder = clip_model.visual
        self.text_encoder = TextEncoder(clip_model)
        self.logit_scale = clip_model.logit_scale
        self.dtype = clip_model.dtype
        self.clip = clip_model
        self.classnames = classnames
        self.cfg = cfg

    def forward(self, image):
        image_features = self.image_encoder(image.type(self.dtype))

        # prompts = self.prompt_learner()
        # tokenized_prompts = self.tokenized_prompts
        prompts, tokenized_prompts = self.prompt_learner()
        text_features = self.text_encoder(prompts, tokenized_prompts)

        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)


        logit_scale = self.logit_scale.exp()
        logits = logit_scale * image_features @ text_features.t()

        return logits, image_features, text_features

    def zero_shot_forward(self, image, device):
        temp = CUSTOM_TEMPLATES[self.cfg.DATASET.NAME]
        prompts = [temp.format(c.replace("_", " ")) for c in self.classnames]
        prompts = torch.cat([clip.tokenize(p) for p in prompts])
        prompts = prompts.to(device)
        with torch.no_grad():
            text_features = self.clip.encode_text(prompts)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        image_features = self.clip.encode_image(image)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        logit_scale = self.clip.logit_scale.exp()
        logits = logit_scale * image_features @ text_features.t()
        return logits, image_features, text_features


    def description_forward(self, image, encodings, device, num_classes, encode_type='class'):
        image = image.to(device)
        image_encodings = self.clip.encode_image(image)
        logit_scale = self.clip.logit_scale.exp() 
        if encode_type == 'class':
            similarity = logit_scale * image_encodings @ encodings.T
        else:
            image_description_similarity = [None] * num_classes
            image_description_similarity_cumulative = [None] * num_classes
            for i, (k,v) in enumerate(encodings.items()):
                dot_product = logit_scale * image_encodings @ v.T
                image_description_similarity[i] = dot_product
                image_description_similarity_cumulative[i] = aggregate_similarity(image_description_similarity[i])
            similarity = torch.stack(image_description_similarity_cumulative, dim=1)
        return similarity
        
    
    def feature_forward(self, image, device):
        image = image.to(device)
        image_encodings = self.clip.encode_image(image)
        return image_encodings.detach().cpu()
    



@TRAINER_REGISTRY.register()
class MYTrainer(TrainerX):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.GCE_loss = GeneralizedCrossEntropy(q=0.5)
        self.gt_label_dict = self.get_gt_label(cfg)
    def check_cfg(self, cfg):
        assert cfg.TRAINER.MYTrainer.PREC in ["fp16", "fp32", "amp"]

    def build_model(self):
        cfg = self.cfg
        classnames = self.dm.dataset.classnames
        print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})")
        clip_model = load_clip_to_cpu(cfg)

        if cfg.TRAINER.MYTrainer.PREC == "fp32" or cfg.TRAINER.MYTrainer.PREC == "amp":
            # CLIP's default precision is fp16
            clip_model.float()

        print("Building custom CLIP")
        self.model = CustomCLIP(cfg, classnames, clip_model)

        print("Turning off gradients in both the image and the text encoder")
        for name, param in self.model.named_parameters():
            if "prompt_learner" not in name:
                param.requires_grad_(False)

        if cfg.MODEL.INIT_WEIGHTS:
            load_pretrained_weights(self.model.prompt_learner, cfg.MODEL.INIT_WEIGHTS)

        self.model.to(self.device)
        # NOTE: only give prompt_learner to the optimizer
        self.optim = build_optimizer(self.model.prompt_learner, cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("prompt_learner", self.model.prompt_learner, self.optim, self.sched)

        self.scaler = GradScaler() if cfg.TRAINER.MYTrainer.PREC == "amp" else None

        # Note that multi-gpu training could be slow because CLIP's size is
        # big, which slows down the copy operation in DataParallel
        device_count = torch.cuda.device_count()
        if device_count > 1:
            print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!")
            self.model = nn.DataParallel(self.model)

    def get_gt_label(self, cfg):
            dataset_map = {
                "MYImageNet":"imagenet",
                "MYCaltech101":'Caltech-101/caltech-101', #"caltech-101",
                "MYOxfordPets":'Pets',#"oxford_pets",
                "MYUCF101":"ucf101",
                "MYOxfordFlowers":'Flowers',#"oxford_flowers",
                "MYStanfordCars":'Cars',#"stanford_cars",
                "MYFGVCAircraft":'Aircraft/fgvc-aircraft-2013b/aircraft',#"fgvc_aircraft",
                "MYDescribableTextures":"DTD/dtd",
                "MYEuroSAT":'Eurosat',#"eurosat",
                "MYFood101":'Food101/food-101',#"food-101",
                "MYSUN397":'sun397',#"sun397"          
            }
            dataset_dir = dataset_map[self.cfg.DATASET.NAME]
            root = os.path.abspath(os.path.expanduser(cfg.DATASET.ROOT))
            dataset_dir = os.path.join(root, dataset_dir)
            gt_labels = os.path.join(dataset_dir, "{}_GTlabels.json".format(self.cfg.DATASET.NAME))
            if os.path.exists(gt_labels):
                with open(gt_labels, "rb") as f:
                    gt_label_dict = json.load(f)
                print("Loaded training GT labels from {}".format(gt_labels))
            else:
                print("Generating training GT labels to {}".format(gt_labels))
                gt_label_dict = {}
                for batch_idx, batch in enumerate(self.train_loader_x):
                    input, label, impath = self.parse_batch_test_with_impath(batch)
                    for l, ip in zip(label, impath):
                        ip = './data/' + ip.split('/data/')[-1]
                        gt_label_dict[ip] = l.item()
                with open(gt_labels, "w") as outfile:
                    json.dump(gt_label_dict, outfile)
            return gt_label_dict

    def forward_backward(self, batch, unlabel_batch=None):
        gt_label_list = []
        image, label, impath = self.parse_batch_test_with_impath(batch)

        if unlabel_batch != None:
            image_u, label_u, impath_u = self.parse_batch_test_with_impath(unlabel_batch)

        # image, label = self.parse_batch_train(batch)
        for ip in impath:
            ip = './data/' + ip.split('/data/')[-1]
            gt_label = self.gt_label_dict[ip]
            gt_label_list.append(gt_label)
        gt_label = torch.tensor(gt_label_list, dtype=label.dtype).to(label.device)
        prec = self.cfg.TRAINER.MYTrainer.PREC
        
        if prec == "amp":
            with autocast():
                output, image_features, text_features = self.model(image)
                # loss = F.cross_entropy(output, label, self.class_weights)
                loss = F.cross_entropy(output, label)
            self.optim.zero_grad()
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optim)
            self.scaler.update()
        elif unlabel_batch != None:
            prob = [self.mixup_prob['./data/'+ip.split('/data/')[-1]] for ip in impath]
            label_p = [self.pred_label['./data/'+ip.split('/data/')[-1]] for ip in impath]

            prob = torch.tensor(prob).cuda()
            label_p = torch.tensor(label_p).cuda()

            alpha = 4
            l = np.random.beta(alpha, alpha)        
            l = max(l, 1-l)

            if len(image) > len(image_u):
                catimage = [image_u for _ in range(int(len(image)/len(image_u)+1))]
                catlabel = [label_u for _ in range(int(len(image)/len(image_u)+1))]
                image_u = torch.cat(catimage, dim=0)[:len(image)]
                label_u = torch.cat(catlabel, dim=0)[:len(image)]




            output, _, _ = self.model(image * l + image_u[:len(image)] * (1-l))

            loss_a = F.cross_entropy(output, label, reduction='none') * prob + (1-prob) * F.cross_entropy(output, label_p, reduction='none')
            loss_b = F.cross_entropy(output, label_u[:len(image)], reduction='none')
            loss = l * loss_a + (1-l) * loss_b
            loss = loss.mean()
            self.model_backward_and_update(loss)

        else:
            output, image_features, text_features = self.model(image)
        
            if self.cfg.TRAINER.MYTrainer.USE_ROBUSTLOSS:
                loss = self.GCE_loss(output, label)
            else:
                loss = F.cross_entropy(output, label)
            self.model_backward_and_update(loss)

        loss_summary = {
            "loss": loss.item(),
            "acc": compute_accuracy(output, gt_label)[0].item(),
        }

        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()

        return loss_summary

    def parse_batch_train(self, batch):
        input = batch["img"]
        label = batch["label"]
        input = input.to(self.device)
        label = label.to(self.device)
        return input, label

    def parse_batch_train_with_impath(self, batch):
        input = batch["img"]
        label = batch["label"]
        impath = batch["impath"]

        input = input.to(self.device)
        label = label.to(self.device)

        return input, label, impath


    

    def load_model(self, directory, epoch=None):
        if not directory:
            print("Note that load_model() is skipped as no pretrained model is given")
            return

        names = self.get_model_names()

        # By default, the best model is loaded
        model_file = "model-best.pth.tar"

        if epoch is not None:
            model_file = "model.pth.tar-" + str(epoch)

        for name in names:
            model_path = osp.join(directory, name, model_file)

            if not osp.exists(model_path):
                raise FileNotFoundError('Model not found at "{}"'.format(model_path))

            checkpoint = load_checkpoint(model_path)
            state_dict = checkpoint["state_dict"]
            epoch = checkpoint["epoch"]

            # Ignore fixed token vectors
            if "token_prefix" in state_dict:
                del state_dict["token_prefix"]

            if "token_suffix" in state_dict:
                del state_dict["token_suffix"]

            print("Loading weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch))
            # set strict=False
            self._models[name].load_state_dict(state_dict, strict=False)

    def load_model_by_id(self, directory, model_id, epoch=None):
        if not directory:
            print(
                'Note that load_model() is skipped as no pretrained model is given'
            )
            return

        names = self.get_model_names()

        # By default, the best model is loaded
        model_file = 'model-best-{}.pth.tar'.format(model_id)

        if epoch is not None:
            model_file = 'model.pth.tar-' + str(epoch)

        for name in names:
            model_path = osp.join(directory, name, model_file)

            if not osp.exists(model_path):
                raise FileNotFoundError(
                    'Model not found at "{}"'.format(model_path)
                )

            checkpoint = load_checkpoint(model_path)
            state_dict = checkpoint['state_dict']
            epoch = checkpoint['epoch']

            # Ignore fixed token vectors
            if 'token_prefix' in state_dict:
                del state_dict['token_prefix']

            if 'token_suffix' in state_dict:
                del state_dict['token_suffix']

            print(
                'Loading weights to {} '
                'from "{}" (epoch = {})'.format(name, model_path, epoch)
            )
            # set strict=False
            self._models[name].load_state_dict(state_dict, strict=False)

    @torch.no_grad()
    def test(self, split=None, trainer_list=None, encode_type='none', clean_type='none'):
        """A generic testing pipeline."""

        self.set_model_mode("eval")
        self.evaluator.reset()

        if self.cfg.TRAINER.MYTrainer.FP_TYPE == 'sym':
            fp_desc = 'fp'+str(self.cfg.TRAINER.MYTrainer.FP)
        else:
            fp_desc = 'fp'+str(self.cfg.TRAINER.MYTrainer.FP)+'_asym'

        save_path = os.path.join(self.cfg.TEST.Analyze_Result_Path, self.cfg.DATASET.NAME, fp_desc, self.cfg.MODEL.BACKBONE.NAME,
        str(self.cfg.OPTIM.MAX_EPOCH)+'_'+str(self.cfg.SEED)+'_'+str(self.cfg.DATASET.NUM_SHOTS)+'_random_init_'+str(self.cfg.TRAINER.MYTrainer.CLASS_TOKEN_POSITION)+'_encode_type_'+str(encode_type)+'_clean_type_'+str(clean_type))

        
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        results_id = 0
        while os.path.exists(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id))):
            results_id += 1
        self.per_image_txt_writer = open(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id)), 'w')
        self.per_class_txt_writer = open(os.path.join(save_path, 'per_class_results_{}_{}.txt'.format(split, results_id)), 'w')

        if split is None:
            split = self.cfg.TEST.SPLIT

        if split == "val" and self.val_loader is not None:
            data_loader = self.val_loader
            print("Do evaluation on {} set".format(split))
        elif split=="novel":
            data_loader = self.test_novel_loader
            print("Do evaluation on test novel set")
        elif split=="base":
            data_loader = self.test_base_loader
            print("Do evaluation on test base set")
        elif split=="all":
            data_loader = self.test_loader
            print("Do evaluation on test set")
        else:
            data_loader = self.test_loader
            print("Do evaluation on test set")

        outputs_all = []
        label_all = []
        image_features_all = []
        text_features_all = []
        for batch_idx, batch in enumerate(data_loader):
            input, label = self.parse_batch_test(batch)
            if trainer_list is None or len(trainer_list)==1:
                output, image_features, text_features = self.model_inference(input)
                image_features_all.append(image_features)
                text_features_all.append(text_features)
            else:
                # ensemble
                outputs = [t.model_inference(input)[0] for t in trainer_list]
                output = sum(outputs) / len(outputs)
            self.evaluator.process(output, label, self.per_image_txt_writer, self.per_class_txt_writer)
            outputs_all.append(output)
            label_all.append(label)
        results = self.evaluator.evaluate()
        if split in ['all', 'train', 'test', 'novel', 'base']:
            if len(outputs_all) != 0:
                outputs_all = torch.cat(outputs_all, dim=0)
                label_all = torch.cat(label_all, dim=0)
                image_features_all = torch.cat(image_features_all, dim=0)
                text_features_all = text_features_all[0]
                torch.save(image_features_all, os.path.join(save_path, '{}_v_features.pt'.format(split)))
                torch.save(image_features_all, os.path.join(save_path, '{}_targets.pt'.format(split)))
                torch.save(outputs_all, os.path.join(save_path, '{}_logits.pt'.format(split)))
                torch.save(text_features_all, os.path.join(save_path, '{}_l_features.pt'.format(split)))


        self.per_image_txt_writer.close()
        self.per_class_txt_writer.close()


        # Confusion Matrix
        label_numpy = label_all.detach().cpu().numpy()
        logit_numpy = outputs_all.max(dim=1)[1].detach().cpu().numpy()
        cmatrix = confusion_matrix(label_numpy, logit_numpy)
        row_sum = cmatrix.sum(axis=1)
        cmatrix = cmatrix / row_sum[:,np.newaxis]
        ax = sns.heatmap(cmatrix)
        plt.savefig(os.path.join(save_path,'test_{}_{}.png'.format(split, results_id)))
        plt.close()


        for k, v in results.items():
            tag = "{}/{}".format(split, k)
            self.write_scalar(tag, v, self.epoch)

        return list(results.values())[0]

    @torch.no_grad()
    def zero_shot_analyze(self, trainer_list=None):
        """A generic predicting pipeline."""
        self.set_model_mode("eval")
        self.model.eval()
        self.evaluator.reset()

        data_loader = self.train_loader_mytrain
        outputs = []
        image_features_list = []
        img_paths = []
        from tqdm import tqdm
        for batch_idx, batch in tqdm(enumerate(data_loader)):
            input, label, impath = self.parse_batch_test_with_impath(batch)
            if trainer_list is None or len(trainer_list)==1:
                output, image_features, text_features = self.model.zero_shot_forward(input, self.device)
            else:
                # ensemble
                outputs = [t.model.zero_shot_forward(input, self.device)[0] for t in trainer_list]
                output = sum(outputs) / len(outputs)
            outputs.append(output)
            image_features_list.append(image_features)
            img_paths.append(impath)
        mytrain_outputs = torch.cat(outputs, dim=0)
        mytrain_img_paths = np.concatenate(img_paths, axis=0)
        image_features = torch.cat(image_features_list, axis=0)
        # text_features = torch.cat(text_features, axis=0)
        print('image_features', image_features.shape)
        print('text_features', text_features.shape)
        predict_label_dict, _ = select_top_k_similarity_per_class(mytrain_outputs, mytrain_img_paths, -1, image_features, True)
        save_outputs(self.train_loader_x, self, predict_label_dict, self.cfg.DATASET.NAME, text_features, backbone_name=self.cfg.MODEL.BACKBONE.NAME)
        caculate_noise_rate_analyze(predict_label_dict, train_loader=self.train_loader_x, trainer=self)
        return predict_label_dict


    def generate_noise_label(self, file_path, model, encode_type='none'):
        model_path = os.path.join(file_path, model)
        info_path = os.path.join(model_path, '{}.json'.format(self.cfg.DATASET.NAME))
        info = json.load(open(info_path))
        items = []
        for c in info:
            for img_path in info[c]:
                item = info[c][img_path]
                items.append([img_path, int(item[3])]) # 路径 序号
        sorted(items, key=(lambda x:x[1]))
        mytrain_img_paths = np.array(items)[:,0]

        if self.cfg.TRAINER.MYTrainer.FP_TYPE == 'sym':
            noise_label = gen_noise_label_sym(  img_paths=mytrain_img_paths,
                                                K=self.cfg.DATASET.NUM_SHOTS,
                                                random_seed=self.cfg.SEED, 
                                                gt_label_dict=self.gt_label_dict,
                                                num_fp=self.cfg.TRAINER.MYTrainer.FP,
                                                data_name = self.cfg.DATASET.NAME)    
        else:
            noise_label = gen_noise_label_asym( img_paths=mytrain_img_paths,
                                                K=self.cfg.DATASET.NUM_SHOTS,
                                                random_seed=self.cfg.SEED, 
                                                gt_label_dict=self.gt_label_dict,
                                                num_fp=self.cfg.TRAINER.MYTrainer.FP,
                                                data_name = self.cfg.DATASET.NAME)   
            
        if encode_type == 'clean':
            noise_label = self.get_clean_label(noise_label)
        
        elif encode_type == 'cleanse':
            noise_label = self.get_cleanse_label(noise_label)

        return noise_label


    def load_from_exist_file(self, file_path, model_names):
        logits = None
        for model in model_names:
            model_path = os.path.join(file_path, model)
            logist_path = os.path.join(model_path, '{}_logits.pt'.format(self.cfg.DATASET.NAME))
            if logits is None:
                logits = torch.load(logist_path)
            else:
                logits += torch.load(logist_path)

            info_path = os.path.join(model_path, '{}.json'.format(self.cfg.DATASET.NAME))
            info = json.load(open(info_path))
            items = []
            for c in info:
                for img_path in info[c]:
                    item = info[c][img_path]
                    items.append([img_path, int(item[3])]) 
            sorted(items, key=(lambda x:x[1]))
            mytrain_img_paths = np.array(items)[:,0]

        logits /= len(model_names)
        predict_label_dict = None
        return predict_label_dict 

    @torch.no_grad()
    def zero_shot_predict(self, trainer_list=None):
        """A generic predicting pipeline."""
        self.set_model_mode("eval")
        self.model.eval()
        self.evaluator.reset()

        save_path = os.path.join(self.cfg.TEST.Analyze_Result_Path, self.cfg.DATASET.NAME,
        str(self.cfg.OPTIM.MAX_EPOCH)+'_'+str(self.cfg.SEED)+'_'+str(self.cfg.DATASET.NUM_SHOTS))
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        data_loader = self.train_loader_mytrain

        outputs = []
        img_paths = []


        for batch_idx, batch in tqdm(enumerate(data_loader)):
            input, label, impath = self.parse_batch_test_with_impath(batch)
            if trainer_list is None or len(trainer_list)==1:
                output, image_features, text_features = self.model.zero_shot_forward(input, self.device)
            else:
                # ensemble
                outputs = [t.model.zero_shot_forward(input, self.device)[0] for t in trainer_list]
                output = sum(outputs) / len(outputs)
            outputs.append(output)
            img_paths.append(impath)


        outputs = torch.cat(outputs, dim=0)
        img_paths = np.concatenate(img_paths, axis=0)


        if self.cfg.DATASET.CLASS_EQULE is True:
            if self.cfg.DATASET.CONF_THRESHOLD > 0:
                predict_label_dict_1, predict_conf_dict_1 = select_top_k_similarity_per_class(outputs, img_paths, K=self.cfg.DATASET.NUM_SHOTS)
                predict_label_dict_2, predict_conf_dict_2 = select_top_by_value(outputs, img_paths, conf_threshold=self.cfg.DATASET.CONF_THRESHOLD)

                print(len(predict_label_dict_1), 'predict_label_dict_1')
                print(len(predict_label_dict_2), 'predict_label_dict_2')

                predict_label_dict = dict(predict_label_dict_1, **predict_label_dict_2)
                predict_conf_dict = dict(predict_conf_dict_1, **predict_conf_dict_2)
                caculate_noise_rate(predict_label_dict, train_loader=self.train_loader_x, trainer=self)
                print('select {} samples'.format(len(predict_label_dict)))

            else:
                print("K {} shots".format(self.cfg.DATASET.NUM_SHOTS))
                predict_label_dict, predict_conf_dict = select_top_k_similarity_per_class(outputs, img_paths, K=self.cfg.DATASET.NUM_SHOTS)
                caculate_noise_rate(predict_label_dict,  train_loader=self.train_loader_x, trainer=self)
                print('select {} samples'.format(len(predict_label_dict)))

        else:
            print("K", self.cfg.DATASET.NUM_SHOTS*text_features.shape[0])
            predict_label_dict, predict_conf_dict = select_top_k_similarity(outputs, img_paths, K=self.cfg.DATASET.NUM_SHOTS*text_features.shape[0])
            caculate_noise_rate(predict_label_dict, train_loader=self.train_loader_x, trainer=self)
            print('select {} samples'.format(len(predict_label_dict)))
        return predict_label_dict, predict_conf_dict

    @torch.no_grad()
    def zero_shot_test(self, split=None, trainer_list=None):
        """A generic predicting pipeline."""

        self.set_model_mode("eval")
        self.evaluator.reset()

        save_path = os.path.join(self.cfg.TEST.Analyze_Result_Path, self.cfg.DATASET.NAME,
        str(self.cfg.OPTIM.MAX_EPOCH)+'_'+str(self.cfg.SEED)+'_'+str(self.cfg.DATASET.NUM_SHOTS))
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        results_id = 0
        while os.path.exists(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id))):
            results_id += 1
        self.per_image_txt_writer = open(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id)), 'w')
        self.per_class_txt_writer = open(os.path.join(save_path, 'per_class_results_{}_{}.txt'.format(split, results_id)), 'w')

        if split is None:
            split = self.cfg.TEST.SPLIT

        save_path = os.path.join(self.cfg.TEST.Analyze_Result_Path, self.cfg.DATASET.NAME,
        str(self.cfg.OPTIM.MAX_EPOCH)+'_'+str(self.cfg.SEED)+'_'+str(self.cfg.DATASET.NUM_SHOTS))
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        if split == "val" and self.val_loader is not None:
            data_loader = self.val_loader
            print("Do evaluation on {} set".format(split))
        elif split=="novel":
            data_loader = self.test_novel_loader
            print("Do evaluation on test novel set")
        elif split=="base":
            data_loader = self.test_base_loader
            print("Do evaluation on test base set")
        elif split=="all":
            data_loader = self.test_loader
            print("Do evaluation on test set")
        elif split=="train":
            data_loader = self.train_loader_x
            print("Do evaluation on train set")
        else:
            data_loader = self.test_loader
            print("Do evaluation on test set")

        for batch_idx, batch in enumerate(data_loader):
            input, label, impath = self.parse_batch_test_with_impath(batch)
            if trainer_list is None or len(trainer_list)==1:
                output, image_features, text_features = self.model.zero_shot_forward(input, self.device)
            else:
                # ensemble
                outputs = [t.model.zero_shot_forward(input, self.device)[0] for t in trainer_list]
                output = sum(outputs) / len(outputs)
            self.evaluator.process(output, label, self.per_image_txt_writer, self.per_class_txt_writer)
        results = self.evaluator.evaluate()

        for k, v in results.items():
            tag = "{}/{}".format(split, k)
            self.write_scalar(tag, v, self.epoch)

        self.per_image_txt_writer.close()
        self.per_class_txt_writer.close()

        return list(results.values())[0]



    def build_data_loader(self):
        """Create essential data-related attributes.

        A re-implementation of this method must create the
        same attributes (except self.dm).
        """
        _, preprocess = clip.load(self.cfg.MODEL.BACKBONE.NAME)
        dm = MYDataManager(self.cfg, custom_tfm_test=preprocess)
        # _, preprocess = clip.load(self.cfg.MODEL.BACKBONE.NAME)
        # dm = MYDataManager(self.cfg, custom_tfm_test=preprocess)

        self.train_loader_x = dm.train_loader_x
        self.train_loader_u = dm.train_loader_u  # optional, can be None
        self.val_loader = dm.val_loader  # optional, can be None
        self.test_loader = dm.test_loader
        self.train_loader_mytrain = dm.train_loader_mytrain
        self.num_classes = dm.num_classes
        self.num_source_domains = dm.num_source_domains
        self.lab2cname = dm.lab2cname  # dict {label: classname}

        if self.cfg.DATALOADER.OPEN_SETTING:
            self.test_novel_loader = dm.test_novel_loader
            self.test_base_loader = dm.test_base_loader


        self.dm = dm

    def mytrain_with_id(self, model_id, train_label, encode_type = 'none', clean_type = 'none'):
        self.mytrain(self.start_epoch, self.max_epoch, train_label, model_id, encode_type=encode_type, clean_type = clean_type)
        
    def mytrain(self, start_epoch, max_epoch, label, model_id, encode_type, clean_type):
        """Generic training loops."""
        self.start_epoch = start_epoch
        self.max_epoch = max_epoch
        self.label = label
        
        self.before_train()
        self.dataset_update(self.label)
        self.num_batches = len(self.train_loader_mytrain)


        self.per_class_heatmap = np.zeros((self.num_classes, self.max_epoch))
        self.dataset = self.cfg.DATASET.NAME.split('MY')[-1]
        os.makedirs(f'./imgs/{self.dataset}/barplots/', exist_ok=True)
        os.makedirs(f'./imgs/{self.dataset}/histograms/', exist_ok=True)
        os.makedirs(f'./imgs/{self.dataset}/heatmaps/', exist_ok=True)
        self.dataset_update(self.label)

        self.mixup=False
        for self.epoch in range(self.start_epoch, self.max_epoch):
            if clean_type == 'none' or clean_type == 'gce':
                pass
            elif clean_type == 'dmix':
                if self.epoch <= 10:
                    pass
                else:
                    self.dataset_update(self.label)
                    label, noisy, prob, pred_label = self.gmm_label()
                    self.dataset_update(label)

                    self.dm.update_mydataloader(self.label)
                    self.dm.update_mydataloader(noisy)
                    self.train_loader_mytrain_noisy = self.dm.train_loader_mytrain
                    self.mixup=True
                    self.mixup_prob = prob
                    self.pred_label = pred_label

            elif clean_type == 'unicon':
                if self.epoch <= 10:
                    pass
                else:
                    self.dataset_update(self.label)
                    label, noisy, prob, pred_label = self.jsd_label()
                    self.dataset_update(label)

                    self.dm.update_mydataloader(self.label)
                    self.dm.update_mydataloader(noisy)
                    self.train_loader_mytrain_noisy = self.dm.train_loader_mytrain
                    self.mixup=True
                    self.mixup_prob = prob
                    self.pred_label = pred_label

                self.dataset_update(label)
            elif clean_type == 'pond':
                if encode_type not in ['class', 'syn', 'description', 'union']:
                    print(f"Wrong encode... Mismatch between PoND and {encode_type}...")
                    exit()
                self.pond_label()
                self.dataset_update(self.trainset)
            # self.measure_per_class_acc()
            self.before_epoch()
            self.run_epoch_with_mytrain()
            self.after_epoch(model_id, encode_type, clean_type)
            
        self.after_train(model_id, encode_type, clean_type)
        # print(self.mean_gap_sel)
        # np.savetxt(f'./ev_check/{self.dataset}.txt',self.mean_gap_sel.int().numpy(), delimiter=',', newline='\n')


    def measure_per_class_acc(self):
        self.set_model_mode("eval")

        path_all = []
        pred_all = []
        gtlabel_all = []
        for batch_idx, batch in enumerate(self.train_loader_mytrain):
            image, label, impath = self.parse_batch_train_with_impath(batch)
            for ip in impath:
                ip = './data/' + ip.split('/data/')[-1]
                path_all.append(ip)
            output, _, _ = self.model(image)
            output = output.max(dim=1)[1].tolist()
            pred_all += output
        
        for ip in path_all:
            gtlabel_all += [self.gt_label_dict[ip]]
        
        pred_all = torch.tensor(pred_all).reshape(-1)
        gtlabel_all = torch.tensor(gtlabel_all).reshape(-1)

        perclass_accuracy = torch.zeros(self.num_classes)
        for cidx in range(self.num_classes):
            pos = torch.where(gtlabel_all == cidx)[0]
            corr = torch.where(pred_all[pos] == cidx)[0]
            perclass_accuracy[cidx] = len(corr) / len(pos) if len(pos) != 0 else np.inf
        
        perclass_accuracy = np.array(perclass_accuracy)
        self.per_class_heatmap[:, self.epoch] = perclass_accuracy 
        
        argsort = np.argsort(self.per_class_heatmap[:,self.epoch])
        sns.heatmap(self.per_class_heatmap[argsort], cmap="crest")
        plt.savefig(f'./imgs/{self.dataset}/heatmaps/per_class_heatmap.png')
        plt.clf()
        if (self.epoch + 1) % 5 == 0:
            sns.barplot(x=np.arange(self.num_classes), y= perclass_accuracy[argsort], color='k', alpha=0.5)
            plt.savefig(f'./imgs/{self.dataset}/barplots/{self.epoch}_accuracy.png')
            plt.clf()
        

    def print_statistics_cleanse(self):
        torch.set_printoptions(threshold=10_000)
        for encode_type in self.encode_types:
            iternum = 0
            for idx, row in self.label_df.iterrows():
                path = row.name
                if path in self.label.keys():
                    if iternum == 0:
                        logits = torch.zeros(len(self.label), len(row[encode_type]))
                        gtlabels = torch.zeros(len(self.label)).long()
                        labels = torch.zeros(len(self.label)).long()
                    logits[iternum, :] = torch.tensor(row[encode_type])
                    gtlabels[iternum] = row['gtlabel']
                    labels[iternum] = row['label']

                    iternum += 1
                    
            for cidx in range(self.num_classes):
                pos = torch.where(cidx == labels)

                loss = F.cross_entropy(logits[pos], labels[pos], reduction='none')
                pos_clean = torch.where(gtlabels[pos] == labels[pos])
                neg_clean = torch.where(gtlabels[pos] != labels[pos])
                loss = loss / max(loss)
                

                loss_all = np.array(loss).reshape(-1,1)
                gmm = GaussianMixture(n_components=2, max_iter = 100, tol=1e-10, reg_covar=1e-5)
                gmm.fit(loss_all)
                prob = gmm.predict_proba(loss_all)
                prob = prob[:, gmm.means_.argmin()]

                regard_clean = np.where(prob>0.5)
                regard_noise = np.where(prob<=0.5)


                print(f'{encode_type}\t{cidx:3d}\t{float(torch.mean(loss[pos_clean])):.4f}\t{float(torch.mean(loss[neg_clean])):.4f}\t{float(np.abs(gmm.means_[0]-gmm.means_[1])):.4f}')
                
                plt.hist(loss[pos_clean], bins=10, range=(0,1), alpha=0.5, density=True, label='clean')
                plt.hist(loss[neg_clean], bins=10, range=(0,1), alpha=0.5, density=True, label='noise')
                plt.legend()
                plt.savefig(f'./imgs/{self.dataset}/histograms/{encode_type}_{cidx:3d}.png')
                plt.close()

                plt.hist(loss[regard_clean], bins=10, range=(0,1), alpha=0.5, density=True, label='clean')
                plt.hist(loss[regard_noise], bins=10, range=(0,1), alpha=0.5, density=True, label='noise')
                plt.legend()
                plt.savefig(f'./imgs/{self.dataset}/histograms/regard_{encode_type}_{cidx:3d}.png')
                plt.close()

        exit()
                
    def pond_dict_update(self):
        self.set_model_mode("eval")
        outputs_all = []
        label_all = []
        image_features_all = []
        text_features_all = []


        self.dataset_update(self.label)
        dataset = self.train_loader_mytrain.dataset
        loader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=8, drop_last=False, pin_memory=True)

        outputs_all = []
        paths_all = []
        pseudo_dict = {}
        logits_dict = {}
        pseudo = []
        logits = []
        
        for batch_idx, batch in enumerate(loader):
            input, _, path = self.parse_batch_test_with_impath(batch)
            output, _, _ = self.model_inference(input)
            outputs_all += output.detach().cpu().tolist()
            paths_all += path

        for pidx in range(len(paths_all)):
            path = './data/' + paths_all[pidx].split('/data/')[-1]
            label = torch.tensor(outputs_all[pidx]).argmax()
            pseudo_dict[path] = int(label)
            logits_dict[path] = outputs_all[pidx]
        
        # Sorting
        for idx, row in self.label_df.iterrows():
            path = row.name
            pseudo += [pseudo_dict[path]]
            logits += [logits_dict[path]]

        self.label_df = self.label_df.assign(pseudo_label=pseudo, pseudo=logits)
        if 'pseudo' not in self.encode_types:
            self.encode_types += ['pseudo']
        
    def pond_label(self):
        
        self.pond_dict_update()

        paths = []
        mean_gap = torch.zeros(len(self.label))
        if self.epoch == 0:
            self.mean_gap_sel = torch.zeros((self.max_epoch, self.num_classes))
        selected_samples = [[] for _ in range(self.num_classes)]
        
        select_num = torch.zeros(len(self.label))
        pseudolabels = torch.zeros(len(self.encode_types), len(self.label)).long()
        losses = torch.zeros(len(self.encode_types), len(self.label)).long()
        
        for eidx, encode_type in enumerate(self.encode_types):
            iternum = 0
            for idx, row in self.label_df.iterrows():
                path = row.name
                paths += [path]
                if path in self.label.keys():
                    if iternum == 0:
                        logits = torch.zeros(len(self.label), len(row[encode_type]))
                        gtlabels = torch.zeros(len(self.label)).long()
                        labels = torch.zeros(len(self.label)).long()
                    logits[iternum, :] = torch.tensor(row[encode_type])
                    gtlabels[iternum] = row['gtlabel']
                    labels[iternum] = row['label']
                    pseudolabels[eidx, iternum] = torch.argmax(torch.tensor(row[encode_type]))
                    
                    iternum += 1

            losses[eidx,:]  = F.cross_entropy(logits, pseudolabels[eidx,:], reduction='none')



            for cidx in range(self.num_classes):
                pos = torch.where(cidx == labels)[0]

                loss = F.cross_entropy(logits[pos], labels[pos], reduction='none')
                pos_clean = torch.where(gtlabels[pos] == labels[pos])
                neg_clean = torch.where(gtlabels[pos] != labels[pos])

                if max(loss) != 0:
                    loss = loss / max(loss)

                loss_all = np.array(loss).reshape(-1,1)

                if np.any(np.isnan(loss_all)):
                    print("Nan error...")
                    exit()

                

                gmm = GaussianMixture(n_components=2, max_iter = 100, tol=1e-10, reg_covar=1e-5)
                gmm.fit(loss_all)
                prob = gmm.predict_proba(loss_all)
                prob = prob[:, gmm.means_.argmin()]
                prob_pos = np.where(prob>0.5)

                mean_gap_diff = float(np.abs(gmm.means_[0] - gmm.means_[1]))

                if mean_gap[cidx] <= mean_gap_diff:
                    mean_gap[cidx] = mean_gap_diff
                    self.mean_gap_sel[self.epoch, cidx] = eidx
                    selected_set = list(pos[prob_pos].reshape(-1))
                    selected_samples[cidx] = selected_set
                    select_num[cidx] = len(selected_set)
        print(self.mean_gap_sel[self.epoch,:])
        selected_sample_list = [j for l in selected_samples for j in l]
        
        self.clean = {}
        self.trainset = {}
        self.dirty = {}

        self.class_check = {}
        self.descr_check = {}
        self.synon_check = {}
        self.pseudo_check = {}

        for idx in range(len(self.label)):
            if idx in selected_sample_list:
                self.clean[paths[idx]] = int(self.label[paths[idx]])
            else:
                given_label = int(self.label[paths[idx]])
                best_encode_type = int(torch.argmin(losses[:,idx])) 
                pseudo_label = pseudolabels[best_encode_type, idx] # 가장 믿을수 있는 description을 쓰자...?
                self.dirty[paths[idx]] = int(pseudo_label)

                self.class_check[paths[idx]] = pseudolabels[0, idx]
                self.descr_check[paths[idx]] = pseudolabels[1, idx]
                self.synon_check[paths[idx]] = pseudolabels[2, idx]
                self.pseudo_check[paths[idx]] = pseudolabels[3, idx]
                
        # self.trainset = self.clean

        # self.dirty_histogram()
        self.dirty_cleanse = self.dirty_cleansing()
        self.trainset = {**self.clean, **self.dirty_cleanse}
        

        print('-'*40)
        self.print_label_quality(self.clean, 'clean', end_print=True)
        self.print_label_quality(self.dirty, 'dirty', end_print=True)
        self.print_label_quality(self.class_check, 'class', end_print=True)
        self.print_label_quality(self.descr_check, 'descr', end_print=True)
        self.print_label_quality(self.synon_check, 'synon', end_print=True)
        self.print_label_quality(self.pseudo_check, 'pseudo', end_print=True)
        self.print_label_quality(self.dirty_cleanse, 'dirty_cleanse', end_print=True)
        print('-'*40)


    def dirty_cleansing(self):
        paths = []
        pseudolabels = torch.zeros(len(self.dirty)).long()
        softmax_all = torch.zeros(len(self.dirty))
        maxsoftmax_all = torch.zeros(len(self.dirty))
        
        
        ret = {}
        encode_type = 'pseudo'
        iternum = 0
        for idx, row in self.label_df.iterrows():
            path = row.name
            if path in self.dirty.keys():
                paths += [path]
                if iternum == 0:
                    logits = torch.zeros(len(self.dirty), len(row[encode_type]))
                    gtlabels = torch.zeros(len(self.dirty)).long()
                    labels = torch.zeros(len(self.dirty)).long()
                logits[iternum, :] = torch.tensor(row[encode_type])
                gtlabels[iternum] = row['gtlabel']
                labels[iternum] = row['label']
                pseudolabels[iternum] = torch.argmax(torch.tensor(row[encode_type]))
                iternum += 1
        maxsoftmax_all = torch.max(torch.softmax(logits,dim=1),dim=1)[0]
        pos = torch.where(maxsoftmax_all > 0.95)[0]

        for idx in pos:
            ret[paths[idx]] = int(pseudolabels[idx])

        return ret
        
    def dirty_histogram(self):
        paths = []
        pseudolabels = torch.zeros(len(self.encode_types), len(self.dirty)).long()
        softmax_all = torch.zeros(len(self.encode_types), len(self.dirty))
        maxsoftmax_all = torch.zeros(len(self.encode_types), len(self.dirty))
        
        for eidx, encode_type in enumerate(self.encode_types):
            iternum = 0
            for idx, row in self.label_df.iterrows():
                path = row.name
                if path in self.dirty.keys():
                    paths += [path]
                    if iternum == 0:
                        logits = torch.zeros(len(self.dirty), len(row[encode_type]))
                        gtlabels = torch.zeros(len(self.dirty)).long()
                        labels = torch.zeros(len(self.dirty)).long()
                    logits[iternum, :] = torch.tensor(row[encode_type])
                    gtlabels[iternum] = row['gtlabel']
                    labels[iternum] = row['label']
                    pseudolabels[eidx, iternum] = torch.argmax(torch.tensor(row[encode_type]))
                    
                    iternum += 1

            softmax_all[eidx,:] = calc_entropy(logits)
            maxsoftmax_all[eidx,:] = torch.max(torch.softmax(logits,dim=1),dim=1)[0]
            pos_corr = torch.where(gtlabels == pseudolabels[eidx,:])[0]
            pos_inco = torch.where(gtlabels != pseudolabels[eidx,:])[0]

            maxval = float(max(softmax_all[eidx,:]))
            minval = float(min(softmax_all[eidx,:]))
            
            plt.hist(softmax_all[eidx, pos_corr], bins=10, range=(minval, maxval), alpha=0.5, density=True, label=f'clean / {encode_type}')
            plt.hist(softmax_all[eidx, pos_inco], bins=10, range=(minval, maxval), alpha=0.5, density=True, label=f'noisy / {encode_type}')
            plt.legend()
            plt.savefig(f'./imgs/{self.dataset}/histograms/entropy_{encode_type}_{self.epoch:03d}.png')
            plt.close()

            maxval = float(max(maxsoftmax_all[eidx,:]))
            minval = float(min(maxsoftmax_all[eidx,:]))
            
            plt.hist(maxsoftmax_all[eidx, pos_corr], bins=10, range=(minval, maxval), alpha=0.5, density=True, label=f'clean / {encode_type}')
            plt.hist(maxsoftmax_all[eidx, pos_inco], bins=10, range=(minval, maxval), alpha=0.5, density=True, label=f'noisy / {encode_type}')
            plt.legend()
            plt.savefig(f'./imgs/{self.dataset}/histograms/maxsoftmax_{encode_type}_{self.epoch:03d}.png')
            plt.close()
        
            

            
            # softmax_all = np.array(softmax_all).reshape(-1,1)
            # gmm = GaussianMixture(n_components=2, max_iter = 100, tol=1e-10, reg_covar=1e-5)
            # gmm.fit(softmax_all)
            # prob = gmm.predict_proba(loss_all)
            # prob = prob[:, gmm.means_.argmin()]
            # prob_pos = np.where(prob>0.5)
                

            


    # def pond_label(self):
    #     # self.print_statistics_cleanse()
    #     class_to_best_descr, softmax = self.get_best_descr()
    #     ret = self.label.copy()

    #     for pidx, path in enumerate(self.label.keys()):
    #         label = self.label[path]
    #         for eidx, descr in enumerate(self.encode_types):
    #             pred = self.label_df[f'{descr}_label'][path]
    #             use_pseudo=False
    #             # if class_to_best_descr[label] == -1:
    #             #     break
    #             if class_to_best_descr[pred] == eidx:
    #                 ret[path] = pred
    #                 use_pseudo=True
    #                 break
    #         if not use_pseudo:
    #             ret[path] = label
    #     return ret

    def get_best_descr(self):
        pred  = torch.zeros(self.num_classes, len(self.encode_types))
        softmax  = torch.zeros(self.num_classes, len(self.encode_types))
        given = torch.zeros(self.num_classes)
        for idx, row in self.label_df.iterrows():
            path = row.name
            if path in self.label.keys():
                given_label = self.label[path]
                given[given_label] += 1
                for encode_idx, encode_type in enumerate(self.encode_types):
                    logits = torch.tensor(row[encode_type])
                    sm = torch.softmax(logits, dim=0)
                    pseudo_label = int(torch.argmax(logits))
                    if given_label == pseudo_label:
                        pred[pseudo_label, encode_idx] += 1
                        softmax[pseudo_label, encode_idx] = sm[pseudo_label] #max(sm)
        softmax /= pred
        softmax = torch.nan_to_num(softmax, nan=0.0)
        for encode_idx in range(len(self.encode_types)):
            pred[:,encode_idx] /= given




        max_pred, ret = self.argmax_tie_break(pred, softmax)
        discard_class = torch.where(max_pred < 1.1 * max(max_pred))
        
        ret[discard_class] = -1
        return ret, softmax
    
    def argmax_tie_break(self, tens, softmax):
        ret_idx = torch.zeros(len(tens))
        ret_val = torch.zeros(len(tens))
        for idx in range(len(tens)):
            mval = torch.max(tens[idx])
            pos = torch.where(mval == tens[idx])[0]
            if len(pos) == 1:
                ival = int(pos)
            else:
                # ival = torch.randperm(len(pos))[0]
                ival = pos[int(torch.argmax(softmax[idx,pos]))]
            
            ret_idx[idx] = ival
            ret_val[idx] = mval

        return ret_val, ret_idx
            
    def jsd_label(self):
        self.set_model_mode("eval")
        output_all = []
        label_all = []
        path_all = []
        pseudo_all = []

        for batch_idx, batch in enumerate(self.train_loader_mytrain):
            image, label, impath = self.parse_batch_train_with_impath(batch)
            for ip in impath:
                ip = './data/' + ip.split('/data/')[-1]
                path_all.append(ip)

            output, _, _ = self.model(image)
            
            
            
            pseudo_all += list(output.detach().cpu().argmax(dim=1).squeeze())
            output_all.append(output.detach().cpu())
            label_all += list(label.detach().cpu())
            pseudo_all += list(output.detach().cpu().argmax(dim=1).squeeze())

        label_all = torch.tensor(label_all)
        output_all = torch.cat(output_all).detach().cpu().float()
        prob = Calculate_JSD(output_all, label_all, self.num_classes)
        threshold = torch.mean(prob)
        if threshold > 0.5:
            threshold = threshold - (threshold-torch.min(prob))/5
            # threshold = threshold - (threshold-torch.min(prob))
        dR = torch.where(prob < threshold)[0]
        filter_rate = len(dR) / len(prob)
        
        print(filter_rate)

        clean_set = []
        dirty_set = []

        
        for cidx in range(self.num_classes):
            pos = torch.where(label_all == cidx)[0]
            sort = torch.argsort(prob[pos], descending=False).squeeze()

            clean_set += pos[sort[:int(len(pos)* filter_rate)]]
            dirty_set += pos[sort[int(len(pos)* filter_rate):]]
        
        ret = {}
        noisy = {}
        prob_ret = {}
        pred_label = {}

        for p in clean_set:
            ret[path_all[int(p)]] = int(label_all[int(p)])
            prob_ret[path_all[int(p)]] = prob.reshape(-1)[int(p)]
            pred_label[path_all[int(p)]] = int(pseudo_all[int(p)])
        for p in dirty_set:
            noisy[path_all[int(p)]] = int(pseudo_all[int(p)])
            pred_label[path_all[int(p)]] = int(pseudo_all[int(p)])


        self.print_label_quality(ret, 'UNICON_clean', end_print=True)
        self.print_label_quality(noisy, 'UNICON_noisy', end_print=True)
        
        return ret, noisy, prob_ret, pred_label

    def gmm_label(self):
        self.set_model_mode("eval")
        loss_all = []
        label_all = []
        path_all = []
        pseudo_all = []
        maxsoftmax_all = []
        for batch_idx, batch in enumerate(self.train_loader_mytrain):
            image, label, impath = self.parse_batch_train_with_impath(batch)
            for ip in impath:
                ip = './data/' + ip.split('/data/')[-1]
                path_all.append(ip)

            output, _, _ = self.model(image)
            loss = F.cross_entropy(output, label, reduction='none')

            loss_all += list(loss.detach().cpu())
            label_all += list(label.detach().cpu())
            pseudo_all += list(output.detach().cpu().argmax(dim=1).squeeze())
            maxsoftmax_all += list(output.detach().cpu().float().softmax(dim=1).squeeze().max(dim=1)[0])
        loss_all = np.array(loss_all).reshape(-1,1)
        maxsoftmax_all = np.array(maxsoftmax_all)
        gmm = GaussianMixture(n_components=2, max_iter = 50, tol=1e-33, reg_covar=1e-5)
        gmm.fit(loss_all)
        prob = gmm.predict_proba(loss_all)
        prob = prob[:, gmm.means_.argmin()]

        pos = np.where(prob>0.5)[0]
        negpos = np.where(prob<=0.5)[0]

        ret = {}
        noisy = {}
        prob_ret = {}
        pred_label = {}

        for p in pos:
            ret[path_all[int(p)]] = int(label_all[int(p)])
            prob_ret[path_all[int(p)]] = prob.reshape(-1)[int(p)]
            pred_label[path_all[int(p)]] = int(pseudo_all[int(p)])
        for p in negpos:
            noisy[path_all[int(p)]] = int(pseudo_all[int(p)])
            # noisy[path_all[int(p)]] = int(self.gt_label_dict[path_all[int(p)]])
            pred_label[path_all[int(p)]] = int(pseudo_all[int(p)])

        print(len(ret))
        print(len(noisy))

        self.print_label_quality(ret, 'DMIX Clean', end_print=True)
        self.print_label_quality(noisy, 'DMIX Noisy', end_print=True)
        return ret, noisy, prob_ret, pred_label

    def dataset_update(self, label):
        self.dm.update_mydataloader(label)
        self.train_loader_mytrain = self.dm.train_loader_mytrain

    def run_epoch_with_mytrain(self):
        self.set_model_mode("train")
        losses = MetricMeter()
        batch_time = AverageMeter()
        data_time = AverageMeter()
        loader_iteration = iter(self.train_loader_mytrain)


        if self.mixup:
            unlabel_iteration = iter(self.train_loader_mytrain_noisy)

        end = time.time()
        # for self.batch_idx, batch in enumerate(self.train_loader_mytrain):
        for self.batch_idx in range(self.num_batches):
            try:
                batch = next(loader_iteration)
            except:
                loader_iteration = iter(self.train_loader_mytrain)
                batch = next(loader_iteration)

            if self.mixup:
                try:
                    batch_unlabel = next(unlabel_iteration)
                except:
                    unlabel_iteration = iter(self.train_loader_mytrain_noisy)
                    batch_unlabel = next(unlabel_iteration)
                
            else:
                batch_unlabel=None
                

            data_time.update(time.time() - end)
            loss_summary = self.forward_backward(batch, unlabel_batch=batch_unlabel)
            batch_time.update(time.time() - end)
            losses.update(loss_summary)

            if (
                self.batch_idx + 1
            ) % self.cfg.TRAIN.PRINT_FREQ == 0 or self.num_batches < self.cfg.TRAIN.PRINT_FREQ:
                nb_remain = 0
                nb_remain += self.num_batches - self.batch_idx - 1
                nb_remain += (
                    self.max_epoch - self.epoch - 1
                ) * self.num_batches
                eta_seconds = batch_time.avg * nb_remain
                eta = str(datetime.timedelta(seconds=int(eta_seconds)))
                print(
                    "epoch [{0}/{1}][{2}/{3}]\t"
                    "time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
                    "data {data_time.val:.3f} ({data_time.avg:.3f})\t"
                    "eta {eta}\t"
                    "{losses}\t"
                    "lr {lr:.6e}".format(
                        self.epoch + 1,
                        self.max_epoch,
                        self.batch_idx + 1,
                        self.num_batches,
                        batch_time=batch_time,
                        data_time=data_time,
                        eta=eta,
                        losses=losses,
                        lr=self.get_current_lr(),
                    )
                )

            n_iter = self.epoch * self.num_batches + self.batch_idx
            for name, meter in losses.meters.items():
                self.write_scalar("train/" + name, meter.avg, n_iter)
            self.write_scalar("train/lr", self.get_current_lr(), n_iter)

            end = time.time()

    def after_epoch(self, model_id, encode_type, clean_type):
        last_epoch = (self.epoch + 1) == self.max_epoch
        do_test = not self.cfg.TEST.NO_TEST
        meet_checkpoint_freq = (
            (self.epoch + 1) % self.cfg.TRAIN.CHECKPOINT_FREQ == 0
            if self.cfg.TRAIN.CHECKPOINT_FREQ > 0 else False
        )
        # if ((self.epoch + 1) % 5) == 0 and self.cfg.DATASET.NAME!="SSImageNet":
        #     curr_result = self.test(split="test")
        if do_test and self.cfg.TEST.FINAL_MODEL == "best_val":
            curr_result = self.test(split="val", encode_type=encode_type, clean_type=clean_type)
            is_best = curr_result > self.best_result
            if is_best:
                self.best_result = curr_result
                self.save_model(
                    self.epoch,
                    self.output_dir,
                    model_name="model-best-{}.pth.tar".format(model_id)
                )

        if meet_checkpoint_freq or last_epoch:
            self.save_model(
                    self.epoch,
                    self.output_dir,
                    model_name="model-best-{}.pth.tar".format(model_id)
                )

    def after_train(self, model_id, encode_type='none', clean_type='none'):
        print("Finished training")

        do_test = not self.cfg.TEST.NO_TEST
        if do_test:
            if self.cfg.TEST.FINAL_MODEL == "best_val":
                print("Deploy the model with the best val performance")
                self.load_model_by_id(self.output_dir, model_id)
            # self.test(split='novel')
            # self.test(split='base')
            # self.test(split='train')
            self.test(split='test', encode_type=encode_type, clean_type= clean_type)


        # Show elapsed time
        elapsed = round(time.time() - self.time_start)
        elapsed = str(datetime.timedelta(seconds=elapsed))
        print("Elapsed: {}".format(elapsed))

        # Close writer
        self.close_writer()
        
    def parse_batch_test(self, batch):
        input = batch["img"]
        label = batch["label"]

        input = input.to(self.device)
        label = label.to(self.device)

        return input, label

    def parse_batch_test_with_impath(self, batch):
        input = batch["img"]
        label = batch["label"]
        impath = batch["impath"]

        input = input.to(self.device)

        label = label.to(self.device)
        # impath = impath.to(self.device)

        return input, label, impath

    @torch.no_grad()
    def test_with_existing_logits(self, logits, split='test', encode_type='none', clean_type='none'):

        if self.cfg.TRAINER.MYTrainer.FP_TYPE == 'sym':
            fp_desc = 'fp'+str(self.cfg.TRAINER.MYTrainer.FP)
        else:
            fp_desc = 'fp'+str(self.cfg.TRAINER.MYTrainer.FP)+'_asym'

        self.set_model_mode("eval")
        self.evaluator.reset()


        
        save_path = os.path.join(self.cfg.TEST.Analyze_Result_Path, self.cfg.DATASET.NAME, fp_desc, self.cfg.MODEL.BACKBONE.NAME,
        str(self.cfg.OPTIM.MAX_EPOCH)+'_'+str(self.cfg.SEED)+'_'+str(self.cfg.DATASET.NUM_SHOTS)+'_random_init_'+str(self.cfg.TRAINER.MYTrainer.CLASS_TOKEN_POSITION)+'_encode_type_'+str(encode_type)+'_clean_type_'+str(clean_type))

        if not os.path.exists(save_path):
            os.makedirs(save_path)

        results_id = 0
        while os.path.exists(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id))):
            results_id += 1
        self.per_image_txt_writer = open(os.path.join(save_path, 'per_image_results_{}_{}.txt'.format(split, results_id)), 'w')
        self.per_class_txt_writer = open(os.path.join(save_path, 'per_class_results_{}_{}.txt'.format(split, results_id)), 'w')

        if split is None:
            split = self.cfg.TEST.SPLIT

        if split == "val" and self.val_loader is not None:
            data_loader = self.val_loader
            print("Do evaluation on {} set".format(split))
        elif split=="novel":
            data_loader = self.test_novel_loader
            print("Do evaluation on test novel set")
        elif split=="base":
            data_loader = self.test_base_loader
            print("Do evaluation on test base set")
        elif split=="all":
            data_loader = self.test_loader
            print("Do evaluation on test set")
        else:
            data_loader = self.test_loader
            print("Do evaluation on test set")

        label_all = []
        for batch_idx, batch in enumerate(data_loader):
            input, label = self.parse_batch_test(batch)
            label_all.append(label)
        label_all = torch.hstack(label_all)
        
        # Confusion Matrix
        label_numpy = label_all.detach().cpu().numpy()
        logit_numpy = logits.max(dim=1)[1].detach().cpu().numpy()
        cmatrix = confusion_matrix(label_numpy, logit_numpy)
        row_sum = cmatrix.sum(axis=1)
        cmatrix = cmatrix / row_sum[:,np.newaxis]
        ax = sns.heatmap(cmatrix)
        plt.savefig(os.path.join(save_path,'test_{}_{}.png'.format(split, results_id)))
        plt.close()


        self.evaluator.process(logits, label_all, self.per_image_txt_writer, self.per_class_txt_writer)
        results = self.evaluator.evaluate()
        for k, v in results.items():
            tag = "{}/{}".format(split, k)
            self.write_scalar(tag, v, self.epoch)

        return results


    def get_self_info(self, label, encode_type):
        self.dataset_update(label)
        dataset = self.train_loader_mytrain.dataset
        gtlabel = self.get_gt_label_dict(label)

        label_dataframe = pd.DataFrame.from_dict(label, orient='index', columns=['label'])
        gtlabel_dataframe = pd.DataFrame.from_dict(gtlabel, orient='index', columns=['gtlabel'])
        if encode_type.lower() in ['class', 'description', 'syn']:
            encode_dict, pseudo_dict = self.description_based_logits(dataset = dataset, encode_type = encode_type)
            
            pd_encode = pd.DataFrame.from_dict(encode_dict, orient='index', columns=[f'{encode_type.lower()}'])
            pd_pseudo = pd.DataFrame.from_dict(pseudo_dict, orient='index', columns=[f'{encode_type.lower()}_label'])

            self.encode_types = [f'{encode_type.lower()}']
            self.label_df = pd.concat([label_dataframe, gtlabel_dataframe, pd_encode, pd_pseudo], axis=1)

            # self.cleanse_label_selection()

        elif encode_type.lower() in ['union']:
            class_dict, class_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'class')
            descr_dict, descr_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'description')
            synon_dict, synon_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'syn')

            pd_class = pd.DataFrame.from_dict(class_dict, orient='index', columns=['class'])
            pd_descr = pd.DataFrame.from_dict(descr_dict, orient='index', columns=['description'])
            pd_synon = pd.DataFrame.from_dict(synon_dict, orient='index', columns=['synonym'])

            pd_class_pseudo = pd.DataFrame.from_dict(class_pseudo, orient='index', columns=['class_label'])
            pd_descr_pseudo = pd.DataFrame.from_dict(descr_pseudo, orient='index', columns=['description_label'])
            pd_synon_pseudo = pd.DataFrame.from_dict(synon_pseudo, orient='index', columns=['synonym_label'])

            self.encode_types = ['class', 'description', 'synonym']
            self.label_df = pd.concat([label_dataframe, gtlabel_dataframe, pd_class, pd_descr, pd_synon, pd_class_pseudo, pd_descr_pseudo, pd_synon_pseudo], axis=1)

            # predict_df = pd.concat([pd_class, pd_descr, pd_synon, label_dataframe, gtlabel_dataframe], axis=1)
            # self.cleanse_label_selection()

        else:
            self.label_df = label_dataframe
            
        if'class' in  encode_type:
            print('class')
            class_dict, class_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'class')
            pd_class = pd.DataFrame.from_dict(class_dict, orient='index', columns=['class'])
            ret = {}
            for key in class_pseudo.keys():
                if gtlabel[key] == label[key]:
                    ret[key] = gtlabel[key] 
                else:
                    nsoftmax = torch.max(torch.softmax(torch.tensor(class_dict[key]).reshape(1,-1), dim=1))
                    nlabel = torch.argmax(torch.softmax(torch.tensor(class_dict[key]).reshape(1,-1), dim=1))
                    # nsoftmax, nlabel= torch.max(torch.softmax(torch.tensor(class_dict[key]).reshape(1,-1), dim=1))
                    if nsoftmax  > 0.95:
                        ret[key] = int(nlabel)

            return ret
        
        if 'descr' in encode_type:
            print('descr')
            descr_dict, descr_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'description')
            pd_descr = pd.DataFrame.from_dict(descr_dict, orient='index', columns=['description'])
            
            ret = {}
            for key in descr_pseudo.keys():
                if gtlabel[key] == label[key]:
                    ret[key] = gtlabel[key] 
                else:
                    nsoftmax = torch.max(torch.softmax(torch.tensor(descr_dict[key]).reshape(1,-1), dim=1))
                    nlabel = torch.argmax(torch.softmax(torch.tensor(descr_dict[key]).reshape(1,-1), dim=1))
                    # nsoftmax, nlabel= torch.max(torch.softmax(torch.tensor(descr_dict[key]).reshape(1,-1), dim=1))
                    if nsoftmax  > 0.95:
                        ret[key] = int(nlabel)

            return ret
        
        if 'synon' in encode_type :
            print('synon')
            synon_dict, synon_pseudo = self.description_based_logits(dataset = dataset, encode_type = 'syn')
            pd_synon = pd.DataFrame.from_dict(synon_dict, orient='index', columns=['synonym'])
            
            ret = {}
            for key in synon_pseudo.keys():
                if gtlabel[key] == label[key]:
                    ret[key] = gtlabel[key] 
                else:
                    nsoftmax = torch.max(torch.softmax(torch.tensor(synon_dict[key]).reshape(1,-1), dim=1))
                    nlabel = torch.argmax(torch.softmax(torch.tensor(synon_dict[key]).reshape(1,-1), dim=1))
                    # nsoftmax, nlabel= torch.max(torch.softmax(torch.tensor(synon_dict[key]).reshape(1,-1), dim=1))
                    if nsoftmax  > 0.95:
                        ret[key] = int(nlabel)

            return ret
        else:
            return label
        
    def description_based_logits(self, dataset, encode_type):

        device_count = torch.cuda.device_count()
        if device_count > 1:
            model = self.model.module
            
        else:
            model = self.model
        encodings, label_to_class = compute_encodings(model.clip, self.device, self.cfg.DATASET.NAME, model.classnames, encode_type=encode_type)
        image_dataset = IndexDataset(dataset)
        dataset_size = len(image_dataset)
        image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=32, num_workers=8, drop_last=False, pin_memory=True)

        for bidx, batch in enumerate(tqdm(image_loader)):
            image, label, path, idx = batch
            batch_sim = model.description_forward(image, encodings, self.device, self.num_classes, encode_type)
            
            if bidx == 0:
                ret_logit = {} 
                ret_pseudo = {}
            for i, img_path in enumerate(path):
                ret_logit['./data/' + img_path.split('/data/')[-1]] = [batch_sim.detach().cpu().float()[i].tolist()]
                ret_pseudo['./data/' + img_path.split('/data/')[-1]] = int(torch.argmax(batch_sim.detach().cpu().float()[i]))
            
        return ret_logit, ret_pseudo


    def print_label_quality(self, label, option='none', end_print=False):
        pred = torch.zeros(self.num_classes)
        corr = torch.zeros(self.num_classes)
        num =  torch.zeros(self.num_classes)
        
        for key, value in label.items():
            given_label = value
            pred[given_label] += 1
            if given_label == self.gt_label_dict[key]:
                corr[given_label] += 1
            num[self.gt_label_dict[key]] += 1

        if not end_print:
            print('-'*40)
            print(f'=> ({option}) Description prediction result...')
            for i in range(self.num_classes):
                print(f'* [Class {i}] GT_Num: {int(num[i])}, Correct/Pred: {int(corr[i])}/{int(pred[i])}')
        
            print(f'=> ({option}) Aggregated average {float(torch.sum(corr) / torch.sum(num)):.4f}')
            print('-'*40)
        else:
            print(f'=> ({option}) Aggregated average {float(torch.sum(corr) / torch.sum(num)):.4f}')
            

    
    def get_test_label(self):
        ret = {}
        for batch_idx, batch in enumerate(self.test_loader):
            _, label, impath = self.parse_batch_test_with_impath(batch)
            for idx in range(len(label)):
                path = './data/' + impath[idx].split('/data/')[-1]
                ret[path] = label[idx].detach().cpu()
        return ret

    def get_gt_label_dict(self, label):
        ret = {}
        for path in  label.keys():
            if path in self.gt_label_dict.keys():
                ret[path] = self.gt_label_dict[path]
        return ret
