import copy
from dataclasses import dataclass, field, fields, asdict
import json
import logging
import pathlib
from typing import Dict, Optional, Sequence, List
import sys
import torch

import transformers
import gc

from PIL import Image
import numpy as np
import os
from qwen_vl_utils import process_vision_info
from qwen_vl_utils import fetch_image, fetch_video

@dataclass
class DexVLADataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    multimodal_processor: transformers.AutoProcessor=None
    computed_type: torch.dtype=None
    tokenizer: transformers.AutoTokenizer=None
    video: bool=False

    # @profile
    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids = [torch.flip(instance['input_ids'].squeeze(0), dims=[0]) for instance in instances]
        attention_mask = [torch.flip(instance['attention_mask'].squeeze(0), dims=[0]) for instance in instances]
        labels = [torch.flip(instance['labels'].squeeze(0), dims=[0]) for instance in instances]
        raw_images = torch.stack([instances['raw_images'] for instances in instances])
        if self.video:
            video_grid_thw = torch.stack([instances['video_grid_thw'] for instances in instances])
            pixel_values_videos = torch.stack([instances['pixel_values_videos'] for instances in instances])
            pixel_values = None
            image_grid_thw=None
        else:
            image_grid_thw = torch.stack([instances['image_grid_thw'] for instances in instances])
            pixel_values = torch.stack([instances['pixel_values'] for instances in instances])
            pixel_values_videos = None
            video_grid_thw = None

        labels = torch.nn.utils.rnn.pad_sequence(labels,
                                                 batch_first=True,
                                                 padding_value=-100)
        labels = torch.flip(labels, dims=[1]) # left padding
        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids,
                                                    batch_first=True,
                                                    padding_value=self.tokenizer.pad_token_id)
        input_ids = torch.flip(input_ids, dims=[1])
        b = input_ids.shape[0]
        if self.video:
            video_grid_thw = video_grid_thw.reshape(b * video_grid_thw.shape[1], video_grid_thw.shape[2])
            pixel_values_videos = pixel_values_videos.reshape(b * pixel_values_videos.shape[1], pixel_values_videos.shape[2])

        else:
            image_grid_thw = image_grid_thw.reshape(b * image_grid_thw.shape[1], image_grid_thw.shape[2])
            pixel_values = pixel_values.reshape(b * pixel_values.shape[1], pixel_values.shape[2])

        attention_mask = input_ids.ne(self.tokenizer.pad_token_id),
        # attention_mask = torch.nn.utils.rnn.pad_sequence(labels,
        #                                          batch_first=True,
        #                                          padding_value=1)

        # max_length = max([each.shape[-1] for each in input_ids])
        # pad_id = self.tokenizer.pad_token_id
        # for idx,_ in enumerate(input_ids):
        #     length = input_ids[idx].shape[-1]
        #     padd = torch.ones((1, max_length-length), dtype=torch.long, device=input_ids[idx].device)
        #     input_ids[idx] = torch.cat((padd*pad_id,input_ids[idx]), dim=-1)
        #     attention_mask[idx] = torch.cat((padd,attention_mask[idx]), dim=-1)
        #     labels[idx] = torch.cat((padd*-100,labels[idx]), dim=-1)
            
        if not isinstance(instances[0]['action'], torch.Tensor):
            actions = torch.tensor(np.array([instance['action'] for instance in instances]))
            states = torch.tensor(np.array([instance['state'] for instance in instances]))
        else:
            actions = torch.stack([instance['action'] for instance in instances])
            states = torch.stack([instance['state'] for instance in instances])

        is_pad_all = torch.stack([instance['is_pad'] for instance in instances])
        
        #print("#"*60)
        #print(attention_mask.shape)
        #exit(0)
        batch = dict(
            input_ids=input_ids,
            # token_type_ids=model_inputs['token_type_ids'],
            raw_images=raw_images,
            attention_mask=attention_mask[0],
            labels=labels,
            image_grid_thw=image_grid_thw,
            pixel_values_videos=pixel_values_videos,
            actions=actions,
            states=states,
            video_grid_thw=video_grid_thw,
            pixel_values=pixel_values,
            is_pad=is_pad_all,
            # attention_mask=input_ids.ne(temp_pad_token_id),
        )
        del input_ids
        del attention_mask
        del labels
        del pixel_values_videos
        del pixel_values
        del actions
        del states
        del video_grid_thw
        del image_grid_thw
        del is_pad_all
        gc.collect()
        torch.cuda.empty_cache()
        return batch


@dataclass
class PaliGemmaVLADataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    multimodal_processor: transformers.AutoProcessor = None
    computed_type: torch.dtype = None

    # @profile
    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:

        prompt = "Task:"
        raw_langs = [prompt + ins['raw_lang'] for ins in instances]

        images = torch.stack([ins['image'] for ins in instances])

        answers = [ins['reasoning'] for ins in instances]
        # answers = ["aaa" ,'bbb asdasda asda']
        model_inputs = self.multimodal_processor(text=raw_langs, suffix=answers, images=images, return_tensors="pt", padding="longest")

        pixel_values = copy.deepcopy(model_inputs['pixel_values'])
        if not isinstance(instances[0]['action'], torch.Tensor):
            actions = torch.tensor(np.array([instance['action'] for instance in instances]))
            states = torch.tensor(np.array([instance['state'] for instance in instances]))
        else:
            actions = torch.stack([instance['action'] for instance in instances])
            states = torch.stack([instance['state'] for instance in instances])

        is_pad_all = torch.stack([instance['is_pad'] for instance in instances])

        batch = dict(
            input_ids=model_inputs['input_ids'],
            token_type_ids=model_inputs['token_type_ids'],
            attention_mask=model_inputs['attention_mask'],
            labels=model_inputs['labels'],
            actions=actions,
            states=states,
            pixel_values=pixel_values,
            is_pad=is_pad_all,
            # attention_mask=input_ids.ne(temp_pad_token_id),
        )

        del model_inputs
        del pixel_values
        del actions
        del states
        del is_pad_all
        gc.collect()
        torch.cuda.empty_cache()
        return batch
