#*****************************************************************************
# Copyright 2022 NVIDIA Corporation. All rights reserved.
#*****************************************************************************

import sys
import os
import gc
import platform
import traceback
from enum import Enum
from pxr import Usd
from pxr import Plug
from pxr import Sdr
from pxr import Sdf
from pxr import Tf
from pxr import UsdShade
from pxr import Gf
from pxr import UsdGeom
from pxr import Kind
from pxr import UsdUI
import numpy        # used to represent Vectors, Matrices, and Colors
import re           # For command line args
from copy import deepcopy # Dictionary merging
import tempfile

# load the binding module
print("MDL python binding about to load ...")
####################
# Omniverse code
from omni.mdl import pymdlsdk
from omni.mdl import pymdl
import omni.client
async def copy_async(texture_path, dst_path):
    result = await omni.client.copy_async(texture_path, dst_path)
    return (result == omni.client.Result.OK)
# Omniverse code
####################
                
# TODO: loading fails for debug builds of the pymdlsdk.pyd (or pymdlsdk.so on unix) module 
print("MDL python binding loaded")

def str2bool(v):
  return str(v).lower() in ("true", "1")

def str2MDL_TO_USD_OUTPUT(v):
    if v == '1': return OutputType.SHADER
    if v == '2': return OutputType.MATERIAL
    if v == '3': return OutputType.MATERIAL_AND_GEOMETRY

#--------------------------------------------------------------------------------------------------
# Command line arguments
#--------------------------------------------------------------------------------------------------
class Args: # pragma: no cover
    display_usage = False
    options = dict()
    def __init__(self):
        self.options["SEARCH_PATH"] = get_examples_search_path()
        self.options["MODULE"] = "::nvidia::sdk_examples::tutorials"
        self.options["PRIM"] = "/example_material" # Used for USD to MDL conversion
        self.options["STAGE"] = "tutorials.usda"
        self.options["OUT"] = "tutorials.usda"
        # Set FOR_OV to False if a Material/Shader hierarchy is required in the output USD stage
        # SHADER = 1
        # MATERIAL = 2
        # MATERIAL_AND_GEOMETRY = 3
        self.options["MDL_TO_USD_OUTPUT"] = 1
        self.options["NESTED_SHADERS"] = False # See mdl_to_usd_output_material_nested_shaders
        self.options["MDL_TO_USD"] = True
        for a in sys.argv:
            if a == "-help" or a == "-h" or a == "-?":
                self.display_usage = True
                break
            match = re.search("(.+)=(.+)", a)
            if match:
                self.options[match.group(1)] = match.group(2)
        # Convert paths
        self.options["SEARCH_PATH"] = os.path.abspath(self.options["SEARCH_PATH"])
        self.options["OUT"] = os.path.abspath(self.options["OUT"])
        # Convert string args to bool
        self.options["MDL_TO_USD"] = str2bool(self.options["MDL_TO_USD"])
        self.options["NESTED_SHADERS"] = str2bool(self.options["NESTED_SHADERS"])
        # Convert string args to int
        self.options["MDL_TO_USD_OUTPUT"] = str2MDL_TO_USD_OUTPUT(self.options["MDL_TO_USD_OUTPUT"])
        
    def usage(self):
        print("usage: " + os.path.basename(sys.argv[0]) + " [-help|-h|-?] [SEARCH_PATH=<path>] [MODULE=<string>] [PRIM=<string>] [STAGE=<string>] [OUT=<string>] [MDL_TO_USD_OUTPUT=<int>] [NESTED_SHADERS=<bool>] [MDL_TO_USD=<bool>]")
        print("\n")
        print("  -help|-h|-?:\t Display usage")
        print("  SEARCH_PATH:\t Use this as MDL search path (default: {})".format(self.options["SEARCH_PATH"]))
        print("  MODULE:\t Module to load (default: {})".format(self.options["MODULE"]))
        print("  PRIM:\t\t USD prim to convert to MDL (default: {})".format(self.options["PRIM"]))
        print("  STAGE:\t USD stage to load and convert (default: {})".format(self.options["STAGE"]))
        print("  OUT:\t\t Output file (default: {})".format(self.options["OUT"]))
        print("  MDL_TO_USD_OUTPUT:\t Output either a shader (1) or material (2) or material and geometry (3) (default: {})".format(self.options["MDL_TO_USD_OUTPUT"]))
        print("  NESTED_SHADERS:\t Output nested shader tree, vs. flat shader tree (default: {})".format(self.options["NESTED_SHADERS"]))
        print("  MDL_TO_USD:\t MDL to USD conversion ortherwise USD to MDL (default: {})".format(self.options["MDL_TO_USD"]))

#--------------------------------------------------------------------------------------------------
# USD and misc utilities
#--------------------------------------------------------------------------------------------------

MdlIValueKindToUSD = {
    pymdlsdk.IValue.Kind.VK_BOOL: Sdf.ValueTypeNames.Bool,
    pymdlsdk.IValue.Kind.VK_INT : Sdf.ValueTypeNames.Int,
    pymdlsdk.IValue.Kind.VK_ENUM : Sdf.ValueTypeNames.Int,
    pymdlsdk.IValue.Kind.VK_FLOAT : Sdf.ValueTypeNames.Float,
    pymdlsdk.IValue.Kind.VK_DOUBLE : Sdf.ValueTypeNames.Double,
    pymdlsdk.IValue.Kind.VK_STRING : Sdf.ValueTypeNames.String,
    pymdlsdk.IValue.Kind.VK_COLOR : Sdf.ValueTypeNames.Color3f,
    pymdlsdk.IValue.Kind.VK_STRUCT : Sdf.ValueTypeNames.Token,
    pymdlsdk.IValue.Kind.VK_TEXTURE : Sdf.ValueTypeNames.Asset,
    pymdlsdk.IValue.Kind.VK_LIGHT_PROFILE : Sdf.ValueTypeNames.Asset,
    pymdlsdk.IValue.Kind.VK_BSDF_MEASUREMENT : Sdf.ValueTypeNames.Asset
}

MdlITypeKindToUSD = {
    pymdlsdk.IType.Kind.TK_BOOL: Sdf.ValueTypeNames.Bool,
    pymdlsdk.IType.Kind.TK_INT : Sdf.ValueTypeNames.Int,
    pymdlsdk.IType.Kind.TK_ENUM : Sdf.ValueTypeNames.Int,
    pymdlsdk.IType.Kind.TK_FLOAT : Sdf.ValueTypeNames.Float,
    pymdlsdk.IType.Kind.TK_DOUBLE : Sdf.ValueTypeNames.Double,
    pymdlsdk.IType.Kind.TK_STRING : Sdf.ValueTypeNames.String,
    pymdlsdk.IType.Kind.TK_COLOR : Sdf.ValueTypeNames.Color3f,
    pymdlsdk.IType.Kind.TK_STRUCT : Sdf.ValueTypeNames.Token,
    pymdlsdk.IType.Kind.TK_TEXTURE : Sdf.ValueTypeNames.Asset,
    pymdlsdk.IType.Kind.TK_LIGHT_PROFILE : Sdf.ValueTypeNames.Asset,
    pymdlsdk.IType.Kind.TK_BSDF_MEASUREMENT : Sdf.ValueTypeNames.Asset
}
# TODO: 
# mi::neuraylib::IType::TK_BSDF
# mi::neuraylib::IType::TK_EDF
# mi::neuraylib::IType::TK_VDF
def python_vector_to_usd_type(dtype, size):
    if dtype == numpy.int32 or dtype == bool:
        if size == 2:
            return Sdf.ValueTypeNames.Int2
        elif size == 3:
            return Sdf.ValueTypeNames.Int3
        elif size == 4:
            return Sdf.ValueTypeNames.Int4
    elif dtype == numpy.float32:
        if size == 2:
            return Sdf.ValueTypeNames.Float2
        elif size == 3:
            return Sdf.ValueTypeNames.Float3
        elif size == 4:
            return Sdf.ValueTypeNames.Float4
    elif dtype == numpy.float64:
        if size == 2:
            return Sdf.ValueTypeNames.Double2
        elif size == 3:
            return Sdf.ValueTypeNames.Double3
        elif size == 4:
            return Sdf.ValueTypeNames.Double4
    return None

def python_vector_to_usd_value(value):
    dtype = value.dtype
    size = value.size
    out = None
    if dtype == numpy.int32 or dtype == bool:
        if size == 2:
            out = Gf.Vec2i(0)
        elif size == 3:
            out = Gf.Vec3i(0)
        elif size == 4:
            out = Gf.Vec4i(0)
        if out != None:
            for i in range(size):
                out[i] = int(value[i][0])
    elif dtype == numpy.float32:
        if size == 2:
            out = Gf.Vec2f(0)
        elif size == 3:
            out = Gf.Vec3f(0)
        elif size == 4:
            out = Gf.Vec4f(0)
        if out != None:
            for i in range(size):
                out[i] = float(value[i][0])
    elif dtype == numpy.float64:
        if size == 2:
            out = Gf.Vec2d(0)
        elif size == 3:
            out = Gf.Vec3d(0)
        elif size == 4:
            out = Gf.Vec4d(0)
        if out != None:
            for i in range(size):
                out[i] = numpy.float64(value[i][0])
    return out

def custom_data_from_python_vector(value):
    dtype = value.dtype
    out = dict()
    if dtype == bool:
        size = value.size
        out = dict({"mdl":{"type" : "bool{}".format(size)}})
    return out

def python_matrix_to_usd_type(dtype, nrow, ncol):
    #     numpyType   Column Row           OutType
    m = { 'float32' : { 2 : { 2 : Sdf.ValueTypeNames.Matrix2d,
                                3 : Sdf.ValueTypeNames.Float3Array,
                                4 : Sdf.ValueTypeNames.Float4Array},
                          3 : { 2 : Sdf.ValueTypeNames.Float2Array,
                                3 : Sdf.ValueTypeNames.Matrix3d,
                                4 : Sdf.ValueTypeNames.Float4Array},
                          4 : { 2 : Sdf.ValueTypeNames.Float2Array,
                                3 : Sdf.ValueTypeNames.Float3Array,
                                4 : Sdf.ValueTypeNames.Matrix4d }},
          'float64' : { 2 : { 2 : Sdf.ValueTypeNames.Matrix2d,
                                3 : Sdf.ValueTypeNames.Double3Array,
                                4 : Sdf.ValueTypeNames.Double4Array},
                          3 : { 2 : Sdf.ValueTypeNames.Double2Array,
                                3 : Sdf.ValueTypeNames.Matrix3d,
                                4 : Sdf.ValueTypeNames.Double4Array},
                          4 : { 2 : Sdf.ValueTypeNames.Double2Array,
                                3 : Sdf.ValueTypeNames.Double3Array,
                                4 : Sdf.ValueTypeNames.Matrix4d }}}
    return m[dtype.name][ncol][nrow]

