# GromacsWrapper: test_example.py
# Copyright (c) 2009 Oliver Beckstein <orbeckst@gmail.com>
# Released under the GNU Public License 3 (or higher, your choice)
# See the file COPYING for details.

import os.path
import pytest

from numpy.testing import assert_almost_equal

import gromacs
import gromacs.setup

from .datafiles import datafile


@pytest.mark.xfail(
    gromacs.release.startswith("2020.6"),
    reason="pdb2gmx 2020.6 fails to build the TIP4P waters",
)
def test_trj_compact_main(tmpdir):
    pdb = datafile("1ake_A.pdb")
    top = tmpdir.mkdir("top")
    mdpfile = "simple.mdp"
    tprfile = "simple.tpr"
    outfile = "compact.pdb"
    with top.as_cwd():
        f = gromacs.setup.topology(struct=pdb, ff="oplsaa", water="tip4p")
        with open(mdpfile, "w") as mdp:
            mdp.write("; empty mdp file\nrcoulomb = 1\nrvdw = 1\nrlist = 1\n")
        gromacs.grompp(f=mdpfile, o=tprfile, c=f["struct"], p=f["top"])
        gromacs.setup.trj_compact_main(
            s=tprfile, f=f["struct"], o=outfile, input=("protein", "system")
        )
        assert os.path.exists(outfile)


@pytest.fixture(scope="session")
def topology(tmpdir_factory, struct=datafile("1ake_A_protein.pdb")):
    # note: use protein-only input 1ake_A_protein.pdb because solvation fails
    #       if crystal waters are included (in 1ake_A.pdb)
    TMPDIR = tmpdir_factory.mktemp("1ake")
    with TMPDIR.as_cwd():
        topol_args = gromacs.setup.topology(struct=struct, ff="oplsaa", water="tip4p")
    return TMPDIR, topol_args


@pytest.fixture(scope="session")
def solvate(topology):
    TMPDIR, topol_args = topology
    with TMPDIR.as_cwd():
        solvate_args = gromacs.setup.solvate(
            concentration=0.15, water="tip4p", **topol_args
        )
    return TMPDIR, solvate_args


@pytest.fixture(scope="session")
def energy_minimize(solvate, low_performance):
    # run short energy minimization with cheapest minimizer
    TMPDIR, solvate_args = solvate
    nt = 2 if low_performance else 0
    with TMPDIR.as_cwd():
        with pytest.warns(
            gromacs.UsageWarning,
            match="Unprocessed mdp option are interpreted as options for grompp",
        ):
            em_args = gromacs.setup.energy_minimize(
                mdrun_args={"nt": nt},
                integrator="steep",
                emtol=5000,
                maxwarn=1,
                **solvate_args,
            )
    return TMPDIR, em_args


def test_topology(topology):
    TMPDIR, topol_args = topology
    top = topol_args["top"]
    struct = topol_args["struct"]
    posres = topol_args["posres"]
    assert os.path.exists(top)
    assert os.path.exists(struct)
    assert os.path.exists(posres)
    # add more tests for content of files!


def test_solvate(solvate):
    TMPDIR, solvate_args = solvate

    assert_almost_equal(solvate_args["qtot"], 0.0)
    assert os.path.exists(solvate_args["struct"])
    assert os.path.exists(solvate_args["ndx"])
    assert solvate_args["mainselection"] == '"Protein"'
    # add more tests for content of files!


def test_energy_minimize(energy_minimize):
    TMPDIR, em_args = energy_minimize
    assert os.path.exists(em_args["struct"])
    assert os.path.exists(em_args["top"])
    assert em_args["mainselection"] == '"Protein"'
    # add more tests for content of files!


def test_energy_minimize_custom_mdp(
    solvate, low_performance, mdp=datafile("custom_em.mdp")
):
    TMPDIR, solvate_args = solvate
    nt = 2 if low_performance else 0
    with TMPDIR.as_cwd():
        with pytest.warns(
            gromacs.UsageWarning,
            match="Unprocessed mdp option are interpreted as options for grompp",
        ):
            try:
                em_args = gromacs.setup.energy_minimize(
                    mdrun_args={"nt": nt}, mdp=mdp, emtol=5000, **solvate_args
                )
            except gromacs.exceptions.GromacsError as err:
                # sometimes the em does not converge at all, e.g. 5.02988e+04 on atom 3277;
                # (happens on Travis Linux with Gromacs 4.6.5 but not locally or on Travis OSX) so we
                # re-run with a ridiculous tolerance so that we can at least test that the whole
                # function can run to completion
                em_args = gromacs.setup.energy_minimize(
                    mdrun_args={"nt": nt}, mdp=mdp, emtol=6e4, **solvate_args
                )
    assert os.path.exists(em_args["struct"])
    assert os.path.exists(em_args["top"])
    assert em_args["mainselection"] == '"Protein"'
    # add more tests for content of files!
