import math
import random
import warnings
import logging

import torch
from torchvision import transforms
import torchvision.transforms.functional as F

from timm.data import transforms as timm_transforms
from timm.data.transforms_factory import transforms_imagenet_eval

_logger = logging.getLogger(__name__)

# Modified from timm
class AutoRandomResizedCropAndInterpolation:
    """Crop the given PIL Image to random size and aspect ratio with random interpolation.

    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
    is finally resized to given size.
    This is popularly used to train the Inception networks.

    Args:
        size: expected output size of each edge
        scale: range of size of the origin size cropped
        ratio: range of aspect ratio of the origin aspect ratio cropped
        interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(self, size, scale=(0.08, 1.0), 
                 ratio_type="auto",
                 ratio=(3. / 4., 4. / 3.),
                 interpolation='bilinear'):
        if isinstance(size, (list, tuple)):
            self.size = tuple(size)
        else:
            self.size = (size, size)
        if ratio_type != "auto" and ((scale[0] > scale[1]) or (ratio[0] > ratio[1])):
            warnings.warn("range should be of kind (min, max)")

        if interpolation == 'random':
            self.interpolation = timm_transforms._RANDOM_INTERPOLATION
        else:
            self.interpolation = timm_transforms.str_to_interp_mode(interpolation)
        self.scale = scale
        self.ratio_type = ratio_type
        if ratio_type == "auto":
            assert isinstance(ratio, float)
        self.ratio = ratio

    @staticmethod
    def get_params(img, scale, ratio_type, ratio):
        """Get parameters for ``crop`` for a random sized crop.

        Args:
            img (PIL Image): Image to be cropped.
            scale (tuple): range of size of the origin size cropped
            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped

        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
                sized crop.
        """
        area = img.size[0] * img.size[1]

        for attempt in range(10):
            target_area = random.uniform(*scale) * area
            if ratio_type == "auto":
                # When we want the aspect ratio to be image dependent
                aspect_ratio = img.size[0] / img.size[1]
                ratio_diff = random.uniform(-1000.*ratio, 1000.*ratio) / 1000.
                aspect_ratio = aspect_ratio + ratio_diff
            else:
                log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
                aspect_ratio = math.exp(random.uniform(*log_ratio))

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if w <= img.size[0] and h <= img.size[1]:
                i = random.randint(0, img.size[1] - h)
                j = random.randint(0, img.size[0] - w)
                return i, j, h, w

        # Fallback to central crop
        in_ratio = img.size[0] / img.size[1]
        if ratio_type != "auto" and in_ratio < min(ratio):
            w = img.size[0]
            h = int(round(w / min(ratio)))
        elif ratio_type != "auto" and in_ratio > max(ratio):
            h = img.size[1]
            w = int(round(h * max(ratio)))
        else:  # whole image
            w = img.size[0]
            h = img.size[1]
        i = (img.size[1] - h) // 2
        j = (img.size[0] - w) // 2
        return i, j, h, w

    def __call__(self, img):
        """
        Args:
            img (PIL Image): Image to be cropped and resized.

        Returns:
            PIL Image: Randomly cropped and resized image.
        """
        i, j, h, w = self.get_params(img, self.scale, self.ratio_type, self.ratio)
        if isinstance(self.interpolation, (tuple, list)):
            interpolation = random.choice(self.interpolation)
        else:
            interpolation = self.interpolation
        return F.resized_crop(img, i, j, h, w, self.size, interpolation)

    def __repr__(self):
        if isinstance(self.interpolation, (tuple, list)):
            interpolate_str = ' '.join([timm_transforms.interp_mode_to_str(x) for x in self.interpolation])
        else:
            interpolate_str = timm_transforms.interp_mode_to_str(self.interpolation)
        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
        format_string += ', ratio_type={0}'.format(self.ratio_type)
        if self.ratio_type == "auto":
            format_string += ', ratio=+-{0}'.format(self.ratio)    
        else:
            format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
        format_string += ', interpolation={0})'.format(interpolate_str)
        return format_string


TRANSFORM_FACTORY = {
    "scale-and-crop": AutoRandomResizedCropAndInterpolation,
    "hflip": transforms.RandomHorizontalFlip,
    "vflip": transforms.RandomVerticalFlip,
    "timm-eval": transforms_imagenet_eval,
    "resize": transforms.Resize
}