def python_matrix_to_usd_value(value):
    dtype = value.dtype
    nrow = value.shape[0]
    ncol = value.shape[1]
    out = None
    if dtype == numpy.float32:
        if ncol == 2:
            if nrow == 2:
                out = Gf.Matrix2d(0)
                out.SetColumn(0, Gf.Vec2d(float(value[0][0]), float(value[0][1])))
                out.SetColumn(1, Gf.Vec2d(float(value[1][0]), float(value[1][1])))
            elif nrow == 3:
                out = [Gf.Vec3f(float(value[0][0]), float(value[1][0]), float(value[2][0])),
                       Gf.Vec3f(float(value[0][1]), float(value[1][1]), float(value[2][1]))]
            elif nrow == 4:
                out = [Gf.Vec4f(float(value[0][0]), float(value[1][0]), float(value[2][0]), float(value[3][0])),
                       Gf.Vec4f(float(value[0][1]), float(value[1][1]), float(value[2][1]), float(value[3][1]))]
        elif ncol == 3:
            if nrow == 2:
                out = [Gf.Vec2f(float(value[0][0]), float(value[1][0])),
                       Gf.Vec2f(float(value[0][1]), float(value[1][1])),
                       Gf.Vec2f(float(value[0][2]), float(value[1][2]))]
            elif nrow == 3:
                out = Gf.Matrix3d(0)
                out.SetColumn(0, Gf.Vec3d(float(value[0][0]), float(value[0][1]), float(value[0][2])))
                out.SetColumn(1, Gf.Vec3d(float(value[1][0]), float(value[1][1]), float(value[1][2])))
                out.SetColumn(2, Gf.Vec3d(float(value[2][0]), float(value[2][1]), float(value[2][2])))
            elif nrow == 4:
                out = [Gf.Vec4f(float(value[0][0]), float(value[1][0]), float(value[2][0]), float(value[3][0])),
                       Gf.Vec4f(float(value[0][1]), float(value[1][1]), float(value[2][1]), float(value[3][1])),
                       Gf.Vec4f(float(value[0][2]), float(value[1][2]), float(value[2][2]), float(value[3][2]))]
        elif ncol == 4:
            if nrow == 2:
                out = [Gf.Vec2f(float(value[0][0]), float(value[1][0])),
                       Gf.Vec2f(float(value[0][1]), float(value[1][1])),
                       Gf.Vec2f(float(value[0][2]), float(value[1][2])),
                       Gf.Vec2f(float(value[0][3]), float(value[1][3]))]
            elif nrow == 3:
                out = [Gf.Vec3f(float(value[0][0]), float(value[1][0]), float(value[2][0])),
                       Gf.Vec3f(float(value[0][1]), float(value[1][1]), float(value[2][1])),
                       Gf.Vec3f(float(value[0][2]), float(value[1][2]), float(value[2][2])),
                       Gf.Vec3f(float(value[0][3]), float(value[1][3]), float(value[2][3]))]
            elif nrow == 4:
                out = Gf.Matrix4d(0)
                out.SetColumn(0, Gf.Vec4d(float(value[0][0]), float(value[0][1]), float(value[0][2]), float(value[0][3])))
                out.SetColumn(1, Gf.Vec4d(float(value[1][0]), float(value[1][1]), float(value[1][2]), float(value[1][3])))
                out.SetColumn(2, Gf.Vec4d(float(value[2][0]), float(value[2][1]), float(value[2][2]), float(value[2][3])))
                out.SetColumn(3, Gf.Vec4d(float(value[3][0]), float(value[3][1]), float(value[3][2]), float(value[3][3])))
    elif dtype == numpy.float64:
        if ncol == 2:
            if nrow == 2:
                out = Gf.Matrix2d(0)
                out.SetColumn(0, Gf.Vec2d(float(value[0][0]), float(value[0][1])))
                out.SetColumn(1, Gf.Vec2d(float(value[1][0]), float(value[1][1])))
            elif nrow == 3:
                out = [Gf.Vec3d(float(value[0][0]), float(value[1][0]), float(value[2][0])),
                       Gf.Vec3d(float(value[0][1]), float(value[1][1]), float(value[2][1]))]
            elif nrow == 4:
                out = [Gf.Vec4d(float(value[0][0]), float(value[1][0]), float(value[2][0]), float(value[3][0])),
                       Gf.Vec4d(float(value[0][1]), float(value[1][1]), float(value[2][1]), float(value[3][1]))]
        elif ncol == 3:
            if nrow == 2:
                out = [Gf.Vec2d(float(value[0][0]), float(value[1][0])),
                       Gf.Vec2d(float(value[0][1]), float(value[1][1])),
                       Gf.Vec2d(float(value[0][2]), float(value[1][2]))]
            elif nrow == 3:
                out = Gf.Matrix3d(0)
                out.SetColumn(0, Gf.Vec3d(float(value[0][0]), float(value[0][1]), float(value[0][2])))
                out.SetColumn(1, Gf.Vec3d(float(value[1][0]), float(value[1][1]), float(value[1][2])))
                out.SetColumn(2, Gf.Vec3d(float(value[2][0]), float(value[2][1]), float(value[2][2])))
            elif nrow == 4:
                out = [Gf.Vec4d(float(value[0][0]), float(value[1][0]), float(value[2][0]), float(value[3][0])),
                       Gf.Vec4d(float(value[0][1]), float(value[1][1]), float(value[2][1]), float(value[3][1])),
                       Gf.Vec4d(float(value[0][2]), float(value[1][2]), float(value[2][2]), float(value[3][2]))]
        elif ncol == 4:
            if nrow == 2:
                out = [Gf.Vec2d(float(value[0][0]), float(value[1][0])),
                       Gf.Vec2d(float(value[0][1]), float(value[1][1])),
                       Gf.Vec2d(float(value[0][2]), float(value[1][2])),
                       Gf.Vec2d(float(value[0][3]), float(value[1][3]))]
            elif nrow == 3:
                out = [Gf.Vec3d(float(value[0][0]), float(value[1][0]), float(value[2][0])),
                       Gf.Vec3d(float(value[0][1]), float(value[1][1]), float(value[2][1])),
                       Gf.Vec3d(float(value[0][2]), float(value[1][2]), float(value[2][2])),
                       Gf.Vec3d(float(value[0][3]), float(value[1][3]), float(value[2][3]))]
            elif nrow == 4:
                out = Gf.Matrix4d(0)
                out.SetColumn(0, Gf.Vec4d(float(value[0][0]), float(value[0][1]), float(value[0][2]), float(value[0][3])))
                out.SetColumn(1, Gf.Vec4d(float(value[1][0]), float(value[1][1]), float(value[1][2]), float(value[1][3])))
                out.SetColumn(2, Gf.Vec4d(float(value[2][0]), float(value[2][1]), float(value[2][2]), float(value[2][3])))
                out.SetColumn(3, Gf.Vec4d(float(value[3][0]), float(value[3][1]), float(value[3][2]), float(value[3][3])))
    return out

def python_array_to_usd_type(value):
    array_simple_conversion = { 
        bool : Sdf.ValueTypeNames.BoolArray,
        int : Sdf.ValueTypeNames.IntArray,
        # TODO: enum?
        # AddArrayOfSimpleConversion(mi::neuraylib::IType::TK_ENUM, SdfValueTypeNames->IntArray);
        float : Sdf.ValueTypeNames.FloatArray,
        # TODO: double ?
        # AddArrayOfSimpleConversion(mi::neuraylib::IType::TK_DOUBLE, SdfValueTypeNames->DoubleArray);
        str : Sdf.ValueTypeNames.StringArray,
        pymdlsdk.IType.Kind.TK_COLOR :  Sdf.ValueTypeNames.Color3fArray
        # TODO: AddArrayOfSimpleConversion(mi::neuraylib::IType::TK_STRUCT, SdfValueTypeNames->TokenArray);
        # TODO: AddArrayOfSimpleConversion(mi::neuraylib::IType::TK_TEXTURE, SdfValueTypeNames->AssetArray);
        }

    array_of_vector_conversion = {
        numpy.bool_ : { 2 : Sdf.ValueTypeNames.Int2Array, 
                        3 : Sdf.ValueTypeNames.Int3Array, 
                        4 : Sdf.ValueTypeNames.Int4Array },
        numpy.int32 : { 2 : Sdf.ValueTypeNames.Int2Array, 
                        3 : Sdf.ValueTypeNames.Int3Array, 
                        4 : Sdf.ValueTypeNames.Int4Array },
        numpy.float32 : { 2 : Sdf.ValueTypeNames.Float2Array, 
                          3 : Sdf.ValueTypeNames.Float3Array, 
                          4 : Sdf.ValueTypeNames.Float4Array },
        numpy.float64 : { 2 : Sdf.ValueTypeNames.Double2Array, 
                          3 : Sdf.ValueTypeNames.Double3Array, 
                          4 : Sdf.ValueTypeNames.Double4Array }
    }

    # // Array of matrix
    #     numpyType   Column Row           OutType
    array_of_matrix = \
        { numpy.float32 : { 2 : { 2 : Sdf.ValueTypeNames.Matrix2dArray,
                                3 : Sdf.ValueTypeNames.FloatArray,
                                4 : Sdf.ValueTypeNames.FloatArray},
                          3 : { 2 : Sdf.ValueTypeNames.FloatArray,
                                3 : Sdf.ValueTypeNames.Matrix3dArray,
                                4 : Sdf.ValueTypeNames.FloatArray},
                          4 : { 2 : Sdf.ValueTypeNames.FloatArray,
                                3 : Sdf.ValueTypeNames.FloatArray,
                                4 : Sdf.ValueTypeNames.Matrix4dArray }},
          numpy.float64 : { 2 : { 2 : Sdf.ValueTypeNames.Matrix2dArray,
                                3 : Sdf.ValueTypeNames.DoubleArray,
                                4 : Sdf.ValueTypeNames.DoubleArray},
                          3 : { 2 : Sdf.ValueTypeNames.DoubleArray,
                                3 : Sdf.ValueTypeNames.Matrix3dArray,
                                4 : Sdf.ValueTypeNames.DoubleArray},
                          4 : { 2 : Sdf.ValueTypeNames.DoubleArray,
                                3 : Sdf.ValueTypeNames.DoubleArray,
                                4 : Sdf.ValueTypeNames.Matrix4dArray }}}


    dtype = type(value[0])
    if dtype in array_simple_conversion:
        return array_simple_conversion[dtype]
    elif dtype == numpy.ndarray:
        shape = value[0].shape
        # Array
        if len(shape) > 1 and shape[1] > 1:
            # Array of matrices
            elemtype = type(value[0][0][0])
            if elemtype in array_of_matrix:
                if shape[1] in array_of_matrix[elemtype]:
                    if shape[0] in array_of_matrix[elemtype][shape[1]]:
                        return array_of_matrix[elemtype][shape[1]][shape[0]]
        elif type(value[0][0]) == numpy.ndarray:
            # Array of vectors
            etype = type(value[0][0][0])
            size = len(value[0])
            if etype in array_of_vector_conversion:
                if size in array_of_vector_conversion[etype]:
                    return array_of_vector_conversion[etype][size]
            else:
                print("Array of vector type not supported for type/size: {}/{}".format(etype, size))
        elif len(value[0]) == 3:
            # Assume this is a color
            return array_simple_conversion[pymdlsdk.IType.Kind.TK_COLOR]
        else:
            print("Array type not supported for type/shape: {}/{}".format(dtype, shape))
    else:
        print("Type not supported for type: {}".format(dtype))

def python_array_to_usd_value(value):
    dtype = type(value[0])
    out = None
    if dtype == bool or \
        dtype == int or \
        dtype == float or \
        dtype == str:
        return value
    elif dtype == numpy.ndarray:
        # Array
        shape = value[0].shape
        if len(shape) > 1 and shape[1] > 1:
            # Array of matrices
            elemtype = type(value[0][0][0])
            nrow = shape[0]
            ncol = shape[1]
            out = None
            arraysize = len(value)
            if ncol == nrow:
                if ncol == 2:
                    mattype = Gf.Matrix2d
                if ncol == 3:
                    mattype = Gf.Matrix3d
                if ncol == 4:
                    mattype = Gf.Matrix4d
                out = []
                for i in range(arraysize):
                    mat = mattype(0)
                    for c in range(ncol):
                        row = []
                        for r in range(nrow):
                            row.append(float(value[i][c][r]))
                        mat.SetColumn(c, row)                         
                    out.append(mat)
            else:
                if elemtype == numpy.float32:
                    otype = float
                elif elemtype == numpy.float64:
                    otype = numpy.float64
                out = []
                for i in range(arraysize):
                    for c in range(ncol):
                        for r in range(nrow):
                            out.append(otype(value[i][r][c]))
        elif type(value[0][0]) == numpy.ndarray:
            # Array of vectors
            etype = type(value[0][0][0])
            size = len(value)
            out = []
            if etype == numpy.bool_ or \
               etype == numpy.int32:
                otype = int
            elif etype == numpy.float32:
                otype = float
            elif etype == numpy.float64:
                otype = numpy.float64
            itemsize = len(value[0])
            for i in range(size):
                ll = []
                for j in range(itemsize):
                    ll.append(otype(value[i][j]))
                out.append(ll)
        elif len(value[0]) == 3:
            # Assume this is a color
            size = len(value)
            out = []
            for i in range(size):
                out.append([value[0][0], value[0][1], value[0][2]])
    return out

def custom_data_from_python_array(value):
    # type = dict({"mdl":{"type" : "array of matrix"}})
    out = dict()
    dtype = type(value[0])
    arraysize = len(value)
    if dtype == numpy.ndarray:
        # Array
        shape = value[0].shape
        if len(shape) > 1 and shape[1] > 1:
            # Array of matrices
            elemtype = type(value[0][0][0])
            nrow = shape[0]
            ncol = shape[1]
            elemtype = type(value[0][0][0])
            if ncol != nrow or elemtype == numpy.float32:
                # Non square matrices
                typename = ""
                if elemtype == numpy.float32:
                    typename = "float"
                elif elemtype == numpy.float64:
                    typename = "double"
                out = dict({"mdl":{"type" : "{}{}x{}[{}]".format(typename, ncol, nrow, arraysize)}})
        elif type(value[0][0]) == numpy.ndarray:
            # Array of vectors
            etype = type(value[0][0][0])
            if etype == numpy.bool_:
                itemsize = len(value[0])
                out = dict({"mdl":{"type" : "bool{}[{}]".format(itemsize, arraysize)}})
    return out

def mdl_type_to_usd_type(val : pymdl.ArgumentConstant):
    if not val:
        return None
    kind = val.type.kind
    if kind in MdlITypeKindToUSD:
        return MdlITypeKindToUSD[kind]
    else:
        # A bit more work is required to derive the type
        if kind == pymdlsdk.IType.Kind.TK_VECTOR:
            return python_vector_to_usd_type(val.value.dtype, val.value.size)
        elif kind == pymdlsdk.IType.Kind.TK_MATRIX:
            return python_matrix_to_usd_type(val.value.dtype, val.value.shape[0], val.value.shape[1])
        elif kind == pymdlsdk.IType.Kind.TK_ARRAY:
            return python_array_to_usd_type(val.value)
    return None

class Stack:
    def __init__(self):
        self.items = []

    def isEmpty(self):
        return self.items == []

    def push(self, item):
        self.items.append(item)

    def pop(self):
        return self.items.pop()

    def top(self):
        return self.items[len(self.items)-1]

    def size(self):
        return len(self.items)

# When converting from MDL to USD user can choose the USD output format
# SHADER: A single shader is output at the root of the USD Stage
# MATERIAL: A single material is output at the root of the USD Stage
# MATERIAL_AND_GEOMETRY: Geometry and Material (bound to the geometry) are output in the USD Stage
class OutputType(Enum):
    SHADER = 1
    MATERIAL = 2
    MATERIAL_AND_GEOMETRY = 3

