# -*- 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 Lesser GNU Public Licence, v2.1 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 pytest
from io import StringIO
import os

import MDAnalysis as mda
import numpy as np
from MDAnalysisTests import make_Universe
from MDAnalysisTests.coordinates.base import _SingleFrameReader, BaseWriterTest
from MDAnalysisTests.datafiles import FHIAIMS
from numpy.testing import (
    assert_equal,
    assert_array_almost_equal,
    assert_almost_equal,
)


@pytest.fixture(scope="module")
def universe():
    return mda.Universe(FHIAIMS, FHIAIMS)


@pytest.fixture(scope="module")
def universe_from_one_file():
    return mda.Universe(FHIAIMS)


@pytest.fixture(scope="module")
def ref():
    return RefFHIAIMS()


@pytest.fixture(scope="class")
def good_input_natural_units():
    buffer = "atom 0.1 0.2 0.3 H\natom 0.2 0.4 0.6 H\natom 0.3 0.6 0.9 H"
    return StringIO(buffer)


@pytest.fixture(scope="class")
def good_input_with_velocity():
    buffer = "atom 0.1 0.1 0.1 H\nvelocity 0.1 0.1 0.1"
    return StringIO(buffer)


class RefFHIAIMS(object):
    from MDAnalysis.coordinates.FHIAIMS import FHIAIMSReader, FHIAIMSWriter

    filename, trajectory, topology = [FHIAIMS] * 3
    reader, writer = FHIAIMSReader, FHIAIMSWriter
    pos_atom1 = np.asarray([6.861735, 2.103823, 37.753513], dtype=np.float32)
    dimensions = np.asarray(
        [18.6, 18.6, 55.8, 90.0, 90.0, 90.0], dtype=np.float32
    )
    n_atoms = 6
    n_frames = 1
    time = 0.0
    ext = ".in"
    prec = 6
    container_format = True
    changing_dimensions = False


class TestFHIAIMSReader(object):

    prec = 6

    @pytest.fixture(scope="class")
    def bad_input_missing_periodicity(self):
        buffer = "lattice_vector 1.0 0.0 0.0\nlattice_vector 0.0 1.0 0.0\natom 0.1 0.1 0.1 H"
        return StringIO(buffer)

    @pytest.fixture(scope="class")
    def bad_input_relative_positions_no_box(self):
        buffer = "atom_frac 0.1 0.1 0.1 H"
        return StringIO(buffer)

    @pytest.fixture(scope="class")
    def bad_input_missing_velocity(self):
        buffer = "atom 0.1 0.1 0.1 H\natom 0.2 0.2 0.2 H\nvelocity 0.1 0.1 0.1"
        return StringIO(buffer)

    @pytest.fixture(scope="class")
    def bad_input_velocity_wrong_position(self):
        buffer = "atom 0.1 0.1 0.1 H\natom 0.2 0.2 0.2 H\nvelocity 0.1 0.1 0.1\nvelocity 0.1 0.1 0.1"
        return StringIO(buffer)

    @pytest.fixture(scope="class")
    def bad_input_wrong_input_line(self):
        buffer = "garbage"
        return StringIO(buffer)

    @pytest.fixture(scope="class")
    def good_input_mixed_units(self):
        buffer = "lattice_vector 1.0 0.0 0.0\nlattice_vector 0.0 2.0 0.0\nlattice_vector 0.0 0.0 3.0\natom 0.1 0.2 0.3 H\natom_frac 0.2 0.2 0.2 H\natom_frac 0.3 0.3 0.3 H"
        return StringIO(buffer)

    def test_single_file(self, universe, universe_from_one_file):
        assert_almost_equal(
            universe.atoms.positions,
            universe_from_one_file.atoms.positions,
            self.prec,
            "FHIAIMSReader failed to load universe from single file",
        )

    def test_uses_FHIAIMSReader(self, universe):
        from MDAnalysis.coordinates.FHIAIMS import FHIAIMSReader

        assert isinstance(
            universe.trajectory, FHIAIMSReader
        ), "failed to choose FHIAIMSReader"

    def test_dimensions(self, ref, universe):
        assert_almost_equal(
            ref.dimensions,
            universe.dimensions,
            self.prec,
            "FHIAIMSReader failed to get unitcell dimensions",
        )

    def test_n_atoms(self, ref, universe):
        assert_equal(
            ref.n_atoms,
            universe.trajectory.n_atoms,
            "FHIAIMSReader failed to get the right number of atoms",
        )

    def test_fhiaims_positions(self, ref, universe):
        # first particle
        assert_almost_equal(
            ref.pos_atom1,
            universe.atoms.positions[0],
            self.prec,
            "FHIAIMSReader failed to read coordinates properly",
        )

    def test_n_frames(self, ref, universe):
        assert_equal(
            ref.n_frames,
            universe.trajectory.n_frames,
            "wrong number of frames",
        )

    def test_time(self, ref, universe):
        assert_equal(
            ref.time, universe.trajectory.time, "wrong time of the frame"
        )

    def test_bad_input_missing_periodicity(
        self, bad_input_missing_periodicity
    ):
        with pytest.raises(ValueError, match="Found partial periodicity"):
            u = mda.Universe(bad_input_missing_periodicity, format="FHIAIMS")

    def test_bad_input_relative_positions_no_box(
        self, bad_input_relative_positions_no_box
    ):
        with pytest.raises(
            ValueError,
            match="Found relative coordinates in FHI-AIMS file without lattice info",
        ):
            u = mda.Universe(
                bad_input_relative_positions_no_box, format="FHIAIMS"
            )

    def test_bad_input_missing_velocity(self, bad_input_missing_velocity):
        with pytest.raises(
            ValueError, match="Found incorrect number of velocity tags"
        ):
            u = mda.Universe(bad_input_missing_velocity, format="FHIAIMS")

    def test_bad_input_velocity_wrong_position(
        self, bad_input_velocity_wrong_position
    ):
        with pytest.raises(
            ValueError, match="Non-conforming line .velocity must follow"
        ):
            u = mda.Universe(
                bad_input_velocity_wrong_position, format="FHIAIMS"
            )

    def test_bad_input_wrong_input_line(self, bad_input_wrong_input_line):
        with pytest.raises(ValueError, match="Non-conforming line"):
            u = mda.Universe(bad_input_wrong_input_line, format="FHIAIMS")

    def test_good_input_with_velocity(self, good_input_with_velocity):
        u = mda.Universe(good_input_with_velocity, format="FHIAIMS")
        assert_almost_equal(
            u.atoms.velocities[0],
            np.asarray([0.1, 0.1, 0.1]),
            self.prec,
            "FHIAIMSReader failed to read velocities properly",
        )

    def test_mixed_units(
        self, good_input_natural_units, good_input_mixed_units
    ):
        u_natural = mda.Universe(good_input_natural_units, format="FHIAIMS")
        u_mixed = mda.Universe(good_input_mixed_units, format="FHIAIMS")
        print(u_natural.atoms.positions)
        print(u_mixed.atoms.positions)
        assert_almost_equal(
            u_natural.atoms.positions,
            u_mixed.atoms.positions,
            self.prec,
            "FHIAIMSReader failed to read positions in lattice units properly",
        )


