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.utils import ANSWER_LIST, SHORT_QUESTION_LIST
from PIL import Image
import copy
import json
from FCLM.datasets.utils import (
    MR_SINGLE_ANSWER_LIST,
    MR_MULTI_ANSWER_LIST,
    ANSWER_LIST,
    DEFAULT_IM_END_TOKEN,
    DEFAULT_IM_START_TOKEN,
    DEFAULT_IMAGE_PATCH_TOKEN,
    DEFAULT_IMAGE_TOKEN,
    EXPLANATORY_QUESTION_LIST,
    LONG_QUESTION_LIST,
    SHORT_QUESTION_LIST,
    EXPAND_LONG_QUESTION_LIST,
)
class MultiReasonSegDataset(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 * 2 * 10,
        precision: str = "fp32",
        image_size: int = 1024,
        num_classes_per_sample: int = 3,
        exclude_val=False,
        reason_seg_data="MultiReasonSeg|train",
        explanatory=0.1,
        num_classes_per_question=1,
        seg_token_num=1,
        pad_train_clip_images=False,
        masks_process_with_clip=False,
        preprocessor_config='',
        use_expand_question_list=True
    ):
        self.exclude_val = exclude_val
        self.reason_seg_data = reason_seg_data
        self.samples_per_epoch = samples_per_epoch
        self.explanatory = explanatory
        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.short_question_list = SHORT_QUESTION_LIST
        self.long_question_list = LONG_QUESTION_LIST
        self.answer_list = ANSWER_LIST
        self.single_answer_list = MR_SINGLE_ANSWER_LIST
        self.multi_answer_list = MR_MULTI_ANSWER_LIST   
        self.seg_token_num = seg_token_num
        self.num_classes_per_question = num_classes_per_question
        
        self.masks_process_with_clip = masks_process_with_clip
        self.pad_train_clip_images = pad_train_clip_images
       
        if use_expand_question_list:
            self.long_question_list.extend(EXPAND_LONG_QUESTION_LIST)
        reason_seg_data, split = reason_seg_data.split("|")
        # import ipdb;ipdb.set_trace()
        self.reason_seg_data = []
        json_file_name = './data/muse/MUSE_train.json'
        with open(json_file_name, 'r') as f:
            reason_file = json.load(f)
            print(
                "\033[1;32mdataset {} has {} samples.\033[0m".format(
                    'MUSE',
                    len(reason_file),
                )   
            )
        self.reason_seg_data.extend(reason_file)

        json_file_name = './data/muse/MUSE_val.json'
        with open(json_file_name, 'r') as f:
            reason_file = json.load(f)
            print(
                "\033[1;32mdataset {} has {} samples.\033[0m".format(
                    'MUSE',
                    len(reason_file),
                )   
            )
        self.reason_seg_data.extend(reason_file)

        json_file_name = './data/muse/MUSE_test_less.json'
        with open(json_file_name, 'r') as f:
            reason_file = json.load(f)
            print(
                "\033[1;32mdataset {} has {} samples.\033[0m".format(
                    'MUSE',
                    len(reason_file),
                )   
            )
        self.reason_seg_data.extend(reason_file)

        json_file_name = './data/muse/MUSE_test_many.json'
        with open(json_file_name, 'r') as f:
            reason_file = json.load(f)
            print(
                "\033[1;32mdataset {} has {} samples.\033[0m".format(
                    'MUSE',
                    len(reason_file),
                )   
            )
        self.reason_seg_data.extend(reason_file)
        images = []
        anns = []
        questions = []
        answers = []
        self.data_args = data_args
        
        
        
    def __len__(self):
        # return self.samples_per_epoch
        return len(self.reason_seg_data)

    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):
        #
        
        # idx = random.randint(0, len(self.reason_seg_data) - 1)
        image_info = self.reason_seg_data[idx]

        if 'file_name' in image_info:
            image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/train2014")
            image_path = os.path.join(image_root, image_info['file_name'])
        else:
            if 'train2017' in image_info['coco_url']:
                image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/train2017")
                image_path = os.path.join(image_root, image_info['coco_url'].split('/')[-1])
            else:
                image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/val2017")
                image_path = os.path.join(image_root, image_info['coco_url'].split('/')[-1])

        anns = image_info['ann_list']
        question = image_info['questions'] if 'questions' in image_info else None
        gt_answer = image_info['answers'] if 'answers' in image_info else None
        if question is not None:
            text_answers = image_info['text_answers'] if 'text_answers' in image_info else [None] * len(gt_answer)
        else:
            text_answers = None
        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 = []
        

        masks = []
        if len(anns) == 0:
            return self[0]

        category_ids = [ann['category_id'] for ann in anns]
        category_ids = list(set(category_ids))
        sampled_num = min(self.num_classes_per_sample, len(category_ids))
        sampled_category_ids = np.random.choice(category_ids, size=sampled_num, replace=False)

        sampled_sents = question
        sampled_answers = gt_answer
        sampled_masks = masks
        sample_text_answers = text_answers

        image_name = image_path.split("/")[-1]
        questions = []
        answers = []
        use_assign_list = []
        seg_token = ["[SEG{}]".format(i) for i in range(self.seg_token_num)]
        seg_token = ' '.join(seg_token)

        if question is not None:
            for text, answer_list, text_answer in zip(sampled_sents, sampled_answers, sample_text_answers):
                # if is_sentence:
                question_template = random.choice(self.long_question_list)
                questions.append(question_template.format(sent=text))
                for answer in answer_list:
                    rle = mask.frPyObjects(answer["segmentation"], image_info["height"], image_info["width"])
                    m = mask.decode(rle)
                    if len(m.shape) > 2:
                        # assert m.shape[-1] == 1, m.shape
                        m = np.sum(m, axis=2)  # so
                    m = m.astype(np.uint8)
                    masks.append(m)
                use_assign = False
                if text_answer is not None:
                    if text_answer.count('{seg}') != len(answer_list):
                        return self[0]
                    try:
                        _text_answer = text_answer.format(seg='[SEG]') if self.seg_token_num == 1 else text_answer.format(seg=seg_token)
                    except:
                        return self.__getitem__(0)
                    answers.append(_text_answer)
                    use_assign_list.append(False)
                else:
                    target_list = [a['rephrased_name'] if (random.random() > 0.1 and 'rephrased_name' in a) else a['category_name'] for a in answer_list ]
                    target_answer = []
                    separate_answer = random.randint(0, 1)
                    _seg = ['[SEG]'] * len(target_list)
                    if len(target_list) > 1:
                        part1 = ', '.join(_seg[:-1])
                        part2 = ' and ' + _seg[-1]
                        _seg = part1 + part2 
                    else:
                        _seg = _seg[0]
                    if separate_answer:
                        choice_list = self.single_answer_list
                        answer_temp = random.choice(choice_list) if self.seg_token_num == 1 else random.choice(choice_list).replace('[SEG]', seg_token)
                        use_assign = False if "{class_name}" in answer_temp else True
                        for i, sampled_cls in enumerate(target_list):
                            _answer_temp = answer_temp.format(class_name=sampled_cls) if "{class_name}" in answer_temp else answer_temp
                            target_answer.append(_answer_temp[:-1])
                        if len(target_answer) > 1:
                            part1 = ', '.join(target_answer[:-1])
                            part2 = ' and ' + target_answer[-1]
                            target_answer = part1 + part2 + '.'
                        else:
                            target_answer = target_answer[0] + '.'
                    else:
                        answer_temp = random.choice(self.multi_answer_list)
                        _answer_temp = answer_temp.format(class_name=', '.join(target_list).lower(), seg=_seg) if "{class_name}" in answer_temp else answer_temp.format(seg=_seg)
                        use_assign = False if "{class_name}" in answer_temp else True
                        _answer_temp = _answer_temp if self.seg_token_num == 1 else _answer_temp.replace('[SEG]', seg_token)
                        target_answer = _answer_temp

                    answers.append(target_answer)
                    use_assign_list.append(use_assign)
        else:
            for sampled_category_id in sampled_category_ids:
                question_template = random.choice(self.instance_question_list)
                category_names = self.lvis_name_dict[str(sampled_category_id)]
                category_name = random.choice(category_names)
                questions.append(question_template.format(class_name=category_name))
                answer_list = [ann for ann in anns if ann['category_id'] == sampled_category_id]
                for answer in answer_list:
                    rle = mask.frPyObjects(answer["segmentation"], image_info["height"], image_info["width"])
                    m = mask.decode(rle)
                    if len(m.shape) > 2:
                        # assert m.shape[-1] == 1, m.shape
                        m = np.sum(m, axis=2)  # so
                    m = m.astype(np.uint8)
                    masks.append(m)

                target_list = [a['rephrased_name'] if random.random() > 0.1 else a['category_name'] for a in answer_list ]
                target_answer = []
                separate_answer = random.randint(0, 1)
                _seg = ['[SEG]'] * len(target_list)
                if len(target_list) > 1:
                    part1 = ', '.join(_seg[:-1])
                    part2 = ' and ' + _seg[-1]
                    _seg = part1 + part2 
                else:
                    _seg = _seg[0]

                separate_answer = random.randint(0, 1)
                # if len(answer_list) == 1 or separate_answer:
                choice_list = self.single_answer_list
                answer_temp = random.choice(choice_list) if self.seg_token_num == 1 else random.choice(choice_list).replace('[SEG]', seg_token)
                use_assign = False if "{class_name}" in answer_temp else True
                for i, sampled_cls in enumerate(target_list):
                    _answer_temp = answer_temp.format(class_name=sampled_cls) if "{class_name}" in answer_temp else answer_temp
                    target_answer.append(_answer_temp[:-1])
                if len(target_answer) > 1:
                    part1 = ', '.join(target_answer[:-1])
                    part2 = ' and ' + target_answer[-1]
                    target_answer = part1 + part2 + '.'
                else:
                    target_answer = target_answer[0] + '.'
                

                answers.append(target_answer)
                use_assign_list.append(use_assign)
        
        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()
        processor=self.data_args.image_processor
        qwen_img_process=Image.open(image_path).convert('RGB')
        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)
            new_inputs=torch.cat((old_data_dict['input_ids'][:,:insert_position],img_padding,old_data_dict['input_ids'][:,insert_position+1:]),dim=1)       
            new_target=torch.cat((old_data_dict['labels'][:,:insert_position],target_padding,old_data_dict['labels'][:,insert_position+1:]),dim=1)
            data_dict_bs['input_ids'].append(new_inputs.squeeze(0))
            data_dict_bs['labels'].append(new_target.squeeze(0))

        try:
            masks = np.stack(masks, axis=0)
        except:
            print(f"empty mask, use another sample")
            return self.__getitem__(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
        N_SEG=0
        
        for each in sources['conversations']:
            N_SEG+=self.tokenizer.encode(each[1]['value']).count(self.tokenizer.encode('[SEG]')[0])
        data_dict_bs['img_metas']['SEG_num'] = N_SEG
        try:
            # 检查分割数量与掩码数量是否匹配
            assert N_SEG == masks.shape[0], \
                f"样本错误 SEG数量 {N_SEG} 与掩码数量 {masks.shape[0]} 不匹配"
        except AssertionError as e:
            print(f"警告: {e}，将使用第一个样本替代")
            return self.__getitem__(0)  # 调用第一个样本
        return data_dict_bs



def transform_mask(masks, size):
    height, width = masks.shape[-2:]
    short, long = (width, height) if width <= height else (height, width)
    requested_new_short = size
    new_short, new_long = requested_new_short, int(requested_new_short * long / short)
    new_shape = (new_long, new_short) if width <= height else (new_short, new_long)
    masks = F.interpolate(masks[None].float(), size=new_shape, mode="nearest")[0].bool()

    orig_height, orig_width = new_shape
    crop_height, crop_width = size, size
    crop_height, crop_width = int(crop_height), int(crop_width)
    top = (orig_height - crop_height) // 2
    bottom = top + crop_height
    # In case size is odd, (image_shape[1] + size[1]) // 2 won't give the proper result.
    left = (orig_width - crop_width) // 2
    right = left + crop_width
    assert top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width
    # if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width:
    masks = masks[..., top:bottom, left:right]

    return masks



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 = MultiReasonSegDataset(tokenizer=tokenizer,base_image_dir='./data')
    # import tqdm
    # for i in tqdm.tqdm(range(len(dataset))):
    #     data = dataset[i]

    dataloader = DataLoader(
        dataset,
        batch_size=1,  # 如果你不需要批处理，可以保持为 1
        shuffle=False,  # 如果你只是遍历，通常不需要打乱
        num_workers=16,  # 多线程数，根据机器配置调整
        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 做进一步处理