class ConverterContext():


    def __init__(self):
        self.neuray = None
        self.transaction = None
        self.stage = None
        self.modules = Stack()
        self.usd_materials = Stack()
        self.parameter_names = Stack()
        self.usd_shaders = Stack()
        self.custom_data = Stack()
        self.type_annotation = True
        self.mdl_to_usd_output = OutputType.SHADER
        self.mdl_to_usd_output_material_nested_shaders = False # Nested/deep or flat shader tree
        self.ov_neuray = None
        GetUniqueNameInStage.s_uniqueID = 0 # Reset the unique name counter to get same output between 2 runs

    def enable_type_annotation(self):
        self.type_annotation = True

    def disable_type_annotation(self):
        self.type_annotation = False

    def is_type_annotation_enabled(self):
        return self.type_annotation == True

    def set_neuray(self, neuray):
        self.neuray = neuray

    def set_transaction(self, transaction):
        self.transaction = transaction

    def set_stage(self, stage):
        self.stage = stage

    def push_module(self, module):
        self.modules.push(module)

    def pop_module(self):
        return self.modules.pop()

    def current_module(self):
        return self.modules.top()

    def push_usd_material(self, material):
        self.usd_materials.push(material)

    def pop_usd_material(self):
        return self.usd_materials.pop()

    def current_usd_material(self):
        return self.usd_materials.top()

    def push_parameter_name(self, parm):
        self.parameter_names.push(parm)

    def pop_parameter_name(self):
        return self.parameter_names.pop()

    def current_parameter_name(self):
        return self.parameter_names.top()

    def push_usd_shader(self, shader):
        self.usd_shaders.push(shader)

    def pop_usd_shader(self):
        return self.usd_shaders.pop()

    def current_usd_shader(self):
        return self.usd_shaders.top()

    def push_custom_data(self, custom_data):
        self.custom_data.push(custom_data)

    def pop_custom_data(self):
        return self.custom_data.pop()

    def current_custom_data(self):
        return self.custom_data.top()

# Example: 'mdl::OmniSurface::OmniSurfaceBase::OmniSurfaceBase'
# Return: 'OmniSurface/OmniSurfaceBase.mdl'
def prototype_to_source_asset(prototype):
    source_asset = prototype
    if source_asset[:5] == "mdl::":
        source_asset = source_asset[3:] # Remove leading "mdl"
    if source_asset[:2] == "::":
        source_asset = source_asset[2:] # Remove leading "::"
    ll = source_asset.split("::")
    source_asset = '/'.join(ll[0:-1]) # Drop last part which is material name
    source_asset = source_asset + ".mdl" # Append ".mdl" extension
    return source_asset

def module_mdl_name_to_source_asset(context, module_mdl_name):

    # Hack to convert builtin module functions
    if module_mdl_name == '::scene' or module_mdl_name == '::state':
        return 'nvidia/support_definitions.mdl'

    unmangle_helper = UnmangleAndFixMDLSearchPath(context)
    (isMangled, wasSuccess, newmodule_mdl_name) = unmangle_helper.unmangle_mdl_module(module_mdl_name)
    if wasSuccess and newmodule_mdl_name != module_mdl_name:
        return newmodule_mdl_name + '.mdl'
    if isMangled and not wasSuccess:
        # Better return empty asset than garbage asset
        return ''

    source_asset = module_mdl_name
    if source_asset[:5] == "mdl::":
        source_asset = source_asset[3:] # Remove leading "mdl"
    if source_asset[:2] == "::":
        source_asset = source_asset[2:] # Remove leading "::"
    source_asset = source_asset.replace("::", "/") # Replace "::" with "/"
    source_asset = source_asset + ".mdl" # Append ".mdl" extension
    return source_asset

def source_asset_to_module_name(source_asset, context):
    module_name = source_asset
    module_name = os.path.normpath(module_name)
    if os.path.isabs(module_name):
        # Absolute path, try to find a mathing search path and remove it
        with context.neuray.get_api_component(pymdlsdk.IMdl_configuration) as cfg:
            if cfg.is_valid_interface():
                for i in range(cfg.get_mdl_paths_length()):
                    sp = cfg.get_mdl_path(i)
                    sp = os.path.normpath(sp.get_c_str())
                    if module_name.find(sp) == 0:
                        module_name = module_name[len(sp):]
                        break
    # Split path
    ll = module_name.split(os.sep)
    # Concat path using '::' separator
    module_name = "::".join(ll)
    if not module_name[0:2] == "::":
        # Prepend '::'
        module_name = "::" + module_name
    if module_name[-4:] == ".mdl":
        # Remove '.mdl' extension
        module_name = module_name[0:-4]
    return module_name

def db_name_to_prim_name(dbname):
    name = dbname
    name = name.split("(")[0] # Remove arguments
    name = name.split("::")[-1] # Remove package, module
    name = name.replace(".", "_") # Replace '.'
    return name

def mdl_name_to_sub_identifier(mdl_name, fct_call=False):
    sub_id = mdl_name
    if sub_id[:5] == "mdl::":
        sub_id = sub_id[3:] # Remove leading "mdl"
    # split the identifier from the arguments
    names = sub_id.split("(")
    # keep only the last name from the identifier
    sub_id = names[0].split("::")[-1]
    if fct_call:
        # concat with args
        if len(names) > 1:
            sub_id = sub_id + "(" + names[1]
    return sub_id

# Return a unique SdfPath for the current stage
class GetUniqueNameInStage:
    s_uniqueID = 0
    @staticmethod
    def get(stage, path):
        unique_path = path
        while stage.GetPrimAtPath(unique_path).IsValid():
            unique_path = Sdf.Path(str(path) + str(GetUniqueNameInStage.s_uniqueID))
            GetUniqueNameInStage.s_uniqueID += 1
        return unique_path

class GetUniqueName:
    s_uniqueID = 0
    @staticmethod
    def get(transaction, prefix_in):
        prefix = prefix_in
        if len(prefix)==0:
            prefix = "elt"
        if not transaction.access_as(pymdlsdk.IInterface, prefix).is_valid_interface():
            return prefix
        while True:
            prefix = prefix_in + "_" + str(GetUniqueName.s_uniqueID)
            GetUniqueName.s_uniqueID += 1
            if not transaction.access_as(pymdlsdk.IInterface, prefix).is_valid_interface():
                return prefix

def create_stage(args):
    stageName = args.options["OUT"]
    stage = Usd.Stage.CreateNew(stageName)
    stage.SetMetadata('comment', 'MDL to USD conversion')
    return stage

def load_stage(stageName):
    stage = None
    try:
        stage = Usd.Stage.Open(stageName)
    except:
        print("Error:\tFailed to load stage: {}".format(stageName))
    return stage

def save_stage(stage):
    stage.GetRootLayer().Save()

def export_stage(stage, filename):
    stage.GetRootLayer().Export(filename)

def material_to_stage_output_material(context, function: pymdl.FunctionDefinition):
    return material_to_stage_output_material_and_geo(context, function, want_geometry = False)

def material_to_stage_output_material_and_geo(context, function: pymdl.FunctionDefinition, want_geometry = True):
    neuray = context.neuray
    transaction = context.transaction
    stage = context.stage

    db_name = function.dbName

    materialName = db_name_to_prim_name(db_name)
    
    # Start at root
    rootPath = Sdf.Path('/')
    spherePrim = None

    if want_geometry:
        xform = UsdGeom.Xform.Define(stage,"/World")
        # Add geometry
        spherePrim = UsdGeom.Sphere.Define(stage, xform.GetPath().AppendChild('Geom'))
        model = Usd.ModelAPI(xform)
        model.SetKind(Kind.Tokens.component)
        # Create Looks scope
        # Scope is the simplest grouping primitive...
        scope = UsdGeom.Scope.Define(stage, xform.GetPath().AppendChild('Looks'))
        rootPath = scope.GetPath()
        model = Usd.ModelAPI(scope)
        model.SetKind(Kind.Tokens.model)
    
    material = None

    # Create material
    material = UsdShade.Material.Define(stage, rootPath.AppendChild(materialName))
    if spherePrim != None:
        # Bind material to geometry
        UsdShade.MaterialBindingAPI(spherePrim).Bind(material)

    # Create surface shader    
    surfaceShader = UsdShade.Shader.Define(stage, material.GetPath().AppendChild("Shader"))
    # Create surface shader output port
    outPort = surfaceShader.CreateOutput('out', Sdf.ValueTypeNames.Token)
    
    # Create material output port
    terminal = material.CreateOutput('mdl:surface', Sdf.ValueTypeNames.Token)
    # Connect material output to shader output
    terminal.ConnectToSource(outPort)
    
    terminal = material.CreateOutput('mdl:volume', Sdf.ValueTypeNames.Token)
    terminal.ConnectToSource(outPort)
    
    terminal = material.CreateOutput('mdl:displacement', Sdf.ValueTypeNames.Token)
    terminal.ConnectToSource(outPort)
    
    # Determine module asset
    module_interface : pymdl.Module
    module_interface = context.current_module()
    module = module_interface.mdlName
    source_asset = module_mdl_name_to_source_asset(context, module)

    # Example:
    # uniform token info:implementationSource = "sourceAsset"
    # uniform asset info:mdl:sourceAsset = @nvidia/core_definitions.mdl@
    # uniform token info:mdl:sourceAsset:subIdentifier = "::nvidia::core_definitions::flex_material"
    surfaceShader.GetImplementationSourceAttr().Set("sourceAsset")
    surfaceShader.SetSourceAsset(Sdf.AssetPath(source_asset), "mdl")
    node_identifier = mdl_name_to_sub_identifier(db_name)
    surfaceShader.GetPrim().CreateAttribute("info:mdl:sourceAsset:subIdentifier", Sdf.ValueTypeNames.Token, False, Sdf.VariabilityUniform).Set(node_identifier)

    context.stage.SetDefaultPrim(material.GetPrim())

    context.push_usd_material(material)
    context.push_usd_shader(surfaceShader)
    context.push_custom_data(dict())

    definition_to_stage(context, function)

    # Annotations
    anno_dict = get_annotations_dict(context, function)
    add_annotations_to_prim(surfaceShader.GetPrim(), anno_dict)
    add_annotations_to_node(surfaceShader, anno_dict)

    context.pop_custom_data()
    context.pop_usd_shader()
    context.pop_usd_material()

def material_to_stage_output_shader(context, function: pymdl.FunctionDefinition):
    neuray = context.neuray
    transaction = context.transaction
    stage = context.stage

    db_name = function.dbName

    materialName = db_name_to_prim_name(db_name)
    
    # Start at root
    rootPath = Sdf.Path('/')
    material = UsdShade.Shader.Define(stage, rootPath.AppendChild(materialName))
    # Create output port
    outPort = material.CreateOutput('out', Sdf.ValueTypeNames.Token).SetRenderType("material")

    # Determine module asset
    module_interface : pymdl.Module
    module_interface = context.current_module()
    module = module_interface.mdlName
    source_asset = module_mdl_name_to_source_asset(context, module)

    # Example:
    # uniform token info:implementationSource = "sourceAsset"
    # uniform asset info:mdl:sourceAsset = @nvidia/core_definitions.mdl@
    # uniform token info:mdl:sourceAsset:subIdentifier = "::nvidia::core_definitions::flex_material"
    material.GetImplementationSourceAttr().Set("sourceAsset")
    material.SetSourceAsset(Sdf.AssetPath(source_asset), "mdl")
    node_identifier = mdl_name_to_sub_identifier(db_name)
    material.GetPrim().CreateAttribute("info:mdl:sourceAsset:subIdentifier", Sdf.ValueTypeNames.Token, False, Sdf.VariabilityUniform).Set(node_identifier)

    context.stage.SetDefaultPrim(material.GetPrim())

    context.push_usd_material(material)
    context.push_usd_shader(material)
    context.push_custom_data(dict())

    definition_to_stage(context, function)

    # Annotations
    anno_dict = get_annotations_dict(context, function)
    add_annotations_to_prim(material.GetPrim(), anno_dict)
    add_annotations_to_node(material, anno_dict)

    context.pop_custom_data()
    context.pop_usd_shader()
    context.pop_usd_material()

def material_to_stage(context, function: pymdl.FunctionDefinition):

    if context.mdl_to_usd_output == OutputType.SHADER:
        return material_to_stage_output_shader(context, function)
    elif context.mdl_to_usd_output == OutputType.MATERIAL:
        return material_to_stage_output_material(context, function)
    elif context.mdl_to_usd_output == OutputType.MATERIAL_AND_GEOMETRY:
        return material_to_stage_output_material_and_geo(context, function)

def module_to_stage(context, module: pymdl.Module):

    context.push_module(module)

    for simple_name, overloads in module.functions.items():
        f: pymdl.FunctionDefinition
        for f in overloads:
            material_to_stage(context, f)

    anno_dict = get_annotations_dict(context, module)
    add_annotations_to_stage(context, context.stage, anno_dict)

    context.pop_module()

def mdl_prim_to_usd(context : ConverterContext, stage: Usd.Stage, prim: Usd.Prim, createNewMaterial: bool = False):
    inst_name = None
    mdl_entity = None
    mdl_entity_snapshot = None
    if context.ov_neuray == None:
        return

    # Get the shader prim from Material
    shader_prim = get_shader_prim(context.stage, prim.GetPath())  

    mdl_entity = context.ov_neuray.createMdlEntity(shader_prim.GetPath().pathString)
    mdl_entity_snapshot = None
    if not mdl_entity.getMdlModule() == None:
        mdl_entity_snapshot = context.ov_neuray.createMdlEntitySnapshot(mdl_entity)
        inst_name = mdl_entity_snapshot.dbName

    source_asset = Sdf.AssetPath("")
    shader = UsdShade.Shader(shader_prim)
    if shader:
        imp_source = shader.GetImplementationSourceAttr().Get()
        if imp_source == "sourceAsset":
            # Retrieve module name
            source_asset = shader.GetSourceAsset("mdl")

    material = UsdShade.Material(prim)

    if createNewMaterial:
        rootPath = Sdf.Path(prim.GetPath().GetParentPath())
        # Create material
        materialName = prim.GetName() + "_converted"
        # Find unique name
        materialName = GetUniqueNameInStage.get(stage, rootPath.AppendChild(materialName))
        material = UsdShade.Material.Define(stage, materialName)

    # Create surface shader
    surfaceShader = UsdShade.Shader.Define(stage, material.GetPath().AppendChild(shader_prim.GetName()))
    # Create output port
    outPort = surfaceShader.CreateOutput('out', Sdf.ValueTypeNames.Token)

    # Create material output port
    terminal = material.CreateOutput('mdl:surface', Sdf.ValueTypeNames.Token)
    # Connect material output to shader output
    terminal.ConnectToSource(outPort)
    
    terminal = material.CreateOutput('mdl:volume', Sdf.ValueTypeNames.Token)
    terminal.ConnectToSource(outPort)
    
    terminal = material.CreateOutput('mdl:displacement', Sdf.ValueTypeNames.Token)
    terminal.ConnectToSource(outPort)

    context.push_usd_material(material)
    context.push_usd_shader(surfaceShader)
    context.push_custom_data(dict())

    fctCall = pymdl.FunctionCall._fetchFromDb(context.transaction, inst_name)
    
    handle_function_call(context, fctCall)

    context.pop_usd_material()
    context.pop_usd_shader()
    context.pop_custom_data()

    if not mdl_entity_snapshot == None:
        context.ov_neuray.destroyMdlEntitySnapshot(mdl_entity_snapshot)
    if not mdl_entity == None:
        context.ov_neuray.destroyMdlEntity(mdl_entity)


