from __future__ import annotations
from . import DataCollection, DataObject, AttributeDataObject, TriangleMesh
from ._data_objects_dict import DataObjectsDict
from ..modifiers import PythonModifier
from ..pipeline import StaticSource, Modifier
from ..nonpublic import PipelineStatus
import ovito
from typing import Optional, Any
import collections.abc

# Helper class used to implement the DataCollection.attributes field.
class _AttributesView(collections.abc.MutableMapping):

    def __init__(self, data_collection):
        """ Constructor that stores away a back-pointer to the owning DataCollection instance. """
        self._collection = data_collection

    def __len__(self):
        count = 0
        for obj in self._collection.objects:
            if isinstance(obj, AttributeDataObject):
                count += 1
        return count

    def __getitem__(self, key):
        if not isinstance(key, str):
            raise TypeError("Attribute key must be a string")
        for obj in self._collection.objects:
            if isinstance(obj, AttributeDataObject) and obj.identifier == key:
                return obj.value
        raise KeyError(f"Attribute '{key}' does not exist in data collection. Available attributes are: {list(self)}")

    def __setitem__(self, key, value):
        if not isinstance(key, str):
            raise TypeError("Attribute key must be a string")
        for obj in self._collection.objects:
            if isinstance(obj, AttributeDataObject) and obj.identifier == key:
                if not value is None:
                    self._collection.make_mutable(obj).value = value
                else:
                    self._collection.objects.remove(obj)
                return
        if not value is None:
            self._collection.objects.append(AttributeDataObject(identifier = key, value = value))

    def __delitem__(self, key):
        """ Removes a global attribute from the data collection. """
        if not isinstance(key, str):
            raise TypeError("Attribute key must be a string")
        for obj in self._collection.objects:
            if isinstance(obj, AttributeDataObject) and obj.identifier == key:
                self._collection.objects.remove(obj)
                return
        raise KeyError(f"Attribute '{key}' does not exist in data collection.")

    def __iter__(self):
        """ Returns an iterator over the names of all global attributes. """
        for obj in self._collection.objects:
            if isinstance(obj, AttributeDataObject):
                yield obj.identifier

    def __repr__(self):
        return repr(dict(self))

# Implementation of the DataCollection.attributes field.
def _DataCollection_attributes(self) -> collections.abc.MutableMapping[str, Any]:
    """
    This field contains a dictionary view with all the *global attributes* currently associated with this data collection.
    Global attributes are key-value pairs that represent small tokens of information, typically simple value types such as ``int``, ``float`` or ``str``.
    Every attribute has a unique identifier such as ``'Timestep'`` or ``'ConstructSurfaceMesh.surface_area'``. This identifier serves as lookup key in the :py:attr:`!attributes` dictionary.
    Identifiers starting with ``'.'`` are hidden in the GUI. Attributes are dynamically generated by modifiers in a data pipeline or come from the data source.
    For example, if the input simulation file contains timestep information, the timestep number is made available by the :py:attr:`~ovito.pipeline.FileSource` as the
    ``'Timestep'`` attribute. It can be retrieved from pipeline's output data collection:

        >>> pipeline = import_file('snapshot_140000.dump')
        >>> pipeline.compute().attributes['Timestep']
        140000

    Some modifiers report their calculation results by adding new attributes to the data collection. See each modifier's
    reference documentation for the list of attributes it generates. For example, the number of clusters identified by the
    :py:class:`~ovito.modifiers.ClusterAnalysisModifier` is available in the pipeline output as an attribute named
    ``ClusterAnalysis.cluster_count``::

        pipeline.modifiers.append(ClusterAnalysisModifier(cutoff = 3.1))
        data = pipeline.compute()
        nclusters = data.attributes["ClusterAnalysis.cluster_count"]

    The :py:func:`ovito.io.export_file` function can be used to output dynamically computed attributes to a text file, possibly as functions of time::

        export_file(pipeline, "data.txt", "txt/attr",
            columns = ["Timestep", "ClusterAnalysis.cluster_count"],
            multiple_frames = True)

    If you are writing your own :ref:`modifier function <writing_custom_modifiers>`, you let it add new attributes to a data collection.
    In the following example, the :py:class:`~ovito.modifiers.CommonNeighborAnalysisModifier` first inserted into the
    pipeline generates the ``'CommonNeighborAnalysis.counts.FCC'`` attribute to report the number of atoms that
    have an FCC-like coordination. To compute an atomic *fraction* from that, we need to divide the count by the total number of
    atoms in the system. To this end, we append a user-defined modifier function
    to the pipeline, which computes the fraction and outputs the value as a new attribute named ``'fcc_fraction'``.

    .. literalinclude:: ../example_snippets/python_modifier_generate_attribute.py
        :lines: 6-

    """
    return _AttributesView(self)
DataCollection.attributes = property(_DataCollection_attributes)

# Implementation of the DataCollection.triangle_meshes attribute.
def _DataCollection_triangle_meshes(self) -> collections.abc.Mapping[str, TriangleMesh]:
    """
    This is a dictionary view providing key-based access to all :py:class:`TriangleMesh` objects currently stored in
    this data collection. Each :py:class:`TriangleMesh` has a unique :py:attr:`~ovito.data.DataObject.identifier` key,
    which can be used to look it up in the dictionary.
    """
    return DataObjectsDict(self, TriangleMesh)
DataCollection.triangle_meshes = property(_DataCollection_triangle_meshes)