class TestFHIAIMSWriter(BaseWriterTest):
    prec = 6
    ext = ".in"

    @pytest.fixture
    def outfile(self, tmpdir):
        return str(
            tmpdir.mkdir("FHIAIMSWriter").join(
                "primitive-fhiaims-writer" + self.ext
            )
        )

    def test_writer(self, universe, outfile):
        """Test writing from a single frame FHIAIMS file to a FHIAIMS file."""
        universe.atoms.write(outfile)
        u = mda.Universe(FHIAIMS, outfile)
        assert_almost_equal(
            u.atoms.positions,
            universe.atoms.positions,
            self.prec,
            err_msg="Writing FHIAIMS file with FHIAIMSWriter "
            "does not reproduce original coordinates",
        )

    def test_writer_with_velocity(self, good_input_with_velocity, outfile):
        """Test writing from a single frame FHIAIMS file to a FHIAIMS file."""
        universe_in = mda.Universe(good_input_with_velocity, format="FHIAIMS")
        universe_in.atoms.write(outfile)
        u = mda.Universe(outfile)
        assert_almost_equal(
            u.atoms.velocities,
            universe_in.atoms.velocities,
            self.prec,
            err_msg="Writing FHIAIMS file with FHIAIMSWriter "
            "does not reproduce original velocities",
        )

    def test_writer_no_atom_names(self, u_no_names, outfile):
        u_no_names.atoms.write(outfile)
        u = mda.Universe(outfile)
        expected = np.array(["X"] * u_no_names.atoms.n_atoms)
        assert_equal(u.atoms.names, expected)

    def test_writer_with_n_atoms_none(self, good_input_natural_units, outfile):
        u = mda.Universe(good_input_natural_units, format="FHIAIMS")
        with mda.Writer(outfile, natoms=None) as w:
            w.write(u.atoms)
            with open(outfile, "r") as fhiaimsfile:
                line = fhiaimsfile.readline().strip()
                assert line.startswith(
                    "atom"
                ), "Line written incorrectly with FHIAIMSWriter"
                assert line.endswith(
                    "H"
                ), "Line written incorrectly with FHIAIMSWriter"
                line = np.asarray(line.split()[1:-1], dtype=np.float32)
                assert_almost_equal(
                    line,
                    [0.1, 0.2, 0.3],
                    self.prec,
                    err_msg="Writing FHIAIMS file with FHIAIMSWriter "
                    "does not reproduce original positions",
                )
