Source code for schrodinger.application.livedesign.io_utils

"""
LiveDesign protocol io utils

Copyright Schrodinger, LLC. All rights reserved.
"""

import csv

from random import choice

import schrodinger

from schrodinger import structure

from schrodinger.structutils import color

ALLOWED_COLOR_SCHEMES = [
    'element_customc_darkblue', 'element_customc_blue',
    'element_customc_aquamarine', 'element_customc_turquoise',
    'element_customc_springgreen_carbons', 'element_customc_darkgreen',
    'element_customc_green', 'element_customc_limegreen',
    'element_customc_yellowgreen', 'element_customc_yellow',
    'element_customc_orange', 'element_customc_maroon', 'element_customc_red',
    'element_customc_pink', 'element_customc_plum', 'element_customc_purple',
    'element_customc_bluepurple', 'element_customc_white'
]

SCH_VER = schrodinger.get_release_name()


def get_structure_ids(infile):
    """
    Create a set of titles for all structures found in the input file

    :param infile: Input SD file
    :type infile: str

    :return: Set containing all structure titles
    :rtype: set
    """

    ids = set()

    for st in structure.StructureReader(infile):
        title = st.title

        if not title:
            raise RuntimeError('Structure found with missing title')

        ids.add(title)

    return ids


def apply_random_color(st):
    """
    Apply a random color scheme to the structure

    :param st: Structure to be colored
    :type st: `schrodinger.structure.Structure`
    """

    color.apply_color_scheme(st, choice(ALLOWED_COLOR_SCHEMES))


def record_failures(results_list, failed_ids, prop_keys, task):
    """
    Create a dictionary with empty values for the corporate ID that failed a
    stage in a LiveDesign protocol and updated the list that is passed in.

    :param results_list: List of dictionaries where each dictionary contains the
        required data to write to the CSV file
    :type results_list: list

    :param failed_ids: Set of corporate IDs that failed a given task
    :type failed_ids: set

    :param prop_keys: List of properties that will be written out to the output
        CSV file
    :type prop_keys: list

    :param task: Name of the task/stage where the failed IDs come from
    :type task: str

    :return: Update list of dictionaries where each dictionary contains the
        required data to write to the output CSV file
    :rtype: list
    """

    for corp_id in failed_ids:
        results_dict = {}
        for i in prop_keys:
            results_dict[i] = ''

        results_dict['Corporate ID'] = corp_id
        results_dict['status'] = f'{SCH_VER}: Dropped by {task}'

        results_list.append(results_dict)

    # Return an updated list
    return results_list


def create_csv(results_list, csv_file):
    """
    Write a CSV file from a list of dictionaries. The output file name
    'results.csv', is required by LiveDesign so the module-level constant will
    always be used

    :param results_list: List of dictionaries where each dictionary contains the
        required data to write to the CSV file. Column headers will be take from
        the keys of the first dict
    :type results_list: list

    :param csv_file: CSV filename
    :type csv_file: str
    """

    fieldnames = results_list[0].keys()

    # Sort the list by the corporate ID to allow for consistent CSV file creation
    results_list = sorted(results_list, key=lambda x: x['Corporate ID'])

    with open(csv_file, 'w') as output_file:
        writer = csv.DictWriter(output_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(results_list)