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 PIL import Image
import copy
import json
from pycocotools.coco import COCO
from detectron2.data import detection_utils as utils
from detectron2.data import transforms as T
from detectron2.data.transforms import TransformGen
from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
from detectron2.structures import BitMasks, Boxes, Instances
from FCLM.datasets.utils import ANSWER_LIST, LONG_QUESTION_LIST
import numpy as np
MY_COCO_CATEGORIES=dict()
for each in COCO_CATEGORIES:
    MY_COCO_CATEGORIES[each['id']]=dict()
    MY_COCO_CATEGORIES[each['id']]['color']=each['color']
    MY_COCO_CATEGORIES[each['id']]['isthing']=each['isthing']
    MY_COCO_CATEGORIES[each['id']]['name']=each['name']

def get_prompt_templates():
    prompt_templates = [
        '{}.',
        'a photo of a {}.',
        'a bad photo of a {}.',
        'a photo of many {}.',
        'a sculpture of a {}.',
        'a photo of the hard to see {}.',
        'a low resolution photo of the {}.',
        'a rendering of a {}.',
        'graffiti of a {}.',
        'a bad photo of the {}.',
        'a cropped photo of the {}.',
        'a tattoo of a {}.',
        'the embroidered {}.',
        'a photo of a hard to see {}.',
        'a bright photo of a {}.',
        'a photo of a clean {}.',
        'a photo of a dirty {}.',
        'a dark photo of the {}.',
        'a drawing of a {}.',
        'a photo of my {}.',
        'the plastic {}.',
        'a photo of the cool {}.',
        'a close-up photo of a {}.',
        'a black and white photo of the {}.',
        'a painting of the {}.',
        'a painting of a {}.',
        'a pixelated photo of the {}.',
        'a sculpture of the {}.',
        'a bright photo of the {}.',
        'a cropped photo of a {}.',
        'a plastic {}.',
        'a photo of the dirty {}.',
        'a jpeg corrupted photo of a {}.',
        'a blurry photo of the {}.',
        'a photo of the {}.',
        'a good photo of the {}.',
        'a rendering of the {}.',
        'a {} in a video game.',
        'a photo of one {}.',
        'a doodle of a {}.',
        'a close-up photo of the {}.',
        'the origami {}.',
        'the {} in a video game.',
        'a sketch of a {}.',
        'a doodle of the {}.',
        'a origami {}.',
        'a low resolution photo of a {}.',
        'the toy {}.',
        'a rendition of the {}.',
        'a photo of the clean {}.',
        'a photo of a large {}.',
        'a rendition of a {}.',
        'a photo of a nice {}.',
        'a photo of a weird {}.',
        'a blurry photo of a {}.',
        'a cartoon {}.',
        'art of a {}.',
        'a sketch of the {}.',
        'a embroidered {}.',
        'a pixelated photo of a {}.',
        'itap of the {}.',
        'a jpeg corrupted photo of the {}.',
        'a good photo of a {}.',
        'a plushie {}.',
        'a photo of the nice {}.',
        'a photo of the small {}.',
        'a photo of the weird {}.',
        'the cartoon {}.',
        'art of the {}.',
        'a drawing of the {}.',
        'a photo of the large {}.',
        'a black and white photo of a {}.',
        'the plushie {}.',
        'a dark photo of a {}.',
        'itap of a {}.',
        'graffiti of the {}.',
        'a toy {}.',
        'itap of my {}.',
        'a photo of a cool {}.',
        'a photo of a small {}.',
        'a tattoo of the {}.',
    ]
    return prompt_templates

def prompt_engineering(classnames, topk=1, suffix='.'):
    prompt_templates = get_prompt_templates()
    temp_idx = np.random.randint(min(len(prompt_templates), topk))

    if isinstance(classnames, list):
        classname = random.choice(classnames)
    else:
        classname = classnames

    return prompt_templates[temp_idx].replace('.', suffix).format(classname.replace(',', '').replace('+', ' '))
    


DEFAULT_IMAGE_TOKEN = '<image>'


