# -*- coding: utf-8 -*-
#
# ProDy: A Python Package for Protein Dynamics Analysis
#
# Copyright (C) 2010-2014 University of Pittsburgh
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#
"""This module defines a class handling normal mode analysis data."""
from past.utils import old_div
import numpy as np
from .mode import Mode
from .modeset import ModeSet
__all__ = ['NMA']
[docs]class NMA(object):
    """A class for handling Normal Mode Analysis (NMA) data."""
[docs]    def __init__(self, title='Unknown'):
        self._title = str(title).strip()
        self._n_modes = 0
        self._cov = None
        self._n_atoms = 0
        self._dof = 0
        self._array = None  # modes/eigenvectors
        self._eigvals = None
        self._vars = None  # evs for PCA, inverse evs for ENM
        self._trace = None
        self._is3d = True  # is set to false for GNM
        self._indices = None 
[docs]    def __len__(self):
        return self._n_modes 
    def __getitem__(self, index):
        """A list or tuple of integers can be used for indexing."""
        if self._n_modes == 0:
            raise ValueError('{0} modes are not calculated, use '
                             'calcModes() method'.format(str(self)))
        if isinstance(index, slice):
            indices = np.arange(*index.indices(len(self)))
            if len(indices) > 0:
                return ModeSet(self, indices)
        elif isinstance(index, (list, tuple)):
            for i in index:
                assert isinstance(i, int), 'all indices must be integers'
            if len(index) == 1:
                return self._getMode(index[0])
            return ModeSet(self, index)
        try:
            index = int(index)
        except Exception:
            raise IndexError('indices must be int, slice, list, or tuple')
        else:
            return self._getMode(index)
    def __iter__(self):
        for i in range(self._n_modes):
            yield self[i]
    def __repr__(self):
        return '<{0}: {1} ({2} modes; {3} atoms)>'.format(
            self.__class__.__name__, self._title, self._n_modes, self._n_atoms)
    def __str__(self):
        return self.__class__.__name__ + ' ' + self._title
    def _reset(self):
        self._n_modes = 0
        self._cov = None
        self._n_atoms = 0
        self._dof = 0
        self._array = None
        self._eigvals = None
        self._vars = None
        self._trace = None
        self._is3d = True
    def _getMode(self, index):
        if self._n_modes == 0:
            raise ValueError('{0} modes are not calculated, use '
                             'calcModes() method'.format(str(self)))
        if index >= self._n_modes or index < -self._n_modes:
            raise IndexError('{0} contains {1} modes, try 0 <= index < '
                             '{1}'.format(str(self), self._n_modes))
        if index < 0:
            index += self._n_modes
        return Mode(self, index)
[docs]    def getTrace(self):
        """Return trace, and emit a warning message if trace is calculated
        using eigenvalues of a subset of variances (eigenvalues or inverse
        eigenvalues)."""
        trace = self._trace
        if trace is None:
            if self._vars is None:
                raise ValueError('variances are not set or calculated')
            trace = self._vars.sum()
            diff = self._dof - self._n_modes
            if self._is3d and diff > 6:
                diff = True
            elif diff > 1:
                diff = True
            else:
                diff = False
            if diff:
                from prody import LOGGER
                LOGGER.warn('Total variance for {0} is calculated using '
                            '{1} available modes out of {2} possible.'.format(
                                str(self), self._n_modes, self._dof))
        return trace 
[docs]    def getModel(self):
        """Return self."""
        return self 
[docs]    def is3d(self):
        """Return **True** if model is 3-dimensional."""
        return self._is3d 
[docs]    def numAtoms(self):
        """Return number of atoms."""
        return self._n_atoms 
[docs]    def numModes(self):
        """Return number of modes in the instance (not necessarily maximum
        number of possible modes)."""
        return self._n_modes 
[docs]    def numDOF(self):
        """Return number of degrees of freedom."""
        return self._dof 
[docs]    def getTitle(self):
        """Return title of the model."""
        return self._title 
[docs]    def setTitle(self, title):
        """Set title of the model."""
        self._title = str(title) 
[docs]    def getEigvals(self):
        """Return eigenvalues.  For :class: `.PCA` and :class: `.EDA` models
        built using coordinate data in Å, unit of eigenvalues is `|A2|`.  For
        :class: `.ANM`, :class: `.GNM`, and :class: `.RTB`, on the other hand,
        eigenvalues are in arbitrary or relative units but they correlate with
        stiffness of the motion along associated eigenvector."""
        if self._eigvals is None:
            return None
        return self._eigvals.copy() 
