#!/usr/bin/env python3

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 IndexRenderTest(RtxTest):

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

    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))

    def get_volumes_path(self, volumeSubPath: Path):
        return Sdf.AssetPath(str(VOLUMES_DIR.joinpath(volumeSubPath)))

    def change_reference(self, prim_path, reference):
        stage = self.ctx.get_stage()
        prim = stage.GetPrimAtPath(prim_path)
        ref = prim.GetReferences()
        ref.ClearReferences()
        ref.AddReference(Sdf.Reference(reference))

    async def run_imge_test(self, usd_file: str, image_name: str, fn = None):
        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)

    #
    # The tests
    #

    async def test_render_torus(self):
        await self.run_imge_test('torus.usda', 'nvindex-torus.png')

    async def test_render_torus_colormap(self):
        await self.run_imge_test('torus-colormap.usda', 'nvindex-torus-colormap-2.png')

        colormap_prim_path = '/World/Volume/mat/Colormap'

        # Switch to another colormap
        await self.run_imge_test(None, 'nvindex-torus-colormap-3.png',
                                 lambda: self.change_reference(colormap_prim_path, './colormap3.usda'))

        # Switch to colormap defined by RGBA control-points
        colormap = self.ctx.get_stage().GetPrimAtPath(colormap_prim_path)
        colormap.CreateAttribute("domain", Sdf.ValueTypeNames.Float2).Set((0.0, 1.0))
        colormap.CreateAttribute("colormapSource", Sdf.ValueTypeNames.Token).Set("rgbaPoints")
        colormap.CreateAttribute("xPoints", Sdf.ValueTypeNames.FloatArray).Set([0.5, 0.7, 1.0])
        colormap.CreateAttribute("rgbaPoints", Sdf.ValueTypeNames.Float4Array).Set([(0, 1, 0, 0.14), (1, 0, 0, 0.25), (1, 1, 0, 0.25)])

        await self.run_imge_test(None, 'nvindex-torus-colormap-rgba-points.png')

    async def test_render_torus_shader(self):
        await self.run_imge_test('torus-shader.usda', 'nvindex-torus-shader.png')

    async def test_render_torus_settings(self):
        await self.run_imge_test('torus-settings.usda', 'nvindex-torus-settings.png')

    async def test_render_torus_create(self):
        PRIM_PATH = "/World/volume"
        open_vdb_asset = UsdVol.OpenVDBAsset.Define(self.ctx.get_stage(), PRIM_PATH + "/OpenVDBAsset")
        open_vdb_asset.GetFilePathAttr().Set(self.get_volumes_path("torus.vdb"))
        # Use default field in the OpenVDB file, instead of setting it explicitly:
        # open_vdb_asset.GetFieldNameAttr().Set("torus_fog")
        volume = UsdVol.Volume.Define(self.ctx.get_stage(), PRIM_PATH)
        volume.CreateFieldRelationship("torus", open_vdb_asset.GetPath())
        set_transform_helper(volume.GetPath(), translate=Gf.Vec3d(0, 150, 0), scale=Gf.Vec3d(20.0, 20.0, 20.0))

        await wait_for_update()
        await self.run_imge_test(None, 'nvindex-torus-create.png')
