"""Suite of tests to exercise the DataWrapper class as a unit."""
import omni.graph.core as og
import omni.kit.test

# The test needs some non-standard imports to do its thing
from omni.graph.core._impl.dtypes import (
    Bool,
    BundleOutput,
    Double,
    Double2,
    Double3,
    Double4,
    Float,
    Float2,
    Float3,
    Float4,
    Half,
    Half2,
    Half3,
    Half4,
    Int,
    Int2,
    Int3,
    Int4,
    Int64,
    Matrix2d,
    Matrix3d,
    Matrix4d,
    Token,
    UChar,
    UInt,
    UInt64,
)


class TestDataWrapper(omni.kit.test.AsyncTestCase):
    """Wrapper for unit tests for the og.DataWrapper class"""

    # ----------------------------------------------------------------------
    async def test_data_wrapper_creation(self):
        """Test creation of the DataWrapper from the raw og.Type"""

        # Test data that enumerates all of the OGN types. Values of the list entries are:
        #   OGN type name
        #   Expected dtype
        #   Expected shape for non-gathered attributes
        #   Expected shape for gathered attributes
        test_data = [
            ("any", Token, None, (2,)),
            ("bool", Bool, None, (2,)),
            ("bool[]", Bool, (0,), ((2, 2))),
            ("bundle", BundleOutput, None, (2,)),
            ("colord[3]", Double3, (3,), ((2,), 3)),
            ("colord[3][]", Double3, (0, 3), (((2, 2)), 3)),
            ("colord[4]", Double4, (4,), ((2,), 4)),
            ("colord[4][]", Double4, (0, 4), ((2, 2), 4)),
            ("colorf[3]", Float3, (3,), ((2,), 3)),
            ("colorf[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("colorf[4]", Float4, (4,), ((2,), 4)),
            ("colorf[4][]", Float4, (0, 4), ((2, 2), 4)),
            ("colorh[3]", Half3, (3,), ((2,), 3)),
            ("colorh[3][]", Half3, (0, 3), ((2, 2), 3)),
            ("colorh[4]", Half4, (4,), ((2,), 4)),
            ("colorh[4][]", Half4, (0, 4), ((2, 2), 4)),
            ("double", Double, None, (2,)),
            ("double[]", Double, (0,), ((2, 2))),
            ("double[2]", Double2, (2,), ((2,), 2)),
            ("double[2][]", Double2, (0, 2), ((2, 2), 2)),
            ("double[3]", Double3, (3,), ((2,), 3)),
            ("double[3][]", Double3, (0, 3), ((2, 2), 3)),
            ("double[4]", Double4, (4,), ((2,), 4)),
            ("double[4][]", Double4, (0, 4), ((2, 2), 4)),
            ("execution", UInt, None, (2,)),
            ("float", Float, None, (2,)),
            ("float[]", Float, (0,), ((2, 2))),
            ("float[2]", Float2, (2,), ((2,), 2)),
            ("float[2][]", Float2, (0, 2), ((2, 2), 2)),
            ("float[3]", Float3, (3,), ((2,), 3)),
            ("float[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("float[4]", Float4, (4,), ((2,), 4)),
            ("float[4][]", Float4, (0, 4), ((2, 2), 4)),
            ("frame[4]", Matrix4d, (4, 4), ((2,), 4, 4)),
            ("frame[4][]", Matrix4d, (0, 4, 4), ((2, 2), 4, 4)),
            ("half", Half, None, (2,)),
            ("half[]", Half, (0,), ((2, 2))),
            ("half[2]", Half2, (2,), ((2,), 2)),
            ("half[2][]", Half2, (0, 2), ((2, 2), 2)),
            ("half[3]", Half3, (3,), ((2,), 3)),
            ("half[3][]", Half3, (0, 3), ((2, 2), 3)),
            ("half[4]", Half4, (4,), ((2,), 4)),
            ("half[4][]", Half4, (0, 4), ((2, 2), 4)),
            ("int", Int, None, (2,)),
            ("int[]", Int, (0,), ((2, 2))),
            ("int[2]", Int2, (2,), ((2,), 2)),
            ("int[2][]", Int2, (0, 2), ((2, 2), 2)),
            ("int[3]", Int3, (3,), ((2,), 3)),
            ("int[3][]", Int3, (0, 3), ((2, 2), 3)),
            ("int[4]", Int4, (4,), ((2,), 4)),
            ("int[4][]", Int4, (0, 4), ((2, 2), 4)),
            ("int64", Int64, None, (2,)),
            ("int64[]", Int64, (0,), ((2, 2))),
            ("matrixd[2]", Matrix2d, (2, 2), ((2,), 2, 2)),
            ("matrixd[2][]", Matrix2d, (0, 2, 2), ((2, 2), 2, 2)),
            ("matrixd[3]", Matrix3d, (3, 3), ((2,), 3, 3)),
            ("matrixd[3][]", Matrix3d, (0, 3, 3), ((2, 2), 3, 3)),
            ("matrixd[4]", Matrix4d, (4, 4), ((2,), 4, 4)),
            ("matrixd[4][]", Matrix4d, (0, 4, 4), ((2, 2), 4, 4)),
            ("normald[3]", Double3, (3,), ((2,), 3)),
            ("normald[3][]", Double3, (0, 3), (((2, 2)), 3)),
            ("normalf[3]", Float3, (3,), ((2,), 3)),
            ("normalf[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("normalh[3]", Half3, (3,), ((2,), 3)),
            ("normalh[3][]", Half3, (0, 3), ((2, 2), 3)),
            ("objectId", UInt64, None, (2,)),
            ("objectId[]", UInt64, (0,), ((2, 2))),
            ("pointd[3]", Double3, (3,), ((2,), 3)),
            ("pointd[3][]", Double3, (0, 3), (((2, 2)), 3)),
            ("pointf[3]", Float3, (3,), ((2,), 3)),
            ("pointf[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("pointh[3]", Half3, (3,), ((2,), 3)),
            ("pointh[3][]", Half3, (0, 3), ((2, 2), 3)),
            ("quatd[4]", Double4, (4,), ((2,), 4)),
            ("quatd[4][]", Double4, (0, 4), (((2, 2)), 4)),
            ("quatf[4]", Float4, (4,), ((2,), 4)),
            ("quatf[4][]", Float4, (0, 4), ((2, 2), 4)),
            ("quath[4]", Half4, (4,), ((2,), 4)),
            ("quath[4][]", Half4, (0, 4), ((2, 2), 4)),
            ("string", UChar, (0,), ((2, 2))),
            ("target", BundleOutput, None, (2,)),
            ("texcoordd[2]", Double2, (2,), ((2,), 2)),
            ("texcoordd[2][]", Double2, (0, 2), (((2, 2)), 2)),
            ("texcoordd[3]", Double3, (3,), ((2,), 3)),
            ("texcoordd[3][]", Double3, (0, 3), ((2, 2), 3)),
            ("texcoordf[2]", Float2, (2,), ((2,), 2)),
            ("texcoordf[2][]", Float2, (0, 2), ((2, 2), 2)),
            ("texcoordf[3]", Float3, (3,), ((2,), 3)),
            ("texcoordf[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("texcoordh[2]", Half2, (2,), ((2,), 2)),
            ("texcoordh[2][]", Half2, (0, 2), ((2, 2), 2)),
            ("texcoordh[3]", Half3, (3,), ((2,), 3)),
            ("texcoordh[3][]", Half3, (0, 3), ((2, 2), 3)),
            ("timecode", Double, None, (2,)),
            ("timecode[]", Double, (0,), ((2, 2))),
            ("token", Token, None, (2,)),
            ("token[]", Token, (0,), ((2, 2))),
            ("uchar", UChar, None, (2,)),
            ("uchar[]", UChar, (0,), ((2, 2))),
            ("uint", UInt, None, (2,)),
            ("uint[]", UInt, (0,), ((2, 2))),
            ("uint64", UInt64, None, (2,)),
            ("uint64[]", UInt64, (0,), ((2, 2))),
            ("vectord[3]", Double3, (3,), ((2,), 3)),
            ("vectord[3][]", Double3, (0, 3), (((2, 2)), 3)),
            ("vectorf[3]", Float3, (3,), ((2,), 3)),
            ("vectorf[3][]", Float3, (0, 3), ((2, 2), 3)),
            ("vectorh[3]", Half3, (3,), ((2,), 3)),
            ("vectorh[3][]", Half3, (0, 3), ((2, 2), 3)),
        ]

        for name, expected_dtype, expected_ungathered_shape, expected_gathered_shape in test_data:
            attribute_type = og.AttributeType.type_from_ogn_type_name(name)
            ungathered_dtype, ungathered_shape = og.data_shape_from_type(attribute_type, is_gathered=False)
            gathered_dtype, gathered_shape = og.data_shape_from_type(attribute_type, is_gathered=True)

            self.assertEqual(ungathered_dtype.__class__, expected_dtype, f"Mismatched type for {name}")
            self.assertEqual(ungathered_shape, expected_ungathered_shape, f"Mismatched shape for {name}")
            self.assertEqual(ungathered_dtype, gathered_dtype, f"Mismatched gathered type for {name}")
            self.assertEqual(gathered_shape, expected_gathered_shape, f"Mismatched gathered shape for {name}")

    # ----------------------------------------------------------------------
    async def test_data_wrapper_indexing(self):
        """Test indexing of the DataWrapper for all of the documented supported types"""

        # Test data drives all of the test cases for expected index results. The data is a list of test
        # configurations consisting of:
        #    Dtype of the parent wrapper
        #    Shape of the parent wrapper
        #    Index to be be taken to get the child
        #    Expected Dtype of the child wrapper (None means to expect an exception raised)
        #    Expected shape of the child wrapper
        #    Expected memory location offset of the child at index 1 from the parent
        test_data = [
            [Float, None, 1, None, None, 0],  # Single value
            [Float, (3,), 1, Float, None, Float.size],  # Array of 3 simple values
            [Float3, (3,), 1, Float, None, Float.size],  # Tuple of 3 simple values
            [Float3, (2, 3), 1, Float3, (3,), Float3.size],  # Array of 2 triple values
            [Matrix2d, (2, 2), 1, Double2, (2,), Double2.size],  # A 2x2 matrix
            [Matrix2d, (3, 2, 2), 1, Matrix2d, (2, 2), Matrix2d.size],  # An array of 3 2x2 matrixes
            [Float3, (3,), -1, None, None, 0],  # Index too low
            [Float3, (3, 3), -1, None, None, 0],
            [Float3, ((3,), 3), -1, None, None, 0],
            [Float3, ((3, 4), 3), -1, None, None, 0],
            [Float, (3,), -1, None, None, 0],
            [Float, ((3,), 3), -1, None, None, 0],
            [Float, ((3, 4), 3), -1, None, None, 0],
            [Matrix2d, (2, 2), -1, None, None, 0],
            [Matrix2d, (3, 2, 2), -1, None, None, 0],
            [Matrix2d, ((3,), 2, 2), -1, None, None, 0],
            [Matrix2d, ((3, 4), 2, 2), -1, None, None, 0],
            [Float3, (3,), 4, None, None, 0],  # Index too high
            [Float3, (3, 3), 4, None, None, 0],
            [Float3, ((3,), 3), 4, None, None, 0],
            [Float3, ((3, 4), 3), 4, None, None, 0],
            [Float, (3,), 4, None, None, 0],
            [Float, ((3,), 3), 4, None, None, 0],
            [Float, ((3, 4), 3), 4, None, None, 0],
            [Matrix2d, (2, 2), 4, None, None, 0],
            [Matrix2d, (3, 2, 2), 4, None, None, 0],
            [Matrix2d, ((3,), 2, 2), 4, None, None, 0],
            [Matrix2d, ((3, 4), 2, 2), 4, None, None, 0],
            [Float, (3, 3), 1, None, None, 0],  # Too many levels of array
            [Float4, (3,), 1, None, None, 0],  # Mismatched shape size for tuple
            [Float4, (3, 3), 1, None, None, 0],  # Mismatched shape size for tuple array
            [Matrix2d, (3, 3), 1, None, None, 0],  # Mismatched shape size for matrix
            [Matrix2d, (2, 3, 3), 1, None, None, 0],  # Mismatched shape size for matrix array
        ]

        for parent_dtype, parent_shape, child_index, child_dtype, child_shape, child_offset in test_data:
            if child_dtype is None:
                with self.assertRaises(ValueError):
                    parent_wrapper = og.DataWrapper(0, parent_dtype, parent_shape, og.Device.cuda)
                    _ = parent_wrapper[child_index]
            else:
                parent_wrapper = og.DataWrapper(0, parent_dtype, parent_shape, og.Device.cuda)
                child_wrapper = parent_wrapper[child_index]
                self.assertEqual(child_wrapper.dtype.tuple_count, child_dtype.tuple_count)
                self.assertEqual(child_wrapper.dtype.base_type, child_dtype.base_type)
                self.assertEqual(child_wrapper.shape, child_shape)
                self.assertEqual(child_wrapper.memory, child_offset)
