# Source code for openforcefield.utils.utils

    #!/usr/bin/env python
    """
    Utility subroutines

    """

    __all__ = [
        'MessageException',
        'inherit_docstrings',
        'all_subclasses',
        'temporary_cd',
        'temporary_directory',
        'get_data_filename',
        'unit_to_string',
        'quantity_to_string',
        'string_to_unit',
        'string_to_quantity',
        'extract_serialized_units_from_dict',
        'attach_units',
        'detach_units',
        'serialize_numpy',
        'deserialize_numpy',
    ]

    #=============================================================================================
    # GLOBAL IMPORTS
    #=============================================================================================

    import contextlib
    from simtk import unit

    #=============================================================================================
    # COMMON EXCEPTION TYPES
    #=============================================================================================

    class MessageException(Exception):
        """A base class for exceptions that print out a string given in their constructor"""
        def __init__(self, msg):
            super().__init__(self, msg)
            self.msg = msg

    # =============================================================================================
    # UTILITY SUBROUTINES
    # =============================================================================================

    def inherit_docstrings(cls):
        """Inherit docstrings from parent class"""
        from inspect import getmembers, isfunction
        for name, func in getmembers(cls, isfunction):
            if func.__doc__: continue
            for parent in cls.__mro__[1:]:
                if hasattr(parent, name):
                    func.__doc__ = getattr(parent, name).__doc__
        return cls

    def all_subclasses(cls):
        """Recursively retrieve all subclasses of the specified class"""
        return cls.__subclasses__() + [
            g for s in cls.__subclasses__() for g in all_subclasses(s)
        ]

    @contextlib.contextmanager
    def temporary_cd(dir_path):
        """Context to temporary change the working directory.

        Parameters
        ----------
        dir_path : str
            The directory path to enter within the context

        Examples
        --------
        >>> dir_path = '/tmp'
        >>> with temporary_cd(dir_path):
        ...     pass  # do something in dir_path

        """
        import os
        prev_dir = os.getcwd()
        os.chdir(os.path.abspath(dir_path))
        try:
            yield
        finally:
            os.chdir(prev_dir)

    @contextlib.contextmanager
    def temporary_directory():
        """Context for safe creation of temporary directories."""

        import tempfile
        tmp_dir = tempfile.mkdtemp()
        try:
            yield tmp_dir
        finally:
            import shutil
            shutil.rmtree(tmp_dir)

    def get_data_filename(relative_path):
        """Get the full path to one of the reference files in testsystems.
        In the source distribution, these files are in ``openforcefield/data/``,
        but on installation, they're moved to somewhere in the user's python
        site-packages directory.
        Parameters
        ----------
        name : str
            Name of the file to load (with respect to the repex folder).
        """

        from pkg_resources import resource_filename
        import os
        fn = resource_filename('openforcefield', os.path.join(
            'data', relative_path))

        if not os.path.exists(fn):
            raise ValueError(
                "Sorry! %s does not exist. If you just added it, you'll have to re-install"
                % fn)

        return fn

    def unit_to_string(input_unit):
        """
        Serialize a simtk.unit.Unit and return it as a string.

        Parameters
        ----------
        input_unit : A simtk.unit
            The unit to serialize

        Returns
        -------
        unit_string : str
            The serialized unit.
        """

        # Decompose output_unit into a tuples of (base_dimension_unit, exponent)
        unit_string = None

        for unit_component in input_unit.iter_base_or_scaled_units():
            unit_component_name = unit_component[0].name
            # Convert, for example "elementary charge" --> "elementary_charge"
            unit_component_name = unit_component_name.replace(' ', '_')
            if unit_component[1] == 1:
                contribution = '{}'.format(unit_component_name)
            else:
                contribution = '{}**{}'.format(unit_component_name, int(unit_component[1]))
            if unit_string is None:
                unit_string = contribution
            else:
                unit_string += ' * {}'.format(contribution)

        return unit_string

    def quantity_to_string(input_quantity):
        """
        Serialize a simtk.unit.Quantity to a string.

        Parameters
        ----------
        input_quantity : simtk.unit.Quantity
            The quantity to serialize

        Returns
        -------
        output_string : str
            The serialized quantity

        """
        import numpy as np
        if input_quantity is None:
            return None
        unitless_value = input_quantity.value_in_unit(input_quantity.unit)
        # The string representaiton of a numpy array doesn't have commas and breaks the
        # parser, thus we convert any arrays to list here
        if isinstance(unitless_value, np.ndarray):
            unitless_value = list(unitless_value)
        unit_string = unit_to_string(input_quantity.unit)
        print(input_quantity, type(unitless_value))
        output_string = '{} * {}'.format(unitless_value, unit_string)
        return output_string

    def _ast_eval(node):
        """
        Performs an algebraic syntax tree evaluation of a unit.

        Parameters
        ----------
        node : An ast parsing tree node
        """
        import ast
        import operator as op

        operators = {ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul,
            ast.Div: op.truediv, ast.Pow: op.pow, ast.BitXor: op.xor,
            ast.USub: op.neg}

        if isinstance(node, ast.Num): # <number>
            return node.n
        elif isinstance(node, ast.BinOp): # <left> <operator> <right>
            return operators[type(node.op)](_ast_eval(node.left), _ast_eval(node.right))
        elif isinstance(node, ast.UnaryOp): # <operator> <operand> e.g., -1
            return operators[type(node.op)](_ast_eval(node.operand))
        elif isinstance(node, ast.Name):
            # see if this is a simtk unit
            b = getattr(unit, node.id)
            return b
        # TODO: This was a quick hack that surprisingly worked. We should validate this further.
        elif isinstance(node, ast.List):
            return ast.literal_eval(node)
        else:
            raise TypeError(node)

    def string_to_unit(unit_string):
        """
        Deserializes a simtk.unit.Quantity from a string representation, for
        example: "kilocalories_per_mole / angstrom ** 2"

        Parameters
        ----------
        unit_string : dict
            Serialized representation of a simtk.unit.Quantity.

        Returns
        -------
        output_unit: simtk.unit.Quantity
            The deserialized unit from the string
        """
        import ast
        output_unit = _ast_eval(ast.parse(unit_string, mode='eval').body)
        return output_unit

        #if (serialized['unitless_value'] is None) and (serialized['unit'] is None):
        #    return None
        #quantity_unit = None
        #for unit_name, power in serialized['unit']:
        #    unit_name = unit_name.replace(
        #        ' ', '_')  # Convert eg. 'elementary charge' to 'elementary_charge'
        #    if quantity_unit is None:
        #        quantity_unit = (getattr(unit, unit_name)**power)
        #    else:
        #        quantity_unit *= (getattr(unit, unit_name)**power)
        #quantity = unit.Quantity(serialized['unitless_value'], quantity_unit)
        #return quantity

    def string_to_quantity(quantity_string):
        """
        Takes a string representation of a quantity and returns a simtk.unit.Quantity

        Parameters
        ----------
        quantity_string : str
            The quantity to deserialize

        Returns
        -------
        output_quantity : simtk.unit.Quantity
            The deserialized quantity
        """
        if quantity_string is None:
            return None
        # This can be the exact same as string_to_unit
        import ast
        output_quantity = _ast_eval(ast.parse(quantity_string, mode='eval').body)
        return output_quantity

    def extract_serialized_units_from_dict(input_dict):
        """
        Create a mapping of (potentially unit-bearing) quantities from a dictionary, where some keys exist in pairs like
        {'length': 8, 'length_unit':'angstrom'}.

        Parameters
        ----------
        input_dict : dict
           Dictionary where some keys are paired like {'X': 1.0, 'X_unit': angstrom}.

        Returns
        -------
        unitless_dict : dict
           input_dict, but with keys ending in ``_unit`` removed.
        attached_units : dict str : simtk.unit.Unit
           ``attached_units[parameter_name]`` is the simtk.unit.Unit combination that should be attached to corresponding
           parameter ``parameter_name``. For example ``attached_units['X'] = simtk.unit.angstrom.

        """

        # TODO: Should this scheme also convert "1" to int(1) and "8.0" to float(8.0)?
        from collections import OrderedDict
        attached_units = OrderedDict()
        unitless_dict = input_dict.copy()
        keys_to_delete = []
        for key in input_dict.keys():
            if key.endswith('_unit'):
                parameter_name = key[:-5]
                parameter_units_string = input_dict[key]
                try:
                    parameter_units = string_to_unit(parameter_units_string)
                except Exception as e:
                    e.msg = "Could not parse units {}\n".format(
                        parameter_units_string) + e.msg
                    raise e
                attached_units[parameter_name] = parameter_units
                # Remember this key and delete it later (we break the dict if we delete a key in the loop)
                keys_to_delete.append(key)
        # Clean out the '*_unit' keys that we processed
        for key in keys_to_delete:
            del unitless_dict[key]

        return unitless_dict, attached_units

    def attach_units(unitless_dict, attached_units):
        """
        Attach units to dict entries for which units are specified.

        Parameters
        ----------
        unitless_dict : dict
           Dictionary, where some items are to have units applied.
        attached_units : dict [str : simtk.unit.Unit]
           ``attached_units[parameter_name]`` is the simtk.unit.Unit combination that should be attached to corresponding
           parameter ``parameter_name``

        Returns
        -------
        unit_bearing_dict : dict
           Updated dict with simtk.unit.Unit units attached to values for which units were specified for their keys

        """
        temp_dict = unitless_dict.copy()
        for parameter_name, units_to_attach in attached_units.items():
            if parameter_name in temp_dict.keys():
                parameter_attrib_string = temp_dict[parameter_name]
                try:
                    temp_dict[parameter_name] = float(parameter_attrib_string) * units_to_attach
                except ValueError as e:
                    e.msg = (
                        "Expected numeric value for parameter '{}',"
                        "instead found '{}' when trying to attach units '{}'\n"
                    ).format(parameter_name, parameter_attrib_string, units_to_attach)
                    raise e

            # Now check for matches like "phase1", "phase2"
            c = 1
            while (parameter_name + str(c)) in temp_dict.keys():
                indexed_parameter_name = parameter_name + str(c)
                parameter_attrib_string = temp_dict[indexed_parameter_name]
                try:
                    temp_dict[indexed_parameter_name] = float(
                              parameter_attrib_string) * units_to_attach
                except ValueError as e:
                    e.msg = "Expected numeric value for parameter '{}', instead found '{}' when trying to attach units '{}'\n".format(
                        indexed_parameter_name, parameter_attrib_string,
                        units_to_attach)
                    raise e
                c += 1
        return temp_dict

    def detach_units(unit_bearing_dict, output_units=None):
        """
        Given a dict which may contain some simtk.unit.Quantity objects, return the same dict with the Quantities
        replaced with unitless values, and a new dict containing entries with the suffix "_unit" added, containing
        the units.

        Parameters
        ----------
        unit_bearing_dict : dict
            A dictionary potentially containing simtk.unit.Quantity objects as values.
        output_units : dict[str : simtk.unit.Unit], optional. Default = None
            A mapping from parameter fields to the output unit its value should be converted to.
            For example, {'length_unit': unit.angstrom}. If no output_unit is defined for a key:value pair in which
            the value is a simtk.unit.Quantity, the output unit will be the Quantity's unit, and this information
            will be included in the unit_dict return value.

        Returns
        -------
        unitless_dict : dict
            The input smirnoff_dict object, with all simtk.unit.Quantity values converted to unitless values.
        unit_dict : dict
            A dictionary in which keys are keys of simtk.unit.Quantity values in unit_bearing_dict,
            but suffixed with "_unit". Values are simtk.unit.Unit .
        """
        from simtk import unit

        if output_units is None:
            output_units = {}

        # initialize dictionaries for outputs
        unit_dict = {}
        unitless_dict = unit_bearing_dict.copy()

        for key, value in unit_bearing_dict.items():
            # If no conversion is needed, skip this item
            if not isinstance(value, unit.Quantity):
                continue

            # If conversion is needed, see if the user has requested an output unit
            unit_key = key + '_unit'

            if unit_key in output_units:
                output_unit = output_units[unit_key]
            else:
                output_unit = value.unit
            if not (output_unit.is_compatible(value.unit)):
                raise ValueError("Requested output unit {} is not compatible with "
                                 "quantity unit {} .".format(output_unit, value.unit))
            unitless_dict[key] = value.value_in_unit(output_unit)
            unit_dict[unit_key] = output_unit

        return unitless_dict, unit_dict

    def serialize_numpy(np_array):
        """
        Serializes a numpy array into a JSON-compatible string. Leverages the numpy.save function,
        thereby preserving the shape of the input array

        from https://stackoverflow.com/questions/30698004/how-can-i-serialize-a-numpy-array-while-preserving-matrix-dimensions#30699208

        Parameters
        ----------
        np_array : A numpy array
            Input numpy array

        Returns
        -------
        serialized : str
            A serialized representation of the numpy array.
        shape : tuple of ints
            The shape of the serialized array
        """

        bigendian_array = np_array.newbyteorder('>')
        serialized = bigendian_array.tobytes()
        shape = np_array.shape
        return serialized, shape

    def deserialize_numpy(serialized_np, shape):
        """
        Deserializes a numpy array from a JSON-compatible string.

        from https://stackoverflow.com/questions/30698004/how-can-i-serialize-a-numpy-array-while-preserving-matrix-dimensions#30699208

        Parameters
        ----------
        serialized_np : str
            A serialized numpy array
        shape : tuple of ints
            The shape of the serialized array
        Returns
        -------
        np_array : numpy.ndarray
            The deserialized numpy array
        """

        import numpy as np
        dt = np.dtype('float')
        dt.newbyteorder('>')  # set to big-endian
        np_array = np.frombuffer(serialized_np, dtype=dt)
        np_array = np_array.reshape(shape)
        return np_array