## Copyright (c) 2018-2019, NVIDIA CORPORATION.  All rights reserved.
##
## NVIDIA CORPORATION and its licensors retain all intellectual property
## and proprietary rights in and to this software, related documentation
## and any modifications thereto.  Any use, reproduction, disclosure or
## distribution of this software and related documentation without an express
## license agreement from NVIDIA CORPORATION is strictly prohibited.
##
"""The utilities for image comparison"""
import pathlib
import sys
import traceback
import math
from enum import Enum
import numpy as np

import carb
import carb.tokens

import omni.appwindow
from omni.kit.test.teamcity import teamcity_publish_image_artifact


class CompareError(Exception):
    pass


class ComparisonMetric(Enum):
    MAX_DIFFERENCE = 0              # the maximum differrence between pixels. 0 is no difference. In the range [0-255].
    PIXEL_COUNT = 1                 # the total number of pixels that is not the same for two images
    MEAN_ERROR_SQUARED = 2          # Mean root square difference between pixels of two images
    PEAK_SIGNAL_TO_NOISE_RATIO = 3


def computeMSE(difference):
    components = len(difference.getpixel((0, 0)))
    errors = np.asarray(difference) / 255
    return sum(np.mean(np.square(errors), axis = (0, 1))) / components


def compare(image1: pathlib.Path, image2: pathlib.Path, image_diffmap: pathlib.Path,
            metric: ComparisonMetric = ComparisonMetric.MAX_DIFFERENCE):
    """
    Compares two images and return a value that indicates the maximum differrence between pixels. 0 is no difference
    in the range [0-255].

    It uses Pillow for image read.

    Args:
        image1, image2: images to compare
        image_diffmap: the difference map image will be saved if there is any difference between given images

    """
    if not image1.exists():
        raise CompareError(f"File image1 {image1} does not exist")
    if not image2.exists():
        raise CompareError(f"File image2 {image2} does not exist")

    if "PIL" not in sys.modules.keys():
        # Checking if we have Pillow imported
        try:
            from PIL import Image
        except ImportError:
            # Install Pillow if it's not installed
            import omni.kit.pipapi

            omni.kit.pipapi.install("Pillow", module="PIL")

    from PIL import Image
    from PIL import ImageChops

    original = Image.open(str(image1))
    contrast = Image.open(str(image2))

    if original.size != contrast.size:
        raise CompareError(
            f"[omni.ui.test] Can't compare different resolutions\n\n"
            f"{image1} {original.size[0]}x{original.size[1]}\n"
            f"{image2} {contrast.size[0]}x{contrast.size[1]}\n\n"
            f"It's possible that your monitor DPI is not 100%.\n\n"
        )
    
    difference = ImageChops.difference(original, contrast)
    
    metric_value = 0.0

    try:
        if metric is ComparisonMetric.MAX_DIFFERENCE:
            metric_value = sum([sum(i) for i in difference.convert("RGB").getextrema()])
        if metric is ComparisonMetric.PIXEL_COUNT:
            metric_value = sum([sum(
                difference.getpixel((j, i))) > 0 for i in range(difference.height) for j in range(difference.width)])
        if metric is ComparisonMetric.MEAN_ERROR_SQUARED:
            metric_value = computeMSE(difference)
        if metric is ComparisonMetric.PEAK_SIGNAL_TO_NOISE_RATIO:
            metric_value = computeMSE(difference)
            if metric_value < sys.float_info.epsilon:
                metric_value = math.inf
            metric_value = 10 * math.log10(1 / metric_value)
    except():
        raise CompareError(f"[omni.ui.test] Can't compute metric {metric} \n\n")
    
    if metric_value > 0:
        # Images are different
        # Multiply image by 255
        difference = difference.convert("RGB").point(lambda i: min(i * 255, 255))
        difference.save(str(image_diffmap))
    return metric_value


def capture(image_name: str, output_img_dir: str, app_window: omni.appwindow.IAppWindow = None, use_log: bool = True):
    """
    Captures frame.

    Args:
        image_name: the image name of the image and golden image.
        output_img_dir: the directory path that the capture will be saved to.
        golden_img_dir: the directory path that stores the golden image. Leave it to None to use default dir.
    """

    image1 = pathlib.Path(output_img_dir).joinpath(image_name)
    if use_log:
        carb.log_info(f"[tests.compare] Capturing {image1}")

    import omni.renderer_capture

    capture_interface = omni.renderer_capture.acquire_renderer_capture_interface()
    capture_interface.capture_next_frame_swapchain(str(image1), app_window)


def finalize_capture_and_compare(image_name: str, threshold: float, output_img_dir: str, golden_img_dir: str,
                                 app_window: omni.appwindow.IAppWindow = None,
                                 metric: ComparisonMetric = ComparisonMetric.MAX_DIFFERENCE):
    """
    Finalizes capture and compares it with the golden image.

    Args:
        image_name: the image name of the image and golden image.
        threshold: the max threshold to collect TC artifacts.
        output_img_dir: the directory path that the capture will be saved to.
        golden_img_dir: the directory path that stores the golden image. Leave it to None to use default dir.
        app_window: IAppWindow instance for which the capture will happen, use None to capture default app window.

    Returns:
        A value that indicates the maximum difference between pixels. 0 is no difference
        in the range [0-255].
    """
    image1 = pathlib.Path(output_img_dir).joinpath(image_name)
    image2 = pathlib.Path(golden_img_dir).joinpath(image_name)

    image_diffmap = pathlib.Path(output_img_dir).joinpath(f"{pathlib.Path(image_name).stem}.diffmap.png")

    import omni.renderer_capture

    capture_interface = omni.renderer_capture.acquire_renderer_capture_interface()
    capture_interface.wait_async_capture(app_window)

    carb.log_info(f"[tests.compare] Comparing {image1} to {image2}")

    try:
        diff = compare(image1, image2, image_diffmap, metric)
        if diff >= threshold:
            teamcity_publish_image_artifact(image2, "golden", "Reference")
            teamcity_publish_image_artifact(image1, "results", "Generated")
            teamcity_publish_image_artifact(image_diffmap, "results", "Diff")
            
        return diff
    except CompareError as e:
        carb.log_error(f"[tests.compare] Failed to compare images for {image_name}. Error: {e}")
        exc = traceback.format_exc()
        carb.log_error(f"[tests.compare] Traceback:\n{exc}")


async def capture_and_compare(image_name: str, threshold, output_img_dir: str, golden_img_dir: str,
                              app_window: omni.appwindow.IAppWindow = None, use_log: bool = True,
                              metric: ComparisonMetric = ComparisonMetric.MAX_DIFFERENCE):
    """
    Captures frame and compares it with the golden image.

    Args:
        image_name: the image name of the image and golden image.
        threshold: the max threshold to collect TC artifacts.
        golden_img_dir: the directory path that stores the golden image. Leave it to None to use default dir.

    Returns:
        A value that indicates the maximum difference between pixels. 0 is no difference
        in the range [0-255].
    """

    capture(image_name, output_img_dir, app_window, use_log=use_log)

    await omni.kit.app.get_app().next_update_async()

    return finalize_capture_and_compare(image_name, threshold, output_img_dir, golden_img_dir, app_window, metric)

