import os, torch

import numpy as np

from PIL import Image

from torch.utils.data import Dataset

from torchvision import transforms

import cv2

import random

import torchio as tio

import slicerio

import nrrd

import monai

import pickle

import nibabel as nib

from scipy.ndimage import zoom

from monai.transforms import OneOf

import einops

from funcs import *

from torchvision.transforms import InterpolationMode

#from .utils.transforms import ResizeLongestSide

class MRI_dataset(Dataset):

    def __init__(self,args, img_folder, mask_folder, img_list,phase='train',sample_num=50,channel_num=1,crop=False,crop_size=1024,targets=['femur','hip'],part_list=['all'],cls=1,if_prompt=True,prompt_type='point',region_type='largest_15',prompt_num=15,delete_empty_masks=False,if_attention_map=None):

        super(MRI_dataset, self).__init__()

        self.img_folder = img_folder

        self.mask_folder = mask_folder

        self.crop = crop

        self.crop_size = crop_size

        self.phase = phase

        self.channel_num=channel_num

        self.targets = targets

        self.segment_names_to_labels = []

        self.args = args

        self.cls = cls

        self.if_prompt = if_prompt

        self.region_type = region_type

        self.prompt_type = prompt_type

        self.prompt_num = prompt_num

        self.if_attention_map = if_attention_map

        for i,tag in enumerate(targets):

            self.segment_names_to_labels.append((tag,i))

        namefiles = open(img_list,'r')

        self.data_list = namefiles.read().split('\n')[:-1]

        if delete_empty_masks=='delete' or delete_empty_masks=='subsample':

            keep_idx = []

            for idx,data in enumerate(self.data_list):

                mask_path = data.split(' ')[1]

                if os.path.exists(os.path.join(self.mask_folder,mask_path)):

                    msk = Image.open(os.path.join(self.mask_folder,mask_path)).convert('L')

                else:

                    msk = Image.open(os.path.join(self.mask_folder.replace('2D-slices','2D-slices-generated'),mask_path)).convert('L')

                if 'all' in self.targets: # combine all targets as single target

                    mask_cls = np.array(np.array(msk,dtype=int)>0,dtype=int)

                else:

                    mask_cls = np.array(msk==self.cls,dtype=int)

                if part_list[0]=='all' and np.sum(mask_cls)>0:

                    keep_idx.append(idx) 

                elif np.sum(mask_cls)>0:

                    if_keep = False

                    for part in part_list:

                        if mask_path.find(part)>=0:

                            if_keep = True

                    if if_keep:

                        keep_idx.append(idx) 

            print('num with non-empty masks',len(keep_idx),'num with all masks',len(self.data_list))  

            if delete_empty_masks=='subsample':

                empty_idx = list(set(range(len(self.data_list)))-set(keep_idx))

                keep_empty_idx = random.sample(empty_idx, int(len(empty_idx)*0.1))

                keep_idx = empty_idx + keep_idx

            self.data_list = [self.data_list[i] for i in keep_idx] # keep the slices that contains target mask

        if phase == 'train':

            self.aug_img = [transforms.RandomEqualize(p=0.1),

                             transforms.ColorJitter(brightness=0.3, contrast=0.3,saturation=0.3,hue=0.3),

                             transforms.RandomAdjustSharpness(0.5, p=0.5),

                             ]

            self.transform_spatial = transforms.Compose([transforms.RandomResizedCrop(crop_size, scale=(0.8, 1.2)),

                     transforms.RandomRotation(45)])

            transform_img = [transforms.ToTensor()]

        else:

            transform_img = [

                         transforms.ToTensor(),

                             ]

        self.transform_img = transforms.Compose(transform_img)

    def __len__(self):

        return len(self.data_list)

    def __getitem__(self,index):

        # load image and the mask

        data = self.data_list[index]

        img_path = data.split(' ')[0]

        mask_path = data.split(' ')[1]

        slice_num = data.split(' ')[3] # total slice num for this object

        #print(img_path,mask_path)

        try:

            if os.path.exists(os.path.join(self.img_folder,img_path)):

                img = Image.open(os.path.join(self.img_folder,img_path)).convert('RGB')

            else:

                img = Image.open(os.path.join(self.img_folder.replace('2D-slices','2D-slices-generated'),img_path)).convert('RGB')

        except:

            # try to load image as numpy file

            img_arr = np.load(os.path.join(self.img_folder,img_path)) 

            img_arr = np.array((img_arr-img_arr.min())/(img_arr.max()-img_arr.min()+1e-8)*255,dtype=np.uint8)

            img_3c = np.tile(img_arr[:, :,None], [1, 1, 3])

            img = Image.fromarray(img_3c, 'RGB')

        if os.path.exists(os.path.join(self.mask_folder,mask_path)):

            msk = Image.open(os.path.join(self.mask_folder,mask_path)).convert('L')

        else:

            msk = Image.open(os.path.join(self.mask_folder.replace('2D-slices','2D-slices-generated'),mask_path)).convert('L')

        if self.if_attention_map:

            slice_id = int(img_path.split('-')[-1].split('.')[0])

            slice_fraction = int(slice_id/int(slice_num)*4)

            img_id = '/'.join(img_path.split('-')[:-1]) +'_'+str(slice_fraction) + '.npy'

            attention_map = torch.tensor(np.load(os.path.join(self.if_attention_map,img_id)))

        else:

            attention_map = torch.zeros((64,64))

        img = transforms.Resize((self.args.image_size,self.args.image_size))(img)

        msk = transforms.Resize((self.args.image_size,self.args.image_size),InterpolationMode.NEAREST)(msk)

        state = torch.get_rng_state()

        if self.crop:

            im_w, im_h = img.size

            diff_w = max(0,self.crop_size-im_w)

            diff_h = max(0,self.crop_size-im_h)

            padding = (diff_w//2, diff_h//2, diff_w-diff_w//2, diff_h-diff_h//2)

            img = transforms.functional.pad(img, padding, 0, 'constant')

            torch.set_rng_state(state)

            t,l,h,w=transforms.RandomCrop.get_params(img,(self.crop_size,self.crop_size))

            img = transforms.functional.crop(img, t, l, h,w) 

            msk = transforms.functional.pad(msk, padding, 0, 'constant')

            msk = transforms.functional.crop(msk, t, l, h,w)

        if self.phase =='train':

            # add random optimazition

            aug_img_fuc = transforms.RandomChoice(self.aug_img)

            img = aug_img_fuc(img)

        img = self.transform_img(img)

        if self.phase == 'train':

            # It will randomly choose one

            random_transform = OneOf([monai.transforms.RandGaussianNoise(prob=0.5, mean=0.0, std=0.1),\

                                      monai.transforms.RandKSpaceSpikeNoise(prob=0.5, intensity_range=None, channel_wise=True),\

                                      monai.transforms.RandBiasField(degree=3),\

                                      monai.transforms.RandGibbsNoise(prob=0.5, alpha=(0.0, 1.0))

                                     ],weights=[0.3,0.3,0.2,0.2])

            img = random_transform(img).as_tensor()

        else:

            if img.mean()<0.05:

                img = min_max_normalize(img)

                img = monai.transforms.AdjustContrast(gamma=0.8)(img)

        if 'all' in self.targets: # combine all targets as single target

            msk = np.array(np.array(msk,dtype=int)>0,dtype=int)

        else:

            msk = np.array(msk,dtype=int)

        mask_cls = np.array(msk==self.cls,dtype=int)

        if self.phase=='train' and (not self.if_attention_map==None):

            mask_cls = np.repeat(mask_cls[np.newaxis,:, :], 3, axis=0)

            both_targets = torch.cat((img.unsqueeze(0), torch.tensor(mask_cls).unsqueeze(0)),0)

            transformed_targets = self.transform_spatial(both_targets)

            img = transformed_targets[0]

            mask_cls = np.array(transformed_targets[1][0].detach(),dtype=int)

        img = (img-img.min())/(img.max()-img.min()+1e-8)

        img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)

        # generate mask and prompt

        if self.if_prompt:

            if self.prompt_type =='point':

                prompt,mask_now = get_first_prompt(mask_cls,region_type=self.region_type,prompt_num=self.prompt_num)

                pc = torch.as_tensor(prompt[:,:2], dtype=torch.float)

                pl = torch.as_tensor(prompt[:, -1], dtype=torch.float)

                msk = torch.unsqueeze(torch.tensor(mask_now,dtype=torch.long),0)

                return {'image':img,

                    'mask':msk,

                    'point_coords': pc,

                    'point_labels':pl,

                    'img_name':img_path,

                    'atten_map':attention_map,

            }

            elif self.prompt_type =='box':

                prompt,mask_now = get_top_boxes(mask_cls,region_type=self.region_type,prompt_num=self.prompt_num)

                box = torch.as_tensor(prompt, dtype=torch.float)

                msk = torch.unsqueeze(torch.tensor(mask_now,dtype=torch.long),0)

                return {'image':img,

                    'mask':msk,

                    'boxes':box,

                    'img_name':img_path,

                    'atten_map':attention_map,

            }

        else:

            msk = torch.unsqueeze(torch.tensor(mask_cls,dtype=torch.long),0)

            return {'image':img,

                'mask':msk,

                'img_name':img_path,

                'atten_map':attention_map,

        }

