# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# 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
#
import os

import MDAnalysis
import MDAnalysis as mda
from MDAnalysis.analysis import rms, align

from numpy.testing import assert_equal, assert_almost_equal

import numpy as np
import pytest

from MDAnalysis.exceptions import SelectionError, NoDataError
from MDAnalysisTests.datafiles import GRO, XTC, rmsfArray, PSF, DCD


class Testrmsd(object):
    shape = (5, 3)
    # vectors with length one
    ones = np.ones(shape) / np.sqrt(3)

    @pytest.fixture()
    def a(self):
        return self.ones * np.arange(1, 6)[:, np.newaxis]

    @pytest.fixture()
    def b(self, a):
        return a + self.ones

    @pytest.fixture()
    def u(self):
        u = mda.Universe(PSF, DCD)
        return u

    @pytest.fixture()
    def u2(self):
        u = mda.Universe(PSF, DCD)
        return u

    @pytest.fixture()
    def p_first(self, u):
        u.trajectory[0]
        return u.select_atoms('protein')

    @pytest.fixture()
    def p_last(self, u2):
        u2.trajectory[-1]
        return u2.select_atoms('protein')

    # internal test
    def test_p_frames(self, p_first, p_last):
        # check that these fixtures are really different
        assert p_first.universe.trajectory.ts.frame != p_last.universe.trajectory.ts.frame
        assert not np.allclose(p_first.positions, p_last.positions)

    def test_no_center(self, a, b):
        rmsd = rms.rmsd(a, b, center=False)
        assert_almost_equal(rmsd, 1.0)

    def test_center(self, a, b):
        rmsd = rms.rmsd(a, b, center=True)
        assert_almost_equal(rmsd, 0.0)

    def test_list(self, a, b):
        rmsd = rms.rmsd(a.tolist(),
                        b.tolist(),
                        center=False)
        assert_almost_equal(rmsd, 1.0)

    def test_superposition(self, a, b, u):
        bb = u.atoms.select_atoms('backbone')
        a = bb.positions.copy()
        u.trajectory[-1]
        b = bb.positions.copy()
        rmsd = rms.rmsd(a, b, superposition=True)
        assert_almost_equal(rmsd, 6.820321761927005)

    def test_weights(self, a, b):
        weights = np.zeros(len(a))
        weights[0] = 1
        weights[1] = 1
        weighted = rms.rmsd(a, b, weights=weights)
        firstCoords = rms.rmsd(a[:2], b[:2])
        assert_almost_equal(weighted, firstCoords)

    def test_weights_and_superposition_1(self, u):
        weights = np.ones(len(u.trajectory[0]))
        weighted = rms.rmsd(u.trajectory[0], u.trajectory[1],
                            weights=weights, superposition=True)
        firstCoords = rms.rmsd(u.trajectory[0], u.trajectory[1],
                               superposition=True)
        assert_almost_equal(weighted, firstCoords, decimal=5)

    def test_weights_and_superposition_2(self, u):
        weights = np.zeros(len(u.trajectory[0]))
        weights[:100] = 1
        weighted = rms.rmsd(u.trajectory[0], u.trajectory[-1],
                            weights=weights, superposition=True)
        firstCoords = rms.rmsd(u.trajectory[0][:100],
                               u.trajectory[-1][:100],
                               superposition=True)
        # very close to zero, change significant decimal places to 5
        assert_almost_equal(weighted, firstCoords, decimal=5)

    def test_unequal_shape(self):
        a = np.ones((4, 3))
        b = np.ones((5, 3))
        with pytest.raises(ValueError):
            rms.rmsd(a, b)

    def test_wrong_weights(self, a, b):
        w = np.ones(2)
        with pytest.raises(ValueError):
            rms.rmsd(a, b, w)

    def test_with_superposition_smaller(self, p_first, p_last):
        A = p_first.positions
        B = p_last.positions
        rmsd = rms.rmsd(A, B)
        rmsd_superposition = rms.rmsd(A, B, center=True, superposition=True)
        # by design the super positioned rmsd is smaller
        assert rmsd > rmsd_superposition

    def test_with_superposition_equal(self, p_first, p_last):
        align.alignto(p_first, p_last)
        A = p_first.positions
        B = p_last.positions
        rmsd = rms.rmsd(A, B)
        rmsd_superposition = rms.rmsd(A, B, center=True, superposition=True)
        assert_almost_equal(rmsd, rmsd_superposition, decimal=6)


