# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import copy
import weakref
from typing import List

import omni.ui as ui
from pxr import Gf, Sdf, Tf, Usd, UsdShade

from .usd_model_base import *


class FloatModel(ui.SimpleFloatModel):
    def __init__(self, parent):
        super().__init__()
        self._parent = weakref.ref(parent)

    def begin_edit(self):
        parent = self._parent()
        parent.begin_edit(None)

    def end_edit(self):
        parent = self._parent()
        parent.end_edit(None)


class IntModel(ui.SimpleIntModel):
    def __init__(self, parent):
        super().__init__()
        self._parent = weakref.ref(parent)

    def begin_edit(self):
        parent = self._parent()
        parent.begin_edit(None)

    def end_edit(self):
        parent = self._parent()
        parent.end_edit(None)


class UsdAttributeModel(ui.AbstractValueModel, UsdBase):
    """The value model that is reimplemented in Python to watch a USD attribute path"""

    def __init__(
        self,
        stage: Usd.Stage,
        attribute_paths: List[Sdf.Path],
        self_refresh: bool,
        metadata: dict,
        change_on_edit_end=False,
        **kwargs,
    ):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata, change_on_edit_end, **kwargs)
        ui.AbstractValueModel.__init__(self)

    def clean(self):
        UsdBase.clean(self)

    def begin_edit(self):
        UsdBase.begin_edit(self)
        ui.AbstractValueModel.begin_edit(self)

    def end_edit(self):
        UsdBase.end_edit(self)
        ui.AbstractValueModel.end_edit(self)

    def get_value_as_string(self, elide_big_array=True) -> str:
        self._update_value()
        if self._value is None:
            return ""
        elif self._is_big_array and elide_big_array:
            return "[...]"
        else:
            return str(self._value)

    def get_value_as_float(self) -> float:
        self._update_value()
        if self._value is None:
            return 0.0
        else:
            if hasattr(self._value, "__len__"):
                return float(self._value[self._channel_index])
            return float(self._value)

    def get_value_as_bool(self) -> bool:
        self._update_value()
        if self._value is None:
            return False
        else:
            if hasattr(self._value, "__len__"):
                return bool(self._value[self._channel_index])
            return bool(self._value)

    def get_value_as_int(self) -> int:
        self._update_value()
        if self._value is None:
            return 0
        else:
            if hasattr(self._value, "__len__"):
                return int(self._value[self._channel_index])
            return int(self._value)

    def set_value(self, value):
        if UsdBase.set_value(self, value):
            self._value_changed()

    def _on_dirty(self):
        self._value_changed()


class GfVecAttributeSingleChannelModel(UsdAttributeModel):
    def __init__(
        self,
        stage: Usd.Stage,
        attribute_paths: List[Sdf.Path],
        channel_index: int,
        self_refresh: bool,
        metadata: dict,
        change_on_edit_end=False,
        **kwargs,
    ):
        self._channel_index = channel_index
        super().__init__(stage, attribute_paths, self_refresh, metadata, change_on_edit_end, **kwargs)

    def get_value_as_string(self, **kwargs) -> str:
        self._update_value()
        if self._value is None:
            return ""
        else:
            return str(self._value[self._channel_index])

    def get_value_as_float(self) -> float:
        self._update_value()
        if self._value is None:
            return 0.0
        else:
            if hasattr(self._value, "__len__"):
                return float(self._value[self._channel_index])
            return float(self._value)

    def get_value_as_bool(self) -> bool:
        self._update_value()
        if self._value is None:
            return False
        else:
            if hasattr(self._value, "__len__"):
                return bool(self._value[self._channel_index])
            return bool(self._value)

    def get_value_as_int(self) -> int:
        self._update_value()
        if self._value is None:
            return 0
        else:
            if hasattr(self._value, "__len__"):
                return int(self._value[self._channel_index])
            return int(self._value)

    def set_value(self, value):
        vec_value = copy.copy(self._value)
        vec_value[self._channel_index] = value

        if UsdBase.set_value(self, vec_value, self._channel_index):
            self._value_changed()

    def is_different_from_default(self):
        """Override to only check channel"""
        if super().is_different_from_default():
            return self._comp_different_from_default[self._channel_index]
        return False


