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
GCG_QUESTIONS = [
    'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.',
    'Can you provide a thorough description of the this image? Please output with interleaved segmentation masks for the corresponding phrases.',
    'Please describe in detail the contents of the image. Please respond with interleaved segmentation masks for the corresponding parts of the answer.',
    'Could you give a comprehensive explanation of what can be found within this picture? Please output with interleaved segmentation masks for the corresponding phrases.',
    'Could you give me an elaborate explanation of this picture? Please respond with interleaved segmentation masks for the corresponding phrases.',
    'Could you provide me with a detailed analysis of this photo? Please output with interleaved segmentation masks for the corresponding parts of the answer.',
]
DEFAULT_IMAGE_TOKEN = '<image>'



def resize_longest_side(image, target_size):
    """
    调整图像大小，保持纵横比，将最长边缩放到 target_size
    
    参数:
        image: PIL.Image 对象
        target_size: 最长边的目标尺寸（整数）
    
    返回:
        调整后的 PIL.Image 对象
    """
    # 获取原始图像尺寸
    width, height = image.size
    
    # 计算缩放比例
    max_dim = max(width, height)
    scale = target_size / max_dim
    
    # 计算新尺寸（保持整数像素）
    new_width = int(round(width * scale))
    new_height = int(round(height * scale))
    
    # 调整图像大小（使用高质量的 Lanczos 滤镜进行缩放）
    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    return resized_image
def resize_mask_longest_side(mask_tensor, target_size):
    """
    调整0-1 mask张量的大小，保持纵横比，使用最近邻插值（纯PyTorch实现）
    """
    # 处理维度：确保输入是 (B, C, H, W) 格式（interpolate需要批量维度）
    if len(mask_tensor.shape) == 2:  # (H, W) → (1, 1, H, W)
        mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
    elif len(mask_tensor.shape) == 3:  # (C, H, W) → (1, C, H, W)
        mask_tensor = mask_tensor.unsqueeze(0)
    
    _, _, height, width = mask_tensor.shape
    max_dim = max(height, width)
    scale = target_size / max_dim
    new_height, new_width = int(round(height * scale)), int(round(width * scale))
    
    # 使用 interpolate 进行缩放，mode指定为'nearest'
    resized_mask = F.interpolate(
        mask_tensor,
        size=(new_height, new_width),
        mode='nearest-exact'  # 精确最近邻，避免中间值
    )
    
    # 恢复原始维度（去除批量维度）
    if resized_mask.shape[0] == 1:
        resized_mask = resized_mask.squeeze(0)
    
    return resized_mask

