Source code for schrodinger.application.livedesign.import_utils

from . import constants
from . import ld_utils

# Structure properties
SD_ID = 's_sd_ID'

# LiveDesign Metadata keys
ROW_INFO_KEY = 'row_infos'
DISPLAY_ID_KEY = 'display_id'
ENTITY_ID_KEY = 'entity_id'


def cache_entity_IDs(sts, ld_id_dict):
    """
    Store the corporate ID (AKA the entity ID) as a structure property for
    each imported structure

    :param sts: List of structures to set properties for
    :type sts: list(structure.Structure)
    :param ld_id_dict: Live Design structure to entity ID dictionary
    :type ld_id_dict: Dict[Structure, str]
    """
    for st in sts:
        corp_id = ld_id_dict.get(st)
        ld_utils.safely_set_property(st, constants.PROPNAME_IMPORT_ENTITY_ID,
                                     corp_id)


def get_st_entity_id_map(ld_client, sts, lr_id):
    """
    Generate a dictionary mapping each compound's structure to its
    entity id for easier access.

    For LD versions >= 8.1:
    Every structure holds multiple LD IDs, out of which, one is the entity
    ID. We use the live report metadata to obtain the correct ID.

    :param ld_client: LiveReport client
    :type ld_client: LDClient
    :param sts: structures to get map for
    :type sts: structure.Structure
    :param lr_id: the live report id
    :type lr_id: str

    :return: dictionary mapping structure to entitiy id
    :rtype: `dict(structure.Structure, str)`
    """
    # Use metadata to get Entity IDs
    lr_results_metadata = ld_client.live_report_results_metadata(lr_id)
    row_info_list = lr_results_metadata[ROW_INFO_KEY]
    display_to_entity_id_dict = {
        row_info_dict[DISPLAY_ID_KEY].strip():
        row_info_dict[ENTITY_ID_KEY].strip() for row_info_dict in row_info_list
    }
    # Use get to access property as as some structures don't have corporate IDs
    st_to_entity_id = {
        st: display_to_entity_id_dict.get(st.property.get(SD_ID)) for st in sts
    }
    return st_to_entity_id