import PIL
from PIL import Image
from dataclasses import dataclass, field
from datasets import load_dataset
import torch
from .data_processing import pil_to_tensor


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

    def __init__(self, dataset_name, **kwargs):
        override_params = {}
        if dataset_name == "DIV2K":
            override_params = {
                "target_image_size": -1,
                "lock_ratio": True,
                "center_crop": False,
                "padding": False,
            }
        if dataset_name == "imagenet":
            override_params = {"center_crop": True, "padding": False}
        if dataset_name == "movie_posters":
            override_params = {"center_crop": True, "padding": False}
        if dataset_name == "high_quality_1024":
            override_params = {"target_image_size": (1024, 1024)}

        self.data_params = {**kwargs, **override_params}

    def __call__(self, instances):
        images = torch.stack(
            [
                pil_to_tensor(instance["image"], **self.data_params)
                for instance in instances
            ],
            dim=0,
        )
        idx = [instance["idx"] for instance in instances]
        return dict(image=images, idx=idx)


class ImagenetDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_name, split_name="test", n_take=None):
        print(dataset_name, split_name)
        ds = load_dataset("huaweilin/VTBench", name=dataset_name, split=split_name if n_take is None else f"{split_name}[:{n_take}]")
        self.image_list = ds["image"]

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

    def __getitem__(self, idx):
        return dict(
            image=self.image_list[idx],
            idx=idx,
        )


def get_dataset(dataset_name, split_name, n_take):
    dataset = ImagenetDataset(dataset_name, split_name, n_take)
    return dataset