class TfTokenAttributeModel(ui.AbstractItemModel, UsdBase):
    class AllowedTokenItem(ui.AbstractItem):
        def __init__(self, item):
            super().__init__()
            self.token = item
            self.model = ui.SimpleStringModel(item)

    def __init__(self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], self_refresh: bool, metadata: dict):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)

        self._allowed_tokens = []

        self._current_index = ui.SimpleIntModel()
        self._current_index.add_value_changed_fn(self._current_index_changed)

        self._updating_value = False

        self._has_index = False
        self._update_value()
        self._has_index = True

    def clean(self):
        UsdBase.clean(self)

    def get_item_children(self, item):
        self._update_value()
        return self._allowed_tokens

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._current_index

        return item.model

    def begin_edit(self, item):
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        UsdBase.end_edit(self)

    def _current_index_changed(self, model):
        if not self._has_index:
            return

        # if we're updating from USD notice change to UI, don't call set_value
        if self._updating_value:
            return

        index = model.as_int
        if self.set_value(self._get_value_from_index(index)):
            self._item_changed(None)

    def _get_allowed_tokens(self, attr):
        return attr.GetMetadata("allowedTokens")

    def _update_allowed_token(self, token_item=AllowedTokenItem):
        self._allowed_tokens = []

        # For multi prim editing, the allowedTokens should all be the same
        attributes = self._get_attributes()
        attr = attributes[0] if len(attributes) > 0 else None
        if attr:
            for t in self._get_allowed_tokens(attr):
                self._allowed_tokens.append(token_item(t))

    def _update_value(self, force=False):
        was_updating_value = self._updating_value
        self._updating_value = True
        if UsdBase._update_value(self, force):
            # TODO don't have to do this every time. Just needed when "allowedTokens" actually changed
            self._update_allowed_token()

            index = -1
            for i in range(0, len(self._allowed_tokens)):
                if self._allowed_tokens[i].token == self._value:
                    index = i

            if index != -1 and self._current_index.as_int != index:
                self._current_index.set_value(index)
                self._item_changed(None)
        self._updating_value = was_updating_value

    def _on_dirty(self):
        self._item_changed(None)

    def get_value_as_token(self):
        index = self._current_index.as_int
        return self._allowed_tokens[index].token

    def is_allowed_token(self, token):
        return token in [allowed.token for allowed in self._allowed_tokens]

    def _get_value_from_index(self, index):
        return self._allowed_tokens[index].token


class MdlEnumAttributeModel(ui.AbstractItemModel, UsdBase):
    def __init__(self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], self_refresh: bool, metadata: dict):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)

        self._options = []

        self._current_index = ui.SimpleIntModel()
        self._current_index.add_value_changed_fn(self._current_index_changed)

        self._updating_value = False

        self._has_index = False
        self._update_value()
        self._has_index = True

    def clean(self):
        UsdBase.clean(self)

    def get_item_children(self, item):
        self._update_value()
        return self._options

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._current_index

        return item.model

    def begin_edit(self, item):
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        UsdBase.end_edit(self)

    def _current_index_changed(self, model):
        if not self._has_index:
            return

        # if we're updating from USD notice change to UI, don't call set_value
        if self._updating_value:
            return

        index = model.as_int
        if self.set_value(self._options[index].value):
            self._item_changed(None)

    def _update_option(self):
        self._options = []

        # For multi prim editing, the options should all be the same
        attributes = self._get_attributes()
        attr = attributes[0] if len(attributes) > 0 else None
        if attr and isinstance(attr, Usd.Attribute):
            shade_input = UsdShade.Input(attr)
            if shade_input and shade_input.HasRenderType() and shade_input.HasSdrMetadataByKey("options"):
                options = shade_input.GetSdrMetadataByKey("options").split("|")

                class OptionItem(ui.AbstractItem):
                    def __init__(self, display_name: str, value: int):
                        super().__init__()
                        self.model = ui.SimpleStringModel(display_name)
                        self.value = value

                for option in options:
                    kv = option.split(":")
                    self._options.append(OptionItem(kv[0], int(kv[1])))

    def _update_value(self, force=False):
        was_updating_value = self._updating_value
        self._updating_value = True
        if UsdBase._update_value(self, force):
            # TODO don't have to do this every time. Just needed when "option" actually changed
            self._update_option()

            index = -1
            for i in range(0, len(self._options)):
                if self._options[i].value == self._value:
                    index = i

            if index != -1 and self._current_index.as_int != index:
                self._current_index.set_value(index)
                self._item_changed(None)
        self._updating_value = was_updating_value

    def _on_dirty(self):
        self._item_changed(None)

    def get_value_as_string(self):
        index = self._current_index.as_int
        return self._options[index].model.as_string

    def is_allowed_enum_string(self, enum_str):
        return enum_str in [allowed.model.as_string for allowed in self._options]

    def set_from_enum_string(self, enum_str):
        if not self.is_allowed_enum_string(enum_str):
            return

        new_index = -1
        for index, option in enumerate(self._options):
            if option.model.as_string == enum_str:
                new_index = index
                break
        if new_index != -1:
            self._current_index.set_value(new_index)


