"""
Wrapper classes and functions for dealing with the MMKV library, which is
used to read & write Epik, LigPrep, and ConfGen input files.
See MMKVSettings documentation for more details.
Copyright Schrodinger, LLC. All rights reserved.
"""
from schrodinger.infra import mm
from schrodinger.infra import mminit
_initializer = mminit.Initializer([mm.mmkv_initialize], [mm.mmkv_terminate])
# The following functions convert the value passed as argument to
# an array of the appropriate type. Since this value will have
# only one element, conversions are straightforward
# NOTE: These functions are also used by mmim.py
def _convert_to_float_array(value):
    try:
        return [float(value)]
    except ValueError:
        # FIXME do we really want to hide any errors this way?
        return []
def _convert_to_string_array(value):
    return [value]
def _convert_to_int_array(value):
    try:
        return [int(value)]
    except ValueError:
        # FIXME do we really want to hide any errors this way?
        return []
def _convert_to_bool_array(value):
    value = str(value).lower()
    if value in ['yes', 'true', '1']:
        return [True]
    else:
        if value not in ['no', 'false', '0']:
            # FIXME We might want to raise an exception for invalid strings
            print(
                'WARNING: mmim.py:_convert_to_bool_array() invalid bool: "%s"' %
                value)
        return [False]
def _convert_to_bool(value):
    """
    This function will convert a boolean represented as string to Python
    "boolean" type.
    """
    try:
        # Treat YES/NO/TRUE/FALSE strings specially:
        value = value.lower()
        if value in ['yes', 'true', '1']:
            value = True
        elif value in ['no', 'false', '0']:
            value = False
        else:
            # FIXME We might want to raise an exception for invalid strings
            print('WARNING: mmim.py:_convert_to_bool() invalid bool: "%s"' %
                  value)
            value = bool(value)
    except AttributeError:
        # <value> is not of string type (may be bool already)
        value = bool(value)
    return value
#
# MMKV stores a large enumerated list of keywords by data type, and these
# are surrounded by MMKV_ARG_*_BEGIN and MMKV_ARG_*_END values. (See the
# mmlibs/mmkv/mmkv_def.h header file if this doesn't make sense to you.)
# In order for this module to call the right functions in setting and
# getting keyword values, the following dict indexes the various functions
# by the MMKV_ARG_*_END value.
#
# This dict stores a tuple of:
#  (
#   beginning MMKV index for the arg type,
#   string representation of the type,
#   converter function,
#   getter function,
#   setter function,
#   array size get function,
#   array size set function,
#  )
#
# The array size get/set functions are None if the type isn't an array type.
#
_mmkv_arg_handlers = {
    mm.MMKV_ARG_BOOL_END: (
        mm.MMKV_ARG_BOOL_BEGIN,
        "boolean",
        _convert_to_bool,
        mm.mmkv_handle_get_arg_bool,
        mm.mmkv_handle_set_arg_bool,
        None,
        None,
    ),
    mm.MMKV_ARG_BOOL_ARRAY_END: (
        mm.MMKV_ARG_BOOL_ARRAY_BEGIN,
        "boolean array",
        _convert_to_bool_array,
        mm.mmkv_handle_get_indexed_arg_bool,
        mm.mmkv_handle_set_indexed_arg_bool,
        mm.mmkv_handle_get_indexed_arg_bool_count,
        mm.mmkv_handle_set_indexed_arg_bool_count,
    ),
    mm.MMKV_ARG_FLOAT_END: (
        mm.MMKV_ARG_FLOAT_BEGIN,
        "float",
        float,
        mm.mmkv_handle_get_arg_float,
        mm.mmkv_handle_set_arg_float,
        None,
        None,
    ),
    mm.MMKV_ARG_FLOAT_ARRAY_END: (
        mm.MMKV_ARG_FLOAT_ARRAY_BEGIN,
        "float array",
        _convert_to_float_array,
        mm.mmkv_handle_get_indexed_arg_float,
        mm.mmkv_handle_set_indexed_arg_float,
        mm.mmkv_handle_get_indexed_arg_float_count,
        mm.mmkv_handle_set_indexed_arg_float_count,
    ),
    mm.MMKV_ARG_INT_END: (
        mm.MMKV_ARG_INT_BEGIN,
        "integer",
        int,
        mm.mmkv_handle_get_arg_int,
        mm.mmkv_handle_set_arg_int,
        None,
        None,
    ),
    mm.MMKV_ARG_INT_ARRAY_END: (
        mm.MMKV_ARG_INT_ARRAY_BEGIN,
        "integer array",
        _convert_to_int_array,
        mm.mmkv_handle_get_indexed_arg_int,
        mm.mmkv_handle_set_indexed_arg_int,
        mm.mmkv_handle_get_indexed_arg_int_count,
        mm.mmkv_handle_set_indexed_arg_int_count,
    ),
    mm.MMKV_ARG_STRING_END: (
        mm.MMKV_ARG_STRING_BEGIN,
        "string",
        lambda x: x,  # Return without modifying
        mm.mmkv_handle_get_arg_string,
        mm.mmkv_handle_set_arg_string,
        None,
        None,
    ),
    mm.MMKV_ARG_STRING_ARRAY_END: (
        mm.MMKV_ARG_STRING_ARRAY_BEGIN,
        "string array",
        _convert_to_string_array,
        mm.mmkv_handle_get_indexed_arg_string,
        mm.mmkv_handle_set_indexed_arg_string,
        mm.mmkv_handle_get_indexed_arg_string_count,
        mm.mmkv_handle_set_indexed_arg_string_count,
    ),
}
[docs]def get_handlers_for_key(key, arg_handlers):
    """
    Find the data for the specified key in the given argument handlers
    dictionary. Returns a tuple of the converter, getter, and setter
    functions.
    Will raise a KeyError if the wrong key value (or type) is specified,
    because the exception is passed onto the getter/setter methods.
    NOTE: Also used in mmim.py
    """
    arg_type_endpoint = None
    for endpoint in sorted(arg_handlers):
        if key < endpoint:
            arg_type_endpoint = endpoint
            break
    if arg_type_endpoint is None:
        raise KeyError("Argument key %d is of an unsupported type." % key)
    handlers = arg_handlers[arg_type_endpoint]
    if key <= handlers[0]:
        raise KeyError("Argument key %d is not in the "
                       "supported %s range (%d, %d)" %
                       (key, handlers[1], arg_type_endpoint, handlers[0]))
    return handlers 
