# -*- coding: utf-8 -*-
"""This module defines :class:`AtomMap` class that allows for pointing atoms in
arbitrary order.

.. _atommaps:

How AtomMap works
===============================================================================

:class:`AtomMap` class adds great flexibility to manipulating atomic data.

First let's see how an instance of :class:`.Selection` (:class:`.Chain`, or
:class:`.Residue`) works.  Below table shows indices for a selection of atoms
in an :class:`~.AtomGroup` and values returned when
:meth:`~.Selection.getNames`, :meth:`~.Selection.getResnames` and
:meth:`~.Selection.getResnums` methods are called.

.. csv-table:: **Atom Subset**
   :header: "Indices", "Names", "Resnames", "Resnums"

   0, N, PHE, 1
   1, CA, PHE, 1
   2, C, PHE, 1
   3, O, PHE, 1
   4, CB, PHE, 1
   5, CG, PHE, 1
   6, CD1, PHE, 1
   7, CD2, PHE, 1
   8, CE1, PHE, 1
   9, CE2, PHE, 1
   10, CZ, PHE, 1

:class:`~.Selection` instances keep indices ordered and do not allow duplicate
values, hence their use is limited. In an :class:`AtomMap`, indices do not need
to be sorted, duplicate indices may exist, even "DUMMY" atoms are allowed.

Let's say we instantiate the following AtomMap::

    amap = AtomMap(atomgroup, indices=[0, 1, 3, 8, 8, 9, 10],
                   mapping=[5, 6, 7, 0, 1, 2, 3])


The size of the :class:`AtomMap` based on this mapping is 8, since the larger
mapping is 7.

Calling the same functions for this AtomMap instance would result in the
following:

.. csv-table:: **Atom Map**
   :header: "Mapping", "Indices", "Names", "Resnames", "Resnums", \
            "MappedFlags", "DummyFlags"

   0, 8, CE1, PHE, 1, 1, 0
   1, 8, CE1, PHE, 1, 1, 0
   2, 9, CE2, PHE, 1, 1, 0
   3, 10, CZ, PHE, 1, 1, 0
   4, , , , 0, 0, 1
   5, 0, N, PHE, 1, 1, 0
   6, 1, CA, PHE, 1, 1, 0
   7, 3, O, PHE, 1, 1, 0

For unmapped atoms, numeric attributes are set to 0, others to empty string,
i.e. ``""``.

.. seealso::
   :class:`AtomMap` are used by :mod:`.proteins` module functions that
   match or map protein chains.  :ref:`pca-xray` and :ref:`pca-dimer`
   examples that make use of these functions and :class:`AtomMap` class.


"""

from numpy import arange, array, ones, zeros, dtype

from prody.utilities import rangeString, copy

from .atom import Atom
from .fields import ATOMIC_FIELDS
from .fields import wrapGetMethod
from .pointer import AtomPointer

__all__ = ['AtomMap']

DUMMY = -1