class GfVecAttributeModel(ui.AbstractItemModel, UsdBase):
    def __init__(
        self,
        stage: Usd.Stage,
        attribute_paths: List[Sdf.Path],
        comp_count: int,
        tf_type: Tf.Type,
        self_refresh: bool,
        metadata: dict,
        **kwargs,
    ):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata, **kwargs)
        ui.AbstractItemModel.__init__(self)
        self._comp_count = comp_count
        self._data_type_name = "Vec" + str(self._comp_count) + tf_type.typeName[-1]
        self._data_type = getattr(Gf, self._data_type_name)

        class UsdVectorItem(ui.AbstractItem):
            def __init__(self, model):
                super().__init__()
                self.model = model

        # Create root model
        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # Create three models per component
        if self._data_type_name.endswith("i"):
            self._items = [UsdVectorItem(IntModel(self)) for i in range(self._comp_count)]
        else:
            self._items = [UsdVectorItem(FloatModel(self)) for i in range(self._comp_count)]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

        self._edit_mode_counter = 0

    def clean(self):
        UsdBase.clean(self)

    def _construct_vector_from_item(self):
        if self._data_type_name.endswith("i"):
            data = [item.model.get_value_as_int() for item in self._items]
        else:
            data = [item.model.get_value_as_float() for item in self._items]
        return self._data_type(data)

    def _on_value_changed(self, item):
        """Called when the submodel is changed"""

        if self._edit_mode_counter > 0:
            vector = self._construct_vector_from_item()
            index = self._items.index(item)
            if vector and self.set_value(vector, index):
                # Read the new value back in case hard range clamped it
                item.model.set_value(self._value[index])
                self._item_changed(item)
            else:
                # If failed to update value in model, revert the value in submodel
                item.model.set_value(self._value[index])

    def _update_value(self, force=False):
        if UsdBase._update_value(self, force):
            if not self._value:
                for i in range(len(self._items)):
                    self._items[i].model.set_value(0.0)
                return
            for i in range(len(self._items)):
                self._items[i].model.set_value(self._value[i])

    def _on_dirty(self):
        self._item_changed(None)

    def get_item_children(self, item):
        """Reimplemented from the base class"""
        self._update_value()
        return self._items

    def get_item_value_model(self, item, column_id):
        """Reimplemented from the base class"""
        if item is None:
            return self._root_model
        return item.model

    def begin_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user starts editing.
        """
        self._edit_mode_counter += 1
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user finishes editing.
        """

        UsdBase.end_edit(self)
        self._edit_mode_counter -= 1


