from __future__ import annotations
import copy
from typing import Any, Optional, Union
from ovito.data import DataCollection, DataObject, ElementType, Property
from . import EditTypesModifier

# The EditTypesModifier stores whole ElementType objects in its 'edited_types' list, and each of them *replaces*
# the corresponding type definition of the upstream data. Constructing such a replacement from scratch would
# silently discard all attributes the caller does not set explicitly (name, radius, mass, ...) and, worse, a plain
# ElementType would never match an upstream ParticleType, because the modifier matches by numeric ID *and* class.
# The helper methods below hide that by always starting from a copy of the upstream type definition.

def _EditTypesModifier_input_data(modifier: EditTypesModifier, data: Optional[DataCollection]) -> DataCollection:
    """Returns the data collection entering the modifier, i.e. the state of the pipeline right before it."""
    if data is not None:
        return data
    nodes = modifier.get_modification_nodes()
    if not nodes:
        raise RuntimeError("This 'Edit types' modifier is not part of a pipeline yet, so its input data is unknown. "
                           "Insert the modifier into a pipeline first, e.g. by calling pipeline.modifiers.append(modifier), "
                           "or pass the upstream data collection to this method as 'data' keyword argument.")
    if len(nodes) > 1:
        raise RuntimeError("This 'Edit types' modifier is used in more than one pipeline, so its input data is ambiguous. "
                           "Pass the upstream data collection to this method as 'data' keyword argument.")
    return nodes[0].input.compute()

def _EditTypesModifier_input_property(modifier: EditTypesModifier, data: DataCollection):
    """Looks up the typed property selected by the modifier in the given data collection.
       Returns the property and the property container it belongs to."""
    if not modifier.property:
        raise RuntimeError("No typed property has been selected in the 'Edit types' modifier. "
                           "Set the modifier's 'property' field, e.g. to 'particles/Particle Type'.")
    path = data.get(DataObject.Ref(Property, modifier.property), require=False, path=True)
    if path is None:
        raise KeyError(f"The input data of the 'Edit types' modifier contains no property '{modifier.property}'.")
    if len(path) < 2:
        raise KeyError(f"The property '{modifier.property}' is not part of a property container.")
    return path[-1], path[-2]

def _EditTypesModifier_find_edited_type(modifier: EditTypesModifier, type: Union[int, str]) -> Optional[ElementType]:
    """Returns the entry from the modifier's 'edited_types' list matching the given numeric ID or name."""
    # Anything that is not a string is taken to be a numeric ID, so that NumPy integers work too.
    for element_type in modifier.edited_types:
        if isinstance(type, str):
            if element_type.name == type:
                return element_type
        elif element_type.id == int(type):
            return element_type
    return None

def _EditTypesModifier_resolve_type_id(modifier: EditTypesModifier, type: Union[int, str], data: Optional[DataCollection]) -> int:
    """Turns a type name into the corresponding numeric type ID. Numeric IDs are passed through unchanged."""
    if not isinstance(type, str):
        return int(type)
    element_type = _EditTypesModifier_find_edited_type(modifier, type)
    if element_type is not None:
        return element_type.id
    input_property, _ = _EditTypesModifier_input_property(modifier, _EditTypesModifier_input_data(modifier, data))
    return input_property.type_by_name(type).id

# Implementation of the EditTypesModifier.edit_type() method.
def _EditTypesModifier_edit_type(self,
                                 type: Union[int, str],
                                 *,
                                 data: Optional[DataCollection] = None,
                                 **params: Any) -> ElementType:
    """
    Modifies the attributes of an existing element type. This is the standard way of changing the display color,
    radius, or name of a particle or bond type from a Python script.

    The method looks up the type in the modifier's input data, adds a copy of its definition to the
    :py:attr:`edited_types` list, and applies the given keyword arguments to that copy. All attributes you do
    *not* specify keep the values they have in the upstream data. Calling the method again for the same type
    returns the same object, so edits accumulate:

    .. literalinclude:: ../example_snippets/edit_types_modifier.py
       :lines: 10-16

    :param type: The numeric :py:attr:`~ovito.data.ElementType.id` or the :py:attr:`~ovito.data.ElementType.name` of the type to edit.
    :param data: The modifier's input data collection. Only needed if the modifier is not part of a pipeline yet
                 or is shared by several pipelines. Otherwise it is determined automatically.
    :param params: The type attributes to change, e.g. ``color=(1.0, 0.0, 0.0)`` or ``radius=1.2``.
                   See :py:class:`~ovito.data.ElementType` and its sub-classes for the available attributes.
    :return: The mutable :py:class:`~ovito.data.ElementType` object stored in the :py:attr:`edited_types` list.
    """
    element_type = _EditTypesModifier_find_edited_type(self, type)
    if element_type is None:
        input_property, _ = _EditTypesModifier_input_property(self, _EditTypesModifier_input_data(self, data))
        if isinstance(type, str):
            original_type = input_property.type_by_name(type)
        else:
            original_type = input_property.type_by_id(int(type))
        # Detach a mutable copy from the input data. It retains the concrete ElementType sub-class (e.g. ParticleType)
        # and all attribute values of the original type definition.
        element_type = copy.copy(original_type)
        self.edited_types.append(element_type)
        # Editing a type that was previously marked for deletion implicitly restores it, as in the GUI.
        deleted_ids = self.deleted_type_ids
        if element_type.id in deleted_ids:
            deleted_ids.discard(element_type.id)
            self.deleted_type_ids = deleted_ids
    for name, value in params.items():
        setattr(element_type, name, value)
    return element_type

