import torchvision.transforms as transforms
from PIL import ImageOps

class PadWhite(object):
    def __init__(self, size):
        assert isinstance(size, (int, tuple))
        if isinstance(size, tuple):
            self.height, self.width = size
        elif isinstance(size, int):
            self.height = self.width = size

    def __call__(self, img):
        if img.size[0] > self.width or img.size[1] > self.height:
            img.thumbnail((self.width, self.height))
        delta_width = self.width - img.size[0]
        delta_height = self.height - img.size[1]
        pad_width = delta_width // 2
        pad_height = delta_height // 2
        padding = (pad_width, pad_height, delta_width -
                   pad_width, delta_height-pad_height)
        return ImageOps.expand(img, padding, fill=255)

def process_image_pos(image):
    target_size = (400, 512)
    # image = Image.open(img_name).convert("L")
    w, h = image.size
    top_padding, left_padding = 0, 0

    if h <= target_size[0] or w <= target_size[1]:
        delta_height = target_size[0] - h
        delta_width = target_size[1] - w
        pad_height = delta_height // 2
        pad_width = delta_width // 2
        # (left, top, right, bottom)
        padding = (pad_width, pad_height, delta_width -
                    pad_width, delta_height-pad_height)
        image = ImageOps.expand(image, padding, fill=255)
    elif h > 400 or w > 500:
        print("Height screwed")
        
    transform = transforms.ToTensor()
    image = transform(image)
    return image

def process_image_vgg(image):
    input_size = (32, 128)
    transform = transforms.Compose([
        PadWhite(input_size),
        transforms.ToTensor(),
        ])
    image = transform(image)
    return image