class SdfAssetPathAttributeModel(UsdAttributeModel):
    """The value model that is reimplemented in Python to watch a USD attribute path"""

    def __init__(self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], self_refresh: bool, metadata: dict):
        super().__init__(stage, attribute_paths, self_refresh, metadata, True, treat_array_entry_as_comp=True)

    def get_value_as_string(self, **kwargs):
        self._update_value()
        return self._get_value_as_string(self._value)

    def is_valid_path(self) -> bool:
        return self._is_valid_path(self._value)

    def get_resolved_path(self):
        self._update_value()
        return self._get_resolved_path(self._value)

    def set_value(self, path, resolved_path: str=""):
        value = path
        if not isinstance(path, Sdf.AssetPath):
            if resolved_path:
                value = Sdf.AssetPath(path, resolved_path)
            else:
                value = Sdf.AssetPath(path)

        if UsdBase.set_value(self, value):
            self._value_changed()

    def _is_prev_same(self):
        # Strip the resolvedPath from the AssetPath for the comparison, since the prev values don't have resolvedPath.
        return [Sdf.AssetPath(value.path) for value in self._real_values] == self._prev_real_values

    def _save_real_values_as_prev(self):
        # Strip the resolvedPath from the AssetPath so that it can be recomputed.
        self._prev_real_values = [Sdf.AssetPath(value.path) for value in self._real_values]

    def _change_property(self, path: Sdf.Path, new_value, old_value):
        if path.name == "info:mdl:sourceAsset" and new_value.path:
            if isinstance(new_value, Sdf.AssetPath):
                # make the path relative to current edit target layer
                relative_path = omni.usd.make_path_relative_to_current_edit_target(new_value.path)
                new_value = Sdf.AssetPath(relative_path, new_value.resolvedPath)

        if path.name == "info:mdl:sourceAsset":
            if isinstance(new_value, Sdf.AssetPath):
                # "info:mdl:sourceAsset" when is changed update "info:mdl:sourceAsset:subIdentifier"
                stage = self._stage
                asset_attr = stage.GetAttributeAtPath(path)
                subid_attr = stage.GetAttributeAtPath(
                    path.GetPrimPath().AppendElementString(".info:mdl:sourceAsset:subIdentifier")
                )
                if asset_attr and subid_attr:
                    mdl_path = new_value.resolvedPath if new_value.resolvedPath else new_value.path
                    if mdl_path:
                        import asyncio

                        baseclass = super()

                        def have_subids(id_list: list):
                            if len(id_list) > 0:
                                with omni.kit.undo.group():
                                    baseclass._change_property(subid_attr.GetPath(), str(id_list[0]), None)
                                    baseclass._change_property(path, new_value, old_value)

                        asyncio.ensure_future(
                            omni.kit.material.library.get_subidentifier_from_mdl(
                                mdl_file=mdl_path, on_complete_fn=have_subids, use_functions=True, show_alert=True
                            )
                        )
                        return
                    else:
                        super()._change_property(subid_attr.GetPath(), "", None)

        super()._change_property(path, new_value, old_value)

    # private functions to be shared with SdfAssetPathArrayAttributeSingleEntryModel
    def _get_value_as_string(self, value, **kwargs):
        return value.path if value else ""

    def _is_valid_path(self, value) -> bool:
        if not value:
            return False
        path = value.path
        if path.lower().startswith("http"):
            return False
        return bool(path)

    def _get_resolved_path(self, value):
        if self._ambiguous:
            return "Mixed"
        else:
            return value.resolvedPath.replace("\\", "/") if value else ""


class SdfAssetPathArrayAttributeSingleEntryModel(SdfAssetPathAttributeModel):
    def __init__(
        self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], index: int, self_refresh: bool, metadata: dict
    ):
        self._index = index
        super().__init__(stage, attribute_paths, self_refresh, metadata)

    @property
    def index(self):
        return self._index

    def get_value_as_string(self, **kwargs):
        self._update_value()
        return self._get_value_as_string(self._value[self._index])

    def get_value(self):
        return super().get_value()[self._index]

    def is_valid_path(self) -> bool:
        self._update_value()
        return self._is_valid_path(self._value[self._index])

    def get_resolved_path(self):
        self._update_value()
        return self._get_resolved_path(self._value[self._index])

    def set_value(self, path, resolved_path: str=""):
        value = path
        if not isinstance(path, Sdf.AssetPath):
            if resolved_path:
                value = Sdf.AssetPath(path, resolved_path)
            else:
                value = Sdf.AssetPath(path)

        vec_value = Sdf.AssetPathArray(self._value)
        vec_value[self._index] = value

        if UsdBase.set_value(self, vec_value, self._index):
            self._value_changed()

    def _is_prev_same(self):
        # Strip the resolvedPath from the AssetPath for the comparison, since the prev values don't have resolvedPath.
        return [
            [Sdf.AssetPath(value.path) for value in values] for values in self._real_values
        ] == self._prev_real_values

    def _save_real_values_as_prev(self):
        # Strip the resolvedPath from the AssetPath so that it can be recomputed.
        self._prev_real_values = [[Sdf.AssetPath(value.path) for value in values] for values in self._real_values]