class Pixel2Cap_DATASET(torch.utils.data.Dataset):
    """
    Dataset Class for Grounded Conversation Generation (GCG) proposed in GLaMM.
    """
    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, data_args, base_image_dir='./data/pix2cap_data', epoch_samples=8000, precision="fp32",
                 image_size=1024, num_classes_per_sample=3, validation=False, random_sampling=True,
                 image_folder='data/glamm_data/images/grandf/train', ann_file='annotations/GranDf_HA_GCG_train.json',data_name='GranDf'):
        self.data_args=data_args
        self.base_image_dir = base_image_dir
        self.epoch_samples = epoch_samples
        self.num_classes_per_sample = num_classes_per_sample
        self.image_size = image_size
        self.tokenizer = tokenizer
        self.precision = precision
        self.transform = ResizeLongestSide(image_size)
        self.validation = validation
        self.random_sampling = random_sampling
        self.answer_list = ANSWER_LIST
        self.question_list = LONG_QUESTION_LIST
        self.begin_str = f"""{DEFAULT_IMAGE_TOKEN}\nThis provides an overview of the picture.\n"""
        self.validation = validation
        self.image_folder=image_folder
        self.ann_file = os.path.join(self.base_image_dir, 'annotations/pix2cap_coco_train.json')
        
        self.data_infos = json.load(open(self.ann_file))
        print('\033[1;32m' + f"dataset {self.data_infos['info']['description']} initialized {len(self.data_infos['images'])} images {len(self.data_infos['annotations'])} annotations" + '\033[0m')
    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 __len__(self):
        # return 8000
        return len(self.data_infos['images'])
    def __getitem__(self, index):

        file_name=self.data_infos['images'][index]['file_name']
        file_name=os.path.join(self.base_image_dir, 'coco/train2017', file_name)
        # image = cv2.imread(file_name)
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = utils.read_image(file_name, format="RGB")
        image_shape = image.shape[:2]  # h, w

        anno=self.data_infos['annotations'][index]
        segments_infos = [each["id"] for each in anno['segments_info']]
        descriptions = [each["description"] for each in anno['segments_info']]
        pan_seg_gt_path=os.path.join(self.base_image_dir, 'coco/panoptic_train2017', self.data_infos['annotations'][index]['file_name'])
        # pan_seg_gt=cv2.imread(pan_seg_gt_path)
        # pan_seg_gt = cv2.cvtColor(pan_seg_gt, cv2.COLOR_BGR2RGB)
        pan_seg_gt = utils.read_image(pan_seg_gt_path, "RGB")
        from panopticapi.utils import rgb2id

        pan_seg_gt = rgb2id(pan_seg_gt)
        classes = []
        masks = []
        scores = []
        densecaptions = []
        

        for segment_info in anno['segments_info']:
            class_id = segment_info["category_id"]
            if not segment_info["iscrowd"]:
                classes.append(class_id)
                masks.append(pan_seg_gt == segment_info["id"])
                scores.append(segment_info['score'])
                if segment_info["description"]=='':
                    densecaptions += [prompt_engineering(MY_COCO_CATEGORIES[segment_info["category_id"]]['name'].lower().replace('-other','').replace('-merged','').replace('-stuff',''), topk=10000, suffix='')]
                else:
                    densecaptions += [segment_info["description"]]#.lower().translate(str.maketrans('', '', string.punctuation))]
        is_things=[]
        for idx in classes:
            is_things.append(MY_COCO_CATEGORIES[idx]['isthing'])
        classes = np.array(classes)
        is_things = np.array(is_things)
        gt_classes = torch.tensor(classes, dtype=torch.int64)
        is_things = torch.tensor(is_things, dtype=torch.int64)

        if len(masks) == 0:
            return self.__getitem__(0)
        else:
            masks = BitMasks(
                torch.stack([torch.from_numpy(np.ascontiguousarray(x.copy())) for x in masks])
            )
        
        gt_masks = masks.tensor.int()
        gt_boxes = masks.get_bounding_boxes()
        gt_dense = densecaptions
        assert len(gt_masks) == len(gt_boxes) == len(gt_dense) , 'masks: {} boxes: {} dense: {}'.format(len(gt_masks), len(gt_boxes), len(gt_dense))
        
        mask_pixel_num = gt_masks.sum(dim=-1).sum(dim=-1)
        _,indices=torch.sort(mask_pixel_num)
        N = len(indices)
        third = N // 3  # 取整数部分，比如 12 / 3 = 4
        indices = indices[third:]

        if len(indices)> self.num_classes_per_sample:
            sampled_inds = np.random.choice(
                    list(range(len(indices))), size=self.num_classes_per_sample, replace=False
            )
        else:
            sampled_inds = indices


        questions = []
        answers = []
        for inds in sampled_inds:
            text = gt_dense[inds].strip()
            question_template = random.choice(self.question_list)
            questions.append(question_template.format(sent=text.lower()))
            answers.append(random.choice(self.answer_list))
        masks=[gt_masks[i] for i in sampled_inds]
        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)


        data_dict_bs=dict()
        data_dict_bs['input_ids']=[]
        data_dict_bs['labels']=[]
        
        processor=self.data_args.image_processor
        qwen_img_process=Image.open(file_name).convert('RGB')
        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))
            #############################
        masks = np.stack(masks, axis=0)
        masks = torch.from_numpy(masks).float()
        sam_image = self.transform.apply_image(cv2.cvtColor(cv2.imread(file_name), 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'] = sample_num
        return data_dict_bs


if __name__ == '__main__':
    from torch.utils.data import DataLoader
    from tqdm import tqdm
    import transformers
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        './Qwen2.5-VL-3B-Instruct',
        padding_side="right",
        use_fast=True,
    )
    mask_tokens = ['<mask>', '<pos>','<p>','</p>']
    seg_tokens = ['[SEG]']
    tokenizer.add_tokens(mask_tokens, special_tokens=True)
    tokenizer.add_tokens(seg_tokens, special_tokens=True)
    from qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
    processor = Qwen2_5_VLProcessor.from_pretrained('Qwen2.5-VL-3B-Instruct', trust_remote_code=True)
    dataset = Pixel2Cap_DATASET(
        tokenizer=tokenizer,
        base_image_dir='./data/pix2cap_data',
        data_args=processor
        )
    # dataset[0]
    # dataset[1]
    # dataset[2]
    # dataset[3]
    dataloader = DataLoader(
        dataset,
        batch_size=1,  # 如果你不需要批处理，可以保持为 1
        shuffle=False,  # 如果你只是遍历，通常不需要打乱
        num_workers=16,  # 多线程数，根据机器配置调整
        pin_memory=False,  # 如果用到 GPU 可以设为 True 加速
    )
    data_len = len(dataset)
    for i, data in enumerate(tqdm(dataloader, total=data_len)):
        pass  # 你可以在这里对 data 做进一步处理
