import torch
from torchvision.datasets.vision import VisionDataset
from torchvision.datasets.folder import make_dataset, IMG_EXTENSIONS, default_loader
import os
import PIL
import torch.nn.functional as F
import random
import torchjpeg.dct


class RstPatchDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, block_size, transform=None, target_transform=None):
        super(RstPatchDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.block_size = block_size
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.target_transform is not None:
            target = self.target_transform(target)

        patches = ()
        for i in range(20):
            r = random.randint(0, 359)
            s = random.randint(90, 110)
            tx = random.randint(-self.block_size // 10, self.block_size // 10)
            ty = random.randint(-self.block_size // 10, self.block_size // 10)
            tr = sample.rotate(r, resample=PIL.Image.BILINEAR, translate=(tx, ty)).resize((sample.width * s // 100, sample.height * s // 100))
            px = (tr.width - self.block_size) // 2
            py = (tr.height - self.block_size) // 2
            tr = tr.crop((px, py, px + self.block_size, py + self.block_size))
            patches = patches + (self.transform(tr),)
        sample = torch.stack(patches)

        return sample, target


class MyDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, transform=None, target_transform=None):
        super(MyDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return sample, target


class MyDataFeed():
    def __init__(self, root: str, loader=default_loader, extensions=IMG_EXTENSIONS, train_transform=None, val_transform=None, target_transform=None):
        self.root = root
        self.train_transform = train_transform
        self.val_transform = val_transform
        self.target_transform = target_transform
        self.loader = loader
        self.extensions = extensions

        classes = [d.name for d in os.scandir(root) if d.is_dir()] + ['./']
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        self.samples = make_dataset(root, class_to_idx, extensions=self.extensions)

    def load(self, subset_ratio: float, train_ratio: float, block_size: int, seed: int = 8152513, logger = print):
        self.subset_ratio = subset_ratio
        self.train_ratio = train_ratio
        self.val_ratio = 1 - self.train_ratio
        self.block_size = block_size
        logger(f'train {100*subset_ratio:.2f}% of {self.root}, train-val: {100*self.train_ratio:.2f}%-{100*self.val_ratio:.2f}%')

        random.seed(seed)
        random.shuffle(self.samples)
        subset = self.samples[:int(len(self.samples) * subset_ratio)]
        train_set = subset[:int(len(subset) * train_ratio)]
        val_set = subset[int(len(subset) * train_ratio):]
        logger(f'train {len(train_set)}, test {len(val_set)} samples')

        train_dataset = MyDataset(self.root, train_set, self.loader, self.train_transform, self.target_transform)
        val_dataset = MyDataset(self.root, val_set, self.loader, self.val_transform, self.target_transform)

        return train_dataset, val_dataset


class PatchDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, block_size, transform=None, target_transform=None):
        super(PatchDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.block_size = block_size
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        sample = sample.unsqueeze(0)
        if sample.shape[2] < self.block_size or sample.shape[3] < self.block_size:
            sample = F.interpolate(sample, size=(self.block_size, self.block_size), mode='bilinear', align_corners=True)
        else:
            sample = torchjpeg.dct.blockify(sample, size=self.block_size).squeeze(0).permute(1, 0, 2, 3).contiguous()

        return sample, target


class PatchDataFeed():
    def __init__(self, root: str, loader=default_loader, extensions=IMG_EXTENSIONS, transform=None, target_transform=None):
        self.root = root
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        self.extensions = extensions

        classes = [d.name for d in os.scandir(root) if d.is_dir()] + ['./']
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        self.samples = make_dataset(root, class_to_idx, extensions=self.extensions)

    def load(self, subset_ratio: float, train_ratio: float, block_size: int, seed: int = 8152513, logger = print):
        self.subset_ratio = subset_ratio
        self.train_ratio = train_ratio
        self.val_ratio = 1 - self.train_ratio
        self.block_size = block_size
        logger(f'train {100*subset_ratio:.2f}% of {self.root}, train-val: {100*self.train_ratio:.2f}%-{100*self.val_ratio:.2f}%')

        random.seed(seed)
        random.shuffle(self.samples)
        subset = self.samples[:int(len(self.samples) * subset_ratio)]
        train_set = subset[:int(len(subset) * train_ratio)]
        val_set = subset[int(len(subset) * train_ratio):]
        logger(f'train {len(train_set)}, test {len(val_set)} samples')

        train_dataset = PatchDataset(self.root, train_set, self.loader, self.block_size, self.transform, self.target_transform)
        val_dataset = PatchDataset(self.root, val_set, self.loader, self.block_size, self.transform, self.target_transform)

        return train_dataset, val_dataset

import torch
from torchvision.datasets.vision import VisionDataset
from torchvision.datasets.folder import make_dataset, IMG_EXTENSIONS, default_loader
import os
import PIL
import torch.nn.functional as F
import random
import torchjpeg.dct


class RstPatchDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, block_size, transform=None, target_transform=None):
        super(RstPatchDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.block_size = block_size
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.target_transform is not None:
            target = self.target_transform(target)

        patches = ()
        for i in range(20):
            r = random.randint(0, 359)
            s = random.randint(90, 110)
            tx = random.randint(-self.block_size // 10, self.block_size // 10)
            ty = random.randint(-self.block_size // 10, self.block_size // 10)
            tr = sample.rotate(r, resample=PIL.Image.BILINEAR, translate=(tx, ty)).resize((sample.width * s // 100, sample.height * s // 100))
            px = (tr.width - self.block_size) // 2
            py = (tr.height - self.block_size) // 2
            tr = tr.crop((px, py, px + self.block_size, py + self.block_size))
            patches = patches + (self.transform(tr),)
        sample = torch.stack(patches)

        return sample, target


class MyDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, transform=None, target_transform=None):
        super(MyDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return sample, target


class MyDataFeed():
    def __init__(self, root: str, loader=default_loader, extensions=IMG_EXTENSIONS, train_transform=None, val_transform=None, target_transform=None):
        self.root = root
        self.train_transform = train_transform
        self.val_transform = val_transform
        self.target_transform = target_transform
        self.loader = loader
        self.extensions = extensions

        classes = [d.name for d in os.scandir(root) if d.is_dir()] + ['./']
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        self.samples = make_dataset(root, class_to_idx, extensions=self.extensions)

    def load(self, subset_ratio: float, train_ratio: float, block_size: int, seed: int = 8152513, logger = print):
        self.subset_ratio = subset_ratio
        self.train_ratio = train_ratio
        self.val_ratio = 1 - self.train_ratio
        self.block_size = block_size
        logger(f'train {100*subset_ratio:.2f}% of {self.root}, train-val: {100*self.train_ratio:.2f}%-{100*self.val_ratio:.2f}%')

        random.seed(seed)
        random.shuffle(self.samples)
        subset = self.samples[:int(len(self.samples) * subset_ratio)]
        train_set = subset[:int(len(subset) * train_ratio)]
        val_set = subset[int(len(subset) * train_ratio):]
        logger(f'train {len(train_set)}, test {len(val_set)} samples')

        train_dataset = MyDataset(self.root, train_set, self.loader, self.train_transform, self.target_transform)
        val_dataset = MyDataset(self.root, val_set, self.loader, self.val_transform, self.target_transform)

        return train_dataset, val_dataset


class PatchDataset(VisionDataset):
    def __init__(self, root: str, samples, loader, block_size, transform=None, target_transform=None):
        super(PatchDataset, self).__init__(root, transform=transform, target_transform=target_transform)

        self.samples = samples
        self.loader = loader
        self.block_size = block_size
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        sample = sample.unsqueeze(0)
        if sample.shape[2] < self.block_size or sample.shape[3] < self.block_size:
            sample = F.interpolate(sample, size=(self.block_size, self.block_size), mode='bilinear', align_corners=True)
        else:
            sample = torchjpeg.dct.blockify(sample, size=self.block_size).squeeze(0).permute(1, 0, 2, 3).contiguous()

        return sample, target


class PatchDataFeed():
    def __init__(self, root: str, loader=default_loader, extensions=IMG_EXTENSIONS, transform=None, target_transform=None):
        self.root = root
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        self.extensions = extensions

        classes = [d.name for d in os.scandir(root) if d.is_dir()] + ['./']
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        self.samples = make_dataset(root, class_to_idx, extensions=self.extensions)

    def load(self, subset_ratio: float, train_ratio: float, block_size: int, seed: int = 8152513, logger = print):
        self.subset_ratio = subset_ratio
        self.train_ratio = train_ratio
        self.val_ratio = 1 - self.train_ratio
        self.block_size = block_size
        logger(f'train {100*subset_ratio:.2f}% of {self.root}, train-val: {100*self.train_ratio:.2f}%-{100*self.val_ratio:.2f}%')

        random.seed(seed)
        random.shuffle(self.samples)
        subset = self.samples[:int(len(self.samples) * subset_ratio)]
        train_set = subset[:int(len(subset) * train_ratio)]
        val_set = subset[int(len(subset) * train_ratio):]
        logger(f'train {len(train_set)}, test {len(val_set)} samples')

        train_dataset = PatchDataset(self.root, train_set, self.loader, self.block_size, self.transform, self.target_transform)
        val_dataset = PatchDataset(self.root, val_set, self.loader, self.block_size, self.transform, self.target_transform)

        return train_dataset, val_dataset