class SdfAssetPathArrayAttributeItemModel(ui.AbstractItemModel):
    class SdfAssetPathItem(ui.AbstractItem):
        """Single item of the model"""

        def __init__(
            self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], index: int, self_refresh: bool, metadata: dict
        ):
            super().__init__()
            self.sdf_asset_path_model = SdfAssetPathArrayAttributeSingleEntryModel(
                stage, attribute_paths, index, self_refresh, metadata
            )

        def destroy(self):
            self.sdf_asset_path_model.clean()

    def __init__(self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], self_refresh: bool, metadata: dict, delegate):
        super().__init__()
        self._delegate = delegate  # keep a reference of the delegate os it's not destroyed
        self._value_model = UsdAttributeModel(stage, attribute_paths, False, metadata)
        value = self._value_model.get_value()
        self._entries = []

        self._repopulate_entries(value)

    def clean(self):
        self._delegate = None

        for entry in self._entries:
            entry.destroy()
        self._entries.clear()

        if self._value_model:
            self._value_model.clean()
            self._value_model = None

    @property
    def value_model(self) -> UsdAttributeModel:
        return self._value_model

    def get_item_children(self, item):
        """Returns all the children when the widget asks it."""
        if item is not None:
            # Since we are doing a flat list, we return the children of root only.
            # If it's not root we return.
            return []

        return self._entries

    def get_item_value_model_count(self, item):
        return 1

    def get_item_value_model(self, item, column_id):
        return (item.sdf_asset_path_model, self._value_model)

    def get_drag_mime_data(self, item):
        return str(item.sdf_asset_path_model.index)

    def drop_accepted(self, target_item, source, drop_location=-1):
        try:
            source_id = self._entries.index(source)
        except ValueError:
            # Not in the list. This is the source from another model.
            return False
 
        return not target_item and drop_location >= 0

    def drop(self, target_item, source, drop_location=-1):
        try:
            source_id = self._entries.index(source)
        except ValueError:
            # Not in the list. This is the source from another model.
            return

        if source_id == drop_location:
            # Nothing to do
            return

        value = list(self._value_model.get_value())
        moved_entry_value = value[source_id]
        del value[source_id]

        if drop_location > len(value):
            # Drop it to the end
            value.append(moved_entry_value)
        else:
            if source_id < drop_location:
                # Because when we removed source, the array became shorter
                drop_location = drop_location - 1

            value.insert(drop_location, moved_entry_value)
        self._value_model.set_value(value)

    def _repopulate_entries(self, value: Sdf.AssetPathArray):
        for entry in self._entries:
            entry.destroy()
        self._entries.clear()

        stage = self._value_model.stage
        metadata = self._value_model.metadata
        attribute_paths = self._value_model.get_attribute_paths()

        for i in range(len(value)):
            model = SdfAssetPathArrayAttributeItemModel.SdfAssetPathItem(stage, attribute_paths, i, False, metadata)
            self._entries.append(model)

        self._item_changed(None)

    def _on_usd_changed(self, *args, **kwargs):
        # forward to all sub-models
        self._value_model._on_usd_changed(*args, **kwargs)
        for entry in self._entries:
            entry.sdf_asset_path_model._on_usd_changed(*args, **kwargs)

    def _set_dirty(self, *args, **kwargs):
        # forward to all sub-models
        self._value_model._set_dirty(*args, **kwargs)

        new_value = self._value_model.get_value()
        if len(new_value) != len(self._entries):
            self._repopulate_entries(new_value)
        else:
            for entry in self._entries:
                entry.sdf_asset_path_model._set_dirty(*args, **kwargs)

    def get_value(self, *args, **kwargs):
        return self._value_model.get_value(*args, **kwargs)

    def set_value(self, *args, **kwargs):
        return self._value_model.set_value(*args, **kwargs)