class TestRMSD(object):
    @pytest.fixture()
    def universe(self):
        return MDAnalysis.Universe(PSF, DCD)

    @pytest.fixture()
    def outfile(self, tmpdir):
        return os.path.join(str(tmpdir), 'rmsd.txt')

    @pytest.fixture()
    def correct_values(self):
        return [[0, 1, 0], [49, 50, 4.68953]]

    @pytest.fixture()
    def correct_values_mass(self):
        return [[0, 1, 0], [49, 50, 4.74920]]

    @pytest.fixture()
    def correct_values_mass_add_ten(self):
        return [[0, 1, 0.0632], [49, 50, 4.7710]]

    @pytest.fixture()
    def correct_values_group(self):
        return [[0, 1, 0, 0, 0],
                [49, 50, 4.7857, 4.7048, 4.6924]]

    @pytest.fixture()
    def correct_values_backbone_group(self):
        return [[0, 1, 0, 0, 0],
                [49, 50, 4.6997, 1.9154, 2.7139]]

    def test_rmsd(self, universe, correct_values):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, select='name CA')
        RMSD.run(step=49)
        assert_almost_equal(RMSD.results.rmsd, correct_values, 4,
                            err_msg="error: rmsd profile should match" +
                            "test values")

    def test_rmsd_frames(self, universe, correct_values):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, select='name CA')
        RMSD.run(frames=[0, 49])
        assert_almost_equal(RMSD.results.rmsd, correct_values, 4,
                            err_msg="error: rmsd profile should match" +
                            "test values")

    def test_rmsd_unicode_selection(self, universe, correct_values):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, select=u'name CA')
        RMSD.run(step=49)
        assert_almost_equal(RMSD.results.rmsd, correct_values, 4,
                            err_msg="error: rmsd profile should match" +
                            "test values")

    def test_rmsd_atomgroup_selections(self, universe):
        # see Issue #1684
        R1 = MDAnalysis.analysis.rms.RMSD(universe.atoms,
                                          select="resid 1-30").run()
        R2 = MDAnalysis.analysis.rms.RMSD(universe.atoms.select_atoms("name CA"),
                                          select="resid 1-30").run()
        assert not np.allclose(R1.results.rmsd[:, 2], R2.results.rmsd[:, 2])

    def test_rmsd_single_frame(self, universe):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, select='name CA',
                                            ).run(start=5, stop=6)
        single_frame = [[5, 6, 0.91544906]]
        assert_almost_equal(RMSD.results.rmsd, single_frame, 4,
                            err_msg="error: rmsd profile should match" +
                            "test values")

    def test_mass_weighted(self, universe, correct_values):
        # mass weighting the CA should give the same answer as weighing
        # equally because all CA have the same mass
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, select='name CA',
                                            weights='mass').run(step=49)

        assert_almost_equal(RMSD.results.rmsd, correct_values, 4,
                            err_msg="error: rmsd profile should match"
                            "test values")

    def test_custom_weighted(self, universe, correct_values_mass):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe, weights="mass").run(step=49)

        assert_almost_equal(RMSD.results.rmsd, correct_values_mass, 4,
                            err_msg="error: rmsd profile should match"
                            "test values")

    def test_weights_mass_is_mass_weighted(self, universe):
        RMSD_mass = MDAnalysis.analysis.rms.RMSD(universe,
                                                 weights="mass").run(step=49)
        RMSD_cust = MDAnalysis.analysis.rms.RMSD(universe,
                                                 weights=universe.atoms.masses).run(step=49)
        assert_almost_equal(RMSD_mass.results.rmsd, RMSD_cust.results.rmsd, 4,
                            err_msg="error: rmsd profiles should match for 'mass' "
                            "and universe.atoms.masses")

    def test_custom_weighted_list(self, universe, correct_values_mass):
        weights = universe.atoms.masses
        RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                            weights=list(weights)).run(step=49)
        assert_almost_equal(RMSD.results.rmsd, correct_values_mass, 4,
                            err_msg="error: rmsd profile should match" +
                            "test values")

    def test_custom_groupselection_weights_applied_1D_array(self, universe):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                            select='backbone',
                                            groupselections=['name CA and resid 1-5', 'name CA and resid 1'],
                                            weights=None,
                                            weights_groupselections=[[1, 0, 0, 0, 0], None]).run(step=49)

        assert_almost_equal(RMSD.results.rmsd.T[3], RMSD.results.rmsd.T[4], 4,
                            err_msg="error: rmsd profile should match "
                            "for applied weight array and selected resid")

    def test_custom_groupselection_weights_applied_mass(self, universe, correct_values_mass):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                            select='backbone',
                                            groupselections=['all', 'all'],
                                            weights=None,
                                            weights_groupselections=['mass',
                                                                     universe.atoms.masses]).run(step=49)

        assert_almost_equal(RMSD.results.rmsd.T[3], RMSD.results.rmsd.T[4], 4,
                            err_msg="error: rmsd profile should match "
                            "between applied mass and universe.atoms.masses")

    def test_rmsd_scalar_weights_raises_ValueError(self, universe):
        with pytest.raises(ValueError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, weights=42)

    def test_rmsd_string_weights_raises_ValueError(self, universe):
        with pytest.raises(ValueError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, weights="Jabberwock")

    def test_rmsd_mismatched_weights_raises_ValueError(self, universe):
        with pytest.raises(ValueError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, weights=universe.atoms.masses[:-1])

    def test_rmsd_misuse_weights_for_groupselection_raises_TypeError(self, universe):
        with pytest.raises(TypeError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, groupselections=['all'],
                weights=[universe.atoms.masses, universe.atoms.masses[:-1]])

    def test_rmsd_mismatched_weights_in_groupselection_raises_ValueError(self, universe):
        with pytest.raises(ValueError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, groupselections=['all'],
                weights=universe.atoms.masses,
                weights_groupselections = [universe.atoms.masses[:-1]])

    def test_rmsd_list_of_weights_wrong_length(self, universe):
        with pytest.raises(ValueError):
            RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, groupselections=['backbone', 'name CA'],
                weights='mass',
                weights_groupselections=[None])

    def test_rmsd_group_selections(self, universe, correct_values_group):
        RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                            groupselections=['backbone', 'name CA']
                                            ).run(step=49)
        assert_almost_equal(RMSD.results.rmsd, correct_values_group, 4,
                            err_msg="error: rmsd profile should match"
                            "test values")

    def test_rmsd_backbone_and_group_selection(self, universe,
                                               correct_values_backbone_group):
        RMSD = MDAnalysis.analysis.rms.RMSD(
            universe,
            reference=universe,
            select="backbone",
            groupselections=['backbone and resid 1:10',
                             'backbone and resid 10:20']).run(step=49)
        assert_almost_equal(
            RMSD.results.rmsd, correct_values_backbone_group, 4,
            err_msg="error: rmsd profile should match test values")

    def test_ref_length_unequal_len(self, universe):
        reference = MDAnalysis.Universe(PSF, DCD)
        reference.atoms = reference.atoms[:-1]
        with pytest.raises(SelectionError):
            RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                                reference=reference)

    def test_mass_mismatches(self, universe):
        reference = MDAnalysis.Universe(PSF, DCD)
        reference.atoms.masses = 10
        with pytest.raises(SelectionError):
            RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                                reference=reference)

    def test_ref_mobile_mass_mismapped(self, universe,correct_values_mass_add_ten):
        reference = MDAnalysis.Universe(PSF, DCD)
        universe.atoms.masses = universe.atoms.masses + 10
        RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                                reference=reference,
                                                select='all',
                                                weights='mass',
                                                tol_mass=100)
        RMSD.run(step=49)
        assert_almost_equal(RMSD.results.rmsd, correct_values_mass_add_ten, 4,
                            err_msg="error: rmsd profile should match "
                            "between true values and calculated values")

    def test_group_selections_unequal_len(self, universe):
        reference = MDAnalysis.Universe(PSF, DCD)
        reference.atoms[0].residue.resname = 'NOTMET'
        with pytest.raises(SelectionError):
            RMSD = MDAnalysis.analysis.rms.RMSD(universe,
                                                reference=reference,
                                                groupselections=['resname MET', 'type NH3'])

    def test_rmsd_attr_warning(self, universe):
        RMSD = MDAnalysis.analysis.rms.RMSD(
                universe, select='name CA').run(stop=2)

        wmsg = "The `rmsd` attribute was deprecated in MDAnalysis 2.0.0"
        with pytest.warns(DeprecationWarning, match=wmsg):
            assert_equal(RMSD.rmsd, RMSD.results.rmsd)