#--------------------------------------------------------------------------------------------------
# Utilities
#--------------------------------------------------------------------------------------------------

def get_examples_search_path():
    """Try to get the example search path or returns 'mdl' sub folder of the current directory if it failed."""

    # get the environment variable that is used in all MDL SDK examples
    example_sp = os.getenv('MDL_SAMPLES_ROOT')

    # fall back to a path relative to this script file
    if example_sp == None or not os.path.exists(example_sp):
        example_sp = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..')

    # go down into the mdl folder
    example_sp = os.path.join(example_sp, 'mdl')

    # fall back to the current folder
    if not os.path.exists(example_sp):
        example_sp = './mdl'

    return os.path.abspath(example_sp)

#--------------------------------------------------------------------------------------------------
# MDL Python Example
#--------------------------------------------------------------------------------------------------

# return usd_value, connection, usd_type
def get_usd_value_and_type(context, val : pymdl.ArgumentConstant):
    if not val:
        return (None, None, None)
    kind = val.type.kind

    usd_type = mdl_type_to_usd_type(val)
    if usd_type == None:
        return (None, None, None)

    usd_val = None
    connection = None
    # -------------------------------- Atomic ---------------------------------
    if  kind == pymdlsdk.IType.Kind.TK_BOOL or \
        kind == pymdlsdk.IType.Kind.TK_INT or \
        kind == pymdlsdk.IType.Kind.TK_FLOAT or \
        kind == pymdlsdk.IType.Kind.TK_DOUBLE or \
        kind == pymdlsdk.IType.Kind.TK_STRING:
        usd_val = val.value
    elif kind == pymdlsdk.IType.Kind.TK_ENUM:
        usd_val = val.value[1]
    # -------------------------------- Compound ---------------------------------
    elif kind == pymdlsdk.IType.Kind.TK_VECTOR:
        usd_val = python_vector_to_usd_value(val.value)
        if context.is_type_annotation_enabled():
            custom_data = custom_data_from_python_vector(val.value)
            context.current_custom_data().update(custom_data)
    elif kind == pymdlsdk.IType.Kind.TK_MATRIX:
        usd_val = python_matrix_to_usd_value(val.value)
    elif kind == pymdlsdk.IType.Kind.TK_COLOR:
        usd_val = Gf.Vec3f(val.value[0], val.value[1], val.value[2])
    elif kind == pymdlsdk.IType.Kind.TK_ARRAY:
        usd_val = python_array_to_usd_value(val.value)
        if context.is_type_annotation_enabled():
            custom_data = custom_data_from_python_array(val.value)
            context.current_custom_data().update(custom_data)
    elif kind == pymdlsdk.IType.Kind.TK_STRUCT:
        parm_name = context.current_parameter_name()
        shader = context.current_usd_shader()
        usd_mat = context.current_usd_shader()
        new_shader = UsdShade.Shader.Define(context.stage,
            GetUniqueNameInStage.get(context.stage, usd_mat.GetPath().AppendChild("ShaderAttachement_" + parm_name)))
        connection = new_shader.CreateOutput('out', Sdf.ValueTypeNames.Token)
        context.push_usd_shader(new_shader)
        for item in val.value:
            field_name = item     
            context.push_parameter_name(field_name)
            context.push_custom_data(dict())
            v = val.value[item]
            handle_value(context, v)
            context.pop_custom_data()
            context.pop_parameter_name()
        context.pop_usd_shader()
    # -------------------------------- Resource ---------------------------------
    elif kind == pymdlsdk.IType.Kind.TK_TEXTURE:
        ftex = val.value[0]
        gamma = val.value[1]
        with context.transaction.access_as(pymdlsdk.ITexture, ftex) as text:
            if text.is_valid_interface():
                with context.transaction.access_as(pymdlsdk.IImage, text.get_image()) as image:
                    if image.is_valid_interface():
                        usd_val = image.get_filename(0,0)
    elif kind == pymdlsdk.IType.Kind.TK_LIGHT_PROFILE:
        v = val.value
        with context.transaction.access_as(pymdlsdk.ILightprofile, v) as obj:
            if obj.is_valid_interface():
                usd_val = obj.get_filename()
    elif kind == pymdlsdk.IType.Kind.TK_BSDF_MEASUREMENT:
        v = val.value
        with context.transaction.access_as(pymdlsdk.IBsdf_measurement, v) as obj:
            if obj.is_valid_interface():
                usd_val = obj.get_filename()
    else:
        print("Error: Unknown IValue Kind for parm/kind: {}/{}".format(context.current_parameter_name(), kind))

    return (usd_val, connection, usd_type)

def dict_of_dicts_merge(x, y):
    z = {}
    overlapping_keys = x.keys() & y.keys()
    for key in overlapping_keys:
        z[key] = dict_of_dicts_merge(x[key], y[key])
    for key in x.keys() - overlapping_keys:
        z[key] = deepcopy(x[key])
    for key in y.keys() - overlapping_keys:
        z[key] = deepcopy(y[key])
    return z

# Return a dictionary of pair [string, value]
#
def get_annotations_dict(context, val):
    main_dict = dict()
    # We do not want to collect type annotations during conversion of annotation expressions
    # Otherwise the type for the annotations is mized with the parameter annotations
    context.disable_type_annotation()
    try:
        # Handle annotations
        if val.annotations:
            anno_dict = dict()
            # print(f"  Annotations:")
            anno: pymdl.Annotation
            for anno in val.annotations:
                # print(f"  - Simple Name: {anno.simpleName}")
                # print(f"    Qualified Name: {anno.name}")
                # // Special case of annotations without expressions, e.g. ::anno::hidden()
                # // We create a boolean VtValue and set it to true
                # // Annotation:
                # //   ::anno::hidden()
                # // becomes:
                # //   bool "::anno::hidden()" = 1
                if len(anno.arguments.items()) == 0:
                    usd_value = True
                    usd_type = bool
                    anno_dict[anno.name] = usd_value
                arg: pymdl.val
                for arg_name, arg in anno.arguments.items():
                    # print(f"    ({arg.type.kind}) {arg_name}: {arg.value}")
                    (usd_value, connection, usd_type) = get_usd_value_and_type(context, arg)
                    # print(f"    ({usd_value}) {connection}: {usd_type}")
                    if usd_value != None and usd_type != None:
                        anno_dict[anno.name+"::"+arg_name] = usd_value
            if len(anno_dict) > 0:
                main_dict.update(dict({"mdl" : {"annotations": anno_dict}}))

        # Handle special values, concat (possibly empty) context dictionary
        main_dict = dict_of_dicts_merge(main_dict, context.current_custom_data())

    except Exception as e:
        print("Error while retrieving annotations")
    finally:
        context.enable_type_annotation()
        return main_dict

def add_annotations_to_stage(context, stage, anno_dict):
    try:
        if "::anno::display_name(string)::name" in anno_dict["mdl"]["annotations"]:
            stage.SetMetadata('comment', 'Conversion from MDL module: ' + anno_dict["mdl"]["annotations"]["::anno::display_name(string)::name"])
        else:
            if context.current_module() != None:
                stage.SetMetadata('comment', 'Conversion from MDL module: ' + context.current_module().mdlName)
                
        # TODO: Add more annotations to Stage as Metadata: version, ...
    except:
        pass

def add_annotations_to_node(node, anno_dict):
    # display name, group and description
    #
    # sdrMetadata = {
    #     dictionary ui = {
    #         string displayGroup = "float2_parm group"
    #     }
    # }
    #
    if "mdl" in anno_dict and "annotations" in anno_dict["mdl"]:
        for anno, key in ( \
            ("::anno::display_name(string)::name", "ui:displayName"), \
            ("::anno::in_group(string)::group", "ui:displayGroup"), \
            ("::anno::in_group(string,string)::group", "ui:displayGroup"), \
            ("::anno::in_group(string,string,string)::group", "ui:displayGroup"), \
            ("::anno::description(string)::description", "ui:description")):
            if anno in anno_dict["mdl"]["annotations"]:
                node.SetSdrMetadataByKey( key, anno_dict["mdl"]["annotations"][anno])
    #
    # See Also:
    #
    # ui = Usd.SceneGraphPrimAPI(minput.GetAttr())
    # a = ui.CreateDisplayNameAttr()
    #
    # TODO: Could also use?
    #
    # minput.GetAttr().SetDisplayName("FOOBAR")
    # minput.GetAttr().SetDisplayGroup("GROUP")
    try:
        if "::anno::display_name(string)::name" in anno_dict["mdl"]["annotations"]:
            node.GetAttr().SetDisplayName(anno_dict["mdl"]["annotations"]["::anno::display_name(string)::name"])
        for anno in ( \
            "::anno::in_group(string)::group", \
            "::anno::in_group(string,string)::group", \
            "::anno::in_group(string,string,string)::group" \
            ):
            if anno in anno_dict["mdl"]["annotations"]:
                node.GetAttr().SetDisplayGroup(anno_dict["mdl"]["annotations"][anno])
    except:
        pass

def add_annotations_to_prim(usd_prim, anno_dict):
    if len(anno_dict) > 0:
        cd = usd_prim.GetCustomData()
        cd.update(anno_dict)
        usd_prim.SetCustomData(cd)

    try:
        ui = UsdUI.UsdUIBackdrop(usd_prim)
        if "::anno::description(string)::description" in anno_dict["mdl"]["annotations"]:
            ui.CreateDescriptionAttr(anno_dict["mdl"]["annotations"]["::anno::description(string)::description"])
    except:
        pass

    try:
        ui = UsdUI.SceneGraphPrimAPI(usd_prim)
        if "::anno::display_name(string)::name" in anno_dict["mdl"]["annotations"]:
            ui.CreateDisplayNameAttr(anno_dict["mdl"]["annotations"]["::anno::display_name(string)::name"])
        for anno in ( \
            "::anno::in_group(string)::group", \
            "::anno::in_group(string,string)::group", \
            "::anno::in_group(string,string,string)::group" \
            ):
            if anno in anno_dict["mdl"]["annotations"]:
                ui.CreateDisplayGroupAttr(anno_dict["mdl"]["annotations"][anno])
    except:
        pass

    # Hidden
    if "mdl" in anno_dict and "annotations" in anno_dict["mdl"]:
        if "::anno::hidden()" in anno_dict["mdl"]["annotations"]:
            if anno_dict["mdl"]["annotations"]["::anno::hidden()"]:
                usd_prim.SetHidden(True)

def handle_value(context, val : pymdl.ArgumentConstant):
    if not val:
        return

    (usd_value, connection, usd_type) = get_usd_value_and_type(context, val)
    if (usd_value == None and connection == None) or usd_type == None:
        return

    parm_name = context.current_parameter_name()
    shader = context.current_usd_shader()
    minput = shader.CreateInput(parm_name, usd_type)
    if usd_value != None:
        minput.Set(usd_value)
    else:
        minput.ConnectToSource(connection)

    anno_dict = get_annotations_dict(context, val)
    add_annotations_to_prim(minput.GetAttr(), anno_dict)
    add_annotations_to_node(minput, anno_dict)
    if val.type.kind == pymdlsdk.IType.Kind.TK_ENUM:
        minput.SetRenderType(val.type.symbol)

def handle_expression_constant(context, expression_constant : pymdl.ArgumentConstant):
    if not expression_constant:
        return
    handle_value(context, expression_constant)

def handle_function_call(context, function_call : pymdl.FunctionCall):
    if not function_call:
        return

    # print("*   Function definition: {}".format(function_call.functionDefinition))
    # print("*   MDL function definition: {}".format(function_call.mdlFunctionDefinition))
    # print("*   Parameter count: {}".format(len(function_call.parameters)))

    fdef_db_name = function_call.functionDefinition
    function_def = pymdl.FunctionDefinition._fetchFromDb(context.transaction, fdef_db_name)
    if not function_def:
        return

    sourceAsset = module_mdl_name_to_source_asset(context, function_def.mdlModuleName)
    node_identifier = mdl_name_to_sub_identifier( function_def.mdlName, True )

    # Determine if function definition is variant
    fctdef = context.transaction.access_as(pymdlsdk.IFunction_definition, function_def.dbName)
    prototype = fctdef.get_prototype()
    if prototype:
        node_identifier = mdl_name_to_sub_identifier(prototype)
        sourceAsset = prototype_to_source_asset(prototype)

    shader = context.current_usd_shader()
    shader.GetImplementationSourceAttr().Set("sourceAsset")
    shader.SetSourceAsset(Sdf.AssetPath(sourceAsset), "mdl")
    shader.GetPrim().CreateAttribute("info:mdl:sourceAsset:subIdentifier", Sdf.ValueTypeNames.Token, False, Sdf.VariabilityUniform).Set(node_identifier)

    for name, argument in function_call.parameters.items():
        # print(f"* Name: {name}")
        # print(f"  Type: {argument.type.kind}")
        # if argument.type.symbol:
        #     print(f"  Type Symbol: {argument.type.symbol}")
        # print(f"  Value: {argument.value}")
        # print(f"  Value is Constant: {isinstance(argument, pymdl.ArgumentConstant)}")
        # print(f"  Value is Attachment: {isinstance(argument, pymdl.ArgumentCall)}")

        context.push_parameter_name(name)
        context.push_custom_data(dict())

        handle_expression(context, argument)
        
        context.pop_custom_data()
        context.pop_parameter_name()

