Source code for schrodinger.application.livedesign.io_utils

"""
LiveDesign protocol io utils

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

import csv

from random import choice

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'
]


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):
        try:
            title = st.title
        except KeyError:
            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 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)