[docs]    def getVariances(self):
        """Return variances.  For :class: `.PCA` and :class: `.EDA` models
        built using coordinate data in Å, unit of variance is `|A2|`.  For
        :class: `.ANM`, :class: `.GNM`, and :class: `.RTB`, on the other hand,
        variance is the inverse of the eigenvalue, so it has arbitrary or
        relative units."""
        if self._vars is None:
            return None
        return self._vars.copy() 
[docs]    def getArray(self):
        """Return a copy of eigenvectors array."""
        if self._array is None:
            return None
        return self._array.copy() 
    getEigvecs = getArray
    def _getArray(self):
        """Return eigenvectors array."""
        if self._array is None:
            return None
        return self._array
[docs]    def getCovariance(self):
        """Return covariance matrix.  If covariance matrix is not set or yet
        calculated, it will be calculated using available modes."""
        if self._cov is None:
            if self._array is None:
                return None
            self._cov = np.dot(self._array,
                               np.dot(np.diag(self._vars), self._array.T))
        return self._cov 
[docs]    def calcModes(self):
        """""" 
[docs]    def addEigenpair(self, vector, value=None):
        """Add eigen *vector* and eigen *value* pair(s) to the instance.
        If eigen *value* is omitted, it will be set to 1.  Inverse
        eigenvalues are set as variances."""
        if self._array is None:
            self.setEigens()
        try:
            ndim, shape = vector.ndim, vector.shape
        except AttributeError:
            raise TypeError('vector must be a numpy array')
        if ndim > 2:
            raise ValueError('vector must be 1- or 2-dimensional')
        elif ndim == 2 and shape[0] < shape[1]:
            raise ValueError('eigenvectors must correspond to vector columns')
        else:
            vector = vector.reshape((shape[0], 1))
        if self.eigvals is None:
            if ndim == 1:
                self.eigvals = np.ones(1)
            else:
                self.eigvals = np.ones(shape[2])
        else:
            try:
                ndim2, shape2 = self.eigvals.ndim, self.eigvals.shape
            except AttributeError:
                try:
                    self.eigvals = np.array([value], float)
                except Exception:
                    raise TypeError('value must be a number or array')
            else:
                if value.ndim > 1:
                    raise ValueError('value must be a 1-dimensional array')
                elif value.shape[0] != value.shape[0]:
                    raise ValueError('number of eigenvectors and eigenvalues '
                                     'must match')
        if vector.shape[0] != self._array.shape[0]:
            raise ValueError('shape of vector do not match shape of '
                             'existing eigenvectors')
        self._array = np.concatenate((self._array, vector), 1)
        self._eigvals = np.concatenate((self._eigvals, value))
        self._n_modes += shape[1]
        self._vars = old_div(1, self._eigvals) 
[docs]    def setEigens(self, vectors, values=None):
        """Set eigen *vectors* and eigen *values*.  If eigen *values* are
        omitted, they will be set to 1.  Inverse eigenvalues are set as
        variances."""
        try:
            ndim, shape = vectors.ndim, vectors.shape
        except AttributeError:
            raise TypeError('vectors must be a numpy array')
        if ndim != 2:
            raise ValueError('vectors must be a 2-dimensional array')
        else:
            dof = shape[0]
            if self._is3d:
                n_atoms = old_div(dof, 3)
            else:
                n_atoms = dof
            if self._n_atoms > 0 and n_atoms != self._n_atoms:
                raise ValueError('vectors do not have the right shape, '
                                 'which is (M,{0})'.format(n_atoms * 3))
            n_modes = vectors.shape[1]
        if values is not None:
            if not isinstance(vectors, np.ndarray):
                raise TypeError(
                    'values must be a numpy.ndarray, not {0}'.format(
                        type(vectors)))
            elif values.ndim != 1:
                raise ValueError('values must be a 1-dimensional array')
            else:
                if values.shape[0] != vectors.shape[1]:
                    raise ValueError(
                        'number of vectors and values do not match')
        else:
            values = np.ones(n_modes)
        self._array = vectors
        self._eigvals = values
        self._dof = dof
        self._n_atoms = n_atoms
        self._n_modes = n_modes
        self._vars = old_div(1, values)