# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#

"""Coordinate fitting and alignment --- :mod:`MDAnalysis.analysis.align`
=====================================================================

:Author: Oliver Beckstein, Joshua Adelman
:Year: 2010--2013
:Copyright: GNU Public License v3

The module contains functions to fit a target structure to a reference
structure. They use the fast QCP algorithm to calculate the root mean
square distance (RMSD) between two coordinate sets [Theobald2005]_ and
the rotation matrix *R* that minimizes the RMSD [Liu2010]_. (Please
cite these references when using this module.).

Typically, one selects a group of atoms (such as the C-alphas),
calculates the RMSD and transformation matrix, and applys the
transformation to the current frame of a trajectory to obtain the
rotated structure. The :func:`alignto` and :class:`AlignTraj`
functions can be used to do this for individual frames and
trajectories respectively.

The :ref:`RMS-fitting-tutorial` shows how to do the individual steps
manually and explains the intermediate steps.
"""

import numpy as np
from   .qcprot import CalcRMSDRotationalMatrix

def rotation_matrix(a, b, weights=None):
    r"""Returns the 3x3 rotation matrix `R` for RMSD fitting coordinate
    sets `a` and `b`.

    The rotation matrix `R` transforms vector `a` to overlap with
    vector `b` (i.e., `b` is the reference structure):

    .. math::
       \mathbf{b} = \mathsf{R} \cdot \mathbf{a}

    Parameters
    ----------
    a : array_like
        coordinates that are to be rotated ("mobile set"); array of N atoms
        of shape N*3 as generated by, e.g.,
        :attr:`MDAnalysis.core.groups.AtomGroup.positions`.
    b : array_like
        reference coordinates; array of N atoms of shape N*3 as generated by,
        e.g., :attr:`MDAnalysis.core.groups.AtomGroup.positions`.
    weights : array_like (optional)
        array of floats of size N for doing weighted RMSD fitting (e.g. the
        masses of the atoms)

    Returns
    -------
    R : ndarray
        rotation matrix
    rmsd : float
        RMSD between `a` and `b` before rotation
    ``(R, rmsd)`` rmsd and rotation matrix *R*

    Example
    -------
    `R` can be used as an argument for
    :meth:`MDAnalysis.core.groups.AtomGroup.rotate` to generate a rotated
    selection, e.g. ::

        >>> from MDAnalysisTests.datafiles import TPR, TRR
        >>> from MDAnalysis.analysis import align
        >>> A = mda.Universe(TPR,TRR)
        >>> B = A.copy()
        >>> R = rotation_matrix(A.select_atoms('backbone').positions,
        ...                     B.select_atoms('backbone').positions)[0]
        >>> A.atoms.rotate(R)
        <AtomGroup with 47681 atoms>
        >>> A.atoms.write("rotated.pdb")

    Notes
    -----
    The function does *not* shift the centers of mass or geometry;
    this needs to be done by the user.

    See Also
    --------
    MDAnalysis.analysis.rms.rmsd: Calculates the RMSD between *a* and *b*.
    alignto: A complete fit of two structures.
    AlignTraj: Fit a whole trajectory.
    """


    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)

    if a.shape != b.shape:
        raise ValueError("'a' and 'b' must have same shape")

    if np.allclose(a, b) and weights is None:
        return np.eye(3, dtype=np.float64), 0.0

    N = b.shape[0]

    if weights is not None:
        # qcp does NOT divide weights relative to the mean
        weights = np.asarray(weights, dtype=np.float64) / np.mean(weights)

    rot = np.zeros(9, dtype=np.float64)

    # Need to transpose coordinates such that the coordinate array is
    # 3xN instead of Nx3. Also qcp requires that the dtype be float64
    # (I think we swapped the position of ref and traj in CalcRMSDRotationalMatrix
    # so that R acts **to the left** and can be broadcasted; we're saving
    # one transpose. [orbeckst])
    rmsd = CalcRMSDRotationalMatrix(a, b, N, rot, weights)
    return rot.reshape(3, 3), rmsd


def _fit_to(mobile_coordinates, ref_coordinates, weights=None):
    r"""Perform an rmsd-fitting to determine rotation matrix and align atoms

    Parameters
    ----------
    mobile_coordinates : ndarray
        Coordinates of atoms to be aligned
    ref_coordinates : ndarray
        Coordinates of atoms to be fit against
    weights : array_like (optional)
       choose weights. With ``None`` weigh each atom equally. If a float array
       of the same length as `mobile_coordinates` is provided, use each element
       of the `array_like` as a weight for the corresponding atom in
       `mobile_coordinates`.

    Returns
    -------
    coordinates : ndarray
        Coordinates of translated and rotated atoms
    min_rmsd : float
        Minimum rmsd of coordinates

    Notes
    -----
    This function assumes that `mobile_coordinates` and `ref_coordinates` have
    already been shifted so that their centers of geometry (or centers of mass,
    depending on `weights`) coincide at the origin. `mobile_com` and `ref_com`
    are the centers *before* this shift.

    1. The rotation matrix :math:`\mathsf{R}` is determined with
       :func:`rotation_matrix` directly from `mobile_coordinates` and
       `ref_coordinates`.
    2. `mobile_atoms` :math:`X` is rotated according to the
       rotation matrix and the centers according to

       .. math::

           X' = \mathsf{R}(X - \bar{X}) + \bar{X}_{\text{ref}}

       where :math:`\bar{X}` is the center.

    """
    dr_mobile_coordinates = np.average(mobile_coordinates, axis=0)
    dr_ref_coordinates    = np.average(ref_coordinates, axis=0)

    R, min_rmsd = rotation_matrix(mobile_coordinates - dr_mobile_coordinates, 
                                  ref_coordinates    - dr_ref_coordinates,
                                  weights=weights)
    
    new_mobile_coordinates = \
            dr_ref_coordinates + (R @ (mobile_coordinates - dr_mobile_coordinates).T).T

    return new_mobile_coordinates, min_rmsd

