# -*- 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
#
from __future__ import absolute_import

import pytest
import MDAnalysis as mda

from MDAnalysisTests.topology.base import ParserBase
from MDAnalysisTests.datafiles import (
    PSF_NAMD_GBIS,
    PRM
)

pmd = pytest.importorskip('parmed')

class BaseTestParmedParser(ParserBase):
    parser = mda.topology.ParmEdParser.ParmEdParser
    expected_attrs = ['ids', 'names', 'types', 'masses',
                      'charges', 'altLocs', 'occupancies',
                      'tempfactors', 'gbscreens', 'solventradii',
                      'nbindices', 'rmins', 'epsilons', 'rmin14s',
                      'epsilon14s', 'elements', 'chainIDs',
                      'resids', 'resnames', 'resnums',
                      'segids',
                      'bonds', 'ureybradleys', 'angles', 
                      'dihedrals', 'impropers', 'cmaps']

    expected_n_atoms = 0
    expected_n_residues = 1
    expected_n_segments = 1
    expected_n_bonds = 0
    expected_n_angles = 0
    expected_n_dihedrals = 0
    expected_n_impropers = 0
    expected_n_cmaps = 0
    expected_n_ureybradleys = 0

    @pytest.fixture
    def filename(self):
        return pmd.load_file(self.ref_filename)

    @pytest.fixture
    def universe(self, filename):
        return mda.Universe(filename)

    def test_creates_universe(self, filename):
        u = mda.Universe(filename)
        assert isinstance(u, mda.Universe)

    def test_bonds_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx)
                      for a in filename.bonds])
        assert len(top.bonds.values) == len(unique)
    
    def test_angles_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx, a.atom3.idx) 
                      for a in filename.angles])
        assert len(top.angles.values) == len(unique)

    def test_dihedrals_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx, a.atom3.idx, a.atom4.idx)
             for a in filename.dihedrals])
        assert len(top.dihedrals.values) == len(unique)
    
    def test_impropers_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx, a.atom3.idx, a.atom4.idx)
                      for a in filename.impropers])
        assert len(top.impropers.values) == len(unique)

    def test_cmaps_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx, a.atom3.idx, 
                       a.atom4.idx, a.atom5.idx)
                      for a in filename.cmaps])
        assert len(top.cmaps.values) == len(unique)
    
    def test_ureybradleys_total_counts(self, top, filename):
        unique = set([(a.atom1.idx, a.atom2.idx)
                      for a in filename.urey_bradleys])
        assert len(top.ureybradleys.values) == len(unique)



class TestParmedParserPSF(BaseTestParmedParser):
    """
    PSF with CMAPs
    """
    
    ref_filename = PSF_NAMD_GBIS
    
    expected_n_atoms = 3341
    expected_n_residues = 214
    expected_n_bonds = 3365
    expected_n_angles = 6123
    expected_n_dihedrals = 8921
    expected_n_impropers = 541
    expected_n_cmaps = 212

    def test_bonds_atom_counts(self, universe):
        assert len(universe.atoms[[0]].bonds) == 4
        assert len(universe.atoms[[42]].bonds) == 1

    @pytest.mark.parametrize('value', (
        (0, 1), 
        (0, 2), 
        (0, 3), 
        (0, 4),
        ))
    def test_bonds_identity(self, top, value):
        vals = top.bonds.values
        assert value in vals or value[::-1] in vals
    
    def test_bond_types(self, universe):
        b1 = universe.bonds[0]

        assert b1.order == 1.0
        assert tuple(b1.indices) == (0, 1)
        assert b1.type.type is None

    def test_angles_atom_counts(self, universe):
        assert len(universe.atoms[[0]].angles), 9
        assert len(universe.atoms[[42]].angles), 2

    @pytest.mark.parametrize('value', (
        (1, 0, 2), 
        (1, 0, 3), 
        (1, 0, 4),
        ))
    def test_angles_identity(self, top, value):
        vals = top.angles.values
        assert value in vals or value[::-1] in vals

    def test_dihedrals_atom_counts(self, universe):
        assert len(universe.atoms[[0]].dihedrals) == 14

    @pytest.mark.parametrize('value', (
        (0, 4, 6, 7), 
        (0, 4, 6, 8), 
        (0, 4, 6, 9), 
        (0, 4, 17, 18),
        ))
    def test_dihedrals_identity(self, top, value):
        vals = top.dihedrals.values
        assert value in vals or value[::-1] in vals
        

    @pytest.mark.parametrize('value', (
        (17, 19, 21, 41, 43),
        (60, 62, 64, 79, 81)
    ))
    def test_cmaps_identity(self, top, value):
        vals = top.cmaps.values
        assert value in vals or value[::-1] in vals
    
class TestParmedParserPRM(BaseTestParmedParser):
    """
    PRM
    """
    
    ref_filename = PRM
    
    expected_n_atoms = 252
    expected_n_residues = 14

    def test_bonds_atom_counts(self, universe):
        assert len(universe.atoms[[0]].bonds) == 4
        assert len(universe.atoms[[42]].bonds) == 1

    @pytest.mark.parametrize('value', (
        (10, 11),
        (10, 12),
        (4, 6),
        (4, 10),
        ))
    def test_bonds_identity(self, top, value):
        vals = top.bonds.values
        assert value in vals or value[::-1] in vals
    
    def test_bond_types(self, universe):
        b1 = universe.bonds[0]

        assert b1.order == 1.0
        assert tuple(b1.indices) == (0, 1)
        assert b1.type.type.k == 434
        assert b1.type.type.req == 1.010

    def test_angles_atom_counts(self, universe):
        assert len(universe.atoms[[0]].angles), 9
        assert len(universe.atoms[[42]].angles), 2

    @pytest.mark.parametrize('value', (
        (11, 10, 12),
        (10, 12, 14),
        (6, 4, 10),
        (4, 10, 11),
        ))
    def test_angles_identity(self, top, value):
        vals = top.angles.values
        assert value in vals or value[::-1] in vals

    def test_dihedrals_atom_counts(self, universe):
        assert len(universe.atoms[[0]].dihedrals) == 14

    @pytest.mark.parametrize('value', (
        (11, 10, 12, 14),
        (10, 12, 14, 16),
        ))
    def test_dihedrals_identity(self, top, value):
        vals = top.dihedrals.values
        assert value in vals or value[::-1] in vals
    
    def test_dihedral_types(self, universe):
        ag = universe.atoms[[10,12,14,16]]
        dih = universe.dihedrals.atomgroup_intersection(ag,
                                                        strict=True)[0]
        assert len(dih.type) == 4
        for i, (phi_k, per) in enumerate((
            (2.0, 1),
            (2.0, 2),
            (0.4, 3),
            (0.0, 4),
        )):
            assert dih.type[i].type.phi_k == phi_k
            assert dih.type[i].type.per == per
