import argparse
import os
import time

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


def get_tensor_dtype(tensor_info):
    dtype = getattr(tensor_info, 'dtype', None)
    if dtype is None:
        dtype = getattr(tensor_info, 'type', None)
    return str(dtype).lower() if dtype is not None else ''


def format_tensor_stats(name, tensor):
    tensor_np = np.asarray(tensor)
    return '{}: dtype={}, shape={}, min={:.6f}, max={:.6f}'.format(
        name,
        tensor_np.dtype,
        tensor_np.shape,
        float(tensor_np.min()),
        float(tensor_np.max()),
    )


def preprocess(image_path, height, width, resize, input_dtype):
    original_image = Image.open(image_path).convert('RGB')
    original_size = original_image.size
    if resize:
        model_image = original_image.resize((width, height), Image.BILINEAR)
    else:
        if original_image.size[0] < width or original_image.size[1] < height:
            raise ValueError('Image is smaller than axmodel input size: {}'.format(image_path))
        model_image = original_image.crop((0, 0, width, height))

    image_np = np.asarray(model_image)
    if 'float' in input_dtype:
        image_np = image_np.astype(np.float32) / 255.0
    else:
        image_np = image_np.astype(np.uint8)
    image_np = image_np.transpose(2, 0, 1)[None, :, :, :]
    return original_image, original_size, image_np


def postprocess(enhanced_np, original_size):
    enhanced_np = enhanced_np[0].transpose(1, 2, 0)
    max_value = float(enhanced_np.max())
    min_value = float(enhanced_np.min())
    if max_value <= 1.5 and min_value >= -0.5:
        enhanced_np = np.clip(enhanced_np, 0.0, 1.0) * 255.0
    else:
        enhanced_np = np.clip(enhanced_np, 0.0, 255.0)
    enhanced_np = enhanced_np.round().astype(np.uint8)
    enhanced_image = Image.fromarray(enhanced_np)
    return enhanced_image.resize(original_size, Image.BILINEAR)


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


def save_compare(original_image, enhanced_image, result_path):
    original_labeled = add_label(original_image, 'Original')
    enhanced_labeled = add_label(enhanced_image, 'Result')
    compare_image = Image.new('RGB', (original_labeled.width + enhanced_labeled.width, original_labeled.height))
    compare_image.paste(original_labeled, (0, 0))
    compare_image.paste(enhanced_labeled, (original_labeled.width, 0))
    result_dir = os.path.dirname(result_path)
    if result_dir and not os.path.exists(result_dir):
        os.makedirs(result_dir)
    compare_image.save(result_path)


def build_axmodel_session(axmodel_path):
    return axe.InferenceSession(axmodel_path, providers=['AxEngineExecutionProvider'])


def infer(config):
    if not os.path.isfile(config.input):
        raise ValueError('Input image does not exist: {}'.format(config.input))

    session = build_axmodel_session(config.axmodel)
    input_info = session.get_inputs()[0]
    output_info = session.get_outputs()[0]
    input_name = input_info.name
    input_shape = input_info.shape
    input_dtype = get_tensor_dtype(input_info)
    output_dtype = get_tensor_dtype(output_info)
    height = int(input_shape[2]) if config.height <= 0 else config.height
    width = int(input_shape[3]) if config.width <= 0 else config.width

    original_image, original_size, input_np = preprocess(
        config.input,
        height,
        width,
        bool(config.resize),
        input_dtype,
    )
    start = time.time()
    axmodel_outputs = session.run(None, {input_name: input_np})
    axmodel_enhanced = axmodel_outputs[0]
    elapsed = time.time() - start

    enhanced_image = postprocess(axmodel_enhanced, original_size)
    save_compare(original_image, enhanced_image, config.output)

    print('Input image:', config.input)
    print('Output image:', config.output)
    print('Model input dtype:', input_dtype or 'unknown')
    print('Model output dtype:', output_dtype or 'unknown')
    print(format_tensor_stats('Prepared input', input_np))
    print(format_tensor_stats('Model output', axmodel_enhanced))
    print('axmodel time:', elapsed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--axmodel', type=str, default='zerodcepp_512_sf8.axmodel')
    parser.add_argument('--input', type=str, default='101_3_.png')
    parser.add_argument('--output', type=str, default='axmodel_res.png')
    parser.add_argument('--height', type=int, default=512)
    parser.add_argument('--width', type=int, default=512)
    parser.add_argument('--resize', type=int, default=1)
    config = parser.parse_args()

    infer(config)