import numpy

from . import Property
from ._ovito_ndarray_adapter import add_ndarray_interface


# Give the Property class a Numpy-like interface.
add_ndarray_interface(Property)


# The generic ndarray interface implements the 'shape', 'ndim' and 'dtype' attributes as well as element
# iteration by delegating to the NumPy array view of the property. String-typed properties cannot be exposed
# as NumPy arrays, which means the delegation would raise an error for them. Yet, these properties do have a
# well-defined shape, dimensionality and element type, and their elements can be visited one by one.
# We therefore override the affected parts of the interface here and implement them on the basis of the
# property's meta-data and its indexed element access whenever the property is of data type 'str'.


def _Property_shape(self):
    if self.data_type == Property.DataType.String:
        # Mirrors the shape convention used for the NumPy views of numeric properties:
        # Scalar properties without named components are exposed as 1-dimensional arrays.
        if self.component_count == 1 and not self.component_names:
            return (self.size,)
        return (self.size, self.component_count)
    return self.__array__().shape


Property.shape = property(_Property_shape)


def _Property_ndim(self):
    if self.data_type == Property.DataType.String:
        return len(self.shape)
    return self.__array__().ndim


Property.ndim = property(_Property_ndim)


def _Property_dtype(self):
    if self.data_type == Property.DataType.String:
        # NumPy's variable-width UTF-8 string dtype is the closest match for OVITO's string properties.
        # It is also one of the dtypes accepted by PropertyContainer.create_property() for creating
        # string properties, which makes the value returned here round-trip.
        return numpy.dtypes.StringDType()
    return self.__array__().dtype


Property.dtype = property(_Property_dtype)


def _Property_iter(self):
    if self.data_type == Property.DataType.String:
        # Fetch all strings in one go using indexed access and iterate over the resulting list.
        # Note: For string properties with more than one vector component, this raises an IndexError,
        # because such properties do not support element-wise access from Python.
        return iter(self[...])
    return iter(self.__array__())


Property.__iter__ = _Property_iter


# Returns a NumPy array wrapper for a property.
# For backward compatibility with OVITO 2.9.0:
def _Property_array(self):
    # This attribute returns a NumPy array, which provides read access to the per-element data stored in this property object.
    return self.__array__()


Property.array = property(_Property_array)


# Returns a NumPy array wrapper for a property with write access.
# For backward compatibility with OVITO 2.9.0:
def _Property_marray(self):
    with self:
        return self.__array__()


# This is needed to enable the augmented assignment operators (+=, -=, etc.) for the 'marray' property.
# For backward compatibility with OVITO 2.9.0:
def _Property_marray_assign(self, other):
    if not hasattr(other, "__array_interface__"):
        raise ValueError(
            "Only objects supporting the array interface can be assigned to the 'marray' property."
        )
    o = other.__array_interface__
    s = self.__array_interface__
    if (
        o["shape"] != s["shape"]
        or o["typestr"] != s["typestr"]
        or o["data"] != s["data"]
    ):
        raise ValueError(
            "Assignment to the 'marray' property is restricted. Left and right-hand side must be identical."
        )
    # Assume that the data has been changed in the meantime.
    self.notify_object_changed()


# For backward compatibility with OVITO 2.9.0:
Property.marray = property(_Property_marray, _Property_marray_assign)
Property.get_type_by_id = lambda self, id: self.type_by_id(id)
Property.get_type_by_name = lambda self, name: self.type_by_name(name)
Property.type_list = property(lambda self: self.types)
