import os.path
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
import os
import os.path
import torch

from ..util.aug_utils import RandomScale, RandomSizeCrop, DivisibleCrop

class SingleImageDataset(Dataset):
    def __init__(self, cfg, src_img=None):
        self.cfg = cfg

        self.base_transforms = transforms.Compose(
            [
                transforms.Lambda(lambda x: transforms.ToTensor()(x).unsqueeze(0)),
                DivisibleCrop(cfg["d_divisible_crops"]),
            ]
        )

        # used to create the internal dataset
        self.input_transforms = transforms.Compose(
            [
                transforms.RandomApply(
                    [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1)],
                    p=cfg["jitter_p"],
                ),
                transforms.RandomHorizontalFlip(p=cfg["flip_p"]),
                RandomScale((cfg["scale_min"], cfg["scale_max"])),
                RandomSizeCrop(cfg["crops_min_cover"]),
                self.base_transforms,
            ]
        )

        # open source image
        self.src_img = src_img

        if cfg["resize_input"] > 0:
            self.src_img = transforms.Resize(cfg["resize_input"])(self.src_img)

        self.step = -1

    def get_img(self):
        return self.base_transforms(self.src_img)

    def __getitem__(self, index):
        self.step += 1
        sample = {"step": self.step}
        if self.step % self.cfg["source_image_every"] == 0:
            sample["input_image"] = self.get_img()

        sample["input_crop"] = self.input_transforms(self.src_img)

        return sample

    def __len__(self):
        return 1

    def update_img(self, src_img: Image.Image):
        """
        PIL.Image.Image: source image, should be converted to RGB first.
        """
        self.src_img = src_img
        if self.cfg["resize_input"] > 0:
            self.src_img = transforms.Resize(self.cfg["resize_input"])(self.src_img)
        self.step = -1
