Source code for schrodinger.application.matsci.maestrocmds

"""
Common Maestro commands.

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

import contextlib
import random

import seaborn as sns

import schrodinger

maestro = schrodinger.get_maestro()

XKCD = 'xkcd'
CRAYON = 'crayon'
PALETTES = {
    XKCD: (sns.xkcd_palette, sns.xkcd_rgb),
    CRAYON: (sns.crayon_palette, sns.crayons)
}


def get_colors(palette=None, seed=None, n_colors=10, normalized=False):
    """
    Return RGB color tuples.

    :type palette: str
    :param palette: if not a palette in PALETTES then this is the
        palette kwarg passed to seaborn.color_palette

    :type seed: None or int
    :param seed: if given randomize the RGB color tuples using the
        given seed

    :type n_colors: int
    :param n_colors: the number of colors to return

    :type normalized: bool
    :param normalized: whether to normalize the RGB tuples

    :rtype: list
    :return: contains RGB color tuples
    """

    pair = PALETTES.get(palette)
    if pair:
        rgb_getter, color_dict = pair
        rgbs = rgb_getter(color_dict.keys())
    else:
        rgbs = sns.color_palette(palette=palette, n_colors=n_colors)
    if seed:
        myrandom = random.Random(seed)
        myrandom.shuffle(rgbs)
    rgbs = rgbs[:n_colors]
    if not normalized:
        rgbs = [tuple(int(255 * item) for item in rgb) for rgb in rgbs]
    return rgbs


def _get_color_set_1(seed=123, n_colors=20):
    """
    Return RGB color tuples for set 1.

    :type seed: None or int
    :param seed: if given randomize the RGB color tuples using the
        given seed

    :type n_colors: int
    :param n_colors: the number of colors to return

    :rtype: list
    :return: contains RGB color tuples
    """

    # the following gives 10 unique colors (this is the maximum number
    # of unique colors avilable with the given options before it starts
    # reusing colors)
    rgbs = get_colors(palette=None, seed=None, n_colors=10)

    n_colors = n_colors - 10
    if n_colors > 0:
        rgbs += get_colors(palette=CRAYON, seed=seed, n_colors=n_colors)

    return rgbs


def normalize_rgb_color(rgb):
    """
    Normalize the given color.

    :type rgb: tuple
    :param rgb: an RGB triple of integers in [0, 255]

    :rtype: tuple
    :return: normalized RGB triple of floats in [0, 1]
    """

    return tuple(x / 255 for x in rgb)


def get_atoms_asl(idxs):
    """
    Return the atoms ASL.

    :type idxs: list(int)
    :param idxs: atom indices

    :rtype: str
    :return: the atoms ASL
    """

    idxs_str = ','.join(map(str, idxs))
    return f'atom.n {idxs_str}'


def execute_cmd(cmd):
    """
    Execute the Maestro command.

    :type cmd: str
    :param cmd: the Maestro command
    """

    if maestro:
        maestro.command(cmd)


def color_atoms(rgb, idxs):
    """
    Color atoms in the Maestro WS.

    :type rgb: tuple
    :param rgb: an RGB triple of integers in [0, 255]

    :type idxs: list(int)
    :param idxs: atom indices
    """

    cmd = 'coloratomrgb red=%s green=%s blue=%s' % rgb
    asl = get_atoms_asl(idxs)
    execute_cmd(f'{cmd} {asl}')


def color_atom_groups(all_idxs, rgbs=None):
    """
    Color groups of atoms in the Maestro WS.

    :type all_idxs: list(list(int))
    :param all_idxs: groups of atom indices

    :type rgbs: list(tuple) or None
    :param rgbs: RGB triples of integers in [0, 255]
        for coloring the different groups of atoms, if
        None then some defaults colors are used
    """

    if not rgbs:
        # seed of 4 gives decent colors
        rgbs = _get_color_set_1(seed=4, n_colors=len(all_idxs))

    assert len(all_idxs) <= len(rgbs)

    for idxs, rgb in zip(all_idxs, rgbs):
        color_atoms(rgb, idxs)


@contextlib.contextmanager
def select_all_atoms():
    """
    Context manager to temporarily select all atoms in
    the Maestro WS.
    """

    execute_cmd('workspaceselectionclear')
    execute_cmd('workspaceselectionreplacenoundo all')
    yield
    execute_cmd('workspaceselectionclearnoundo')
    execute_cmd('workspaceselectionclear')


def select_atoms(index):
    """
    Select specified atoms in Maestro WS

    :type index: list
    :param index: list of atom indexes
    """
    execute_cmd('workspaceselectionreplace atom.n ' + ','.join(map(str, index)))


def color_by_element():
    """
    Color all atoms in the Maestro WS by element.
    """

    with select_all_atoms():
        execute_cmd('colorscheme scheme="Element" at.selected')


def set_ballnstick_repr(idxs):
    """
    Set the representation of atoms in the Maestro WS to ball-and-stick.

    :type idxs: list(int)
    :param idxs: atom indices
    """

    maestro.command('repatombonds all')
    cmd = 'repatom  rep=ballnstick'
    asl = get_atoms_asl(idxs)
    execute_cmd(f'{cmd} {asl}')


def set_included(entry_id):
    """
    Set the included entry to that of the given entry ID.

    :type entry_id: str
    :param entry_id: the entry ID
    """

    execute_cmd(f'entrywsincludeonly entry {entry_id}')