import copy
import random
import os
import numpy as np
import torch
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
LIMIT = " Answer the questions using Chinese phrases."
QUESTIONS =  [
    'Give me a short description of <region>.',
    'Can you give me a short description of <region>?',
    'Can you provide me with a short description of the region in the picture marked by <region>?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in few words?",
    'What can you tell me about the region indicated by <region> in the image in few words?',
    "I'd like to know more about the area in the photo labeled <region>. Can you give me a concise description?",
    'Could you describe the region shown as <region> in the picture concisely?',
    'What can you give me about the region outlined by <region> in the photo?',
    'Please provide me with a brief description of the region marked with <region> in the image.',
    'Can you give me a brief introduction of the region labeled as <region> in the picture?',
    "I'm interested in knowing the region represented by <region> in the photo. Can you describe it in several words?",
    'What is the region outlined by <region> in the picture like? Could you give me a streamlined description?',
    'Can you provide me with a brief description of the region in the picture marked by <region>, please?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in few words, please?",
    'What can you tell me about the region indicated by <region> in the image?',
    "I'd like to know more about the area in the photo labeled <region>, please. Can you give me a simple description?",
    'Could you describe the region shown as <region> in the picture in several words?',
    'Please provide me with a simple description of the region marked with <region> in the image, please.',
    "I'm interested in learning more about the region represented by <region> in the photo. Can you describe it in few words, please?",
    'What is the region outlined by <region> in the picture like, please? Could you give me a simple and clear description?',
    'Please describe the region <region> in the image concisely.',
    'Can you offer a simple analysis of the region <region> in the image?',
    'Could tell me something about the region highlighted by <region> in the picture briefly?',
    'Can you share a simple rundown of the region denoted by <region> in the presented image?'
]

class SA1BDATA(Dataset):
    def __init__(self,
                 tokenizer=None,
                 data_args=None,
                 parquet_file=None,
                 ann_file=None,
                 img_prefix=None,
                 max_gt_per_img=5,
                 ):

        self.data_args = data_args
        self.tokenizer = tokenizer
        self.ann_file = ann_file
        self.img_prefix = img_prefix
        self.max_gt_per_img = max_gt_per_img
        self.begin_str = """<image>\nThis provides an overview of the picture.\n"""
        self.dataset_pd = pd.read_parquet(parquet_file).sample(n=300000, random_state=42)  #.head(1000) #.sample(n=1000, random_state=42)
        super().__init__()
    def __len__(self):
        return len(self.dataset_pd )
    def __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        data_dict = self.process_text(data_item=data_item)
        return data_dict
    def get_data_item(self, idx):
        data_dict={}
        rows = self.dataset_pd.iloc[idx]
        img_path=self.img_prefix + '/' + rows['opensource_url'].split('/')[-1]
        image = Image.open(img_path).convert('RGB')
        mask_ls=[each.split('/')[-1] for each in rows['matched_local_urls']]
        for i,each in enumerate(mask_ls):
            each=each.split('.')[0]+'.npz'
            mask_ls[i]=self.ann_file+'/'+each
        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'] = [np.load(each)['mask'] for each in mask_ls]
        data_dict['global_caption'] = rows['global_caption']
        data_dict['local_captions'] = ast.literal_eval(rows['local_captions'])
        data_dict['gt_labels'] = ast.literal_eval(rows['matched_local_captions'])
        data_dict['img_path']=img_path
        assert len(data_dict['gt_masks']) == len(data_dict['gt_labels']) 
        return data_dict
    def process_text(self, data_item):
        image = data_item['img']
        img_path = data_item['img_path']
        
        ori_labels = data_item['gt_labels']
        ori_masks = np.array(data_item['gt_masks'])
        ori_masks = torch.from_numpy(ori_masks) 

        shuffle_ids = torch.randperm(len(ori_labels))
        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]

        sources = dict()

        sources['conversations'] = []
        
        for i in range(len(ori_labels)):
            question = random.choice(QUESTIONS).strip()
            question = question.replace('<region>', '<mask><pos>')
            if i == 0:
                question = self.begin_str + question
            question += LIMIT
            answer = ori_labels[i]
            sources['conversations'].append(
                {'from': 'human', 'value': question})
            sources['conversations'].append({'from': 'gpt', 'value': answer})
        
        
        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')
        W,H=qwen_img_process.size
        qwen_img_process=qwen_img_process.resize((W//2,H//2))
        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
        #############################
        # get single
        # if isinstance(i, int):
        #     data_dict = dict(input_ids=data_dict['input_ids'][0],
        #                      labels=data_dict['labels'][0])
        H,W=image.shape[1],image.shape[2]
        image=torch.nn.functional.interpolate(image.unsqueeze(0),
                                                size=(H//2,W//2),
                                                mode='bilinear',
                                                align_corners=False).squeeze(0)
        data_dict['image'] = image
        ori_masks=torch.nn.functional.interpolate(ori_masks.float().unsqueeze(0),
                                                size=(H//2,W//2),
                                                mode='nearest'
                                                ).squeeze(0)
        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 = SA1BDATA(
        parquet_file='./data/SA1B/data/train-0000-of-0013.parquet',
        ann_file='./data/SA1B/data/mask_output',
        img_prefix="./data/SA1B/data/img")
    dataset[2]
    dataloader = DataLoader(dataset, batch_size=32, num_workers=4)  # 多进程加速