def handle_material_instance(context, minst):
    if not minst.is_valid_interface():
        return
    pass

def handle_expression_call(context, expression_call : pymdl.ArgumentCall):
    if not expression_call:
        return
    neuray = context.neuray
    transaction = context.transaction

    # NEW
    fcall : pymdl.FunctionCall
    fcall = pymdl.FunctionCall._fetchFromDb(transaction, expression_call.value)
    if fcall:
        handle_function_call(context, fcall)
        return
    # TODO?
    # minst = transaction.access_as(pymdlsdk.IFunction_call, expression_call.get_call())
    # if minst.is_valid_interface():
    #     handle_material_instance(context, minst)

def handle_expression_parameter(context, expression_parameter):
    if not expression_parameter.is_valid_interface():
        return
    pass

def handle_expression_direct_call(context, expression_direct_call):
    if not expression_direct_call.is_valid_interface():
        return
    pass

def handle_expression_temporary(context, expression_temporary):
    if not expression_temporary.is_valid_interface():
        return
    pass

def handle_expression(context, argument : pymdl.Argument):
    if not argument:
        return 

    if isinstance(argument, pymdl.ArgumentConstant):
        # A constant expression. See #mi::neuraylib::IExpression_constant.
        # argument_cst = pymdl.ArgumentConstant(argument) # FAILS
        argument_cst : pymdl.ArgumentConstant
        argument_cst = argument
        handle_expression_constant(context, argument_cst)
    elif isinstance(argument, pymdl.ArgumentCall):
        # Need to:
        # - create a parameter of the given type
        # - create a new shader in the current material
        # - Connect the parameter to the new shader
        # - Evaluate the expression
        expression_kind = argument.type.kind
        # If the exact return type can not be determined (e.g. vectors, where to get the size?) set the parm type to Token by default
        parm_type = Sdf.ValueTypeNames.Token
        parm_name = context.current_parameter_name()
        if expression_kind in MdlITypeKindToUSD:
            parm_type = MdlITypeKindToUSD[expression_kind]
        else:
            print("Warning: Indirect call expression parm/kind : {}/{}".format(parm_name,argument.type.kind))

        # - create a parameter of the given type
        usd_shader = context.current_usd_shader()
        minput = usd_shader.CreateInput(parm_name, parm_type)
        # By default shaders are not nested, they are created immediately under the material
        usd_mat = context.current_usd_material()
        if context.mdl_to_usd_output_material_nested_shaders == True:
            # Nested shaders case, create under current shader
            usd_mat = context.current_usd_shader()
        # - create a new shader in thecurrent material (or current shader in the nested case)
        shName = GetUniqueNameInStage.get(context.stage, usd_mat.GetPath().AppendChild(parm_name))
        shader = UsdShade.Shader.Define(context.stage, shName)
        outPort = shader.CreateOutput('out', Sdf.ValueTypeNames.Token)
        # - Connect the parameter to the new shader
        minput.ConnectToSource(outPort)

        context.push_usd_shader(shader)

        # - Evaluate the expression
        argument_call : pymdl.ArgumentCall
        argument_call = argument
        handle_expression_call(context, argument_call)
        
        # Annotations
        anno_dict = get_annotations_dict(context, argument_call)
        add_annotations_to_prim(minput.GetAttr(), anno_dict)
        add_annotations_to_node(minput, anno_dict)

        context.pop_usd_shader()

    # TODO
    # elif kind == pymdlsdk.IExpression.Kind.EK_PARAMETER:
    #     # A parameter reference expression. See #mi::neuraylib::IExpression_parameter.
    #     # handle_expression_parameter(context, expression.get_interface(pymdlsdk.IExpression_parameter))
    #     pass
    # elif kind == pymdlsdk.IExpression.Kind.EK_DIRECT_CALL:
    #     # A direct call expression. See #mi::neuraylib::IExpression_direct_call.
    #     # handle_expression_direct_call(context, expression.get_interface(pymdlsdk.IExpression_direct_call))
    #     pass
    # elif kind == pymdlsdk.IExpression.Kind.EK_TEMPORARY:
    #     # A temporary reference expression. See #mi::neuraylib::IExpression_temporary.
    #     # handle_expression_temporary(context, expression.get_interface(pymdlsdk.IExpression_temporary))
    #     pass

def definition_to_stage(context, function: pymdl.FunctionDefinition):
    neuray = context.neuray
    transaction = context.transaction

    argument: pymdl.Argument
    for name, argument in function.parameters.items():
        # print(f"* Name: {name}")
        # print(f"  Type: {argument.type.kind}")
        # if argument.type.symbol:
        #     print(f"  Type Symbol: {argument.type.symbol}")
        # print(f"  Value: {argument.value}")
        # print(f"  Value is Constant: {isinstance(argument, pymdl.ArgumentConstant)}")
        # print(f"  Value is Attachment: {isinstance(argument, pymdl.ArgumentCall)}")

        context.push_parameter_name(name)
        context.push_custom_data(dict())

        handle_expression(context, argument)
        
        context.pop_custom_data()
        context.pop_parameter_name()

def get_db_module_name(neuray, module_mdl_name):
    """Return the db name of the given module."""

    module_db_name = None
    # When the module is loaded we can access it and all its definitions by accessing the DB
    # for that we need to get a the database name of the module using the factory
    with neuray.get_api_component(pymdlsdk.IMdl_factory) as factory:
        with factory.get_db_module_name(module_mdl_name) as istring:
            if istring.is_valid_interface():
                module_db_name = istring.get_c_str() # note, even though this name simple and could
                                                     # be constructed by string operations, use the
                                                     # factory to be save in case of unicode encodings
                                                     # and potential upcoming changes in the future

        # # shortcut for the function above
        # # this chaining is a bit special. In the C++ interface it's not possible without
        # # leaking memory. Here we create the smart pointer automatically. However, the IString
        # # which is created temporay here is released at the end the `load_module` function, right?
        # # This might be unexpected, especially when we rely on the RAII pattern and that
        # # objects are disposed at certain points in time (usually before committing a transaction)
        # module_db_name_2 = factory.get_db_module_name(module_mdl_name).get_c_str()
        # # note, we plan to map compatible types to python. E.g. the IString class may disappear
        # return module_db_name_2
    return module_db_name

#--------------------------------------------------------------------------------------------------

def get_db_definition_name(neuray, function_mdl_name):
    """Return the db name of the given function definition."""

    with neuray.get_api_component(pymdlsdk.IMdl_factory) as factory:
        with factory.get_db_definition_name(function_mdl_name) as istring:
            if istring.is_valid_interface():
                return istring.get_c_str()
    return None

#--------------------------------------------------------------------------------------------------
def load_module(context, source_asset):
    success = False
    module_db_name = ""
    if not context.ov_neuray == None:
        # We are in OV

        if context.transaction:
            context.transaction.commit()
            context.transaction = None

        module_db_name = load_module_from_ov(context.ov_neuray, source_asset.path)

        dbScopeName = "rtx_scope"
        rtxTransactionReadHandle = context.ov_neuray.createReadingTransaction(dbScopeName)
        context.transaction: pymdlsdk.ITransaction = pymdlsdk.attach_itransaction(rtxTransactionReadHandle)

        module = pymdl.Module._fetchFromDb(context.transaction, module_db_name)
        success = (not module == None)
    else:
        module_mdl_name = source_asset_to_module_name(source_asset.path, context)
        success = load_module_from_neuray(context.neuray, context.transaction, module_mdl_name)
        if success:
            module_db_name = get_db_module_name(context.neuray, module_mdl_name)

    module = pymdl.Module._fetchFromDb(context.transaction, module_db_name)
    print(f'Loaded module (from omni.mdl.neuraylib={context.ov_neuray} / source asset={source_asset}): {module.filename}')

    return (success, module_db_name)

#--------------------------------------------------------------------------------------------------

def load_module_from_neuray(neuray, transaction, module_mdl_name):
    """Load the module given its name.
    Returns true if the module is loaded to database"""
    
    with neuray.get_api_component(pymdlsdk.IMdl_impexp_api) as imp_exp:
        with neuray.get_api_component(pymdlsdk.IMdl_factory) as factory:
            # for illustation, we don't use a `with` block for the `context`, instead we release manually
            context = factory.create_execution_context()

        res = imp_exp.load_module(transaction, module_mdl_name, context)
    context.release()
    return res >= 0

#--------------------------------------------------------------------------------------------------

def load_module_from_ov(rtxneuray, module_path):
    """Load the module given its name.
    Returns rtxModule.dbName if the module is loaded to database, otherwise return empty string """

    rtxModule = None
    try:    
        rtxModule = rtxneuray.createMdlModule(module_path)
    except:
        return ""
    return rtxModule.dbName

#--------------------------------------------------------------------------------------------------
class ConvertVectorAndArray:
    def create_type_or_value(self, usd_type, factory):
        if usd_type in (\
            Sdf.ValueTypeNames.Float2,\
            Sdf.ValueTypeNames.Float2Array,\
            Sdf.ValueTypeNames.Float3,\
            Sdf.ValueTypeNames.Float3Array,\
            Sdf.ValueTypeNames.Float4,\
            Sdf.ValueTypeNames.Float4Array,\
            Sdf.ValueTypeNames.FloatArray):
           return factory.create_float()
        elif usd_type in (\
            Sdf.ValueTypeNames.Int2,\
            Sdf.ValueTypeNames.Int2Array,\
            Sdf.ValueTypeNames.Int3,\
            Sdf.ValueTypeNames.Int3Array,\
            Sdf.ValueTypeNames.Int4,\
            Sdf.ValueTypeNames.Int4Array,\
            Sdf.ValueTypeNames.IntArray):
           return factory.create_int()
        elif usd_type in (\
            Sdf.ValueTypeNames.Double2,\
            Sdf.ValueTypeNames.Double2Array,\
            Sdf.ValueTypeNames.Double3,\
            Sdf.ValueTypeNames.Double3Array,\
            Sdf.ValueTypeNames.Double4,\
            Sdf.ValueTypeNames.Double4Array,\
            Sdf.ValueTypeNames.DoubleArray,\
            Sdf.ValueTypeNames.Matrix2d,\
            Sdf.ValueTypeNames.Matrix3d,\
            Sdf.ValueTypeNames.Matrix4d):
           return factory.create_double()
        elif usd_type == Sdf.ValueTypeNames.BoolArray:
           return factory.create_bool()
        elif usd_type == Sdf.ValueTypeNames.StringArray:
           return factory.create_string()
        return None
    def convert(self, input, usd_type, type_factory, value_factory, expression_factory):
        value = input.Get()
        if value == None:
            return None
        if usd_type == Sdf.ValueTypeNames.Color3fArray:
            value = input.Get()
            array_size = len(value)
            vector_size = len(value[0])
            assert(vector_size == 3)
            atomic_type = type_factory.create_color()
            if not atomic_type.is_valid_interface():
                return None
            array_type = type_factory.create_immediate_sized_array(atomic_type, array_size)
            mdl_value_array = value_factory.create_array(array_type)
            mdl_value_array.set_size(array_size)
            for index in range(array_size):
                usd_value = value[index]
                atomic_value = value_factory.create_color(usd_value[0],usd_value[1],usd_value[2])
                mdl_value_array.set_value(index, atomic_value)
            expression = expression_factory.create_constant(mdl_value_array)
            return expression
        elif usd_type == Sdf.ValueTypeNames.StringArray:
            value = input.Get()
            array_size = len(value)
            atomic_type = type_factory.create_string()
            if not atomic_type.is_valid_interface():
                return None
            array_type = type_factory.create_immediate_sized_array(atomic_type, array_size)
            mdl_value_array = value_factory.create_array(array_type)
            mdl_value_array.set_size(array_size)
            for index in range(array_size):
                usd_value = value[index]
                atomic_value = value_factory.create_string()
                atomic_value.set_value(usd_value)
                mdl_value_array.set_value(index, atomic_value)
            expression = expression_factory.create_constant(mdl_value_array)
            return expression
        elif numpy.isscalar(value[0]):
            # Vector
            vector_size = len(value)
            type = self.create_type_or_value(usd_type, type_factory)
            if not type.is_valid_interface():
                return None
            vector_type = type_factory.create_vector(type, vector_size)
            if not vector_type.is_valid_interface():
                return None
            mdl_value_vector = value_factory.create_vector(vector_type)
            for index in range(vector_size):
                usd_value = value[index]
                atomic_value = self.create_type_or_value(usd_type, value_factory)
                atomic_value.set_value(usd_value)
                mdl_value_vector.set_value(index, atomic_value)
            expression = expression_factory.create_constant(mdl_value_vector)
            return expression
        else:
            # Array
            value = input.Get()
            array_size = len(value)
            vector_size = len(value[0])
            type = self.create_type_or_value(usd_type, type_factory)
            if not type.is_valid_interface():
                return None
            vector_type = type_factory.create_vector(type, vector_size)
            if not vector_type.is_valid_interface():
                return None
            array_type = type_factory.create_immediate_sized_array(vector_type, array_size)
            mdl_value_array = value_factory.create_array(array_type)
            mdl_value_array.set_size(array_size)
            for index in range(array_size):
                mdl_value_vector = value_factory.create_vector(vector_type)
                usd_value = value[index]
                for index2 in range(vector_size):
                    atomic_value = self.create_type_or_value(usd_type, value_factory)
                    atomic_value.set_value(usd_value[index2])
                    mdl_value_vector.set_value(index2, atomic_value)
                mdl_value_array.set_value(index, mdl_value_vector)
            expression = expression_factory.create_constant(mdl_value_array)
            return expression

