import os
import omni.kit.test
import omni.graph.core as og
import omni.graph.core.tests as ogts
from omni.graph.core.tests.omnigraph_test_utils import _TestGraphAndNode
from omni.graph.core.tests.omnigraph_test_utils import _test_clear_scene
from omni.graph.core.tests.omnigraph_test_utils import _test_setup_scene
from omni.graph.core.tests.omnigraph_test_utils import _test_verify_scene


class TestOgn(ogts.OmniGraphTestCase):

    TEST_DATA = [
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 6649909271, False],
                ['inputs:min', {'type': 'uint', 'value': 0}, False],
                ['inputs:max', {'type': 'uint', 'value': 4294967295}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint', 'value': 0}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 6159018942, False],
                ['inputs:min', {'type': 'uint[]', 'value': [0, 100]}, False],
                ['inputs:max', {'type': 'uint', 'value': 4294967295}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint[]', 'value': [2147483648, 2160101208]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 3280530163, False],
                ['inputs:min', {'type': 'uint', 'value': 0}, False],
                ['inputs:max', {'type': 'uint[]', 'value': [4294967295, 199]}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint[]', 'value': [4294967295, 19]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 6159018942, False],
                ['inputs:min', {'type': 'int', 'value': -2147483648}, False],
                ['inputs:max', {'type': 'int', 'value': 2147483647}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'int', 'value': 0}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 6649909271, False],
                ['inputs:min', {'type': 'int[2]', 'value': [-2147483648, -100]}, False],
                ['inputs:max', {'type': 'int', 'value': 2147483647}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'int[2]', 'value': [-2147483648, 1629773655]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 3280530163, False],
                ['inputs:min', {'type': 'int', 'value': -2147483648}, False],
                ['inputs:max', {'type': 'int[2]', 'value': [2147483647, 99]}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'int[2]', 'value': [2147483647, -2146948710]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 8280086, False],
                ['inputs:min', {'type': 'float', 'value': 0}, False],
                ['inputs:max', {'type': 'float', 'value': 1}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'float', 'value': 0}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 17972581, False],
                ['inputs:min', {'type': 'float[]', 'value': [0, -10]}, False],
                ['inputs:max', {'type': 'float', 'value': 1}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'float[]', 'value': [0.5, -6.7663326]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 15115159, False],
                ['inputs:min', {'type': 'float', 'value': 0}, False],
                ['inputs:max', {'type': 'float[]', 'value': [1, 10]}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'float[]', 'value': [0.9999999403953552, 4.0452986]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 14092058508772706262, False],
                ['inputs:min', {'type': 'uint64', 'value': 0}, False],
                ['inputs:max', {'type': 'uint64', 'value': 18446744073709551615}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint64', 'value': 0}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 9302349107990861236, False],
                ['inputs:min', {'type': 'uint64', 'value': 0}, False],
                ['inputs:max', {'type': 'uint64', 'value': 18446744073709551615}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint64', 'value': 9223372036854775808}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 1955209015103813879, False],
                ['inputs:min', {'type': 'uint64', 'value': 0}, False],
                ['inputs:max', {'type': 'uint64', 'value': 18446744073709551615}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint64', 'value': 18446744073709551615}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 123456789, False],
                ['inputs:min', {'type': 'uint64', 'value': 1099511627776}, False],
                ['inputs:max', {'type': 'uint64', 'value': 1125899906842624}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uint64', 'value': 923489197424953}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 1955209015103813879, False],
                ['inputs:min', {'type': 'double[2][]', 'value': [[0, -10], [10, 0]]}, False],
                ['inputs:max', {'type': 'double', 'value': 1}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'double[2][]', 'value': [[0, 0.28955788], [7.98645811, 0.09353537]]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 123456789, False],
                ['inputs:min', {'type': 'half[]', 'value': [0, -100]}, False],
                ['inputs:max', {'type': 'half[]', 'value': [1, 100]}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'half[]', 'value': [0.17993164, -76.375]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 0, False],
                ['inputs:min', {'type': 'uchar[]', 'value': [0, 100]}, False],
                ['inputs:max', {'type': 'uchar[]', 'value': [255, 200]}, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'uchar[]', 'value': [153, 175]}, False],
                ['outputs:execOut', 1, False],
            ],
        },
        {
            'inputs': [
                ['inputs:useSeed', True, False],
                ['inputs:seed', 9302349107990861236, False],
                ['inputs:execIn', 1, False],
                ['inputs:isNoise', True, False],
            ],
            'outputs': [
                ['outputs:random', {'type': 'double', 'value': 0.5}, False],
                ['outputs:execOut', 1, False],
            ],
        },
    ]

    async def test_generated(self):
        test_info = _TestGraphAndNode()
        controller = og.Controller()
        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            test_info = await _test_setup_scene(self, controller, "/TestGraph", "TestNode_omni_graph_nodes_RandomNumeric", "omni.graph.nodes.RandomNumeric", test_run, test_info)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.RandomNumeric User test case #{i+1}")

    async def test_vectorized_generated(self):
        test_info = _TestGraphAndNode()
        controller = og.Controller()
        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            test_info = await _test_setup_scene(self, controller, "/TestGraph", "TestNode_omni_graph_nodes_RandomNumeric","omni.graph.nodes.RandomNumeric", test_run, test_info, 16)
            await controller.evaluate(test_info.graph)
            _test_verify_scene(self, controller, test_run, test_info, f"omni.graph.nodes.RandomNumeric User test case #{i+1}", 16)

    async def test_thread_safety(self):
        import omni.kit
        # Generate multiple instances of the test setup to run them concurrently
        instance_setup = dict()
        for n in range(24):
            instance_setup[f"/TestGraph_{n}"] = _TestGraphAndNode()

        for i, test_run in enumerate(self.TEST_DATA):
            await _test_clear_scene(self, test_run)
            for (key, test_info) in instance_setup.items():
                test_info = await _test_setup_scene(self, og.Controller(allow_exists_prim=True), key, "TestNode_omni_graph_nodes_RandomNumeric", "omni.graph.nodes.RandomNumeric", test_run, test_info)
            self.assertEqual(len(og.get_all_graphs()), 24)

            # We want to evaluate all graphs concurrently. Kick them all.
            # Evaluate multiple times to skip 2 serial frames and increase chances for a race condition.
            for _ in range(10):
                await omni.kit.app.get_app().next_update_async()

            for (key, test_instance) in instance_setup.items():
                _test_verify_scene(self, og.Controller(), test_run, test_info, f"omni.graph.nodes.RandomNumeric User test case #{i+1}, instance{key}")

async def test_data_access(self):
    from omni.graph.nodes.ogn.OgnRandomNumericDatabase import OgnRandomNumericDatabase
    test_file_name = "OgnRandomNumericTemplate.usda"
    usd_path = os.path.join(os.path.dirname(__file__), "usd", test_file_name)
    if not os.path.exists(usd_path):
        self.assertTrue(False, f"{usd_path} not found for loading test")
    (result, error) = await ogts.load_test_file(usd_path)
    self.assertTrue(result, f'{error} on {usd_path}')
    test_node = og.Controller.node("/TestGraph/Template_omni_graph_nodes_RandomNumeric")
    database = OgnRandomNumericDatabase(test_node)
    self.assertTrue(test_node.is_valid())
    node_type_name = test_node.get_type_name()
    self.assertEqual(og.GraphRegistry().get_node_type_version(node_type_name), 1)

    def _attr_error(attribute: og.Attribute, usd_test: bool) -> str:
        test_type = "USD Load" if usd_test else "Database Access"
        return f"{node_type_name} {test_type} Test - {attribute.get_name()} value error"


    self.assertTrue(test_node.get_attribute_exists("inputs:execIn"))
    attribute = test_node.get_attribute("inputs:execIn")
    db_value = database.inputs.execIn

    self.assertTrue(test_node.get_attribute_exists("inputs:isNoise"))
    attribute = test_node.get_attribute("inputs:isNoise")
    db_value = database.inputs.isNoise
    expected_value = False
    actual_value = og.Controller.get(attribute)
    ogts.verify_values(expected_value, actual_value, _attr_error(attribute, True))
    ogts.verify_values(expected_value, db_value, _attr_error(attribute, False))

    self.assertTrue(test_node.get_attribute_exists("inputs:useSeed"))
    attribute = test_node.get_attribute("inputs:useSeed")
    db_value = database.inputs.useSeed
    expected_value = False
    actual_value = og.Controller.get(attribute)
    ogts.verify_values(expected_value, actual_value, _attr_error(attribute, True))
    ogts.verify_values(expected_value, db_value, _attr_error(attribute, False))

    self.assertTrue(test_node.get_attribute_exists("outputs:execOut"))
    attribute = test_node.get_attribute("outputs:execOut")
    db_value = database.outputs.execOut

    self.assertTrue(test_node.get_attribute_exists("state:gen"))
    attribute = test_node.get_attribute("state:gen")
    db_value = database.state.gen
