import torch
import os
import random
import json
import math
from warnings import warn
import os.path as osp
import numpy as np
from PIL import Image, ImageDraw
from torch.utils.data import Dataset
import torchvision.transforms as T

class CustomDataset(Dataset):
    def __init__(
        self, 
        json_file: list = ['dataset.json'],
        condition_size: int = 512,
        target_size: int = 512,
        condition_type = ["subject"],
        return_pil_image: bool = False,
        ):
        """
        初始化数据集。

        参数:
            json_file (list): JSON 文件的路径。
            transform (callable, optional): 可选的图像变换操作。
        """
        # 加载 JSON 文件
        if isinstance(json_file, list):
            merged_data = []
            for f in json_file:
                with open(f, 'r') as fp:
                    merged_data.extend(json.load(fp))
            self.data = merged_data
        else:
            with open(json_file, 'r') as f:
                self.data = json.load(f)

    def __len__(self):
        """返回数据集的大小。"""
        return len(self.data)

    def __getitem__(self, idx):
        """根据索引获取一个样本。"""
        item = self.data[idx]
        return_dict = {
            'image_path': item['image']['image_path'],
            'ref_image_path': item['ref_images'][0]['image_path_with_box'],
            'prompt': item['decoupled_prompt'],
            "subject": item["metadata"]["object_name"],
        }
        return return_dict

class TestDataset(CustomDataset):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __getitem__(self, idx):
        return self.data[idx]

class LatentDataset(Dataset):
    def __init__(
        self,
        json_file: str = 'dataset.json',
        condition_type = ["subject"],
        condition_size: int = 512,
        target_size: int = 512,
        ):
        """
        初始化数据集。

        参数:
            json_file (list): JSON 文件的路径。
        """
        self.condition_type = condition_type
        self.data_root = os.path.dirname(json_file)
        self.transform = T.Compose([
            T.Resize((condition_size, condition_size)),  # 调整大小
            T.ToTensor()           # 转换为张量
        ])
        with open(json_file, 'r') as f:
            self.data = json.load(f)

    def __len__(self):
        """返回数据集的大小。"""
        return len(self.data)

    def __getitem__(self, idx):
        """根据索引获取一个样本。"""
        item = self.data[idx]

        if item['ref_image_path']:
            ref_image_path = item['ref_image_path']
            ref_image = Image.open(ref_image_path).convert('RGB')

        # 应用图像变换
        if self.transform:
            ref_image = self.transform(ref_image)


        prompt_embed = torch.load(
            os.path.join(self.data_root, item["prompt_embed_path"]),
            map_location="cpu",
            weights_only=True,
        )
        pooled_prompt_embeds = torch.load(
            os.path.join(
                self.data_root, item["pooled_prompt_embeds_path"]
            ),
            map_location="cpu",
            weights_only=True,
        )
        text_ids = torch.load(
            os.path.join(
                self.data_root, item["text_ids_path"]
            ),
            map_location="cpu",
            weights_only=True,
        )
        subject = item["subject"]

        return prompt_embed, pooled_prompt_embeds, text_ids, ref_image, item['prompt'], item["original_image_path"], item["ref_image_path"], subject