import argparse
import os

import axengine as axe
import numpy as np
from PIL import Image, ImageDraw


LABEL_HEIGHT = 32


def parse_args():
    parser = argparse.ArgumentParser(description='Run Zero-DCE axmodel inference.')
    parser.add_argument('--image', type=str, default='./10.jpg')
    parser.add_argument('--axmodel', type=str, default='Zero-DCE.axmodel')
    parser.add_argument('--output', type=str, default='axmodel_result.png')
    parser.add_argument('--height', type=int, default=256)
    parser.add_argument('--width', type=int, default=256)
    return parser.parse_args()


def load_image(image_path, height, width):
    org_image = Image.open(image_path).convert('RGB')
    input_image = org_image.resize((width, height), Image.BILINEAR)
    image_array = np.asarray(input_image).astype(np.uint8)
    input_array = np.transpose(image_array, (2, 0, 1))
    input_array = np.expand_dims(input_array, axis=0)
    return org_image, input_array


def run_axmodel(input_array, axmodel_path):
    session = axe.InferenceSession(axmodel_path, providers=['AxEngineExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    output = session.run([output_name], {input_name: input_array})[0]
    return output


def output_to_pil(output_array, org_size):
    output = np.squeeze(output_array, axis=0)
    output = np.transpose(output, (1, 2, 0))
    output = np.clip(output, 0.0, 1.0)
    output = (output * 255.0).astype(np.uint8)
    res_image = Image.fromarray(output)
    return res_image.resize(org_size, Image.BILINEAR)


def add_label(image, text):
    canvas = Image.new('RGB', (image.width, image.height + LABEL_HEIGHT), color=(255, 255, 255))
    canvas.paste(image, (0, LABEL_HEIGHT))
    draw = ImageDraw.Draw(canvas)
    draw.text((10, 8), text, fill=(255, 0, 0))
    return canvas


def save_compare_image(org_image, res_image, output_path):
    org_labeled = add_label(org_image, 'org')
    res_labeled = add_label(res_image, 'res')
    compare = Image.new('RGB', (org_labeled.width + res_labeled.width, org_labeled.height), color=(255, 255, 255))
    compare.paste(org_labeled, (0, 0))
    compare.paste(res_labeled, (org_labeled.width, 0))
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    compare.save(output_path)


def main():
    args = parse_args()

    org_image, input_array = load_image(args.image, args.height, args.width)
    axmodel_output = run_axmodel(input_array, args.axmodel)
    res_image = output_to_pil(axmodel_output, org_image.size)

    print('image:', args.image)
    print('axmodel:', args.axmodel)
    print('input shape:', input_array.shape)
    print('axmodel output shape:', axmodel_output.shape)
    print('org size:', org_image.size)

    save_compare_image(org_image, res_image, args.output)
    print('Saved result image:', args.output)
    print('Comparison layout: org | res')


if __name__ == '__main__':
    main()