# Source code for mdtraj.utils.unit

    ##############################################################################
    # MDTraj: A Python Library for Loading, Saving, and Manipulating
    #         Molecular Dynamics Trajectories.
    # Copyright 2012-2013 Stanford University and the Authors
    #
    # Authors: Robert McGibbon
    # Contributors:
    #
    # MDTraj is free software: you can redistribute it and/or modify
    # it under the terms of the GNU Lesser General Public License as
    # published by the Free Software Foundation, either version 2.1
    # of the License, or (at your option) any later version.
    #
    # This library is distributed in the hope that it will be useful,
    # but WITHOUT ANY WARRANTY; without even the implied warranty of
    # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    # GNU Lesser General Public License for more details.
    #
    # You should have received a copy of the GNU Lesser General Public
    # License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
    ##############################################################################
    """

    Unit processing for MDTraj. This subpackage is a port of openmm.unit from
    OpenMM. Unlike in openmm, the MDTraj library **does not pass around
    "united quantities"**

    The only publicly facing API from this package, for the purpose of MDTraj,
    is "in_units_of", which does unit conversion of numbers or numpy arrays
    where the input and output units are passed as strings.

    """

    import ast
    import sys

    import numpy as np

    from mdtraj.utils import import_
    from mdtraj.utils.unit import unit_definitions
    from mdtraj.utils.unit.quantity import Quantity

    UNIT_DEFINITIONS = unit_definitions

    try:
        import openmm.unit as openmm_unit
    except ImportError:
        pass

    __all__ = ["in_units_of"]

    class _UnitContext(ast.NodeTransformer):
        """Node transformer for an AST hack that turns raw strings into
        complex simt.unit.Unit expressions. See _str_to_unit for how this
        is used -- it's not really meant to stand on its own
        """

        # we want to do some validation to ensure that the AST only
        # contains "safe" operations. These are operations that can reasonably
        # appear in unit expressions
        allowed_ops = [
            ast.Expression,
            ast.BinOp,
            ast.Name,
            ast.Attribute,
            ast.Pow,
            ast.Div,
            ast.Mult,
            ast.Constant,
        ]

        def visit(self, node):
            if not any(isinstance(node, a) for a in self.allowed_ops):
                raise ValueError(
                    "Invalid unit expression. Contains dissallowed operation %s" % node.__class__.__name__,
                )
            return super().visit(node)

        def visit_Name(self, node):
            # we want to prefix all names to look like unit.nanometers instead
            # of just "nanometers", because I don't want to import * from
            # units into this module
            if not hasattr(unit_definitions, node.id):
                # also, let's take this opporunity to check that the node.id
                # (which supposed to be the name of the unit, like "nanometers")
                # is actually an attribute in openmm.unit
                raise ValueError("%s is not a valid unit" % node.id)

            return ast.Attribute(
                value=ast.Name(id="unit_definitions", ctx=ast.Load()),
                attr=node.id,
                ctx=ast.Load(),
            )

    _unit_context = _UnitContext()  # global instance of the visitor

    def _str_to_unit(unit_string, openmm=False):
        """eval() based transformer that extracts a openmm.unit object
        from a string description.

        Parameters
        ----------
        unit_string : str
            string description of a unit. this may contain expressions with
            multiplication, division, powers, etc.

        Examples
        --------
        >>> type(_str_to_unit('nanometers**2/meters*gigajoules'))
        <class 'openmm.unit.unit.Unit'>
        >>> str(_str_to_unit('nanometers**2/meters*gigajoules'))
        'nanometer**2*gigajoule/meter'

        """
        # parse the string with the ast, and then run out unit context
        # visitor on it, which will basically change bare names like
        # "nanometers" into "unit.nanometers" and simulataniously check that
        # there's no nefarious stuff in the expression.

        assert isinstance(unit_string, str)
        unit_definitions = UNIT_DEFINITIONS
        if openmm:
            unit_definitions = openmm_unit.unit_definitions
        parsed = ast.parse(unit_string, mode="eval")
        node = _unit_context.visit(parsed)
        fixed_node = ast.fix_missing_locations(node)
        output = eval(compile(fixed_node, "<string>", mode="eval"), {}, locals())
        return output

    def in_units_of(quantity, units_in, units_out, inplace=False):
        """Convert a numerical quantity between unit systems.

        Parameters
        ----------
        quantity : {number, np.ndarray, openmm.unit.Quantity}
            quantity can either be a unitted quantity -- i.e. instance of
            openmm.unit.Quantity, or just a bare number or numpy array
        units_in : str
            If you supply a quantity that's not a openmm.unit.Quantity, you should
            tell me what units it is in. If you don't, i'm just going to echo you
            back your quantity without doing any unit checking.
        units_out : str
            A string description of the units you want out. This should look
            like "nanometers/picosecond" or "nanometers**3" or whatever
        inplace : bool
            Attempt to do the transformation inplace, by mutating the `quantity`
            argument and avoiding a copy. This is only possible if `quantity` is a
            writable numpy array.

        Returns
        -------
        rquantity : {number, np.ndarray}
            The resulting quantity, in the new unit system. If the function was
            called with `inplace=True` and `quantity` was a writable numpy array,
            `rquantity` will alias the same memory as the input `quantity`, which
            will have been changed inplace. Otherwise, if a copy was required,
            `rquantity` will point to new memory.

        Examples
        --------
        >>> in_units_of(1, 'meter**2/second', 'nanometers**2/picosecond')
        1000000.0
        """
        if quantity is None:
            return quantity

        if "openmm.unit" in sys.modules and isinstance(quantity, openmm_unit.Quantity):
            units_in = quantity.unit
            units_out = _str_to_unit(units_out, openmm=True)
            quantity = quantity._value
        elif isinstance(quantity, Quantity):
            units_in = quantity.unit
            units_out = _str_to_unit(units_out)
            quantity = quantity._value
        else:
            if units_in is None:
                return quantity
            units_in = _str_to_unit(units_in)
            units_out = _str_to_unit(units_out)

        if not units_in.is_compatible(units_out):
            raise TypeError(f'Unit "{units_in}" is not compatible with Unit "{units_out}".')

        factor = units_in.conversion_factor_to(units_out)
        if inplace and (isinstance(quantity, np.ndarray) and quantity.flags["WRITEABLE"]):
            quantity *= factor
            return quantity
        return quantity * factor