from dataclasses import dataclass
import torch
import transformers
from torch.utils.data import ConcatDataset
import json
from FCLM.constants import IGNORE_INDEX

from .stage2_data import COCODataset, RefCOCO, RefCOCOP
from .vcr import VCRDataset
from .vg import VGDATA
from .stage2_data import PascalPart
from .stage2_data import PartImagenet
from .osprey_724k import OspreyDetailedDescription, OspreyConversations, OspreyShortForm, OspreyPartLevel, OspreyLVISPosNeg
from .sa1b import SA1BDATA
from .reason_seg import REASON_SEG
from .refer_seg_dataset import ReferSegDataset
from .DAM_data import DAM_DATA,PACO_LVIS_Object_Class_dataset
from .GCG_data import GCGBaseDataset,GranDfDataset,OpenPsgGCGDataset,Flickr30kGCGDataset,RefCOCOgGCGDataset
from .muse import MultiReasonSegDataset
from .Pixel2cap import Pixel2Cap_DATASET
@dataclass
class DataCollatorForDetDataset(object):

    tokenizer: transformers.PreTrainedTokenizer
    def __call__(self, instances):
        #### Qwen
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            print("Fixed Qwen Pad Token ID:", self.tokenizer.pad_token_id)
        
        if isinstance(instances[0], tuple):
            ins_DLC=instances[0][0]
            ins_DLRES=instances[0][1]
            instances=[ins_DLC,ins_DLRES]
            input_ids, labels, img_metas, masks, pixel_values, image_grid_thw = tuple([instance.get(key,None) for instance in instances]
                                  for key in ('input_ids',
                                              'labels',
                                              'img_metas',
                                              'masks',
                                              'pixel_values',
                                              'image_grid_thw'))
        else:
            input_ids, labels, img_metas, masks, pixel_values, image_grid_thw = tuple([instance.get(key,None) for instance in instances]
                                    for key in ('input_ids',
                                                'labels',
                                                'img_metas',
                                                'masks',
                                                'pixel_values',
                                                'image_grid_thw'))
        input_ids_bs=[]
        for each in input_ids:
            input_ids_bs.extend(each)
        labels_bs=[]
        for each in labels:
            labels_bs.extend(each)
        # img_metas_bs=[]
        img_metas_bs=dict()
        img_metas_bs['image']=[]
        img_metas_bs['sample_num']=[]
        img_metas_bs['SEG_num']=[]
        for each in img_metas:
            if each is not None:
                img_metas_bs['image'].extend(each['image'])
                img_metas_bs['sample_num'].append(each['sample_num'])
                img_metas_bs['SEG_num'].append(each['SEG_num'])
        masks_bs=[]
        # 
        for each in masks:
            # masks_bs.extend(list(torch.unbind(each, dim=0))) #grounding
            masks_bs.append(each)
        pixel_values_bs=[]
        for each in pixel_values:
            if len(each.shape)==3: #grounding
                # pixel_values_bs.extend(list(torch.unbind(each, dim=0)))  
                pixel_values_bs.append(each[0]) #一个图只保留一个 pixel
            if len(each.shape)==2: #caption  
                pixel_values_bs.append(each)  
        image_grid_thw_bs=[]
        for each in image_grid_thw:
            image_grid_thw_bs.append(each[0])
        input_ids_bs = torch.nn.utils.rnn.pad_sequence(
            input_ids_bs,
            batch_first=True,
            padding_value=self.tokenizer.pad_token_id)
        labels_bs = torch.nn.utils.rnn.pad_sequence(labels_bs,
                                                 batch_first=True,
                                                 padding_value=IGNORE_INDEX)
        batch = dict(
            input_ids=input_ids_bs,
            labels=labels_bs,
            attention_mask=input_ids_bs.ne(self.tokenizer.pad_token_id),
            img_metas=img_metas_bs,
            masks = masks_bs,
            pixel_values = pixel_values_bs,
            image_grid_thw = image_grid_thw_bs
        )
        if 'image' in instances[0]:
            images = [instance['image'] for instance in instances]
            # if all(x is not None and x.shape == images[0].shape for x in images):
            #     batch['images'] = torch.cat(images,dim=0)
            # else:
            batch['images'] = images
        return batch

