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 DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    computed_type: torch.dtype=None
    tokenizer: transformers.AutoTokenizer=None

    # @profile
    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids = [instance['input_ids'].squeeze(0) for instance in instances]
        pixel_values = torch.stack([instances['pixel_values'] for instances in instances])

        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids,
                                                    batch_first=True,
                                                    padding_value=self.tokenizer.pad_token_id)

        attention_mask = input_ids.ne(self.tokenizer.pad_token_id),
            
        if not isinstance(instances[0]['actions'], torch.Tensor):
            actions = torch.tensor(np.array([instance['actions'] for instance in instances]))
            states = torch.tensor(np.array([instance['states'] for instance in instances]))
        else:
            actions = torch.stack([instance['actions'] for instance in instances])
            states = torch.stack([instance['states'] for instance in instances])

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

        batch = dict(
            input_ids=input_ids,
            attention_mask=attention_mask[0],
            actions=actions,
            states=states,
            pixel_values=pixel_values,
            is_pad=is_pad_all,
        )
        del input_ids
        del attention_mask
        del pixel_values
        del actions
        del states
        del is_pad_all
        gc.collect()
        torch.cuda.empty_cache()
        return batch