from __future__ import absolute_import

from torchvision.transforms import *
from PIL import Image
import random
import numpy as np

class Random2DTranslation(object):
    """
    With a probability, first increase image size to (1 + 1/8), and then perform random crop.

    Args:
        height (int): target height.
        width (int): target width.
        p (float): probability of performing this transformation. Default: 0.5.
    """
    def __init__(self, height, width, p=0.5, interpolation=Image.BILINEAR):
        self.height = height
        self.width = width
        self.p = p
        self.interpolation = interpolation

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

        Returns:
            PIL Image: Cropped image.
        """
        if random.random() < self.p:
            return img.resize((self.width, self.height), self.interpolation)
        new_width, new_height = int(round(self.width * 1.125)), int(round(self.height * 1.125))
        resized_img = img.resize((new_width, new_height), self.interpolation)
        x_maxrange = new_width - self.width
        y_maxrange = new_height - self.height
        x1 = int(round(random.uniform(0, x_maxrange)))
        y1 = int(round(random.uniform(0, y_maxrange)))
        croped_img = resized_img.crop((x1, y1, x1 + self.width, y1 + self.height))
        return croped_img

if __name__ == '__main__':
    """import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-impath', type=str)
    parser.add_argument('-nlevel', type=float, default=0.1)
    args = parser.parse_args()

    RC = RandomOcclusion(nlevel=args.nlevel, p=1)
    im = Image.open(args.impath)
    transformed_im = RC(im)

    basename = osp.basename(args.impath)
    save_name = osp.splitext(basename)[0] + '_nlevel_' + str(args.nlevel) + osp.splitext(basename)[1]
    transformed_im.save(save_name)"""
    pass