import carb
import os.path

from pxr import Gf, UsdGeom, UsdLux, Sdf

import omni.hydratexture
import omni.kit.test
from omni.syntheticdata import SyntheticData, SyntheticDataStage

from ..utils import add_semantics

# Test the instance mapping update Fabric flag
class TestInstanceMappingUpdate(omni.kit.test.AsyncTestCase):

    def __init__(self, methodName: str) -> None:
        super().__init__(methodName=methodName)
        # Dictionnary containing the pair (file_path , reference_data). If the reference data is None only the existence of the file is validated.
        self._golden_references = {} 

    def _texture_render_product_path(self, hydra_texture) -> str:
        '''Return a string to the UsdRender.Product used by the texture'''
        render_product = hydra_texture.get_render_product_path()
        if render_product and (not render_product.startswith('/')):
            render_product = '/Render/RenderProduct_' + render_product
        return render_product

    def _assert_count_equal(self, counter_template_name, count):
        count_output = SyntheticData.Get().get_node_attributes(
            counter_template_name,
            ["outputs:count"],
            self._render_product_path
        )
        assert "outputs:count" in count_output
        assert count_output["outputs:count"] == count

    def _activate_fabric_time_range(self) -> None:
        sdg_iface = SyntheticData.Get()
        if not sdg_iface.is_node_template_registered("TestSimFabricTimeRange"):
            sdg_iface.register_node_template(
                    SyntheticData.NodeTemplate(
                        SyntheticDataStage.ON_DEMAND,
                        "omni.syntheticdata.SdTestSimFabricTimeRange"
                    ),
                    template_name="TestSimFabricTimeRange"
                )
        sdg_iface.activate_node_template(
            "TestSimFabricTimeRange", 
            attributes={"inputs:timeRangeName":"testFabricTimeRangeTrigger"}
        )
        if not sdg_iface.is_node_template_registered("TestPostRenderFabricTimeRange"):
            sdg_iface.register_node_template(
                SyntheticData.NodeTemplate(
                    SyntheticDataStage.POST_RENDER,
                    "omni.syntheticdata.SdFabricTimeRangeExecution",
                    [
                        SyntheticData.NodeConnectionTemplate(
                            SyntheticData.renderer_template_name(),
                            attributes_mapping=
                            {
                                "outputs:rp": "inputs:renderResults",
                                "outputs:gpu": "inputs:gpu"
                            }
                        )
                    ]
                ),
                template_name="TestPostRenderFabricTimeRange"
            )
        sdg_iface.activate_node_template(
            "TestPostRenderFabricTimeRange", 
            0,
            [self._render_product_path],
            attributes={"inputs:timeRangeName":"testFabricTimeRangeTrigger"}
        )
        if not sdg_iface.is_node_template_registered("TestPostProcessFabricTimeRange"):
            sdg_iface.register_node_template(
                SyntheticData.NodeTemplate(
                    SyntheticDataStage.ON_DEMAND,
                    "omni.syntheticdata.SdFabricTimeRangeExecution",
                    [
                        SyntheticData.NodeConnectionTemplate("PostProcessDispatch"),
                        SyntheticData.NodeConnectionTemplate("TestPostRenderFabricTimeRange")
                    ]
                ),
                template_name="TestPostProcessFabricTimeRange"
            )
        sdg_iface.activate_node_template(
            "TestPostProcessFabricTimeRange", 
            0,
            [self._render_product_path],
            attributes={"inputs:timeRangeName":"testFabricTimeRangeTrigger"}
        )
        if not sdg_iface.is_node_template_registered("TestPostProcessFabricTimeRangeCounter"):
            sdg_iface.register_node_template(
                SyntheticData.NodeTemplate(
                    SyntheticDataStage.ON_DEMAND,
                    "omni.graph.action.Counter",
                    [
                        SyntheticData.NodeConnectionTemplate(
                            "TestPostProcessFabricTimeRange", 
                            attributes_mapping={"outputs:exec": "inputs:execIn"}
                        )
                    ]
                ),
                template_name="TestPostProcessFabricTimeRangeCounter"
            )
        sdg_iface.activate_node_template(
            "TestPostProcessFabricTimeRangeCounter", 
            0,
            [self._render_product_path]
        )

    def _activate_instance_mapping_update(self) -> None:
        sdg_iface = SyntheticData.Get()
        if not sdg_iface.is_node_template_registered("TestPostProcessInstanceMappingUpdate"):
            sdg_iface.register_node_template(
                SyntheticData.NodeTemplate(
                    SyntheticDataStage.ON_DEMAND,
                    "omni.syntheticdata.SdTimeChangeExecution",
                    [
                        SyntheticData.NodeConnectionTemplate("InstanceMappingPtr"),
                        SyntheticData.NodeConnectionTemplate("PostProcessDispatch")
                    ]
                ),
                template_name="TestPostProcessInstanceMappingUpdate"
            )
        if not sdg_iface.is_node_template_registered("TestPostProcessInstanceMappingUpdateCounter"):
            sdg_iface.register_node_template(
                SyntheticData.NodeTemplate(
                    SyntheticDataStage.ON_DEMAND,
                    "omni.graph.action.Counter",
                    [
                        SyntheticData.NodeConnectionTemplate(
                            "TestPostProcessInstanceMappingUpdate", 
                            attributes_mapping={"outputs:exec": "inputs:execIn"}
                        )
                    ]
                ),
                template_name="TestPostProcessInstanceMappingUpdateCounter"
            )
        sdg_iface.activate_node_template(
            "TestPostProcessInstanceMappingUpdateCounter", 
            0, 
            [self._render_product_path]
        )

    async def _request_fabric_time_range_trigger(self, number_of_frames=1):
        sdg_iface = SyntheticData.Get()
        sdg_iface.set_node_attributes("TestSimFabricTimeRange",{"inputs:numberOfFrames":number_of_frames})
        sdg_iface.request_node_execution("TestSimFabricTimeRange")
        await omni.kit.app.get_app().next_update_async()

    async def setUp(self):
        """Called at the begining of every tests"""
        self._settings = carb.settings.acquire_settings_interface()
        self._hydra_texture_factory = omni.hydratexture.acquire_hydra_texture_factory_interface()

        self._usd_context_name = ''
        self._usd_context = omni.usd.get_context(self._usd_context_name)
        await self._usd_context.new_stage_async()

        # renderer
        renderer = "rtx"
        if renderer not in self._usd_context.get_attached_hydra_engine_names():
            omni.usd.add_hydra_engine(renderer, self._usd_context)

        # create the hydra textures
        self._hydra_texture_0 = self._hydra_texture_factory.create_hydra_texture(
            "TEX0",
            1920,
            1080,
            self._usd_context_name,
            hydra_engine_name=renderer,
            is_async=self._settings.get("/app/asyncRendering")
        )

        self._hydra_texture_rendered_counter = 0
        def on_hydra_texture_0(event: carb.events.IEvent):
            self._hydra_texture_rendered_counter += 1
        self._hydra_texture_rendered_counter_sub = self._hydra_texture_0.get_event_stream().create_subscription_to_push_by_type(
            omni.hydratexture.EVENT_TYPE_DRAWABLE_CHANGED,
            on_hydra_texture_0,
            name='async rendering test drawable update',
        )

        stage = omni.usd.get_context().get_stage()
        world_prim = UsdGeom.Xform.Define(stage,"/World")
        UsdGeom.Xformable(world_prim).AddTranslateOp().Set((0, 0, 0))
        UsdGeom.Xformable(world_prim).AddRotateXYZOp().Set((0, 0, 0))

        self._render_product_path = self._texture_render_product_path(self._hydra_texture_0)
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path)
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path)
        
    async def tearDown(self):
        """Called at the end of every tests"""
        self._hydra_texture_rendered_counter_sub = None
        self._hydra_texture_0 = None

        self._usd_context.close_stage()
        omni.usd.release_all_hydra_engines(self._usd_context)

        self._hydra_texture_factory = None
        self._settings = None

        wait_iterations = 6
        for _ in range(wait_iterations):
            await omni.kit.app.get_app().next_update_async()

    async def test_case_0(self):
        """Test case 0 : no time range"""
        self._activate_fabric_time_range()
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 11)
        self._assert_count_equal("TestPostProcessFabricTimeRangeCounter", 0)

    async def test_case_1(self):
        """Test case 1 : setup a time range of 5 frames"""
        self._activate_fabric_time_range()
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path)
        await self._request_fabric_time_range_trigger(5)
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 11)
        self._assert_count_equal("TestPostProcessFabricTimeRangeCounter", 5)

    async def test_case_2(self):
        """Test case 2 : initial instance mapping setup"""
        self._activate_instance_mapping_update()
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 11)
        self._assert_count_equal("TestPostProcessInstanceMappingUpdateCounter", 1)

    async def test_case_3(self):
        """Test case 3 : setup an instance mapping with 1, 2, 3, 4 changes"""
        stage = omni.usd.get_context().get_stage()
        self._activate_instance_mapping_update()
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 1)
        self._assert_count_equal("TestPostProcessInstanceMappingUpdateCounter", 1)
        sphere_prim = stage.DefinePrim("/World/Sphere", "Sphere")
        add_semantics(sphere_prim, "sphere")
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 3)
        self._assert_count_equal("TestPostProcessInstanceMappingUpdateCounter", 2)
        sub_sphere_prim = stage.DefinePrim("/World/Sphere/Sphere", "Sphere")
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 5)
        self._assert_count_equal("TestPostProcessInstanceMappingUpdateCounter", 3)
        add_semantics(sub_sphere_prim, "sphere")
        await omni.syntheticdata.sensors.next_render_simulation_async(self._render_product_path, 1)
        self._assert_count_equal("TestPostProcessInstanceMappingUpdateCounter", 4)