import io
import os
import copy
import json
import logging
import torch
import transformers
from typing import List, Optional, Tuple, Union, Dict, Sequence
from torch.utils.data import Dataset
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
from vary.utils.constants import *

"""
这段代码定义了一个名为 BaseDataset 的基础数据集类，它继承自 PyTorch 的 Dataset 类。这个类主要用于处理和存储数据，以便于后续的模型训练。
"""

class BaseDataset(Dataset):
    def __init__(
        self, 
        datasets: str, # 是数据集的路径
        tokenizer: transformers.PreTrainedTokenizer,  # 用于分词的tokenizer
        multimodal_cfg: dict # 多模态配置字典，包含不同模态的配置信息
    ):  
        """
        初始化BaseDataset类的实例。
        
        参数:
        - datasets: str，数据集的路径。
        - tokenizer: transformers.PreTrainedTokenizer，用于对文本进行分词的预训练tokenizer。
        - multimodal_cfg: dict，包含多种模态配置的字典，例如图像和文本的配置。
        """
        
        super(BaseDataset, self).__init__()
        self.tokenizer = tokenizer
        self.multimodal_cfg = multimodal_cfg

         # 记录使用多少个tokens来表示图像
        logging.warning(f"Using {multimodal_cfg['image_token_len']} tokens for representing image")

    def image_processor(self, image):
        """
        image_processor 方法是用来处理图像的。这个方法接收一个图像作为输入，然后根据 multimodal_cfg 中的配置信息来处理这个图像。
        首先，获取两个图像处理器：processor 和 processor_high。然后，根据 multimodal_cfg 中的 image_aspect_ratio 配置来决定如何处理图像。
        如果 image_aspect_ratio 是 "keep"，那么会保持图像的长宽比；如果 image_aspect_ratio 是 "pad"，那么会将图像扩展到正方形；
        如果 image_aspect_ratio 是其他值，那么会直接处理图像。最后，使用 processor_high 来处理复制的图像，并返回处理后的两个图像。
        """
        """
        处理图像，通过两个处理器对图像进行处理：第一个处理器通常为预训练的模型（如CLIP、VIT），第二个处理器通常为设计的图像编码器（如SAM、SWIN、CNN）。

        参数:
        - image: 图像对象，待处理的图像。

        返回值:
        - image: 经过第一个处理器处理后的图像张量。
        - image_high: 经过第二个处理器处理后的图像张量。
        """
        processor = self.multimodal_cfg['image_processor']  # the first processor, usually is the clip pretrained model (vit)
        processor_high = self.multimodal_cfg['image_processor_high'] # the second processor, usually is the designed image encoder (sam/swin/cnn)
        image_high = image.copy()
        # TODO the 'keep', 'padding' only used for the first processor
        # 根据配置决定如何处理图像以保持纵横比或填充到固定大小
        if self.multimodal_cfg['image_aspect_ratio'] == 'keep':
            # 保持纵横比处理图像
            max_hw, min_hw = max(image.size), min(image.size)
            aspect_ratio = max_hw / min_hw
            max_len, min_len = 448, 224
            shortest_edge = int(min(max_len / aspect_ratio, min_len))
            # 将图像转换为张量
            image = processor.preprocess(image, return_tensors='pt', do_center_crop=False, size={"shortest_edge": shortest_edge})['pixel_values'][0]
        elif self.multimodal_cfg['image_aspect_ratio'] == 'pad':
             # 通过填充方式处理图像，保持为正方形
            def expand2square(pil_img, background_color):
                width, height = pil_img.size
                if width == height:
                    return pil_img
                elif width > height:
                    result = Image.new(pil_img.mode, (width, width), background_color)
                    result.paste(pil_img) # for simpler box processing
                    return result
                else:
                    result = Image.new(pil_img.mode, (height, height), background_color)
                    result.paste(pil_img) # for simpler box processing
                    return result
            image = expand2square(image, tuple(int(x*255) for x in processor.image_mean))
            image = processor.preprocess(image, return_tensors='pt', do_center_crop=False, size={"shortest_edge": 224})['pixel_values'][0]
        else:
            image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
         # 使用第二个处理器处理图像
        image_high = processor_high(image_high)

        return image, image_high
    


    def __len__(self):
        return len(self.list_data_dict)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        pass