#!/usr/bin/env python3

import carb
import omni.kit.commands
import omni.kit.test
import omni.usd
from omni.rtx.tests import RtxTest, testSettings, postLoadTestSettings
from omni.rtx.tests.test_common import set_transform_helper, wait_for_update
from omni.kit.test_helpers_gfx.compare_utils import ComparisonMetric

from pxr import Sdf, Gf, UsdVol
from pathlib import Path

EXTENSION_DIR = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))
TESTS_DIR = EXTENSION_DIR.joinpath('data', 'tests')
USD_DIR = TESTS_DIR.joinpath('usd')
GOLDEN_IMAGES_DIR = TESTS_DIR.joinpath('golden')
VOLUMES_DIR = TESTS_DIR.joinpath('volumes')
OUTPUTS_DIR = Path(omni.kit.test.get_test_output_path())

# This class is auto-discoverable by omni.kit.test
class IndexRtxCompositeRenderTest(RtxTest):

    WINDOW_SIZE = (512, 512)
    THRESHOLD = 1e-6

    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        self.ctx.new_stage()
        self.add_dir_light()
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    # Overridden with custom paths
    async def capture_and_compare(self, img_subdir: Path = "", golden_img_name=None, threshold=THRESHOLD,
                                  metric: ComparisonMetric = ComparisonMetric.MEAN_ERROR_SQUARED):
        golden_img_dir = GOLDEN_IMAGES_DIR.joinpath(img_subdir)
        output_img_dir = OUTPUTS_DIR.joinpath(img_subdir)

        if not golden_img_name:
            golden_img_name = f"{self.__test_name}.png"

        return await self._capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir, metric)

    # Overridden with custom paths
    def open_usd(self, usdSubpath: Path):
        path = USD_DIR.joinpath(usdSubpath)
        self.ctx.open_stage(str(path))

    async def run_imge_test(self, usd_file: str, image_name: str, fn = None, threshold=THRESHOLD):
        if usd_file:
            self.open_usd(usd_file)

        if fn:
            fn()

        await wait_for_update()
        await self.capture_and_compare(golden_img_name=image_name, threshold=threshold)

    #
    # The tests
    #

    async def test_render_torus_volume_with_geometry(self):
        self.open_usd('torus-volume-with-geometry.usda')

        settings = carb.settings.get_settings()
        for depth_mode in [0, 1, 2, 3]:
            settings.set("/rtx/index/compositeDepthMode", depth_mode)
            await wait_for_update()

            # Non-stable depth mode is more noisy
            threshold = (1e-2 if depth_mode == 2 else IndexRtxCompositeRenderTest.THRESHOLD)

            await self.run_imge_test(None, f"nvindex-torus-composite-depth-mode-{depth_mode}.png", threshold=threshold)
