import copy
import random
import os
import numpy as np
import torch
import json
import ast
from FCLM.train.train import preprocess, preprocess_multimodal
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from torchvision import transforms
from pycocotools import mask
from copy import deepcopy
from typing import Tuple
class ResizeLongestSide:
    """
    Resizes images to the longest side 'target_length', as well as provides
    methods for resizing coordinates and boxes. Provides methods for
    transforming both numpy array and batched torch tensors.
    """

    def __init__(self, target_length: int) -> None:
        self.target_length = target_length

    def apply_image(self, image: np.ndarray) -> np.ndarray:
        """
        Expects a numpy array with shape HxWxC in uint8 format.
        """
        target_size = self.get_preprocess_shape(
            image.shape[0], image.shape[1], self.target_length
        )
        return np.array(resize(to_pil_image(image), target_size))

    def apply_coords(
        self, coords: np.ndarray, original_size: Tuple[int, ...]
    ) -> np.ndarray:
        """
        Expects a numpy array of length 2 in the final dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).astype(float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

    def apply_boxes(
        self, boxes: np.ndarray, original_size: Tuple[int, ...]
    ) -> np.ndarray:
        """
        Expects a numpy array shape Bx4. Requires the original image size
        in (H, W) format.
        """
        boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)
        return boxes.reshape(-1, 4)

    def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
        """
        Expects batched images with shape BxCxHxW and float format. This
        transformation may not exactly match apply_image. apply_image is
        the transformation expected by the model.
        """
        # Expects an image in BCHW format. May not exactly match apply_image.
        target_size = self.get_preprocess_shape(
            image.shape[0], image.shape[1], self.target_length
        )
        return F.interpolate(
            image, target_size, mode="bilinear", align_corners=False, antialias=True
        )

    def apply_coords_torch(
        self, coords: torch.Tensor, original_size: Tuple[int, ...]
    ) -> torch.Tensor:
        """
        Expects a torch tensor with length 2 in the last dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).to(torch.float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

    def apply_boxes_torch(
        self, boxes: torch.Tensor, original_size: Tuple[int, ...]
    ) -> torch.Tensor:
        """
        Expects a torch tensor with shape Bx4. Requires the original image
        size in (H, W) format.
        """
        boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)
        return boxes.reshape(-1, 4)

    @staticmethod
    def get_preprocess_shape(
        oldh: int, oldw: int, long_side_length: int
    ) -> Tuple[int, int]:
        """
        Compute the output size given input size and target long side length.
        """
        scale = long_side_length * 1.0 / max(oldh, oldw)
        newh, neww = oldh * scale, oldw * scale
        neww = int(neww + 0.5)
        newh = int(newh + 0.5)
        return (newh, neww)