class TestRMSF(object):
    @pytest.fixture()
    def universe(self):
        return mda.Universe(GRO, XTC)

    def test_rmsf(self, universe):
        rmsfs = rms.RMSF(universe.select_atoms('name CA'))
        rmsfs.run()
        test_rmsfs = np.load(rmsfArray)

        assert_almost_equal(rmsfs.results.rmsf, test_rmsfs, 5,
                            err_msg="error: rmsf profile should match test "
                            "values")

    def test_rmsf_single_frame(self, universe):
        rmsfs = rms.RMSF(universe.select_atoms('name CA')).run(start=5, stop=6)

        assert_almost_equal(rmsfs.results.rmsf, 0, 5,
                            err_msg="error: rmsfs should all be zero")

    def test_rmsf_identical_frames(self, universe, tmpdir):

        outfile = os.path.join(str(tmpdir), 'rmsf.xtc')

        # write a dummy trajectory of all the same frame
        with mda.Writer(outfile, universe.atoms.n_atoms) as W:
            for _ in range(universe.trajectory.n_frames):
                W.write(universe)

        universe = mda.Universe(GRO, outfile)
        rmsfs = rms.RMSF(universe.select_atoms('name CA'))
        rmsfs.run()
        assert_almost_equal(rmsfs.results.rmsf, 0, 5,
                            err_msg="error: rmsfs should all be 0")

    def test_rmsf_attr_warning(self, universe):
        rmsfs = rms.RMSF(universe.select_atoms('name CA')).run(stop=2)

        wmsg = "The `rmsf` attribute was deprecated in MDAnalysis 2.0.0"
        with pytest.warns(DeprecationWarning, match=wmsg):
            assert_equal(rmsfs.rmsf, rmsfs.results.rmsf)
