from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms
from typing import List, Tuple
import imgaug.augmenters as iaa
import numpy as np
from sklearn.model_selection import train_test_split

class UnNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        """
        Args:
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
        Returns:
            Tensor: Normalized image.
        """
        for t, m, s in zip(tensor, self.mean, self.std):
            t.mul_(s).add_(m)
            # The normalize code -> t.sub_(m).div_(s)
        return tensor

class DocCleanDataset(Dataset):

    @staticmethod
    def prepareDataset(dataset:str, shuffle=True):
        # imgs_dir = "dataset/raw_data/imgs_Trainblocks"
        with open(f"{dataset}/train_block_names.txt") as train_block_names_file:
            image_names = train_block_names_file.read().splitlines()
            train_img_names, eval_img_names, _, _ = train_test_split(
                image_names, image_names, test_size=0.2, random_state=1, shuffle=shuffle)
            return train_img_names, eval_img_names, dataset

    def __init__(self, img_names: List[str], imgs_dir: str, normalized_tuple: Tuple[List[float], List[float]] = None, dev=False, img_aug=False):
        if dev:
            num = int(len(img_names) * 0.01)
            img_names = img_names[0:num]
        self.img_names = img_names
        self.imgs_dir = imgs_dir
        if normalized_tuple:
            mean, std = normalized_tuple
            self.normalized = transforms.Compose([
                transforms.ToTensor(), 
                transforms.Normalize(mean=mean, std=std)
            ])
            self.aug_seq = iaa.Sometimes(0.7, iaa.OneOf([
                    iaa.SaltAndPepper(p=(0.0, 0.05)),
                    iaa.imgcorruptlike.MotionBlur(severity=2),
                    iaa.SigmoidContrast(gain=(3, 10), cutoff=(0.4, 0.6)),
                    iaa.imgcorruptlike.JpegCompression(severity=2),
                    iaa.GammaContrast((0.5, 2.0)),
                    iaa.LogContrast(gain=(0.5, 0.9)),
                    iaa.GaussianBlur(sigma=(0, 1)),
                    iaa.imgcorruptlike.SpeckleNoise(severity=1),
                    iaa.AdditiveGaussianNoise(scale=(0.03*255, 0.2*255), per_channel=True),
                    iaa.Add((-20, 20), per_channel=0.5),
                    iaa.AddToBrightness((-30, 30))
                ]))
        self.img_aug = img_aug
        self.toTensor = transforms.ToTensor()

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

    def __getitem__(self, index):
        img = Image.open(f"{self.imgs_dir}/{self.img_names[index]}")
        gt = Image.open(f"{self.imgs_dir}/gt{self.img_names[index]}")
        if hasattr(self, 'normalized'):
            img_np = np.array(img)
            if self.img_aug == True:
                img_np = self.aug_seq.augment_images([np.array(img)])[0]
            normalized_img = self.normalized(img_np)
            img = self.toTensor(img_np)
        else:
            img = self.toTensor(img)
            normalized_img = img
        return img, normalized_img, self.toTensor(gt)