def make_multitask_data_module(tokenizer,
                                data_args) :
    """Make dataset and collator for supervised fine-tuning."""

    if data_args.dataset_config is not None:
        dataset_config = json.load(open(data_args.dataset_config))

    train_dataset_caption, train_dataset_grounding = build_osprey_dataset(dataset_config,
                            tokenizer=tokenizer,
                            data_args=data_args)
    import numpy as np
    caption_len=int(np.array([len(each) for each in train_dataset_caption]).sum())
    grounding_len=int(np.array([len(each) for each in train_dataset_grounding]).sum())
    train_dataset = ConcatDataset(train_dataset_caption + train_dataset_grounding)
    data_collator = DataCollatorForDetDataset(tokenizer=tokenizer)
    data_len=[caption_len,grounding_len]
    return dict(train_dataset=train_dataset,
                eval_dataset=None,
                data_collator=data_collator,
                data_len=data_len)

def build_osprey_dataset(dataset_config,
                  tokenizer=None,
                  data_args=None,
                  **kwargs):
    if isinstance(dataset_config, list):
        datasets_caption = []
        datasets_grounding = []
        for cfg in dataset_config:
            if cfg['type'] in ["refer_seg","reason_seg",'GCG_Grandf','GCG_PSG','GCG_Flickr30k','GCG_RefCOCOg','MUSE','pixel2cap']:
                temp_dataset=build_osprey_dataset(cfg, tokenizer=tokenizer, data_args=data_args, **kwargs)
                datasets_grounding.append(temp_dataset)
            else:
                temp_dataset = build_osprey_dataset(cfg, tokenizer=tokenizer, data_args=data_args, **kwargs)
                datasets_caption.append(temp_dataset)

        for dataset in datasets_grounding:
            print('grouding dataset ', type(dataset), f'len = {len(dataset)}')
        for dataset in datasets_caption:
            print('caption dataset ', type(dataset), f'len = {len(dataset)}')

        # return ConcatDataset(datasets)
        return datasets_caption, datasets_grounding
    dataset_type = dataset_config.pop('type')

    if dataset_type == 'coco_data':
        dataset = COCODataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    
    elif dataset_type == 'vcr':
        dataset = VCRDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'VGDATA':
        dataset = VGDATA(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'RefCOCO':
        dataset = RefCOCO(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'RefCOCOP':
        dataset = RefCOCOP(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'PascalPart':
        dataset = PascalPart(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'PartImagenet':
        dataset = PartImagenet(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'OspreyDetailedDescription':
        dataset = OspreyDetailedDescription(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'OspreyConversations':
        dataset = OspreyConversations(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'OspreyShortForm':
        dataset = OspreyShortForm(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'OspreyPartLevel':
        dataset = OspreyPartLevel(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'OspreyLVISPosNeg':
        dataset = OspreyLVISPosNeg(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'SA1BDATA':
        dataset = SA1BDATA(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'reason_seg':
        dataset = REASON_SEG(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type == 'refer_seg':
        dataset = ReferSegDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='DAM_DATA':
        dataset = DAM_DATA(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='GCG_Grandf':
        dataset = GranDfDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='GCG_PSG':
        dataset = OpenPsgGCGDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='GCG_Flickr30k':
        dataset = Flickr30kGCGDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='GCG_RefCOCOg':
        dataset = RefCOCOgGCGDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='MUSE':
        dataset = MultiReasonSegDataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='pixel2cap':
        dataset = Pixel2Cap_DATASET(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    elif dataset_type =='paco_lvis_obj':
        dataset = PACO_LVIS_Object_Class_dataset(
            **dataset_config,
            tokenizer=tokenizer,
            data_args=data_args,
            **kwargs,
        )
    else:
        raise NotImplementedError

    return dataset



class ConcatDataset(ConcatDataset):
    def __init__(self, datasets):
        super().__init__(datasets)

    def collater(self, samples):

        all_keys = set()
        for s in samples:
            all_keys.update(s)

        shared_keys = all_keys
        for s in samples:
            shared_keys = shared_keys & set(s.keys())

        samples_shared_keys = []
        for s in samples:
            samples_shared_keys.append({k: s[k] for k in s.keys() if k in shared_keys})

        return self.datasets[0].collater(samples_shared_keys)

