import torch
import torchvision.transforms as transforms
from torchvision.datasets import DatasetFolder
from torchvision.datasets.folder import (
    make_dataset,
    pil_loader,
    default_loader,
    IMG_EXTENSIONS,
)
from torchvision.transforms.functional import rotate


class DatasetFolderMultiTransform(DatasetFolder):
    """Adapts class DatasetFolder of PyTorch in a way that one sample is transformed several times.
    Args:
        n_transforms (int): number of transformations per sample
        all others: see torchvision.datasets.DatasetFolder
    """

    def __init__(
        self,
        conf: dict,
        root,
        loader,
        extensions=None,
        transform=None,
        target_transform=None,
        is_valid_file=None,
        n_transforms=1,
    ):
        super().__init__(
            root,
            loader,
            extensions=extensions,
            transform=transform,
            target_transform=target_transform,
        )
        classes, class_to_idx = self._find_classes(self.root)
        if is_valid_file is not None:
            extensions = None
        self.samples = make_dataset(self.root, class_to_idx, extensions)
        self.n_transforms = n_transforms
        self.get_fixed = False  # set to true if the rotations should be fixed and regularly over 360 degrees
        self.fixed_degrees = [i * 360.0 / n_transforms for i in range(n_transforms)]
        self.conf = conf

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            samples = list()
            for i in range(self.n_transforms):
                if self.get_fixed:
                    # print(f"i={i}: calling fixed_rotation({sample}, {self.fixed_degrees[i]}))")
                    samples.append(self.fixed_rotation(sample, self.fixed_degrees[i]))
                else:
                    new = self.transform(sample)
                    samples.append(new)
                    # print(f"i={i}: calling transform({sample})")
            samples = torch.stack(samples, dim=0)
        if self.target_transform is not None:
            # print(f"calling target_transform({target})")
            target = self.target_transform(target)
        return samples, target

    def fixed_rotation(self, sample, degrees):
        cust_rot = lambda x: rotate(x, degrees, False, False, None)
        augmentative_transforms = [cust_rot]
        if (
            self.conf.get("transf_brightness") > 0.0
            or self.conf.get("transf_contrast") > 0.0
            or self.conf.get("transf_saturation") > 0.0
        ):
            augmentative_transforms += [
                transforms.ColorJitter(
                    brightness=self.conf.get("transf_brightness"),
                    contrast=self.conf.get("transf_contrast"),
                    saturation=self.conf.get("transf_saturation"),
                )
            ]
        tfs = (
            [transforms.Resize(self.conf.get("img_size"))]
            + augmentative_transforms
            + [
                transforms.ToTensor(),
                transforms.Normalize(
                    self.conf.get("norm_mean"), self.conf.get("norm_std")
                ),
            ]
        )
        return transforms.Compose(tfs)(sample)


class ImageFolderMultiTransform(DatasetFolderMultiTransform):
    """Adapts class ImageFolder of PyTorch in a way that one sample can be transformed several times.
    Args:
        n_transforms (int): number of transformations per sample
        all others: see ImageFolder
    """

    def __init__(
        self,
        conf: dict,
        root,
        transform=None,
        target_transform=None,
        loader=default_loader,
        is_valid_file=None,
        n_transforms=None,
    ):
        self.conf = conf
        if n_transforms is None:
            n_transforms = self.conf.get("n_transforms")

        super().__init__(
            self.conf,
            root,
            loader,
            IMG_EXTENSIONS,
            transform=transform,
            target_transform=target_transform,
            is_valid_file=is_valid_file,
            n_transforms=n_transforms,
        )
        self.imgs = self.samples