class GCGBaseDataset(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, 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.INPUT_PIXEL=768
        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.question_templates = GCG_QUESTIONS
        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, ann_file)
        self.data_infos = self._load_annotations(self.ann_file)
        print('\033[1;32m' + f"dataset GCG: {data_name} dataset initialized {len(self.data_infos)} res:{self.INPUT_PIXEL}" + '\033[0m')
        
    def _load_annotations(self, ann_file):
        with open(ann_file, 'r') as f:
            data_infos = json.load(f)
        data_infos = data_infos[0: 1000] if self.validation else data_infos
        return data_infos

    def _parse_annotations(self, ann_info):
        image_path = os.path.join(self.image_folder, ann_info['file_name'])
        annotations = {'labels': [], 'caption': [], 'masks': [], 'tokens_positive': [],
                       'file_name': ann_info['file_name']}
        width, height = Image.open(image_path).size
        annotations['caption'] = ann_info['caption'].strip('"').strip()

        for word, grounding in ann_info["groundings"].items():
            annotations['labels'].append(word)
            annotations['tokens_positive'].append(grounding["token_positives"])

            # Convert segmentation to binary mask
            binary_mask = np.zeros((height, width), dtype=np.uint8)
            for rle in grounding["rle_masks"]:
                m = mask.decode(rle).astype(np.uint8)
                binary_mask += m.squeeze()
            annotations['masks'].append(binary_mask)

        return annotations

    def __getitem__(self, index):
        while True:
            ann_info = self.data_infos[index%(len(self.data_infos))]
            # Parse annotation info
            ann = self._parse_annotations(ann_info)
            image_path = os.path.join(self.image_folder, ann['file_name'])
            if len(ann['labels']) > 0:
                break
            else:
                index = random.randint(0, len(self.data_infos) - 1)
        data_item = {"image_path": image_path, "filename": ann['file_name'], "caption": ann['caption'],
            "labels": ann['labels'], "masks": ann['masks'], "tokens_positive": ann['tokens_positive']}
        return self.process_data(data_item)

    def __len__(self):
        return len(self.data_infos)



    def create_conversations(self, caption, tokens_positive):
        question = random.choice(self.question_templates).strip()

        # Prepare caption with tags
        def tag_caption(caption, tokens):
            for start, end in sorted(tokens, key=lambda x: x[0], reverse=True):
                caption = f"{caption[:start]}<p> {caption[start:end]} </p> [SEG]{caption[end:]}"
            return caption

        detailed_answer = tag_caption(caption, tokens_positive)
        sources = dict()
        sources['conversations'] = []
        cur_sample=[]
        cur_sample.append(
                {'from': 'human', 'value': self.begin_str + question})
        cur_sample.append({'from': 'gpt', 'value': detailed_answer})
        sources['conversations'].append(cur_sample)
        return sources
    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 process_data(self, data_item):
        data_dict_bs=dict()
        data_dict_bs['input_ids']=[]
        data_dict_bs['labels']=[]
        data_labels = data_item['labels']
        masks = data_item['masks']
        caption = data_item['caption']
        tokens_positive = data_item['tokens_positive']
        image_path = data_item['image_path']

        # Function to sort elements based on the start index of each phrase
        def sort_by_start_index(items, order):
            return [items[i] for i in order]

        # Sort phrases based on their appearance in the sentence
        phrase_order = sorted(range(len(tokens_positive)), key=lambda x: tokens_positive[x][0])
        masks = sort_by_start_index(masks, phrase_order)
        data_labels = sort_by_start_index(data_labels, phrase_order)
        tokens_positive = sort_by_start_index(tokens_positive, phrase_order)
        sources = self.create_conversations(caption, tokens_positive)
        
        sources=preprocess_multimodal(
                copy.deepcopy([sources['conversations']][0]),
                self.data_args,
                0,if_qwen=True)
        text = sources[0][0]['value']
        processor=self.data_args.image_processor
        # processor = self.data_args
        qwen_img=Image.open(image_path).convert('RGB')
        W,H=qwen_img.size
        qwen_img_process=resize_longest_side(qwen_img,self.INPUT_PIXEL)
        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,
            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))
        masks = np.stack(masks, axis=0)
        masks = torch.from_numpy(masks)
        masks=resize_mask_longest_side(masks,self.INPUT_PIXEL)
        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())
        data_dict_bs['image'] = sam_image.unsqueeze(0)
        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)
        data_dict_bs['img_metas']['sample_num'] = 1
        N_SEG=self.tokenizer.encode(sources[0][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

class GranDfDataset(GCGBaseDataset):
    """
    Human annotated dataset proposed in GLaMM as part of GranDf dataset.
    """
    def __init__(self, tokenizer, data_args, base_image_dir, epoch_samples=8000, precision="fp32",
                 image_size=1024, num_classes_per_sample=3, validation=False, random_sampling=True,
                 image_folder='', json_path=''):

        image_folder=os.path.join(base_image_dir, 'images/grandf/train')
        ann_file='annotations/GranDf_HA_GCG_train.json'
        super().__init__(
            tokenizer, data_args, base_image_dir, epoch_samples, precision, image_size, num_classes_per_sample,
            validation, random_sampling, image_folder, ann_file, data_name='GranDf')

class OpenPsgGCGDataset(GCGBaseDataset):
    """
    Human annotated dataset proposed in GLaMM as part of GranDf dataset.
    """
    def __init__(self, tokenizer, data_args, base_image_dir, epoch_samples=8000, precision="fp32",
                 image_size=1024, num_classes_per_sample=3, validation=False, random_sampling=True,
                 image_folder='', ann_file=''):

        image_folder=os.path.join(base_image_dir, 'images/coco2017')
        ann_file='annotations/OpenPsgGCG_train.json'
        super().__init__(
            tokenizer, data_args, base_image_dir, epoch_samples, precision, image_size, num_classes_per_sample,
            validation, random_sampling, image_folder, ann_file, data_name='OpenPsg')


class Flickr30kGCGDataset(GCGBaseDataset):
    def __init__(self, tokenizer, data_args, base_image_dir, epoch_samples=8000, precision="fp32",
                 image_size=1024, num_classes_per_sample=3, validation=False, random_sampling=True,
                 image_folder='', ann_file=''):
        image_folder = os.path.join(base_image_dir, 'images/flickr30k/Flickr30K/')
        ann_file = 'annotations/flickr_mergedGT_GCG_train.json'
        super().__init__(
            tokenizer, data_args, base_image_dir, epoch_samples, precision, image_size, num_classes_per_sample,
            validation, random_sampling, image_folder, ann_file, data_name='Flickr30k')
        # Filter out images smaller than the minimum size
        self.data_infos = [self.data_infos[i] for i in self._filter_images(min_size=32)]

    def _load_annotations(self, ann_file):
        # Load annotations and filter out images with very short captions
        self.coco = COCO(ann_file)
        self.image_ids = self.coco.getImgIds()
        data_infos = []
        total_ann_ids = []
        removed_img_count = 0
        for img_id in self.image_ids:
            if len(data_infos) == 1000 and self.validation:
                # Only limited images for validation
                break
            info = self.coco.loadImgs([img_id])[0]
            if len(info['caption'].split(' ')) < 3:
                removed_img_count += 1
                continue
            info['filename'] = info['file_name'].split('_')[-1]
            info['height'] = int(info['height'])
            info['width'] = int(info['width'])
            data_infos.append(info)
            ann_ids = self.coco.getAnnIds(imgIds=[img_id])
            total_ann_ids.extend(ann_ids)
        assert len(set(total_ann_ids)) == len(total_ann_ids), f"Non-unique annotation IDs in '{ann_file}'!"
        print(f'Removed {removed_img_count} images.')
        return data_infos

    def _filter_images(self, min_size):
        return [i for i, info in enumerate(self.data_infos) if min(info['width'], info['height']) >= min_size]

    def _parse_annotations(self, img_info, ann_info):
        annotations = {'bboxes': [], 'labels': [], 'bboxes_ignore': [], 'caption': img_info['caption'], 'masks': [],
                       'tokens_positive': []}
        for ann in ann_info:
            if ann.get('ignore', False):
                continue
            x1, y1, w, h = ann['bbox']
            inter_w = max(0, min(x1 + w, img_info['width']) - max(x1, 0))
            inter_h = max(0, min(y1 + h, img_info['height']) - max(y1, 0))
            if inter_w * inter_h == 0 or ann['area'] <= 0 or w < 1 or h < 1:
                continue
            bbox = [x1, y1, x1 + w, y1 + h]
            annotations['bboxes'].append(bbox)
            tokens_positive = ann['tokens_positive']
            gt_label = [img_info['caption'][span[0]:span[1]] for span in tokens_positive]
            annotations['labels'].append(gt_label[0])
            annotations['tokens_positive'].append(tokens_positive[0])

            rle = ann['sam_mask']
            mask_decoded = mask.decode(rle).astype(np.uint8)
            annotations['masks'].append(mask_decoded)

        # Convert bounding boxes to numpy arrays
        annotations['bboxes'] = np.array(annotations['bboxes'], dtype=np.float32) if annotations[
            'bboxes'] else np.zeros((0, 4), dtype=np.float32)
        annotations['bboxes_ignore'] = np.array(annotations['bboxes_ignore'], dtype=np.float32) if annotations[
            'bboxes_ignore'] else np.zeros((0, 4), dtype=np.float32)

        return annotations

    def __getitem__(self, index):
        img_info = self.data_infos[index] if (self.validation or not self.random_sampling) \
            else self.data_infos[random.randint(0, len(self.data_infos) - 1)]
        ann_ids = self.coco.getAnnIds(imgIds=img_info['id'])
        ann_info = self.coco.loadAnns(ann_ids)
        image_path = os.path.join(self.image_folder, img_info['file_name'])
        # Parse annotation info
        ann = self._parse_annotations(img_info, ann_info)
        data_item = {"image_path": image_path, "filename": img_info['file_name'], "width": img_info['width'],
                     "height": img_info['height'], "bbox": ann['bboxes'], "caption": ann['caption'],
                     "labels": ann['labels'], "masks": ann['masks'], "tokens_positive": ann['tokens_positive']}
        return self.process_data(data_item)


class RefCOCOgGCGDataset(GCGBaseDataset):
    def __init__(self, tokenizer, data_args, base_image_dir, epoch_samples=8000, precision="fp32",
                 image_size=1024, num_classes_per_sample=3, validation=False, random_sampling=True,
                 image_folder='', ann_file=''):

        image_folder=os.path.join(base_image_dir, 'images/coco2014/train2014/')
        ann_file='annotations/RefCOCOg_GCG_train.json'
        super().__init__(
            tokenizer, data_args, base_image_dir, epoch_samples, precision, image_size, num_classes_per_sample,
            validation, random_sampling, image_folder, ann_file, data_name='RefCOCOg')

    def _parse_annotations(self, ann_info):
        image_path = os.path.join(self.image_folder, ann_info['img_file_name'])
        annotations = {'labels': [], 'caption': [], 'masks': [], 'tokens_positive': [],
                       'file_name': ann_info['img_file_name']}
        width, height = Image.open(image_path).size
        orig_caption = ann_info['caption'].strip('"').strip()
        annotations['caption'] = orig_caption.lower()

        for detail in ann_info['refs']:
            phrase = detail['sentence']
            if phrase.lower() in annotations['caption']:
                annotations['labels'].append(phrase)
                index = annotations['caption'].find(phrase)
                end_index = index + len(phrase) if index != -1 else -1
                annotations['tokens_positive'].append([index, end_index])

                # Convert segmentation to binary mask
                binary_mask = np.zeros((height, width), dtype=np.uint8)
                for seg in detail["segmentation"]:
                    rles = mask.frPyObjects([seg], height, width)
                    m = mask.decode(rles)
                    m = m.astype(np.uint8)
                    binary_mask += m.squeeze()
                annotations['masks'].append(binary_mask)

        # Sort tokens_positive and corresponding lists
        tokens_positive = annotations['tokens_positive']
        sorted_indices = sorted(range(len(tokens_positive)), key=lambda i: tokens_positive[i][0])
        annotations['tokens_positive'] = [tokens_positive[i] for i in sorted_indices]
        annotations['masks'] = [annotations['masks'][i] for i in sorted_indices]
        annotations['labels'] = [annotations['labels'][i] for i in sorted_indices]

        # Trimming overlapping intervals
        for i in range(len(tokens_positive)):
            for j in range(i + 1, len(tokens_positive)):
                # If there is overlap
                if tokens_positive[i][1] >= tokens_positive[j][0]:
                    # Modify the end index of phrase i to be one less than the start index of phrase j
                    tokens_positive[i][1] = tokens_positive[j][0] - 1
                    # Modify the phrases to reflect the change in indices
                    annotations['labels'][i] = orig_caption[tokens_positive[i][0]:tokens_positive[i][1] + 1]
                    break  # Exit inner loop since i was modified

        return annotations

    def __getitem__(self, index):
        while True:
            ann_dict = self.data_infos[index] if (self.validation or not self.random_sampling) \
                else self.data_infos[random.randint(0, len(self.data_infos) - 1)]
            ann_info = next(iter(ann_dict.values()))
            # Parse annotation info
            ann = self._parse_annotations(ann_info)
            image_path = os.path.join(self.image_folder, ann['file_name'])
            # Check if len(gt_phrases) > 0 and if True, break the loop
            if len(ann['labels']) > 0:
                break
            else:
                index = random.randint(0, len(self.data_infos) - 1)
        data_item = {"image_path": image_path, "filename": ann['file_name'], "caption": ann['caption'],
                     "labels": ann['labels'], "masks": ann['masks'], "tokens_positive": ann['tokens_positive']}

        return self.process_data(data_item)

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 = RefCOCOgGCGDataset(
        tokenizer=tokenizer,
        base_image_dir='./data/glamm_data',
        data_args=processor
        )
    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 做进一步处理