# Implementation of the DataCollection.triangle_meshes_ attribute.
def _DataCollection_triangle_meshes_mutable(self) -> collections.abc.Mapping[str, TriangleMesh]:
    return DataObjectsDict(self, TriangleMesh, always_mutable=True)
DataCollection.triangle_meshes_ = property(_DataCollection_triangle_meshes_mutable)

# Implementation of the DataCollection.apply() method:
def __DataCollection_apply(self, modifier: Modifier, frame: Optional[int] = None):
    """ This method applies a :py:class:`~ovito.pipeline.Modifier` function to the data stored in this collection to modify it in place.

        :param modifier: The :py:class:`~ovito.pipeline.Modifier` instance that should alter the contents of this data collection in place.
        :param frame: Optional animation frame number to be passed to the modifier function, which may use it for time-dependent modifications.

        The method allows modifying a data collection with one of OVITO's modifiers directly without the need to build up a complete
        :py:class:`~ovito.pipeline.Pipeline` first. In contrast to a :ref:`data pipeline <modifiers_overview>`, the :py:meth:`!apply()` method
        executes the modifier function immediately and alters the data in place. In other words, the original data in this :py:class:`DataCollection`
        gets replaced by the output produced by the invoked modifier function. It is possible to first create a copy of
        the original data using the :py:meth:`.clone` method if needed. The following code example
        demonstrates how to use :py:meth:`!apply()` to successively modify a dataset:

        .. literalinclude:: ../example_snippets/data_collection_apply.py
            :lines: 4-10

        Note that it is typically possible to achieve the same result by first populating a :py:class:`~ovito.pipeline.Pipeline` with the modifiers and then calling its
        :py:meth:`~ovito.pipeline.Pipeline.compute` method at the very end:

        .. literalinclude:: ../example_snippets/data_collection_apply.py
            :lines: 15-19

        An important use case of the :py:meth:`!apply()` method is in the implementation of a :ref:`user-defined modifier function <writing_custom_modifiers>`,
        making it possible to invoke other modifiers as sub-routines:

        .. literalinclude:: ../example_snippets/data_collection_apply.py
            :lines: 24-34
    """

    # This method expects a Modifier object as argument.
    if not isinstance(modifier, Modifier):
        if isinstance(modifier, ovito.pipeline.ModifierInterface):
            # Automatically wrap ModifierInterface objects in a PythonModifier object.
            modifier = PythonModifier(delegate=modifier)
        elif callable(modifier):
            # Automatically wrap freestanding Python methods in a PythonModifier object.
            modifier = PythonModifier(function=modifier)
        else:
            raise TypeError("Expected a modifier as argument")

    # Perform an early check that this DataCollection is modifiable and is not shared with other owners.
    # This is required so we can safely replace the contents of this DataCollection with the output of the modifier.
    if not self.is_safe_to_modify:
        raise RuntimeError("This DataCollection has shared ownership and cannot be modified in place. You may have to create a copy first using the clone() method.")

    # Build an ad-hoc pipeline by creating a ModificationNode for the Modifier,
    # which receives the input DataCollection from a StaticSource.
    node = modifier.create_modification_node()
    # Note: We pass a new copy of input data collection to the pipeline, because we cannot be sure that the
    # pipeline system releases all references in time before we modify the data collection in place below.
    node.input = StaticSource(data=self.clone())
    node.modifier = modifier

    # Initialize the modifier within the pipeline, e.g. parameters that depend on the modifier's input.
    modifier.initialize_modifier(node, frame)

    # Evaluate the ad-hoc pipeline.
    state = node._evaluate(frame)
    if state.status.type == PipelineStatus.Type.Error:
        raise RuntimeError(f"Modifier evaluation failed: {state.status.text}")

    # The DataCollection.apply() method is supposed to modify the DataCollection in place.
    # To implement this behavior, move the data objects from the pipeline output collection
    # over to this collection, replacing the original state.
    self._assign_objects(state.data)
DataCollection.apply = __DataCollection_apply

# Implementation of the DataCollection.clone() method:
def __DataCollection_clone(self) -> DataCollection:
    """
        Returns a shallow copy of this :py:class:`DataCollection` containing the same data objects as the original.

        The method can be used to retain a copy of the original data before modifying a data collection in place,
        for example, using the :py:meth:`.apply` method:

        .. literalinclude:: ../example_snippets/data_collection_clone.py
            :lines: 8-12

        Note that the :py:meth:`!clone()` method performs an inexpensive shallow copy, meaning that the newly created collection
        still shares the data objects with the original collection.

        Keep in mind that data objects shared by two or more data collections are implicitly protected against modifications
        to avoid unexpected side effects. Thus, in order to subsequently modify the objects in either the original or the
        copy of the data collection, you have to use the underscore notation or the :py:meth:`DataObject.make_mutable` method
        to make a deep copy of the particular data object(s) you want to modify. For example:

        .. literalinclude:: ../example_snippets/data_collection_clone.py
            :lines: 17-28

        .. tip::

           The :py:meth:`!clone()` method is equivalent to the standard Python function :py:func:`python:copy.copy` when
           applied to a :py:class:`DataCollection`. In fact, most OVITO object types can be shallow-copied with Python's
           :py:func:`python:copy.copy` function and deep-copied with the :py:func:`python:copy.deepcopy` function.
    """
    return self.__copy__()
DataCollection.clone = __DataCollection_clone