class AtomMap(AtomPointer):

    """A class for mapping atomic data."""

    __slots__ = ['_ag', '_indices', '_acsi', '_mapping', '_dummies', '_title',
                 '_len', '_idarray']

    def __init__(self, ag, indices, acsi=None, **kwargs):
        """Instantiate an atom map.

        :arg ag: AtomGroup instance from which atoms are mapped
        :arg indices: indices of mapped atoms
        :arg acsi: active coordinate set index, defaults is that of *ag*
        :arg mapping: mapping of atom *indices*
        :arg dummies: dummy atom indices
        :arg title: title of the instance, default is 'Unknown'

        *mapping* and *dummies* arrays must be provided together.  Length of
        *mapping* must be equal to length of *indices*.  Elements of *mapping*
        must be an ordered in ascending order.  When dummy atoms are present,
        number of atoms is the sum of lengths of *mapping* and *dummies*.

        Following built-in functions are customized for this class:

        * :func:`len` returns the number of atoms in the instance.
        * :func:`iter` yields :class:`.Atom` instances.
        * Indexing returns an :class:`.Atom` or an :class:`.AtomMap` instance
          depending on the type and value of the index."""

        AtomPointer.__init__(self, ag, acsi)

        self._dummies = self._mapping = None
        mapping = kwargs.get('mapping', None)
        dummies = kwargs.get('dummies', False)

        try:
            len(dummies)
        except TypeError:
            dummy_array = False
        else:
            dummy_array = True

        if mapping is None:
            if not kwargs.get('intarrays'):
                indices = array(indices, int)
            self._len = len(indices)

            if dummy_array:
                raise TypeError('mapping and dummies must be provided '
                                'together')

            if dummies:
                dummies = (indices == DUMMY).nonzero()[0]
                if len(dummies):
                    self._dummies = dummies
                    self._mapping = (indices > DUMMY).nonzero()[0]
                    self._indices = indices[self._mapping]
                    self._idarray = indices
                else:
                    self._indices = self._idarray = indices
            else:
                self._indices = self._idarray = indices
        else:
            if dummies is None:
                raise TypeError('mapping and dummies must be provided '
                                'together')

            if len(mapping) != len(indices):
                raise ValueError('indices and mapping arrays must have the '
                                 'same length')
            if not kwargs.get('intarrays'):
                indices = array(indices, int)
                mapping = array(mapping, int)
                if dummy_array:
                    dummies = array(dummies, int)
            #if any(mapping[1:] - mapping[:-1] < 0):
            #    raise ValueError('mapping must be an ordered array')
            self._len = len(indices)
            if dummy_array:
                self._indices = indices
                self._mapping = mapping
                self._dummies = dummies
                self._len += len(dummies)
                self._idarray = idarray = zeros(self._len, int)
                idarray[self._mapping] = self._indices
                idarray[self._dummies] = DUMMY
            else:
                self._indices = self._idarray = indices[mapping]
        self._title = str(kwargs.get('title', 'Unknown'))

    def __repr__(self):

        rep = '<AtomMap: {0} from {1} ({2} atoms'.format(
                self._title, self._ag.getTitle(), self._len)
        if self.numDummies():
            rep += ', {0} mapped, {1} dummy'.format(self.numMapped(),
                                                        self.numDummies())

        n_csets = self._ag.numCoordsets()
        if n_csets > 1:
            rep += '; active #{0} of {1} coordsets)>'.format(
                    self.getACSIndex(), n_csets)
        elif n_csets == 0:
            rep += '; no coordinates'
        return rep + ')>'

    def __str__(self):

        return 'AtomMap {0}'.format(self._title)

    def __len__(self):

        return self._len

    def __getitem__(self, index):

        indices = self._idarray[index]
        try:
            len(indices)
        except TypeError:
            if indices != DUMMY:
                return self._ag[indices]
        else:
            return AtomMap(self._ag, indices, self._acsi,
                       title='({0})[{1}]'.format(self._title, repr(index)),
                       intarrays=True, dummies=self.numDummies())

    def getTitle(self):
        """Returns title of the instance."""

        return self._title

    def setTitle(self, title):
        """Set title of the instance."""

        self._title = str(title)

    def numAtoms(self, flag=None):
        """Returns number of atoms."""

        return len(self._getSubset(flag)) if flag else self._len

    def iterAtoms(self):
        """Yield atoms, and **None** for dummies."""

        ag = self._ag
        acsi = self.getACSIndex()
        for index in self.getIndices():
            yield Atom(ag, index, acsi) if index > DUMMY else None

    __iter__ = iterAtoms

    def getCoords(self):
        """Returns a copy of coordinates from the active coordinate set."""

        coords = self._ag._getCoordsets()
        if coords is not None:
            if self._mapping is None:
                xyz = coords[self.getACSIndex(), self._indices]
            else:
                xyz = zeros((self._len, 3), float)
                xyz[self._mapping] = coords[self.getACSIndex(), self._indices]
            return xyz

    _getCoords = getCoords

    def setCoords(self, coords):
        """Set coordinates of atoms in the active coordinate set."""

        coordsets = self._ag._getCoordsets()
        if coordsets is not None:
            if self._mapping is None:
                coordsets[self.getACSIndex(), self._indices] = coords
            elif self._dummies is None:
                coordsets[self.getACSIndex(), self._indices] = coords


    def getCoordsets(self, indices=None):
        """Returns coordinate set(s) at given *indices*, which may be an integer
        or a list/array of integers."""

        coords = self._ag._getCoordsets()
        if coords is not None:
            if indices is None:
                coords = coords[:, self._indices]
            else:
                try:
                    len(indices)
                except TypeError:
                    coords = coords[indices, self._indices]
                else:
                    coords = coords[indices][:, self._indices]

            if self._mapping is None:
                return coords
            else:
                csets = zeros(coords.shape[:-2] + (self._len, 3))
                csets[:, self._mapping] = coords
                return csets

    _getCoordsets = getCoordsets

    def iterCoordsets(self):
        """Yield copies of coordinate sets."""

        coords = self._ag._getCoordsets()
        if coords is not None:
            mapping = self._mapping
            n_atoms = self._len
            indices = self._indices
            for i in range(self._ag.numCoordsets()):
                xyz = zeros((n_atoms, 3), float)
                xyz[mapping] = coords[i, indices]
                yield xyz

    _iterCoordsets = iterCoordsets

    def getData(self, label):
        """Returns a copy of data associated with *label*, if it is present."""

        data = self._ag._getData(label)
        if data is not None:
            result = zeros((self._len,) + data.shape[1:], data.dtype)
            result[self._mapping] = data[self._indices]
            return result

    _getData = getData

    def getFlags(self, label):
        """Returns a copy of atom flags for given *label*, or **None** when
        flags for *label* is not set."""

        if label == 'dummy':
            flags = zeros(self._len, bool)
            if self._dummies is not None:
                flags[self._dummies] = True
        elif label == 'mapped':
            flags = ones(self._len, bool)
            if self._dummies is not None:
                flags[self._dummies] = False
        else:
            flags = None
            agflags = self._ag._getFlags(label)
            if agflags is not None:
                flags = zeros(self._len, bool)
                flags[self._mapping] = agflags[self._indices]
        return flags

    _getFlags = getFlags

    def _getSubset(self, label):

        return self._idarray[self._getFlags(label)]

    def getIndices(self):
        """Returns a copy of indices of atoms, with maximum integer value
        dummies."""

        return copy(self._idarray)

    def _getIndices(self):
        """Returns indices of atoms, with maximum integer value dummies."""

        return self._idarray

    def getMapping(self):
        """Returns a copy of mapping of indices."""

        mapping = self._mapping
        return arange(self._len) if mapping is None else mapping.copy()

    def _getMapping(self):
        """Returns mapping of indices."""

        mapping = self._mapping
        return arange(self._len) if mapping is None else mapping

    def numMapped(self):
        """Returns number of mapped atoms."""

        return len(self._indices)

    def numDummies(self):
        """Returns number of dummy atoms."""

        return 0 if self._dummies is None else len(self._dummies)

    def getSelstr(self):
        """Returns selection string that selects mapped atoms."""

        return 'index ' + rangeString(self._indices)

    def getHierView(self, **kwargs):
        """Returns a hierarchical view of the this chain."""

        return HierView(self, **kwargs)


for fname, field in ATOMIC_FIELDS.items():

    if field.private:
        continue

    meth = field.meth_pl
    getMeth = 'get' + meth

    def getData(self, meth=field.meth_pl, dtype=field.dtype):
        data = getattr(self._ag, '_get' + meth)()
        if data is not None:
            if self._mapping is None:
                return data[self._indices]
            else:
                result = zeros((self._len,) + data.shape[1:], dtype)
                result[self._mapping] = data[self._indices]
                return result

    getData = wrapGetMethod(getData)
    getData.__name__ = getMeth
    getData.__doc__ = (field.getDocstr('get', selex=False) +
                       ' Entries for dummy atoms will be ``{0}``.'
                       .format(repr(dtype(field.dtype).type())))
    setattr(AtomMap, getMeth, getData)
    setattr(AtomMap, '_' + getMeth, getData)

del getData