[docs]class MMKVArgList:
    """
    A class to provide list-like access to MMKV/MMIM array properties.
    Like an actual Python list, indices start at 0, not 1. Unlike a Python
    list, at this time negative indices and slices are not supported.
    Note that while this class requires an MMKV/MMIM handle, it neither
    initiates nor terminates the mmkv library.
    """
[docs]    def __init__(self, handle, key, converter, getter, setter, get_len,
                 set_len):
        """
        Create an instance from the MMKV/MMIM handle and key (int or string).
        Takes functions to get and set indexed values, and a function to
        determine the length. If these are not provided, they are looked up
        based on the key provided.
        get_len and set_len methods should be defined for array types only.
        """
        # FIXME add ability to initialize from a Python list of strings, ints,
        # floats, or booleans.
        self.handle = handle
        self.key = key
        self._converter = converter
        self._getter = getter
        self._setter = setter
        self._get_len = get_len
        self._set_len = set_len 
    def __getitem__(self, index):
        try:
            return self._getter(self.handle, self.key, index + 1)
        except mm.MmException:
            # If an exception is raised, check to see if it was due to an
            # index error. If it wasn't, just reraise the exception.
            if index < 0 or index >= len(self):
                raise IndexError("List index out of range (index "
                                 "starts at zero)")
            else:
                raise
    def __setitem__(self, index, value):
        # Ev:67229 Attempt to convert the value to float in case it was
        # specified as a string if setting indexed_arg_float type:
        value = self._converter(value)[0]  # Take the first element of the array
        try:
            return self._setter(self.handle, self.key, index + 1, value)
        except mm.MmException:
            # If an exception is raised, check to see if it was due to an
            # index error. If it wasn't, just re-raise the exception.
            if index < 0 or index >= len(self):
                raise IndexError("List index out of range (index "
                                 "starts at zero)")
            else:
                raise
[docs]    def __len__(self):
        return self._get_len(self.handle, self.key) 
[docs]    def setSize(self, size):
        """
        Set the size of the MMKV argument list.
        """
        set_count_function = self._set_len
        return set_count_function(self.handle, self.key, size) 
[docs]    def append(self, value):
        """
        Append a new value to the current MMKV arglist.
        :param value: The new value to append.
        :type  value: str, int, bool, or float depending on the MMKV arglist
                type
        """
        size = len(self) + 1
        self.setSize(size)
        self.__setitem__(size - 1, value) 
    def __iter__(self):
        list_len = len(self)
        for i in range(list_len):
            yield self.__getitem__(i)