class GfQuatAttributeModel(ui.AbstractItemModel, UsdBase):
    def __init__(
        self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], tf_type: Tf.Type, self_refresh: bool, metadata: dict
    ):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)
        data_type_name = "Quat" + tf_type.typeName[-1]
        self._data_type = getattr(Gf, data_type_name)

        class UsdQuatItem(ui.AbstractItem):
            def __init__(self, model):
                super().__init__()
                self.model = model

        # Create root model
        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # Create four models per component
        self._items = [UsdQuatItem(FloatModel(self)) for i in range(4)]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

        self._edit_mode_counter = 0

    def clean(self):
        UsdBase.clean(self)

    def _on_value_changed(self, item):
        """Called when the submodel is chaged"""
        if self._edit_mode_counter > 0:
            quat = self._construct_quat_from_item()
            if quat and self.set_value(quat, self._items.index(item)):
                self._item_changed(item)

    def _update_value(self, force=False):
        if UsdBase._update_value(self, force):
            for i in range(len(self._items)):
                if i == 0:
                    self._items[i].model.set_value(self._value.real)
                else:
                    self._items[i].model.set_value(self._value.imaginary[i - 1])

    def _on_dirty(self):
        self._item_changed(None)

    def get_item_children(self, item):
        """Reimplemented from the base class"""
        self._update_value()
        return self._items

    def get_item_value_model(self, item, column_id):
        """Reimplemented from the base class"""
        if item is None:
            return self._root_model
        return item.model

    def begin_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user starts editing.
        """
        self._edit_mode_counter += 1
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user finishes editing.
        """
        UsdBase.end_edit(self)
        self._edit_mode_counter -= 1

    def _construct_quat_from_item(self):
        data = [item.model.get_value_as_float() for item in self._items]

        return self._data_type(data[0], data[1], data[2], data[3])


# A data model for display orient(quaternion) as rotate(Euler)
class GfQuatEulerAttributeModel(ui.AbstractItemModel, UsdBase):
    axes = [Gf.Vec3d(1, 0, 0), Gf.Vec3d(0, 1, 0), Gf.Vec3d(0, 0, 1)]

    def __init__(
        self, stage: Usd.Stage, attribute_paths: List[Sdf.Path], tf_type: Tf.Type, self_refresh: bool, metadata: dict
    ):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)

        # The underline USD data is still a quaternion
        data_type_name = "Quat" + tf_type.typeName[-1]
        self._data_type = getattr(Gf, data_type_name)

        class UsdFloatItem(ui.AbstractItem):
            def __init__(self, model):
                super().__init__()
                self.model = model

        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # Create three models per component
        self._items = [UsdFloatItem(FloatModel(self)) for _ in range(3)]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

        self._edit_mode_counter = 0

    def clean(self):
        UsdBase.clean(self)

    # Should be used to set the self._items to USD attribute
    def _on_value_changed(self, item):
        if self._edit_mode_counter > 0:
            quat = self._compose([item.model.as_float for item in self._items])
            if quat and self.set_value(quat):
                self._item_changed(item)

    # Unlike normal Vec3 or Quat, we have to update all four values all together
    def update_to_submodels(self):
        if self._edit_mode_counter != 0:
            return
        if self._value is not None:
            e = self._decompose(self._value)
            for i, item in enumerate(self._items):
                item.model.set_value(e[i])

    def _update_value(self, force=False):
        if UsdBase._update_value(self, force):
            self.update_to_submodels()

    def _on_dirty(self):
        self._item_changed(None)

    def get_item_children(self, item):
        self._update_value()
        return self._items

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._root_model
        return item.model

    def begin_edit(self, item):
        self._edit_mode_counter += 1
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        UsdBase.end_edit(self)
        self._edit_mode_counter -= 1

    def _compose(self, eulers):
        nrs = [Gf.Rotation(axis, eulers[i]) for i, axis in enumerate(self.axes)]
        nr = nrs[2] * nrs[1] * nrs[0]
        result = self._data_type(nr.GetQuat())
        return result

    def _decompose(self, value):
        rot = Gf.Rotation(value)
        eulers = rot.Decompose(*self.axes)
        # round epsilons from decompose
        eulers = [round(angle + 1e-4, 3) for angle in eulers]
        return eulers

    def _get_comp_num(self):
        return 3

    def _get_value_by_comp(self, value, comp: int):
        e = self._decompose(value)
        return e[comp]

    def _update_value_by_comp(self, value, comp: int):
        self_e = self._decompose(self._value)
        e = self._decompose(value)
        e[comp] = self_e[comp]
        value = self._compose(e)
        return value

    def _compare_value_by_comp(self, val1, val2, comp: int):
        return Gf.IsClose(self._get_value_by_comp(val1, comp), self._get_value_by_comp(val2, comp), 1e-6)


