import os
import random

import cv2
cv2.setNumThreads(0)
import numpy as np
import torch
import torch.nn.functional as F
from pycocotools import mask
from transformers import CLIPImageProcessor
from FCLM.train.train import preprocess, preprocess_multimodal
# from model.llava import conversation as conversation_lib
from FCLM.model.language_model.segment_anything.utils.transforms import ResizeLongestSide
from FCLM.datasets.grefer import G_REFER
from FCLM.datasets.refer import REFER
from FCLM.datasets.utils import ANSWER_LIST, SHORT_QUESTION_LIST
from PIL import Image
import copy
class ReferSegDataset(torch.utils.data.Dataset):
    pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    img_size = 1024
    ignore_label = 255

    def __init__(
        self,
        tokenizer=None,
        data_args=None,
        base_image_dir=None,
        samples_per_epoch=500 * 8 * 3 * 10,
        # samples_per_epoch=300, #overfit
        precision: str = "fp32",
        image_size: int = 1024,
        num_classes_per_sample: int = 3,
        exclude_val=False,
        # refer_seg_data="refclef||refcoco||refcoco+||refcocog",
        # refer_seg_data="refcoco||refcoco+||refcocog",
        refer_seg_data="refclef||grefcoco||refcoco||refcoco+||refcocog",
        # refer_seg_data="refclef",
    ):
        self.data_args=data_args
        self.exclude_val = exclude_val
        self.samples_per_epoch = samples_per_epoch
        self.num_classes_per_sample = num_classes_per_sample

        self.base_image_dir = base_image_dir
        self.image_size = image_size
        self.tokenizer = tokenizer
        self.precision = precision
        self.transform = ResizeLongestSide(image_size)
        # self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

        self.short_question_list = SHORT_QUESTION_LIST
        self.answer_list = ANSWER_LIST
        
        # DATA_DIR = os.path.join(base_image_dir, "refer_seg")
        DATA_DIR = base_image_dir
        self.refer_seg_ds_list = refer_seg_data.split(
            "||"
        )  # ['refclef', 'refcoco', 'refcoco+', 'refcocog']
        self.refer_seg_data = {}
        for ds in self.refer_seg_ds_list:
            if ds == "refcocog":
                splitBy = "umd"
            else:
                splitBy = "unc"
            if ds == "grefcoco":
                refer_api = G_REFER(DATA_DIR, ds, splitBy)
            else:
                refer_api = REFER(DATA_DIR, ds, splitBy)
            ref_ids_train = refer_api.getRefIds(split="train")
            images_ids_train = refer_api.getImgIds(ref_ids=ref_ids_train)
            refs_train = refer_api.loadRefs(ref_ids=ref_ids_train)

            refer_seg_ds = {}
            refer_seg_ds["images"] = []
            loaded_images = refer_api.loadImgs(image_ids=images_ids_train)
            for item in loaded_images:
                item = item.copy()
                if ds == "refclef":
                    item["file_name"] = os.path.join(
                        DATA_DIR, "images/saiapr_tc-12", item["file_name"]
                    )
                else:
                    item["file_name"] = os.path.join(
                        DATA_DIR, "images/coco_2014/train2014", item["file_name"]
                    )
                refer_seg_ds["images"].append(item)
            refer_seg_ds["annotations"] = refer_api.Anns  # anns_train

            # print(
            #     "dataset {} (refs {}) (train split) has {} images and {} annotations.".format(
            #         ds,
            #         splitBy,
            #         len(refer_seg_ds["images"]),
            #         len(refer_seg_ds["annotations"]),
            #     )
            # )
            print(
                "\033[1;32mdataset {} (refs {}) (train split) has {} images and {} annotations.\033[0m".format(
                    ds,
                    splitBy,
                    len(refer_seg_ds["images"]),
                    len(refer_seg_ds["annotations"]),
                )
            )
            img2refs = {}
            for ref in refs_train:
                image_id = ref["image_id"]
                img2refs[image_id] = img2refs.get(image_id, []) + [
                    ref,
                ]
            refer_seg_ds["img2refs"] = img2refs
            self.refer_seg_data[ds] = refer_seg_ds

    def __len__(self):
        return self.samples_per_epoch

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def __getitem__(self, idx):
        
        ds = random.randint(0, len(self.refer_seg_ds_list) - 1)
        # ds = 0
        ds = self.refer_seg_ds_list[ds]
        refer_seg_ds = self.refer_seg_data[ds]

        images = refer_seg_ds["images"]
        annotations = refer_seg_ds["annotations"]
        img2refs = refer_seg_ds["img2refs"]
        idx = random.randint(0, len(images) - 1)
        # idx = 1
        image_info = images[idx]
        image_path = image_info["file_name"]
        image_id = image_info["id"]
        refs = img2refs[image_id]
        if len(refs) == 0:
            return self.__getitem__(0)

        sents = []
        ann_ids = []
        for ref in refs:
            # import ipdb;ipdb.set_trace()
            for sent in ref["sentences"]:
                text = sent["sent"]
                sents.append(text)
                ann_ids.append(ref["ann_id"])
        if len(sents) >= self.num_classes_per_sample:
            sampled_inds = np.random.choice(
                list(range(len(sents))), size=self.num_classes_per_sample, replace=False
            )
            # sampled_inds = np.array([0])
        else:
            sampled_inds = list(range(len(sents)))
        # print('cur sampling:', len(sampled_inds))
        sampled_sents = np.vectorize(sents.__getitem__)(sampled_inds).tolist()
        # sampled_ann_ids = np.vectorize(ann_ids.__getitem__)(sampled_inds).tolist()
        sampled_ann_ids = [ann_ids[ind] for ind in sampled_inds]
        sampled_classes = sampled_sents
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        
        image = self.transform.apply_image(image)  # preprocess image for sam
        resize = image.shape[:2]

        questions = []
        answers = []
        for text in sampled_classes:
            text = text.strip()
            assert len(text.split("||")) == 1
            question_template = random.choice(self.short_question_list)
            questions.append(question_template.format(class_name=text.lower()))
            answers.append(random.choice(self.answer_list))

        
        sources = dict()
        sources['conversations'] = []
        for i in range(len(questions)):
            cur_sample=[]
            cur_sample.append(
                    {'from': 'human', 'value': questions[i]})
            cur_sample.append({'from': 'gpt', 'value': answers[i]})
            sources['conversations'].append(cur_sample)
        processor=self.data_args.image_processor
        qwen_img_process=Image.open(image_path).convert('RGB')
        # print('referdata size',qwen_img_process.size)
        data_dict_bs=dict()
        data_dict_bs['input_ids']=[]
        data_dict_bs['labels']=[]
        for idx in range(len(questions)):
            sources_each = preprocess_multimodal(
                copy.deepcopy([sources['conversations'][idx]]),
                self.data_args,
                0,if_qwen=True)
            
            text = sources_each[0][0]['value']
            inputs = processor(
                text=[text],
                images=qwen_img_process,
                videos=None,
                padding=True,
                return_tensors="pt",
            )
            tmp=inputs['input_ids']
            N=len(tmp[tmp==self.tokenizer('<|image_pad|>')['input_ids'][0]]) 
            data_dict = preprocess(
                sources_each,
                self.tokenizer,
                has_image=True
                )
            ####repeat <|image_pad|> * N
            old_data_dict=data_dict
            insert_position=torch.where(old_data_dict['input_ids']==self.tokenizer('<|image_pad|>')['input_ids'][0])[1]
            img_padding=torch.tensor([self.tokenizer('<|image_pad|>')['input_ids'][0]]*N).unsqueeze(0)
            target_padding=torch.tensor([-100]*N).unsqueeze(0)
            # (1,N) <|image_pad|>
            new_inputs=torch.cat((old_data_dict['input_ids'][:,:insert_position],img_padding,old_data_dict['input_ids'][:,insert_position+1:]),dim=1)       
            # [1, 435] 345 个 image_pad
            new_target=torch.cat((old_data_dict['labels'][:,:insert_position],target_padding,old_data_dict['labels'][:,insert_position+1:]),dim=1)
            # data_dict['input_ids']=new_inputs #[1, 281]
            # data_dict['labels']=new_target #[1, 281]
            data_dict_bs['input_ids'].append(new_inputs.squeeze(0))
            data_dict_bs['labels'].append(new_target.squeeze(0))
            #############################
        flag = False
        masks = []
        for ann_id in sampled_ann_ids:
            if isinstance(ann_id, list):
                flag = True
                if -1 in ann_id:
                    assert len(ann_id) == 1
                    m = np.zeros((image_info["height"], image_info["width"])).astype(
                        np.uint8
                    )
                else:
                    m_final = np.zeros(
                        (image_info["height"], image_info["width"])
                    ).astype(np.uint8)
                    for ann_id_i in ann_id:
                        ann = annotations[ann_id_i]

                        if len(ann["segmentation"]) == 0:
                            m = np.zeros(
                                (image_info["height"], image_info["width"])
                            ).astype(np.uint8)
                        else:
                            try:
                                if type(ann["segmentation"][0]) == list:  # polygon
                                    rle = mask.frPyObjects(
                                        ann["segmentation"],
                                        image_info["height"],
                                        image_info["width"],
                                    )
                                else:
                                    rle = ann["segmentation"]
                                    for i in range(len(rle)):
                                        if not isinstance(rle[i]["counts"], bytes):
                                            rle[i]["counts"] = rle[i]["counts"].encode()
                            except:
                                print('error',ds,idx)
                                return self.__getitem__(0)
                            m = mask.decode(rle)
                            m = np.sum(
                                m, axis=2
                            )  # sometimes there are multiple binary map (corresponding to multiple segs)
                            m = m.astype(np.uint8)  # convert to np.uint8
                        m_final = m_final | m
                    m = m_final
                masks.append(m)
                continue

            ann = annotations[ann_id]

            if len(ann["segmentation"]) == 0:
                m = np.zeros((image_info["height"], image_info["width"])).astype(
                    np.uint8
                )
                masks.append(m)
                continue

            if type(ann["segmentation"][0]) == list:  # polygon
                rle = mask.frPyObjects(
                    ann["segmentation"], image_info["height"], image_info["width"]
                )
            else:
                rle = ann["segmentation"]
                for i in range(len(rle)):
                    if not isinstance(rle[i]["counts"], bytes):
                        rle[i]["counts"] = rle[i]["counts"].encode()
            m = mask.decode(rle)
            m = np.sum(
                m, axis=2
            )  # sometimes there are multiple binary map (corresponding to multiple segs)
            m = m.astype(np.uint8)  # convert to np.uint8
            masks.append(m)

        masks = np.stack(masks, axis=0)

  
        masks = torch.from_numpy(masks).float()
        sam_image = self.transform.apply_image(cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB))
        resize = sam_image.shape[:2]
        sam_image = self.preprocess(torch.from_numpy(sam_image).permute(2, 0, 1).contiguous())
        sample_num=len(questions)
        data_dict_bs['image'] = sam_image.unsqueeze(0).repeat(sample_num,1,1,1)
        data_dict_bs['masks']=masks
        data_dict_bs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
        data_dict_bs['image_grid_thw'] = inputs['image_grid_thw']
        data_dict_bs['img_metas'] = dict()
        data_dict_bs['img_metas']['image'] = torch.tensor(resize).unsqueeze(0).repeat(sample_num,1)
        data_dict_bs['img_metas']['sample_num'] = sample_num
        # data_dict_bs['img_metas']['SEG_num'] = 1
        data_dict_bs['img_metas']['SEG_num'] = sample_num
        return data_dict_bs
if __name__ == '__main__':
    from torch.utils.data import DataLoader
    from tqdm import tqdm
    dataset = ReferSegDataset(base_image_dir='./data/refer_seg')
    dataloader = DataLoader(
        dataset,
        batch_size=1,  # 如果你不需要批处理，可以保持为 1
        shuffle=False,  # 如果你只是遍历，通常不需要打乱
        num_workers=8,  # 多线程数，根据机器配置调整
        pin_memory=False,  # 如果用到 GPU 可以设为 True 加速
    )
    data_len = len(dataset)
    # 使用 tqdm 显示进度条，遍历 DataLoader
    for i, data in enumerate(tqdm(dataloader, total=data_len)):
        # 这里的 data 就是 dataset[i] 返回的内容
        # 如果你原来的 dataset[i] 有特定处理，可以在这里进行
        pass  # 你可以在这里对 data 做进一步处理