#--------------------------------------------------------------------------------------------------
def convert_value(context, input, mdl_type):
    neuray = context.neuray
    transaction = context.transaction
    factory = neuray.get_api_component(pymdlsdk.IMdl_factory)
    expression_factory = factory.create_expression_factory(transaction)
    value_factory = factory.create_value_factory(transaction)
    type_factory = factory.create_type_factory(transaction)

    name = input.GetBaseName()
    type = input.GetTypeName()
    # Try to follow the connection
    # TODO: not sure this is properly implemented, need further tests
    expression = None
    if input.HasConnectedSource():
        (source, sourceName, sourceType) = input.GetConnectedSource()
        if UsdShade.Shader(source):
            prim = source.GetPrim()
            inst_name = create_material_instance(context, prim.GetPath(), None) # TODO Test None expression_list
            if not inst_name == None:
                expression = expression_factory.create_call(inst_name)
        else:
            expression = convert_value(context, UsdShade.Input(input.GetValueProducingAttribute()[0]), mdl_type)

    if expression == None:
        # No connections
        mdlvalue = None
        if type == Sdf.ValueTypeNames.Int:
            if not input.GetRenderType() == None:
                # Handle Enum
                tf = factory.create_type_factory(transaction)
                if tf and tf.is_valid_interface():
                    et = tf.create_enum(input.GetRenderType())
                    if et and et.is_valid_interface():
                        mdlvalue = value_factory.create_enum(et)
            if mdlvalue == None:
                # Fallback
                mdlvalue = value_factory.create_int()
        elif type == Sdf.ValueTypeNames.Float:
            mdlvalue = value_factory.create_float()
        elif type == Sdf.ValueTypeNames.Bool:
            mdlvalue = value_factory.create_bool()
        elif type == Sdf.ValueTypeNames.Double:
            mdlvalue = value_factory.create_double()
        elif type == Sdf.ValueTypeNames.String:
            mdlvalue = value_factory.create_string()
        elif type == Sdf.ValueTypeNames.Color3f:
            value = input.Get()
            if value:
                mdlvalue = value_factory.create_color(value[0], value[1], value[2])
            else:
                mdlvalue = value_factory.create_color(0, 0, 0)
            expression = expression_factory.create_constant(mdlvalue)
            return expression
        elif type == Sdf.ValueTypeNames.Asset:
            is_texture = mdl_type.get_kind() == pymdlsdk.IType.Kind.TK_TEXTURE
            if not is_texture and mdl_type.get_kind() == pymdlsdk.IType.Kind.TK_ALIAS:
                type_alias = mdl_type.get_interface(pymdlsdk.IType_alias)
                aliased_type = type_alias.get_aliased_type()
                is_texture = aliased_type.get_kind() == pymdlsdk.IType.Kind.TK_TEXTURE
                if is_texture:
                    # Shape
                    # TS_2D           = 0,
                    # TS_3D           = 1,
                    # TS_CUBE         = 2,
                    # TS_PTEX         = 3,
                    # TS_BSDF_DATA    = 4,
                    tt = aliased_type.get_interface(pymdlsdk.IType_texture)
                    shape = tt.get_shape()
                    texture_name = GetUniqueName.get(transaction, "texture")
                    # TODO_PA: Has to specify argc and argv otherwise failure
                    tex = transaction.create("Texture", 0, None)
                    transaction.store(tex, texture_name)
                    value = input.Get()
                    if value:
                        texture = transaction.edit_as(pymdlsdk.ITexture, texture_name)
                        texture.set_image(value.path)
                    tex_type = type_factory.create_texture(shape)
                    mdlvalue = value_factory.create_texture(tex_type, texture_name)
                    if mdlvalue:
                        expression = expression_factory.create_constant(mdlvalue)
                        return expression
        elif type in (
            Sdf.ValueTypeNames.Float2,\
            Sdf.ValueTypeNames.Float2Array,\
            Sdf.ValueTypeNames.Float3,\
            Sdf.ValueTypeNames.Float3Array,\
            Sdf.ValueTypeNames.Float4,\
            Sdf.ValueTypeNames.Float4Array,\
            Sdf.ValueTypeNames.Int2,\
            Sdf.ValueTypeNames.Int2Array,\
            Sdf.ValueTypeNames.Int3,\
            Sdf.ValueTypeNames.Int3Array,\
            Sdf.ValueTypeNames.Int4,\
            Sdf.ValueTypeNames.Int4Array,\
            Sdf.ValueTypeNames.Double2,\
            Sdf.ValueTypeNames.Double2Array,\
            Sdf.ValueTypeNames.Double3,\
            Sdf.ValueTypeNames.Double3Array,\
            Sdf.ValueTypeNames.Double4,\
            Sdf.ValueTypeNames.Double4Array,\
            Sdf.ValueTypeNames.BoolArray,\
            Sdf.ValueTypeNames.IntArray,\
            Sdf.ValueTypeNames.FloatArray,\
            Sdf.ValueTypeNames.DoubleArray,\
            Sdf.ValueTypeNames.Color3fArray,\
            Sdf.ValueTypeNames.StringArray,\
            Sdf.ValueTypeNames.Matrix2d,\
            Sdf.ValueTypeNames.Matrix3d,\
            Sdf.ValueTypeNames.Matrix4d\
            ):
            converter = ConvertVectorAndArray()
            return converter.convert(input, type, type_factory, value_factory, expression_factory)
        # TODO Convert structures
        # elif type == Sdf.ValueTypeNames.Token:
        #     pass
        else:
            print("Error: Unsupported input (parm/type): {}/{}".format(name, type))
        if mdlvalue:
            value = input.Get()
            rtn = 0
            if not value == None:
                rtn = mdlvalue.set_value(value)
            if rtn and not rtn == 0:
                print("set_value() error ({}) for (parm/type/value): {}/{}/{}".format(rtn, name, type, value))
            expression = expression_factory.create_constant(mdlvalue)
            return expression
    return expression

#--------------------------------------------------------------------------------------------------
def set_default_value(context, parm_name, definition):
    defaults = definition.get_defaults()
    i = defaults.get_index(parm_name) 
    if i != -1:
        return defaults.get_expression(i)
    return None

#--------------------------------------------------------------------------------------------------
def convert_parameter(context, definition, input):
    name = input.GetBaseName()
    index = definition.get_parameter_index(name)
    if index >= 0:
        types = definition.get_parameter_types()
        if types:
            type = types.get_type(index)
            if type:
                expression = convert_value(context, input, type)
                if expression and expression.is_valid_interface():
                    return expression
                else:
                    print("Error: Invalid expression for parameter: {}".format(name))
    return None

#--------------------------------------------------------------------------------------------------
def edit_instance(transaction, inst_name):
    return transaction.edit_as(pymdlsdk.IFunction_call, inst_name)

#--------------------------------------------------------------------------------------------------
def dump_expression_list(expression_list):
    print("===== Dump expression list =====")
    for index in range(expression_list.get_size()):
        name = expression_list.get_name(index)
        expr = expression_list.get_expression(index)
        print("Expression name/kind/type kind: {} / {} / {}".format(name, expr.get_kind(), expr.get_type().get_kind()))
    print("===== End dump expression list =====")

#--------------------------------------------------------------------------------------------------
def dump_parameter_types(fctdef, neuray, transaction):
    print("===== Dump parameter types =====")
    name = fctdef.get_mdl_simple_name()
    print("Function name: {}".format(name))
    pt = fctdef.get_parameter_types()

    factory = neuray.get_api_component(pymdlsdk.IMdl_factory)
    type_factory = factory.create_type_factory(transaction)
    print(type_factory.dump(pt).get_c_str())

    # for i in range(pt.get_size()):
    #     t = pt.get_type(i)
    #     k = t.get_kind()
    #     if k == pymdlsdk.IType.Kind.TK_ALIAS:
    #         a = t.get_interface(pymdlsdk.IType_alias)
    #         k = a.get_aliased_type().get_kind()
    #     pn = fctdef.get_parameter_name(i)
    #     print("Parameter name/kind: {} / {}".format(pn, k))
    print("===== End dump parameter types =====")

#--------------------------------------------------------------------------------------------------
def convert_parameters(context, prim_path):
    definition = get_definition(context, prim_path)
    if not definition:
        return None
    stage = context.stage
    prim = get_shader_prim(stage, prim_path)
    # expression_list = definition.get_defaults()
    # expression_list = pymdlsdk.IExpression_list()
    neuray = context.neuray
    factory = neuray.get_api_component(pymdlsdk.IMdl_factory)
    transaction = context.transaction
    expression_factory = factory.create_expression_factory(transaction)    
    expression_list = expression_factory.create_expression_list()
    if prim:
        shader = UsdShade.Shader(prim)
        if shader:
            inputs = shader.GetInputs()
            for input in inputs:
                name = input.GetBaseName()
                expression = convert_parameter(context, definition, input)
                if expression:
                    rtn = expression_list.add_expression(name, expression)
                    if rtn != 0:
                        print("Error")
    # dump_expression_list(expression_list)
    # Handle parameters which are not set 
    defaults = definition.get_defaults()
    parameter_types = definition.get_parameter_types()
    value_factory = factory.create_value_factory(transaction)
    for i in range(definition.get_parameter_count()):
        name = definition.get_parameter_name(i)
        if expression_list.get_index(name) == -1:
            type = parameter_types.get_type(i)
            value = value_factory.create(type)
            expr = expression_factory.create_constant(value)
            expression_list.add_expression(name, expr)

    return expression_list

#--------------------------------------------------------------------------------------------------
def dump_material_instance(context, inst_name):
    transaction = context.transaction
    inst = edit_instance(transaction, inst_name)
    if not inst.is_valid_interface():
        return
    print("============ MDL Dump =============")
    print(inst_name)
    print("")
    neuray = context.neuray
    factory = neuray.get_api_component(pymdlsdk.IMdl_factory)
    expression_factory = factory.create_expression_factory(transaction)
    count = inst.get_parameter_count()
    arguments = inst.get_arguments()
    for index in range(count):
        argument = arguments.get_expression(index)
        name = inst.get_parameter_name(index)
        argument_text = expression_factory.dump(argument, name, 1)
        print(argument_text.get_c_str())
    print("============ MDL Dump =============")

#--------------------------------------------------------------------------------------------------
# Return a shader prim corresponding to the input prim_path. 
# If prim_path is a shader, return the corresponding prim.
# If prim_path is a material, return the first shader found in the list of children.
def get_shader_prim(stage, prim_path):
    prim = stage.GetPrimAtPath(Sdf.Path(prim_path))
    if UsdShade.Shader(prim):
        return prim
    if UsdShade.Material(prim):
        mat = UsdShade.Material(prim)
        if mat.GetOutput("mdl:surface") and mat.GetOutput("mdl:surface").HasConnectedSource():
            (src, name, t) = mat.GetOutput("mdl:surface").GetConnectedSource()
            if UsdShade.Shader(src):
                return src.GetPrim()
        for prim in prim.GetChildren():
            if UsdShade.Shader(prim):
                return prim
    if UsdShade.NodeGraph(prim):
        ng = UsdShade.NodeGraph(prim)
        if ng.GetOutput("out") and ng.GetOutput("out").HasConnectedSource():
            (src, name, t) = ng.GetOutput("out").GetConnectedSource()
            if UsdShade.Shader(src):
                return src.GetPrim()
        for prim in prim.GetChildren():
            if UsdShade.Shader(prim):
                return prim
    return None

#--------------------------------------------------------------------------------------------------
def get_sub_identifier_from_prim(prim):
    sub_id = prim.GetAttribute("info:mdl:sourceAsset:subIdentifier").Get()
    if not sub_id:
         sub_id = prim.GetAttribute("info:sourceAsset:subIdentifier").Get()
    return sub_id

#--------------------------------------------------------------------------------------------------
def get_definition(context, prim_path):
    stage = context.stage
    prim = get_shader_prim(stage, prim_path)
    if not prim:
        print("Error:\tCould not find any shader in prim '{}' from stage '{}'".format(prim_path, stage.GetRootLayer().GetDisplayName()))
    else:
        shader = UsdShade.Shader(prim)
        if shader:
            imp_source = shader.GetImplementationSourceAttr().Get()
            if imp_source == "sourceAsset":
                # Retrieve module name
                source_asset = shader.GetSourceAsset("mdl")
                neuray = context.neuray
                (success, module_db_name) = load_module(context, source_asset)
                if not success:
                    print("Error: Failed to load module: {}".format(source_asset))
                else:
                    sub_id = get_sub_identifier_from_prim(prim)
                    if sub_id:
                        transaction = context.transaction
                        module = pymdl.Module._fetchFromDb(transaction, module_db_name)
                        if module:
                            fct = find_function(module, sub_id)
                            if fct:
                                fctdef = transaction.access_as(pymdlsdk.IFunction_definition, fct.dbName)
                                if fctdef.is_valid_interface():
                                    return fctdef
    return None

#--------------------------------------------------------------------------------------------------
def list_functions(module: pymdl.Module):
    for k in module.functions.keys():
        for index in range(len(module.functions[k])):
            print(module.functions[k][index].mdlName.split("::")[-1])