class GfMatrixAttributeModel(ui.AbstractItemModel, UsdBase):
    def __init__(
        self,
        stage: Usd.Stage,
        attribute_paths: List[Sdf.Path],
        comp_count: int,
        tf_type: Tf.Type,
        self_refresh: bool,
        metadata: dict,
    ):
        UsdBase.__init__(self, stage, attribute_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)
        self._comp_count = comp_count
        data_type_name = "Matrix" + str(self._comp_count) + tf_type.typeName[-1]
        self._data_type = getattr(Gf, data_type_name)

        class UsdMatrixItem(ui.AbstractItem):
            def __init__(self, model):
                super().__init__()
                self.model = model

        # Create root model
        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # Create three models per component
        self._items = [UsdMatrixItem(FloatModel(self)) for i in range(self._comp_count * self._comp_count)]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

        self._edit_mode_counter = 0

    def clean(self):
        UsdBase.clean(self)

    def _on_value_changed(self, item):
        """Called when the submodel is chaged"""

        if self._edit_mode_counter > 0:
            matrix = self._construct_matrix_from_item()
            if matrix and self.set_value(matrix, self._items.index(item)):
                self._item_changed(item)

    def _update_value(self, force=False):
        if UsdBase._update_value(self, force):
            for i in range(len(self._items)):
                self._items[i].model.set_value(self._value[i // self._comp_count][i % self._comp_count])

    def _on_dirty(self):
        self._item_changed(None)
        # it's still better to call _value_changed for all child items
        for child in self._items:
            child.model._value_changed()

    def get_item_children(self, item):
        """Reimplemented from the base class."""
        self._update_value()
        return self._items

    def get_item_value_model(self, item, column_id):
        """Reimplemented from the base class."""
        if item is None:
            return self._root_model
        return item.model

    def begin_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user starts editing.
        """
        self._edit_mode_counter += 1
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        """
        Reimplemented from the base class.
        Called when the user finishes editing.
        """
        UsdBase.end_edit(self)
        self._edit_mode_counter -= 1

    def _construct_matrix_from_item(self):
        data = [item.model.get_value_as_float() for item in self._items]
        matrix = []
        for i in range(self._comp_count):
            matrix_row = []
            for j in range(self._comp_count):
                matrix_row.append(data[i * self._comp_count + j])
            matrix.append(matrix_row)
        return self._data_type(matrix)


class UsdAttributeInvertedModel(UsdAttributeModel):
    def get_value_as_bool(self) -> bool:
        return not super().get_value_as_bool()

    def get_value_as_string(self, **kwargs) -> str:
        return str(self.get_value_as_bool())

    def set_value(self, value):
        super().set_value(not value)

    def get_value(self):
        return not super().get_value()


class SdfTimeCodeModel(UsdAttributeModel):
    def _save_real_values_as_prev(self):
        # SdfTimeCode cannot be inited from another SdfTimeCode, only from float (double in C++)..
        self._prev_real_values = [Sdf.TimeCode(float(value)) for value in self._real_values]
