# Source code for MDAnalysis.converters.OpenMM

    # -*- 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 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
    #

    """OpenMM structure I/O --- :mod:`MDAnalysis.converters.OpenMM`
    ================================================================

    Read coordinates data from a
    `OpenMM <http://docs.openmm.org/latest/api-python/generated/openmm.app.simulation.Simulation.html#openmm.app.simulation.Simulation>`_
    :class:`openmm.app.simulation.Simulation` with :class:`OpenMMReader`
    into a MDAnalysis Universe.

    Also converts other objects within the
    `OpenMM Application Layer <http://docs.openmm.org/latest/api-python/app.html>`_:

    - `openmm.app.pdbfile.PDBFile <http://docs.openmm.org/latest/api-python/generated/openmm.app.pdbfile.PDBFile.html#openmm.app.pdbfile.PDBFile>`_
    - `openmm.app.modeller.Modeller <http://docs.openmm.org/latest/api-python/generated/openmm.app.modeller.Modeller.html#openmm.app.modeller.Modeller>`_
    - `openmm.app.pdbxfile.PDBxFile <http://docs.openmm.org/latest/api-python/generated/openmm.app.pdbxfile.PDBxFile.html#openmm.app.pdbxfile.PDBxFile>`_

    Example
    -------
    OpenMM can read various file formats into OpenMM objects.
    MDAnalysis can then convert some of these OpenMM objects into MDAnalysis Universe objects.

        >>> import openmm.app as app
        >>> import MDAnalysis as mda
        >>> from MDAnalysis.tests.datafiles import PDBX
        >>> pdbxfile = app.PDBxFile(PDBX)
        >>> mda.Universe(pdbxfile)
        <Universe with 60 atoms>

    Classes
    -------

    .. autoclass:: OpenMMSimulationReader
       :members:

    .. autoclass:: OpenMMAppReader
       :members:

    """

    import numpy as np

    from ..coordinates import base

    class OpenMMSimulationReader(base.SingleFrameReaderBase):
        """Reader for OpenMM Simulation objects

        .. versionadded:: 2.0.0
        """

        format = "OPENMMSIMULATION"
        units = {
            "time": "ps",
            "length": "nm",
            "velocity": "nm/ps",
            "force": "kJ/(mol*nm)",
            "energy": "kJ/mol",
        }

        @staticmethod
        def _format_hint(thing):
            """Can this reader read *thing*?"""
            try:
                from openmm.app import Simulation
            except ImportError:
                try:  # pragma: no cover
                    from simtk.openmm.app import Simulation
                except ImportError:
                    return False
            else:
                return isinstance(thing, Simulation)

        def _read_first_frame(self):
            self.n_atoms = self.filename.topology.getNumAtoms()

            self.ts = self._mda_timestep_from_omm_context()

            if self.convert_units:
                self.convert_pos_from_native(self.ts._pos)
                self.ts.triclinic_dimensions = self.convert_pos_from_native(
                    self.ts.triclinic_dimensions, inplace=False
                )
                self.ts.dimensions[3:] = _sanitize_box_angles(
                    self.ts.dimensions[3:]
                )
                self.convert_velocities_from_native(self.ts._velocities)
                self.convert_forces_from_native(self.ts._forces)
                self.convert_time_from_native(self.ts.dt)

        def _mda_timestep_from_omm_context(self):
            """Construct Timestep object from OpenMM context"""
            try:
                import openmm.unit as u
            except ImportError:  # pragma: no cover
                import simtk.unit as u

            state = self.filename.context.getState(
                -1, getVelocities=True, getForces=True, getEnergy=True
            )

            n_atoms = self.filename.context.getSystem().getNumParticles()

            ts = self._Timestep(n_atoms, **self._ts_kwargs)
            ts.frame = 0
            ts.data["time"] = state.getTime()._value
            ts.data["potential_energy"] = (
                state.getPotentialEnergy().in_units_of(u.kilojoule / u.mole)._value
            )
            ts.data["kinetic_energy"] = (
                state.getKineticEnergy().in_units_of(u.kilojoule / u.mole)._value
            )
            ts.triclinic_dimensions = state.getPeriodicBoxVectors(
                asNumpy=True
            )._value
            ts.dimensions[3:] = _sanitize_box_angles(ts.dimensions[3:])
            ts.positions = state.getPositions(asNumpy=True)._value
            ts.velocities = state.getVelocities(asNumpy=True)._value
            ts.forces = state.getForces(asNumpy=True)._value

            return ts

    class OpenMMAppReader(base.SingleFrameReaderBase):
        """Reader for OpenMM Application layer objects

        See also `the object definition in the OpenMM Application layer <http://docs.openmm.org/latest/api-python/generated/openmm.app.simulation.Simulation.html#openmm.app.simulation.Simulation>`_

        .. versionadded:: 2.0.0
        """

        format = "OPENMMAPP"
        units = {"time": "ps", "length": "nm"}

        @staticmethod
        def _format_hint(thing):
            """Can this reader read *thing*?"""
            try:
                from openmm import app
            except ImportError:
                try:  # pragma: no cover
                    from simtk.openmm import app
                except ImportError:
                    return False
            else:
                return isinstance(thing, (app.PDBFile, app.Modeller, app.PDBxFile))

        def _read_first_frame(self):
            self.n_atoms = self.filename.topology.getNumAtoms()

            self.ts = self._mda_timestep_from_omm_app()

            if self.convert_units:
                self.convert_pos_from_native(self.ts._pos)
                if self.ts.dimensions is not None:
                    self.ts.triclinic_dimensions = self.convert_pos_from_native(
                        self.ts.triclinic_dimensions, inplace=False
                    )
                    self.ts.dimensions[3:] = _sanitize_box_angles(
                        self.ts.dimensions[3:]
                    )

        def _mda_timestep_from_omm_app(self):
            """Construct Timestep object from OpenMM Application object"""

            omm_object = self.filename
            n_atoms = omm_object.topology.getNumAtoms()

            ts = self._Timestep(n_atoms, **self._ts_kwargs)
            ts.frame = 0
            if omm_object.topology.getPeriodicBoxVectors() is not None:
                ts.triclinic_dimensions = np.array(
                    omm_object.topology.getPeriodicBoxVectors()._value
                )
                ts.dimensions[3:] = _sanitize_box_angles(ts.dimensions[3:])
            ts.positions = np.array(omm_object.getPositions()._value)

            return ts

    def _sanitize_box_angles(angles):
        """Ensure box angles correspond to first quadrant

        See `discussion on unitcell angles <https://github.com/MDAnalysis/mdanalysis/pull/2917/files#r620558575>`_
        """
        inverted = 180 - angles

        return np.min(np.array([angles, inverted]), axis=0)