#--------------------------------------------------------------------------------------------------
# return (module_name, simple_name, parameters)
# module_name without trailing '::'
# e.g. if sub_id = '::nvidia::foo::bar(float,color)'
# module_name = '::nvidia::foo'
# simple_name = 'bar'
# parameters = 'float,color'
def split_function_name(sub_id: str):
    module_name = ""
    simple_name = ""
    parameters = ""
    lf = sub_id.split('(')
    name_and_module = lf[0]
    lm = name_and_module.split('::')
    if len(lm) > 1:
        module_name = '::'.join(lm[0:-1])
    simple_name = lm[-1]
    if len(lf) > 1:
        parameters = lf[1].split(')')[0]
    return (module_name, simple_name, parameters)

#--------------------------------------------------------------------------------------------------
def find_function(module: pymdl.Module, sub_id: str):
    (module_name, simple_name, parameters) = split_function_name(sub_id)
    if simple_name in module.functions:
        if sub_id == simple_name:
            # Exact match, no overload, return the function
            return module.functions[simple_name][0]
        # Access the function definition, enumerate the overloads
        for index in range(len(module.functions[simple_name])):
            fct = module.functions[simple_name][index]
            (candidate_module_name, candidate_simple_name, candidate_parameters) = split_function_name(fct.mdlName)
            if parameters == candidate_parameters:
                return module.functions[simple_name][index]
    print("Error: Function not found in module (function/module): {} / {}".format(sub_id, module.filename))
    # list_functions(module)
    return None

#--------------------------------------------------------------------------------------------------
def create_material_instance(context, prim_path, expression_list):
    stage = context.stage
    prim = get_shader_prim(stage, prim_path)
    if prim:
        shader = UsdShade.Shader(prim)
        if shader:
            imp_source = shader.GetImplementationSourceAttr().Get()
            if imp_source == "sourceAsset":
                # Retrieve module name
                source_asset = shader.GetSourceAsset("mdl")
                neuray = context.neuray
                (success, module_db_name) = load_module(context, source_asset)
                if success:
                    sub_id = get_sub_identifier_from_prim(prim)
                    if sub_id:
                        transaction = context.transaction
                        module = pymdl.Module._fetchFromDb(transaction, module_db_name)
                        if module:
                            fct = find_function(module, sub_id)
                            if fct:
                                with transaction.access_as(pymdlsdk.IFunction_definition, fct.dbName) as fctdef:
                                    inst = None
                                    if fctdef.is_valid_interface():
                                        if expression_list == None:
                                            expression_list = convert_parameters(context, prim_path)
                                        # dump_parameter_types(fctdef, context.neuray, context.transaction)
                                        (inst, ret) = fctdef.create_function_call_with_ret(expression_list)
                                        fname = "create_material_instance()"
                                        error = {0: "Success.",
                                                -1: "An argument for a non-existing parameter was provided in 'arguments'.",
                                                -2: "The type of an argument in 'arguments' does not have the correct type, see #get_parameter_types().",
                                                -3: "A parameter that has no default was not provided with an argument value.",
                                                -4: "The definition can not be instantiated because it is not exported.",
                                                -5: "A parameter type is uniform, but the corresponding argument has a varying return type.",
                                                -6: "An argument expression is not a constant nor a call.",
                                                -8: "One of the parameter types is uniform, but the corresponding argument or default is a \
                                                    call expression and the return type of the called function definition is effectively varying since the \
                                                    function definition itself is varying.",
                                                -9: "The function definition is invalid due to a module reload, see #is_valid() for diagnostics."}
                                    if inst:
                                        if inst.is_valid_interface():
                                            (module_name, simple_name, parameters) = split_function_name(fct.dbName)
                                            inst_name = module_name + "::" + simple_name + "::converted"
                                            inst_name = GetUniqueName.get(transaction, inst_name)
                                            transaction.store(inst, inst_name)
                                            return inst_name
                                        else:
                                            print("Error: Invalid instance for function: {}".format(fct.mdlSimpleName))
                                            print("{} error {}: {}".format(fname, ret, error[ret]))
                                            if ret == -2:
                                                dump_parameter_types(fctdef, context.neuray, context.transaction)
    return None

#--------------------------------------------------------------------------------------------------
async def convert_usd_to_mdl(context, prim_path, output_filename):
    inst_name = None
    mdl_entity = None
    mdl_entity_snapshot = None
    if not context.ov_neuray == None:
        # We are in OV

        # Get the shader prim from Material
        shader_prim_path = get_shader_prim(context.stage, prim_path)
        
        if shader_prim_path == None:
            print(f"Error: can not access prim: '{prim_path}'")
            return None

        mdl_entity = context.ov_neuray.createMdlEntity(shader_prim_path.GetPath().pathString)
        mdl_entity_snapshot = None
        if not mdl_entity.getMdlModule() == None:
            mdl_entity_snapshot = context.ov_neuray.createMdlEntitySnapshot(mdl_entity)
            inst_name = mdl_entity_snapshot.dbName
        else:
            print(f"Error: can not access entity: '{prim_path}'")
    else:
        # Outside of OV, conversion is done here
        expression_list = convert_parameters(context, prim_path)
        inst_name = create_material_instance(context, prim_path, expression_list)

    if  inst_name is not None:
        dump_material_instance(context, inst_name)
        if output_filename:
            # we use a temporary folder for resources from the server
            with tempfile.TemporaryDirectory() as temp_dir:
                success  = await save_instance_to_module(context, inst_name, prim_path, temp_dir, output_filename)
                if not success:
                    inst_name = None

    if not context.ov_neuray == None:
        # We are in OV
        if not mdl_entity_snapshot == None:
            context.ov_neuray.destroyMdlEntitySnapshot(mdl_entity_snapshot)
        if not mdl_entity == None:
            context.ov_neuray.destroyMdlEntity(mdl_entity)

    return inst_name

#--------------------------------------------------------------------------------------------------
# Split the input string or list of strings using given the split char
# Return a list of strings
def split_list(list_of_strings, split_char):
    res = []
    if type(list_of_strings) == list:
        for s in list_of_strings:
            res += s.split(split_char)
    else:
        res = list_of_strings.split(split_char)
    return res

#--------------------------------------------------------------------------------------------------
# Remove all occurences of element from the list
def remove_all_occurences_of_element_from_list(element, alist):
    try:
        while True:
            alist.remove(element)
    except:
        pass
    finally: 
        return alist

#--------------------------------------------------------------------------------------------------
# Split the input string or list of strings using all the input split chars
def multi_split_list(list_of_strings, list_of_splits):
    res = list_of_strings
    for split in list_of_splits:
        res = split_list(res, split)
    return res

#--------------------------------------------------------------------------------------------------
# Identify a line containing an MD import declaration
# From the MDL 1.7 specs:
# import : import qualified_import {, qualified_import} ;
#        | [export] using import_path
#           import ( * | simple_name {, simple_name} ) ;
def is_import_line(token_list):
    if 'import' in token_list:
        if token_list[0] == 'import' or token_list[0] == 'export' or token_list[0] == 'using':
            return True
    return False

#--------------------------------------------------------------------------------------------------
# Construct a set or strings found in an MDL module
# Split the input strings using several separators (' ',';','(', ')')
# if filter_import_lines is set to True, then only process lines beginning with 'import'
def find_tokens(lines, filter_import_lines = False):
    all_ids = set()
    input_list = lines
    if not type(lines) == list:
        input_list = [lines]
    for l in input_list:
        ids = multi_split_list(l, [' ',';','(', ')'])
        remove_all_occurences_of_element_from_list('', ids)
        if filter_import_lines:
            if not is_import_line(ids):
                continue
        for id in ids:
            if not id.find('::') == -1:
                tokens = id.split('::')
                for t in tokens:
                    if len(t) > 0:
                        all_ids.add(t)
    return all_ids

#--------------------------------------------------------------------------------------------------
# Find MDL identifier in this line of text if text is an import statement
def find_import(line):
    ids = multi_split_list(line, [' ',';','(', ')'])
    remove_all_occurences_of_element_from_list('', ids)
    if is_import_line(ids):
        for id in ids:
            if not id.find('::') == -1:
                return id
    return ''

#--------------------------------------------------------------------------------------------------
# Find MDL identifier in this line of text
def find_identifiers(line):
    rtn_identifiers = set()
    ids = multi_split_list(line, [' ', ';', '(', ')', ',', '\n'])
    remove_all_occurences_of_element_from_list('', ids)
    for id in ids:
        if not id.find('::') == -1:
            rtn_identifiers.add(id)
    return rtn_identifiers

#--------------------------------------------------------------------------------------------------
# Substitute source strings with target strings found in the input table in the input lines
# Input lines is a string or list of strings
def replace_tokens(lines, table):
    if type(lines) == list:
        for i in range(len(lines)):
            for k in table.keys():
                lines[i] = lines[i].replace(k, table[k])
        return lines
    else:
        for k in table.keys():
            lines = lines.replace(k, table[k])
    return lines

#--------------------------------------------------------------------------------------------------
# Helper to perfom unmangling on MDL module files and MDL identifiers
# Un-mangling is only done in the framework of OV if an un-mangling routine exists
class Unmangle:
    unmangle_routine = None
    def __init__(self, context):
        if context.ov_neuray != None:
            if "_unmangleMdlModulePath" in dir(context.ov_neuray):
                self.unmangle_routine = context.ov_neuray._unmangleMdlModulePath

    # Un-mangle a single token
    # token should not contain any '::'
    def _unmangle_token(self, token):
        assert(token.find("::") == -1)

        # Un-mangle the input token
        # Add prefix otherwise unmangling has not effect
        unmangled_token = self.unmangle_routine("::" + token)

        # Remove any '.mdl' extension from unmangled token
        if len(unmangled_token) > 4 and unmangled_token[-4:] == '.mdl':
            unmangled_token = unmangled_token[:-4]

        unmangled = (unmangled_token != token)
            
        return (unmangled, unmangled_token)

    # Build an un-mangling mapping table from the input list of tokens
    def build_mapping_table(self, token_list):
        table = dict()
        for token in token_list:
            (unmangled_flag, unmangled_token) = self._unmangle_token(token)
            if unmangled_flag:
                table[token] = unmangled_token
        return table

    # Un-mangled an MDL identifier
    # Return a tuple (True if unmagling was done, unmangled value)
    def unmangle_mdl_identifier(self, source):
        if self.unmangle_routine == None:
            return (False, source)
        token_list = find_tokens(source)
        table = self.build_mapping_table(token_list)
        unmangled_str = replace_tokens(source, table)
        return (unmangled_str != source, unmangled_str)

    # Un-mangle an MDL file and create a new un-mangled output file
    def unmangle_mdl_file(self, infile, outfile):
        if self.unmangle_routine == None:
            return
        # Only consider 'import' statements
        filter_import_lines = True
        with open(infile, 'r') as f_in:
            lines = f_in.readlines()
            token_list = find_tokens(lines, filter_import_lines)
            mapping = self.build_mapping_table(token_list)
            lines = replace_tokens(lines, mapping)
            with open(outfile, 'w') as f_out:
                f_out.writelines(lines)

#--------------------------------------------------------------------------------------------------
def parse_alias(pattern, line):
    match = pattern.match(line)
    if match == None:
        return(False, '', '')
    alias = match.group("alias")
    path_segment = match.group("value")

    # print(f"alias: '{alias}' path_segment: '{path_segment}'")
    return(True, alias, path_segment)

