import argparse
from math import ceil
from time import time

import pyvips
from affine import Affine
from pyvips import Image, Kernel

pyvips.cache_set_max(0)

parser = argparse.ArgumentParser()
parser.add_argument('filename')
args = parser.parse_args()

scale_x = 6.435381355932203
scale_y = 6.434928631402183
rotation = 34.051919898386

data_img = Image.new_from_file(args.filename)
if data_img.bands == 3:
    data_img = data_img.bandjoin(0)
photo_data = data_img.cast('uchar').write_to_memory()


def transform_image(early_resize: bool):
    bands = 4
    img_format = 'uchar'

    photo_width = data_img.width
    photo_height = data_img.height
    graphic_width = ceil(data_img.width * scale_x)
    graphic_height = ceil(data_img.height * scale_y)

    vips_img = Image.new_from_memory(photo_data, photo_width, photo_height, bands, img_format)

    # Resize operation
    resize_kernel = Kernel.LANCZOS3 if (scale_x <= 1 or scale_y <= 1) else Kernel.MITCHELL
    vips_img = vips_img.resize(scale_x, kernel=resize_kernel, vscale=scale_y)

    if early_resize:
        # Force doing the resize before running affine()
        vips_img = Image.new_from_memory(vips_img.write_to_memory(), vips_img.width, vips_img.height, bands, img_format)

    # Affine transform
    affine_transform = Affine.identity()
    affine_transform = affine_transform.rotation(rotation)
    vips_transform_matrix = (affine_transform.a, affine_transform.b, affine_transform.d, affine_transform.e)
    vips_img = vips_img.premultiply().affine(
        vips_transform_matrix,
        oarea=[0, 0, int(round(graphic_width)), int(round(graphic_height))],
        idx=-int(round(vips_img.width / 2)),
        idy=-int(round(vips_img.height / 2)),
        odx=int(round(graphic_width / 2)),
        ody=int(round(graphic_height / 2)),
    ).unpremultiply()

    # Write to memory to actually do the affine transform
    vips_img.write_to_memory()
    vips_img.write_to_file('out.png')


if __name__ == '__main__':
    start_time = time()
    transform_image(True)
    print('With early resize took %f' % (time() - start_time))

    start_time = time()
    transform_image(False)
    print('Without early resize took %f' % (time() - start_time))