class MRI_dataset_multicls(Dataset):

    def __init__(self, args, img_folder, mask_folder, img_list, phase='train', sample_num=50, channel_num=1,

                 crop=False, crop_size=1024, targets=['combine_all'], part_list=['all'], if_prompt=True, 

                 prompt_type='point', if_spatial = True, region_type='largest_20', prompt_num=20, delete_empty_masks=False, 

                 label_mapping=None, reference_slice_num=0, if_attention_map=None,label_frequency_path=None):

        super(MRI_dataset_multicls, self).__init__()

        self.initialize_parameters(args, img_folder, mask_folder, img_list, phase, sample_num, channel_num,

                                   crop, crop_size, targets, part_list, if_prompt, prompt_type, if_spatial, region_type,

                                   prompt_num, delete_empty_masks, label_mapping, reference_slice_num, if_attention_map,label_frequency_path)

        self.load_label_mapping()

        self.prepare_data_list()

        self.filter_data_list()

        if phase == 'train':

            self.setup_transformations_train(crop_size)

        else:

            self.setup_transformations_other()

    def initialize_parameters(self, args, img_folder, mask_folder, img_list, phase, sample_num, channel_num,

                              crop, crop_size, targets, part_list, if_prompt, prompt_type, if_spatial, region_type,

                              prompt_num, delete_empty_masks, label_mapping, reference_slice_num, if_attention_map,label_frequency_path):

        self.args = args

        self.img_folder = img_folder

        self.mask_folder = mask_folder

        self.img_list = img_list

        self.phase = phase

        self.sample_num = sample_num

        self.channel_num = channel_num

        self.crop = crop

        self.crop_size = crop_size

        self.targets = targets

        self.part_list = part_list

        self.if_prompt = if_prompt

        self.prompt_type = prompt_type

        self.if_spatial = if_spatial

        self.region_type = region_type

        self.prompt_num = prompt_num

        self.delete_empty_masks = delete_empty_masks

        self.label_mapping = label_mapping

        self.reference_slice_num = reference_slice_num

        self.if_attention_map = if_attention_map

        self.label_dic = {}

        self.label_frequency_path = label_frequency_path

    def load_label_mapping(self):

        # Load the basic label mappings from a pickle file

        if self.label_mapping:

            with open(self.label_mapping, 'rb') as handle:

                self.segment_names_to_labels = pickle.load(handle)

            self.label_dic = {seg[1]: seg[0] for seg in self.segment_names_to_labels}

            self.label_name_list = [seg[0] for seg in self.segment_names_to_labels]

            print(self.label_dic)

        else:

            self.label_dic = {value: 'all' for value in range(1, 256)}

        # Load frequency data and remap classes if required

        if 'remap_frequency' in self.targets:

            self.load_and_remap_classes_based_on_frequency()

    def load_and_remap_classes_based_on_frequency(self):

        if self.label_frequency_path:

            with open(self.label_frequency_path, 'r') as file:

                all_label_frequencies = json.load(file)

            all_label_frequencies = all_label_frequencies['train']

            # Example to select the target region dynamically based on some condition or configuration

            target_region = self.part_list[0] 

            if target_region in all_label_frequencies:

                label_frequencies = all_label_frequencies[target_region]

                self.label_frequencies = label_frequencies

                #print(label_frequencies)

                self.remap_classes_based_on_frequency(label_frequencies)

            else:

                print(f"Warning: No frequency data found for the target region '{target_region}'. No remapping applied.")

    def remap_classes_based_on_frequency(self, label_frequencies):

        # Determine the frequency threshold for high vs. low frequency classes

        total = max(label_frequencies.values())

        high_freq_threshold = total * 0.5  # Adjust this threshold as needed

        # Initialize dictionaries to hold new class mappings

        high_freq_classes = {}

        low_freq_classes = {}

        # Assign classes to high or low frequency based on the threshold

        for label, freq in label_frequencies.items():

            if freq >= high_freq_threshold:

                high_freq_classes[label] = freq

            else:

                low_freq_classes[label] = freq

        # Update label dictionary based on the frequency classification

        #self.label_dic: {old_cls: old_name}

        new_label_dic = {}

        for cls, name in self.label_dic.items():

            if name in high_freq_classes:

                new_label_dic[cls] = name  # Retain original name for high frequency classes

            elif name in low_freq_classes:

                new_label_dic[cls] = 'combined_low_freq'  # Combine low frequency classes into one

        self.updated_label_dic = new_label_dic

        #new_label_dic: {old_cls: new_name}

        #print("Updated label dictionary with frequency remapping:", new_label_dic)

        #print('new_label_dic:',new_label_dic)

        # Sort high frequency keys by their frequency in descending order

        sorted_high_freq_labels = sorted(high_freq_classes.items(), key=lambda item: item[1], reverse=True)

        # Create a mapping for high frequency classes based on the sorted order

        original_to_new = {label: idx + 1 for idx, (label, _) in enumerate(sorted_high_freq_labels)}

        combined_low_freq_class_id = len(original_to_new) + 1

        # Ensure combined low frequency class is mapped correctly

        if 'combined_low_freq' in new_label_dic.values():

            for cls in low_freq_classes.keys():

                original_to_new[cls] = combined_low_freq_class_id

        # orignal_to_new {old_name:new_cls} 

        #print('original_to_new:',original_to_new)

        # Create additional dictionaries

        self.old_name_to_new_name = {self.label_dic[cls]: new_label for cls, new_label in new_label_dic.items()}

        self.old_cls_to_new_cls = {cls: original_to_new[self.label_dic[cls]] for cls in self.label_dic.keys() if self.label_dic[cls] in original_to_new}

        print('remapped label dic:',self.old_name_to_new_name)

        print('remapped cls dic:',self.old_cls_to_new_cls)

    def prepare_data_list(self):

        with open(self.img_list, 'r') as namefiles:

            self.data_list = namefiles.read().split('\n')[:-1]

        self.sp_symbol = ',' if ',' in self.data_list[0] else ' '

    def filter_data_list(self):

        keep_idx = []

        for idx, data in enumerate(self.data_list):

            img_path, mask_path = self.extract_paths(data)

            msk = Image.open(os.path.join(self.mask_folder, mask_path)).convert('L')

            mask_cls = self.determine_mask_class(msk)

            if self.should_keep(mask_cls, mask_path):

                keep_idx.append(idx)

                if self.reference_slice_num > 1:

                    self.add_reference_slice(img_path, mask_path, data)

        self.data_list = [self.data_list[i] for i in keep_idx]

        print('num with non-empty masks', len(keep_idx), 'num with all masks', len(self.data_list))

    def extract_paths(self, data):

        img_path = data.split(self.sp_symbol)[0]

        mask_path = data.split(self.sp_symbol)[1]

        return img_path.lstrip('/'), mask_path.lstrip('/')

    def determine_mask_class(self, msk):

        if 'combine_all' in self.targets:

            return np.array(msk, dtype=int) > 0

        elif self.targets[0] in self.label_name_list:

            return np.array(msk, dtype=int) == self.cls

        return np.array(msk, dtype=int)

    def should_keep(self, mask_cls, mask_path):

        if self.delete_empty_masks:

            has_mask = np.any(mask_cls > 0)

            if has_mask:

                if self.part_list[0] == 'all':

                    return True

                return any(mask_path.find(part) >= 0 for part in self.part_list)

            return False

        return True

    def add_reference_slice(self, img_path, mask_path, data):

        volume_name = ''.join(img_path.split('-')[:-1])  # get volume name

        slice_num = data.split(self.sp_symbol)[2]

        if volume_name not in self.reference_slices:

            self.reference_slices[volume_name] = []

        self.reference_slices[volume_name].append((img_path, mask_path, slice_num))

    def setup_transformations_train(self, crop_size):

        self.transform_img = transforms.Compose([

            transforms.ToTensor(),

            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        ])

        self.aug_img = transforms.RandomChoice([

            transforms.RandomEqualize(p=0.1),

            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3),

            transforms.RandomAdjustSharpness(0.5, p=0.5),

        ])

        if self.if_spatial:

                self.transform_spatial = transforms.Compose([transforms.RandomResizedCrop(self.crop_size, scale=(0.5, 1.5), interpolation=InterpolationMode.NEAREST),

                         transforms.RandomRotation(45, interpolation=InterpolationMode.NEAREST)])

    def setup_transformations_other(self):

        self.transform_img = transforms.Compose([

            transforms.ToTensor(),

            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        ])

    def __len__(self):

        return len(self.data_list)

    def __getitem__(self, index):

        # Load image and mask, handle missing files

        data = self.data_list[index]

        img, msk, img_path, mask_path, slice_num = self.load_image_and_mask(data)

        # Optional: Load attention map

        attention_map = self.load_attention_map(img_path, slice_num) if self.if_attention_map else torch.zeros((64, 64))

        # Handle reference slices if necessary

        if self.reference_slice_num > 1:

            img, msk = self.handle_reference_slices(img_path, mask_path, slice_num)

        # Apply transformations

        img, msk = self.apply_transformations(img, msk)

        # Generate and process masks and prompts

        output_dict = self.prepare_output(img, msk, img_path, mask_path,attention_map)

        return output_dict

    def load_image_and_mask(self, data):

        img_path, mask_path = self.extract_paths(data)

        slice_num = data.split(self.sp_symbol)[3]  # Extract total slice number for this object

        img_folder = self.img_folder

        msk_folder = self.mask_folder

        img = Image.open(os.path.join(img_folder, img_path)).convert('RGB')

        msk = Image.open(os.path.join(msk_folder, mask_path)).convert('L')

        # Resize images for processing

        img = transforms.Resize((self.args.image_size, self.args.image_size))(img)

        msk = transforms.Resize((self.args.image_size, self.args.image_size), InterpolationMode.NEAREST)(msk)

        return img, msk, img_path, mask_path, int(slice_num)

    def load_attention_map(self, img_path, slice_num):

        slice_id = int(img_path.split('-')[-1].split('.')[0])

        slice_fraction = int(slice_id / slice_num * 4)

        img_id = '/'.join(img_path.split('-')[:-1]) + '_' + str(slice_fraction) + '.npy'

        attention_map = torch.tensor(np.load(os.path.join(self.if_attention_map, img_id)))

        return attention_map

    def apply_crop(self, img, msk):

        im_w, im_h = img.size

        diff_w = max(0, self.crop_size - im_w)

        diff_h = max(0, self.crop_size - im_h)

        padding = (diff_w // 2, diff_h // 2, diff_w - diff_w // 2, diff_h - diff_h // 2)

        img = transforms.functional.pad(img, padding, 0, 'constant')

        msk = transforms.functional.pad(msk, padding, 0, 'constant')

        t, l, h, w = transforms.RandomCrop.get_params(img, (self.crop_size, self.crop_size))

        img = transforms.functional.crop(img, t, l, h, w)

        msk = transforms.functional.crop(msk, t, l, h, w)

        return img, msk

    def apply_transformations(self, img, msk):

        if self.crop:

            img, msk = self.apply_crop(img, msk)

        if self.phase == 'train':

            img = self.aug_img(img)

        img = self.transform_img(img)

        if self.phase =='train' and self.if_spatial:

            mask_cls = np.array(msk,dtype=int)

            mask_cls = np.repeat(mask_cls[np.newaxis,:, :], 3, axis=0)

            both_targets = torch.cat((img.unsqueeze(0), torch.tensor(mask_cls).unsqueeze(0)),0)

            transformed_targets = self.transform_spatial(both_targets)

            img = transformed_targets[0]

            mask_cls = np.array(transformed_targets[1][0].detach(),dtype=int)

            msk = torch.tensor(mask_cls)

        return img, msk

    def handle_reference_slices(self, img_path, mask_path, slice_num):

        volume_name = ''.join(img_path.split('-')[:-1])

        ref_slices, ref_msks = [], []

        reference_slices = self.reference_slices.get(volume_name, [])

        for ref_slice in reference_slices:

            ref_img_path, ref_msk_path, _ = ref_slice

            ref_img = Image.open(os.path.join(self.img_folder, ref_img_path)).convert('RGB')

            ref_img = transforms.Resize((self.args.image_size, self.args.image_size))(ref_img)

            ref_img = self.transform_img(ref_img)

            ref_img = torch.unsqueeze(ref_img, 0)

            ref_msk = Image.open(os.path.join(self.mask_folder, ref_msk_path)).convert('L')

            ref_msk = transforms.Resize((self.args.image_size, self.args.image_size), InterpolationMode.NEAREST)(ref_msk)

            ref_msk = torch.tensor(ref_msk, dtype=torch.long)

            ref_msks.append(torch.unsqueeze(ref_msk, 0))

        img = torch.cat(ref_slices, dim=0)

        msk = torch.cat(ref_msks, dim=0)

        return img, msk

    def remap_classes_sequentially(self, mask, label_frequencies):

        # Apply the mapping to the mask

        remapped_mask = mask.copy()

        for old_cls, new_cls in  self.old_cls_to_new_cls.items():

            remapped_mask[mask == old_cls] = new_cls

        return remapped_mask

    def prepare_output(self, img, msk, img_path, mask_path, attention_map):

        # Normalize the image

        img = (img - img.min()) / (img.max() - img.min() + 1e-8)

        img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)

        msk = np.array(msk, dtype=int)

        #print('ori_msk:',np.unique(msk))

        if self.label_frequency_path:

            msk = self.remap_classes_sequentially(msk,self.label_frequencies)  # Assuming msk is already using updated IDs

            #print('new_msk------------------------:',self.old_cls_to_new_cls)

            # Prepare one-hot encoding for the remapped classes

        unique_classes = np.unique(msk).tolist()

        if 0 in unique_classes:

            unique_classes.remove(0)

        if len(unique_classes) > 0:

            selected_dic = {k: self.label_dic[k] for k in unique_classes if k in self.label_dic}

        else:

            selected_dic = {}

        if self.targets[0] == 'random':

            mask_cls, selected_label, cls_one_hot = self.handle_random_target(msk, unique_classes, selected_dic)

        elif self.targets[0] in self.label_name_list:

            selected_label = self.targets[0]

            mask_cls = np.array(msk == self.cls, dtype=int)

            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)

            cls_one_hot[self.cls - 1] = 1

        else:

            selected_label = self.targets[0]

            mask_cls = msk

            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)

        # Handle prompts

        if self.if_prompt:

            prompt, mask_now, mask_cls = self.generate_prompt(mask_cls)

            ref_msk,_ = torch.max(mask_now>0,dim=0)

            return_dict = {'image': img, 'mask': mask_now, 'selected_label_name': selected_label,

                           'cls_one_hot': cls_one_hot, 'prompt': prompt, 'img_name': img_path,

                           'mask_ori': msk, 'mask_cls': mask_cls, 'all_label_dic': selected_dic,'ref_mask':ref_msk}

        else:

            if len(mask_cls.shape)==2:

                msk = torch.unsqueeze(torch.tensor(mask_cls,dtype=torch.long),0)

            elif len(mask_cls.shape)==4:

                msk = torch.squeeze(torch.tensor(mask_cls,dtype=torch.long))

            else:

                msk = torch.tensor(mask_cls,dtype=torch.long)

            ref_msk,_ = torch.max(msk>0,dim=0)

            #print('unique mask values:',msk.unique())

            return_dict = {'image': img, 'mask': msk, 'selected_label_name': selected_label,

                           'cls_one_hot': cls_one_hot, 'img_name': img_path, 'mask_ori': msk,'ref_mask':ref_msk}

        return return_dict

    def generate_prompt(self, mask_cls):

        if self.prompt_type == 'point':

            prompt, mask_now = get_first_prompt(mask_cls, region_type=self.region_type, prompt_num=self.prompt_num)

        elif self.prompt_type == 'box':

            prompt, mask_now = get_top_boxes(mask_cls, region_type=self.region_type, prompt_num=self.prompt_num)

        else:

            prompt = mask_now = None

        # Handling the shape of mask_now for return

        if mask_now is not None:

            if len(mask_now.shape) == 2:

                mask_now = torch.unsqueeze(torch.tensor(mask_now, dtype=torch.long), 0)

                mask_cls = torch.unsqueeze(torch.tensor(mask_cls, dtype=torch.long), 0)

            elif len(mask_now.shape) == 4:

                mask_now = torch.squeeze(torch.tensor(mask_now, dtype=torch.long))

            else:

                mask_now = torch.tensor(mask_now, dtype=torch.long)

                mask_cls = torch.tensor(mask_cls, dtype=torch.long)

        return prompt, mask_now, mask_cls

    def handle_random_target(self, msk, unique_classes, selected_dic):

        if len(unique_classes) > 0:

            random_selected_cls = random.choice(unique_classes)

            selected_label = selected_dic[random_selected_cls]

            mask_cls = np.array(msk == random_selected_cls, dtype=int)

            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)

            cls_one_hot[random_selected_cls - 1] = 1

        else:

            selected_label = None

            mask_cls = torch.zeros_like(msk)  # assuming msk is already a numpy array

            cls_one_hot = torch.zeros(len(self.label_dic), dtype=torch.long)

        return mask_cls, selected_label, cls_one_hot