#--------------------------------------------------------------------------------------------------
# This unmangling class uses MDL search path to trim down MDL identifiers
# It uses also MDL aliases to perform substitutions in the MDL identifier
class UnmangleAndFixMDLSearchPath(Unmangle):
    alias_pattern = None
    def __init__(self, context):
        # compile pattern only once as it is constant
        self.alias_pattern = re.compile(r"""\s*using\s* # whitespace, using, whitespace
                         (?P<alias>.*?)                 # alias name
                         \s*=\s*                        # whitespace, equal, whitespace
                         \"(?P<value>.*?)\"             # value in quotes, 
                         \s*;""", re.VERBOSE)           # semicolon, whitespace

        # Build MDL search path list
        self.mdl_search_paths = list()
        with context.neuray.get_api_component(pymdlsdk.IMdl_configuration) as cfg:
            if cfg.is_valid_interface():
                for i in range(cfg.get_mdl_paths_length()):
                    sp = cfg.get_mdl_path(i)
                    sp = os.path.normpath(sp.get_c_str())
                    self.mdl_search_paths.append(sp)

        super(UnmangleAndFixMDLSearchPath, self).__init__(context)

    # Determine if this line is defining an alias
    def __is_alias(self, line):
        (is_alias, alias, path_segment) = parse_alias(self.alias_pattern, line)
        if is_alias:
            return (True, path_segment, alias)
        return (False, '', '')

    # Determine if this identifier is mangled
    def __is_mangled_identifier(self, id):
        token_list = id.split("::")
        remove_all_occurences_of_element_from_list('', token_list)
        for token in token_list:
            (unmangled_flag, unmangled_token) = self._unmangle_token(token)
            if unmangled_flag:
                return unmangled_flag
        return False

    # Determine if this string contains a mangled identifier
    def __get_mangled_identifiers(self, line):
        identifiers = find_identifiers(line)
        rtn_identifiers = set()
        for id in identifiers:
            if len(id) > 0:
                if self.__is_mangled_identifier(id):
                    rtn_identifiers.add(id)
        return rtn_identifiers

    # Build mapping table for a given mangled identifier
    # Using the stored MDL search paths and the MDL aliases
    def __build_table_from_mangled_id(self, id, aliases, makeRelativeIfFailure: bool = False):
        key = ''
        value = ''
        # unmangle using _unmangleMdlModulePath
        if not id.find('::') == 0:
            unmangled_id = self.unmangle_routine("::" + id)
        else:
            unmangled_id = self.unmangle_routine(id)
        unmangled_id = re.sub('^file:/', '', unmangled_id)

        # Replace aliases
        for alias in aliases.keys():
            unmangled_id = unmangled_id.replace(aliases[alias], alias)
        unmangled_id = os.path.normpath(unmangled_id)

        # Remove MDL seach path
        search_path_removed = False
        mapping_table = dict()
        for sp in self.mdl_search_paths:
            temp_id = unmangled_id
            temp_sp = sp
            if platform.system() == 'Windows':
                temp_id = unmangled_id.lower()
                temp_sp = sp.lower()
            # Remove the search path only when the path is found at the head of id
            if temp_id.startswith(temp_sp):
                unmangled_id = unmangled_id[len(sp):]
                search_path_removed = True
                break
        # Conversion is done only if we did find search path prefix in the identifier
        if search_path_removed:
            replaced = unmangled_id
            replaced = replaced.replace('\\', '::')
            replaced = replaced.replace('.mdl::', '::')
            what_to_replace = id
            mapping_table[what_to_replace] = replaced
            return (True, mapping_table)
        else:
            if makeRelativeIfFailure:
                # Turn absolute path to relative
                unmangled_id = os.path.basename(unmangled_id)
                replaced = unmangled_id
                replaced = replaced.replace('.mdl::', '::')
                what_to_replace = id
                mapping_table[what_to_replace] = replaced
                return (True, mapping_table)
        return (False, mapping_table)

    # Un-mangle an MDL file and create a new un-mangled output file
    def unmangle_mdl_file(self, infile, outfile):
        if self.unmangle_routine == None:
            return

        # Read file
        mapping_table = dict()
        with open(infile, 'r') as f_in:
            input_lines = f_in.readlines()
            if not type(input_lines) == list:
                # handle single line
                input_lines = [input_lines]
            # aliases mapping table
            aliases = dict()
            for line in input_lines:
                (alias, key, value) = self.__is_alias(line)
                if alias:
                    aliases[key] = value
                    continue
                ids = self.__get_mangled_identifiers(line)
                for id in ids:
                    (success, new_mapping_table) = self.__build_table_from_mangled_id(id, aliases)
                    if success:
                        # Mege new mapping in
                        mapping_table.update(new_mapping_table)

        # Replace in file and write out result
        with open(infile, 'r') as f_in:
            lines = f_in.readlines()
            lines = replace_tokens(lines, mapping_table)
            with open(outfile, 'w') as f_out:
                f_out.writelines(lines)

    # Un-mangle an MDL file and create a new un-mangled output file
    # Return (is identifier mangled, demangling and removing search path succeeded, un-mangled module name)
    def unmangle_mdl_module(self, module):
        if self.unmangle_routine == None:
            return (False, True, module)

        rtn = module
        isMangled = self.__is_mangled_identifier(module)
        (success, mapping_table) = self.__build_table_from_mangled_id(module, dict(), makeRelativeIfFailure = True)
        if success:
            rtn = mapping_table[module]
            if rtn.find('::') == 0:
                rtn = rtn[2:]
            rtn = rtn.replace("::", "/")
        return (isMangled, success, rtn)

#--------------------------------------------------------------------------------------------------

def is_reserved_word(name : str):
    # there will be an API function for this soon
    reserved_words = ['annotation', 'auto', 'bool', 'bool2', 'bool3', 'bool4', 'break', 'bsdf', 'bsdf_measurement',
                      'case', 'cast', 'color', 'const', 'continue', 'default', 'do', 'double', 'double2', 'double2x2',
                      'double2x3', 'double3', 'double3x2', 'double3x3', 'double3x4', 'double4', 'double4x3', 'double4x4',
                      'double4x2', 'double2x4', 'edf', 'else', 'enum', 'export', 'false', 'float', 'float2', 'float2x2',
                      'float2x3', 'float3', 'float3x2', 'float3x3', 'float3x4', 'float4', 'float4x3', 'float4x4',
                      'float4x2', 'float2x4', 'for', 'hair_bsdf', 'if', 'import', 'in', 'int', 'int2', 'int3', 'int4',
                      'intensity_mode', 'intensity_power', 'intensity_radiant_exitance', 'let', 'light_profile', 'material',
                      'material_emission', 'material_geometry', 'material_surface', 'material_volume', 'mdl', 'module',
                      'package', 'return', 'string', 'struct', 'switch', 'texture_2d', 'texture_3d', 'texture_cube',
                      'texture_ptex', 'true', 'typedef', 'uniform', 'using', 'varying', 'vdf', 'while', 'catch', 'char',
                      'class', 'const_cast', 'delete', 'dynamic_cast', 'explicit', 'extern', 'external', 'foreach',
                      'friend', 'goto', 'graph', 'half', 'half2', 'half2x2', 'half2x3', 'half3', 'half3x2', 'half3x3',
                      'half3x4', 'half4', 'half4x3', 'half4x4', 'half4x2', 'half2x4', 'inline', 'inout', 'lambda',
                      'long', 'mutable', 'namespace', 'native', 'new', 'operator', 'out', 'phenomenon', 'private',
                      'protected', 'public', 'reinterpret_cast', 'sampler', 'shader', 'short', 'signed', 'sizeof',
                      'static', 'static_cast', 'technique', 'template', 'this', 'throw', 'try', 'typeid', 'typename',
                      'union', 'unsigned', 'virtual', 'void', 'volatile', 'wchar_t']
    return name in reserved_words

#--------------------------------------------------------------------------------------------------

# the db name of a texture is constructed here:
# rendering\include\rtx\neuraylib\NeurayLibUtils.h
# computeSceneIdentifierTexture
def parse_ov_resource_name(texture_db_name : str):
    # drop the db prefix
    texture_db_name = texture_db_name[5:]
    undescore = texture_db_name.rfind('_')
    texture_db_name = texture_db_name[:undescore]
    undescore = texture_db_name.rfind('_')
    return texture_db_name[:undescore]

async def prepare_resources_for_export(transaction : pymdlsdk.ITransaction, inst_name, temp_dir):

    functionCall = pymdl.FunctionCall._fetchFromDb(transaction, inst_name)

    for name, param in functionCall.parameters.items():
        # name
        if param.type.kind == pymdlsdk.IType.Kind.TK_TEXTURE and param.value[0] != None:
            # print(f"* Name: {name}")
            # print(f"  Type Kind: {param.type.kind}")
            texture_path = parse_ov_resource_name(param.value[0])
            # print(f"    texture_path: {texture_path}")
            # print(f"    db_name: {param.value[0]}")
            # print(f"    gamma: {param.value[1]}")

            # special handling for resources on the server
            if texture_path.startswith("omniverse:") or \
                texture_path.startswith("file:") or \
                texture_path.startswith("ftp:") or \
                texture_path.startswith("http:") or \
                texture_path.startswith("https:"):
                basename = os.path.basename(texture_path)
                filename, fileext = os.path.splitext(basename)
                id = 0
                dst_path = os.path.join(temp_dir, basename)
                while (os.path.exists(dst_path)):
                    dst_path = os.path.join(temp_dir, f"{filename}_{id}{fileext}")
                    id = id + 1

                result = await copy_async(texture_path, dst_path)
                if not result:
                    print(f"Cannot copy from {texture_path} to {dst_path}, error code: {result}.")
                    continue
                else:
                    print(f"Downloaded resource from {texture_path} to {dst_path}. Will be deleted after export.")
                # use the resource in the temp folder
                texture_path = dst_path

            # create a new image
            new_image_db_name = "mdlexp::" + texture_path
            with transaction.create("Image", 0, None) as new_interface:
                 with new_interface.get_interface(pymdlsdk.IImage) as new_image:
                    new_image.reset_file(texture_path)
                    transaction.store(new_image, new_image_db_name)
                    transaction.remove(new_image_db_name) # drop after transaction is closed

            # assign the image to the texture
            # todo make sure the transaction is aborted, we don't want this change visible to other components
            with transaction.edit_as(pymdlsdk.ITexture, param.value[0]) as itexture:
                itexture.set_image(new_image_db_name)
                itexture.set_gamma(param.value[1]) # TODO check if srgb and linear work with floating precission here

        if isinstance(param, pymdl.ArgumentCall):
            # print(f"* Name: {name}")
            # print(f"  Type Kind: {param.type.kind}")
            # print(f"    dbName: {param.value}")
            await prepare_resources_for_export(transaction, param.value, temp_dir)


# ExportHelper: Helps to save files to Nucleus server.
# If we want to export a file to Nucleus, this helper:
#   1- creates a temp folder,
#   2- substitutes the original Nucleus folder with the temp folder,
#   3- files are saved/exported to the temp folder,
#   4- when finalize() is called, the content of the temp folder is saved to Nucleus server
#      and temp folder is deleted
class ExportHelper:
    def __init__(self, out_filename):
        self.original_filename = out_filename
        self.out_filename = out_filename
        self.temp_dir = None
        # If out_filename is on Nucleus, then change it to temp filename
        if out_filename.startswith("omniverse:") or \
            out_filename.startswith("file:") or \
            out_filename.startswith("ftp:"):
            # create a temp folder
            self.temp_dir = tempfile.TemporaryDirectory()
            # substitutes the original Nucleus folder with the temp folder
            self.out_filename = os.path.join(self.temp_dir.name, os.path.basename(out_filename))
    async def finalize(self, copy_files: bool = True):
        # If out_filename is different from the original one, copy all files
        # from the temp folder to the Nucleus server
        if self.temp_dir != None:
            if copy_files:
                from_dir = os.path.dirname(self.out_filename)
                to_dir = os.path.dirname(self.original_filename)
                for ld in os.listdir(from_dir):
                    f = os.path.basename(ld)
                    result = await copy_async(os.path.join(from_dir, f), to_dir + '/' + f)
                    if not result:
                        print(f"Cannot copy from {f} to {to_dir}, error code: {result}.")
                    else:
                        print(f"Copied resource {f} to {to_dir}.")
            # Delete temp folder
            self.temp_dir.cleanup()


async def save_instance_to_module(context, inst_name, usd_prim, temp_dir, out_filename):
    print("============ MDL Save Instance to Module =============")
    neuray = context.neuray
    transaction = context.transaction
    stage = context.stage

    # access shader prim
    prim = get_shader_prim(stage, usd_prim)
    prim_name = usd_prim.pathString.split('/')[-1]

    if out_filename == None or os.path.isdir(out_filename) or os.path.splitext(out_filename)[1] != '.mdl':
        print(f"Error: output filename is invalid, expected an .mdl file: '{out_filename}'")
        return False

    # Handle Nucleus: OM-46539 Graph to MDL fails to copy Nucleus textures to local filesystem when exporting
    export_helper = ExportHelper(out_filename)
    out_filename = export_helper.out_filename

    # Sanity checks
    with transaction.access_as(pymdlsdk.IFunction_call, inst_name) as inst:
        if inst == None or not inst.is_valid_interface():
            print("Error: Invalid instance name (IFunction_call): {}".format(inst_name))
            return False
        with transaction.access_as(pymdlsdk.IFunction_definition, inst.get_function_definition()) as fd:
            if fd == None or not fd.is_valid_interface():
                print("Error: Invalid IFunction_definition: {}".format(inst.get_function_definition()))
                return False
            module = pymdl.Module._fetchFromDb(transaction, fd.get_module())
            if module == None:
                print("Error: Invalid Module: {}".format(fd.get_module()))
                return False

    # since resourcse are handled by the renderer we need to pass them to neuray before exporting
    await prepare_resources_for_export(transaction, inst_name, temp_dir)

    factory: pymdlsdk.IMdl_factory = neuray.get_api_component(pymdlsdk.IMdl_factory)
    execution_context = factory.create_execution_context()

    # Create the module builder.
    module_name = "mdl::new_module_28f8c871b7034e55a0541be81655ffb1"
    module_builder = factory.create_module_builder(
        transaction,
        module_name,
        pymdlsdk.MDL_VERSION_1_6, # In order to use aliases
        pymdlsdk.MDL_VERSION_LATEST,
        execution_context)

    if not module_builder.is_valid_interface():
        print("Error: Failed to create module builder")
        return False

    module_builder.clear_module(execution_context)
    
    # Create a variant
    success = False
    with transaction.access_as(pymdlsdk.IFunction_call, inst_name) as inst, \
        factory.create_expression_factory(transaction) as ef, \
        ef.create_annotation_block() as empty_anno_block:
        if inst.is_valid_interface():
            
            if is_reserved_word(prim_name):
                variant_name = "Main"
            else:
                variant_name = prim_name

            result = module_builder.add_variant(
                variant_name,
                inst.get_function_definition(),
                inst.get_arguments(),
                empty_anno_block,
                empty_anno_block,
                True,
                execution_context)
            if result != 0:
                print("Error: Failed to add variant to module builder:\n\t'{}'\n\t'{}'\n\t'{}'".format(prim.GetName(), inst.get_function_definition(), inst.get_arguments()))
                for i in range(execution_context.get_messages_count()):
                    print(execution_context.get_message(i).get_string())

            if result == 0:
                with neuray.get_api_component(pymdlsdk.IMdl_impexp_api) as imp_exp:
                    rtn = -1
                    execution_context.set_option("bundle_resources", True)
                    rtn = imp_exp.export_module(transaction, module_name, out_filename, execution_context)

                    if rtn >= 0:
                        unmangle_helper = UnmangleAndFixMDLSearchPath(context)
                        unmangle_helper.unmangle_mdl_file(out_filename, out_filename)
                        print(f"Success: Material '{prim_name}' exported to module '{out_filename}'")
                        # Copy files to Nucleus if needed
                        await export_helper.finalize()
                        success = True
                    else:
                        print(f"Failure: Could not export Material '{prim_name}' to module '{out_filename}'")
                        await export_helper.finalize(copy_files = False)

    module_builder.clear_module(execution_context)
    return success