# Implementation of the EditTypesModifier.add_type() method.
def _EditTypesModifier_add_type(self,
                                id: Optional[int] = None,
                                *,
                                name: str = '',
                                data: Optional[DataCollection] = None,
                                **params: Any) -> ElementType:
    """
    Adds a new element type definition to the selected typed property.

    Note that this only defines the new type; it does not assign it to any elements. Use the
    :py:class:`ComputePropertyModifier` or a :ref:`user-defined modifier function <writing_custom_modifiers>`
    downstream of this modifier to give elements the new type's numeric ID.

    .. literalinclude:: ../example_snippets/edit_types_modifier.py
       :lines: 18-21

    :param id: The numeric ID of the new type. If ``None``, the smallest ID not yet in use is picked.
    :param name: The human-readable name of the new type. If it matches a chemical element symbol, the type's
                 color, radius and mass are preset from OVITO's internal database of chemical elements.
    :param data: The modifier's input data collection. Only needed if the modifier is not part of a pipeline yet
                 or is shared by several pipelines. Otherwise it is determined automatically.
    :param params: Additional type attributes to set, e.g. ``color=(1.0, 0.0, 0.0)``.
    :return: The mutable :py:class:`~ovito.data.ElementType` object stored in the :py:attr:`edited_types` list.
    """
    input_property, container = _EditTypesModifier_input_property(self, _EditTypesModifier_input_data(self, data))
    if id is None:
        id = input_property.unique_type_id()
        # Also skip over IDs that are only defined by this modifier and thus not known to the input property.
        used_ids = {t.id for t in self.edited_types}
        while id in used_ids:
            id += 1
    elif input_property.type_by_id(id, require=False) is not None or _EditTypesModifier_find_edited_type(self, id) is not None:
        raise ValueError(f"A type with numeric ID {id} is already defined for property '{self.property}'. "
                         "Use edit_type() to modify an existing type.")
    # Let the property container decide which ElementType sub-class to instantiate and how to preset its
    # attributes from the type's name. Working on a private copy of the property keeps the input data untouched.
    property_copy = copy.copy(input_property)
    element_type = copy.copy(property_copy.add_type_id(id, container, name))
    self.edited_types.append(element_type)
    for param_name, value in params.items():
        setattr(element_type, param_name, value)
    return element_type

# Implementation of the EditTypesModifier.delete_type() method.
def _EditTypesModifier_delete_type(self, type: Union[int, str], *, data: Optional[DataCollection] = None) -> None:
    """
    Removes an element type definition from the selected typed property.

    Note that this deletes the type *definition* only. Elements still carrying the deleted numeric type ID keep it,
    they just no longer have a type descriptor providing a name and display color. Typically you want to delete
    those elements too, e.g. with the :py:class:`ExpressionSelectionModifier` followed by the
    :py:class:`DeleteSelectedModifier`, inserted upstream of this modifier.

    .. literalinclude:: ../example_snippets/edit_types_modifier.py
       :lines: 23-24

    :param type: The numeric :py:attr:`~ovito.data.ElementType.id` or the :py:attr:`~ovito.data.ElementType.name` of the type to delete.
    :param data: The modifier's input data collection. Only needed if the modifier is not part of a pipeline yet
                 or is shared by several pipelines. Otherwise it is determined automatically.
    """
    type_id = _EditTypesModifier_resolve_type_id(self, type, data)
    deleted_ids = self.deleted_type_ids
    deleted_ids.add(type_id)
    self.deleted_type_ids = deleted_ids

# Implementation of the EditTypesModifier.restore_type() method.
def _EditTypesModifier_restore_type(self, type: Union[int, str], *, data: Optional[DataCollection] = None) -> None:
    """
    Reverts the modifications this modifier makes to the given element type.

    If the type is currently marked for deletion, the deletion is undone. Otherwise the type's edited attributes
    are discarded, so that it is passed through from the upstream data unchanged again.

    :param type: The numeric :py:attr:`~ovito.data.ElementType.id` or the :py:attr:`~ovito.data.ElementType.name` of the type to restore.
    :param data: The modifier's input data collection. Only needed if the modifier is not part of a pipeline yet
                 or is shared by several pipelines. Otherwise it is determined automatically.
    """
    type_id = _EditTypesModifier_resolve_type_id(self, type, data)
    deleted_ids = self.deleted_type_ids
    if type_id in deleted_ids:
        deleted_ids.discard(type_id)
        self.deleted_type_ids = deleted_ids
    else:
        for index, element_type in enumerate(self.edited_types):
            if element_type.id == type_id:
                del self.edited_types[index]
                break

EditTypesModifier.edit_type = _EditTypesModifier_edit_type
EditTypesModifier.add_type = _EditTypesModifier_add_type
EditTypesModifier.delete_type = _EditTypesModifier_delete_type
EditTypesModifier.restore_type = _EditTypesModifier_restore_type