def resize_longest_side(image, target_size):
    """
    调整图像大小，保持纵横比，将最长边缩放到 target_size
    
    参数:
        image: PIL.Image 对象
        target_size: 最长边的目标尺寸（整数）
    
    返回:
        调整后的 PIL.Image 对象
    """
    # 获取原始图像尺寸
    width, height = image.sresize_longest_sideize
    
    # 计算缩放比例
    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 PACO(Dataset):
    def __init__(self,
                 tokenizer=None,
                 data_args=None,
                 data_path="./data/DAM_data",
                 max_gt_per_img=5,
                 ):

        self.data_args = data_args
        self.tokenizer = tokenizer
        self.ann_file_ls = dict()
        self.img_prefix_ls = []
        # self.data_type=['PACO']
        self.data_type=['LVIS']
        for each in self.data_type:
            cur_anns_json=os.path.join(data_path,each,'annotations.json')
            cur_dict=json.load(open(cur_anns_json)) 
            prefixed_dict = {f"{each}_{key}": value for key, value in cur_dict.items()}
            # 检查当前字典与已有字典的重复键
            duplicate_keys = self.ann_file_ls.keys() & prefixed_dict.keys()  # 取交集
            
            if duplicate_keys:
                print(f"警告：数据集 {each} 与已有数据存在重复键：{duplicate_keys}")
            self.ann_file_ls |= prefixed_dict  
            print(f'DAM DATA:{each} len:{len(prefixed_dict.keys())}')
            cur_img=[]
            for img in prefixed_dict.keys():
                img=img.split('_',1)[1]
                cur_img.append(os.path.join(data_path,each,'images',img)+'.jpg')
            self.img_prefix_ls.extend(cur_img)
        assert len(self.img_prefix_ls)==len(self.ann_file_ls.keys())
        self.max_gt_per_img = max_gt_per_img
        self.begin_str = """<image>\nThis provides an overview of the picture.\n"""
        super().__init__()
    def __len__(self):
        return len(self.img_prefix_ls)
    def __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        # return data_item
        data_dict = self.process_text(data_item=data_item)
        return data_dict
    def get_data_item(self, idx):
        
        data_dict={}
        img_path=self.img_prefix_ls[idx]
        image = Image.open(img_path).convert('RGB')
        # 获取图片当前尺寸
        width, height = image.size
        
        # max_size = 1024
        # image = resize_longest_side(image, max_size)
        ann_key=img_path.split('/')[-3]+'_'+img_path.split('/')[-1].split('.')[0]
        annos=self.ann_file_ls[ann_key]
        mask_ls=[]
        for anno in annos:
            m=mask.decode(anno['mask_rle'])
            mask_ls.append(m)
        normalize = transforms.Compose([
            transforms.ToTensor(),                      # 转为Tensor并自动归一化到[0,1]
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],            # ImageNet均值
                std=[0.229, 0.224, 0.225]              # ImageNet标准差
            )
        ])
        image = normalize(image) 
        data_dict['img'] = image 
        data_dict['gt_masks'] = mask_ls
        data_dict['category'] =  [anno['category'] for anno in annos]
        data_dict['caption'] = [anno['caption'] for anno in annos]
        data_dict['img_path']=img_path
        assert len(data_dict['gt_masks']) == len(data_dict['caption']) 
        return data_dict
    def process_text(self, data_item):
        image = data_item['img']
        img_path = data_item['img_path']
        
        # ori_labels = data_item['caption']
        ori_masks = np.array(data_item['gt_masks'])
        ori_masks = torch.from_numpy(ori_masks) 
        ori_class = data_item['category']

        shuffle_ids = torch.randperm(len(ori_class))
        if len(shuffle_ids) > self.max_gt_per_img:
            shuffle_ids = shuffle_ids[:self.max_gt_per_img]
        ori_masks = ori_masks[shuffle_ids]
        # ori_labels = [ori_labels[i] for i in shuffle_ids]
        ori_class = [ori_class[i] for i in shuffle_ids]
        sources = dict()

        sources['conversations'] = []
        
        for i in range(len(ori_class)):
            question = 'Please identify the labels of each marked <region> in the image.'
            question = question.replace('<region>', '<mask><pos>')
            if i == 0:
                question = self.begin_str + question
            answer = ori_class[i]
            sources['conversations'].append(
                {'from': 'human', 'value': question})
            sources['conversations'].append({'from': 'gpt', 'value': answer})

        # print(image.shape)
        sources = preprocess_multimodal(
            copy.deepcopy([sources['conversations']]),
            self.data_args,
            0,if_qwen=True)
        text=sources[0][0]['value']
        processor=self.data_args.image_processor
        qwen_img_process=Image.open(img_path).convert('RGB')
        # 获取图片当前尺寸
        width, height = qwen_img_process.size
        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)
        # (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
        data_dict['labels']=new_target
        ############################
        data_dict['image'] = image
        data_dict['masks'] = ori_masks
        data_dict['pixel_values'] = inputs['pixel_values']
        data_dict['image_grid_thw'] = inputs['image_grid_thw']
        return data_dict
if __name__ == '__main__':
    dataset = PACO()
    total_samples = len(dataset)
    print(f"开始遍历数据集，共 {total_samples} 个样本...")
    dataloader = DataLoader(dataset, batch_size=1, num_workers=4)  # 多进程加速
    for i in range(total_samples):
        item = dataset[i]