"""
Utilities to load Maestro projects into PyMOL.
Example::
    >>> cmd = PymolInstance(["/path/to/pymol"])
    >>> pt = maestro.project_table_get()
    >>> process_prj(cmd, pt)
Copyright Schrodinger LLC, All rights reserved.
Author: Thomas Holder
"""
import math
import os
import re
import shlex
import subprocess
import sys
import tempfile
import warnings
import schrodinger
from schrodinger import project
from schrodinger.application.phase import pt_hypothesis
from schrodinger.infra import mm
from schrodinger.infra import mmproj
from schrodinger.infra import mmsurf
from schrodinger.structutils import analyze
from schrodinger.structutils.color import get_rgb_from_color_index
from schrodinger.utils import log
logger = log.get_output_logger("pymol4maestro")
maestro = schrodinger.get_maestro()
# prefixes for pymol names
PREFIX_MAPS = ''
PREFIX_SURFACES = ''
PREFIX_MEASUREMENTS = ''
PREFIX_CALLOUTS = ''
[docs]class Mapping:
    """
    Mappings from Maestro codes to PyMOL representations and settings
    """
    surface_cmd = {
        mmsurf.MMSURF_STYLE_SOLID: "isosurface",
        mmsurf.MMSURF_STYLE_MESH: "isomesh",
        mmsurf.MMSURF_STYLE_DOT: "isodot",
    }
    ramp_colors = {
        mmsurf.MMSURF_COLOR_RAMP_REDWHITEBLUE: '[red, white, blue]',
        mmsurf.MMSURF_COLOR_RAMP_WHITEBLUE: '[white, blue]',
        mmsurf.MMSURF_COLOR_RAMP_WHITERED: '[white, red]',
        mmsurf.MMSURF_COLOR_RAMP_RAINBOW: '[red, yellow, green, cyan, blue, magenta]',
    }
    stereomethods = {
        'hardware': 'quadbuffer',
        'crosseyed': 'crosseye',
        'walleyed': 'walleye',
        'interlaced': 'byrow',
        'anaglyph': 'anaglyph',
        'chromadepth': 'off',
    } 
def _shlex_list(s) -> list:
    """
    Convert `s` to an argument list for subprocess calls
    :type s: str (deprecated) or iterable
    """
    if isinstance(s, str):
        warnings.warn("type str is deprecated, use a list or tuple of strings",
                      stacklevel=3)
        return shlex.split(s)
    return list(s)
[docs]class PymolInstance:
    """
    Represents a remote PyMOL instance (controlled via a one-way pipe)
    Acts like a proxy to the cmd module, without return values on function
    calls (which would actually be very usefull).
    See also: PyMOL XMLRPC server (pymol -R)
    """
[docs]    def __init__(self, pymol_command=("pymol",)):
        """
        :type pymol_command: list or tuple
        :param pymol_command: path to pymol executable
        """
        self._pymol_command = _shlex_list(pymol_command)
        self._initPipe()
        self.set('ignore_case', 0)
        self._hangingStdinWorkaround()
        self._used_names = set()
        self._group_names = {}
        self.row_pymol_names = {}
        self.sendVersionCheck() 
    def _initPipe(self):
        """
        Set up self._pipe
        """
        command = self._pymol_command + ["-pq"]
        self._pipe = subprocess.Popen(command,
                                      env=self._getEnviron(),
                                      stdin=subprocess.PIPE).stdin
    def _hangingStdinWorkaround(self):
        """
        workaround for PYMOL-234 PYMOL-246 PYMOL-508 PYMOL-510
        """
        self.set('suspend_updates')
        self.pseudoatom('_p', elem='C')
        self.label('_p', 'text_type')
        self.delete('_p')
        self.set('suspend_updates', 0)
    def _getEnviron(self):
        """
        The SCHRODINGER environment may be incompatible with PyMOL. This
        method provides a cleaned environment dictionary for the pymol
        subprocess.
        :rtype: dict
        """
        env = os.environ.copy()
        for env_var in [
                'PYTHONHOME',
                'PYMOL_PATH',
                'PYMOL_EXEC',
                'LD_LIBRARY_PATH',
                'DYLD_LIBRARY_PATH',
                'MACOSX_DEPLOYMENT_TARGET',
        ]:
            env.pop(env_var, None)
        # sort Schrodinger dirs to the end of PATH
        sch = env.get('SCHRODINGER', '')
        path = env.get('PATH', '').split(os.pathsep)
        path = sorted(path, key=lambda p: sch in p)
        env['PATH'] = os.pathsep.join(path)
        return env
    def __getattr__(self, name):
        def wrapper(*args, **kwargs):
            args = list(map(repr, args))
            kwargs = ['{}={}'.format(k, repr(v)) for (k, v) in kwargs.items()]
            return self.do('_ /cmd.{}({})'.format(name,
                                                  ', '.join(args + kwargs)))
        return wrapper
    _re_illegal = re.compile(r'[^-.\w]')
[docs]    def get_legal_name(self, name):
        """
        Replacement for cmd.get_legal_name
        :type name: str
        :param name: name candidate
        :return: legal PyMOL object name
        :rtype: str
        """
        return self._re_illegal.sub('_', name) 
[docs]    def get_unused_name(self, name, alwaysnumber=1):
        """
        Replacement for cmd.get_unused_name, does not talk back to PyMOL
        but maintains it's own set of already used names.
        This is only necessary because the the pipe cannot return values.
        :type name: str
        :param name: name candidate
        :type alwaysnumber: bool
        :param alwaysnumber: if False, only append a number if name already exists
        :return: unused legal PyMOL object name
        :rtype: str
        """
        name = self.get_legal_name(name)
        r, i = name, 1
        if alwaysnumber:
            r = name + '01'
        while r.lower() in self._used_names:
            r = name + '%02d' % i
            i += 1
        self._used_names.add(r.lower())
        return r 
    def _get_unused_group_name(self, group):
        """
        Special function to get a unique name for a group. In Maestro, the
        group "title" is diplayed and doesn't have to be unique.
        :type group: `schrodinger.project.EntryGroup`
        :param group: group instance or None
        """
        if not group:
            return None
        uniquekey = group.name
        try:
            return self._group_names[uniquekey]
        except KeyError:
            pass
        name = group.title or group.name
        name = self.get_unused_name(name, 0)
        self._group_names[uniquekey] = name
        return name
[docs]    def sendVersionCheck(self):
        """
        Print a warning on the PyMOL log window if PyMOL version is too old.
        """
        self.do(
            "_ /if cmd.get_version()[1]<1.61: "
            "print('Warning: PyMOL4Maestro requires PyMOL version 1.6.1 or later.')"
        ) 
[docs]    def do(self, cmmd):
        """
        Send command to PyMOL
        :type cmmd: str
        :param cmmd: PyMOL command
        :return: True on success and False on error
        :rtype: bool
        """
        if not isinstance(cmmd, bytes):
            cmmd = cmmd.encode('utf-8')
        try:
            self._pipe.write(cmmd + b"\n")
            self._pipe.flush()
        except OSError:
            return False
        return True 
[docs]    def close(self):
        """
        Quit PyMOL
        """
        self.do("quit")
        self._pipe.close()  
[docs]class PymolScriptInstance(PymolInstance):
    """
    Represents a PyMOL script for deferred execution.
    """
    def _initPipe(self):
        self._pipe = tempfile.NamedTemporaryFile(delete=False, suffix=".pml")
[docs]    def close(self, args=('-cqk',)):
        """
        Close file handle and execute script in PyMOL
        :type args: list or tuple
        :param args: extra command line arguments for pymol
        """
        self._pipe.close()
        args = _shlex_list(args)
        command = self._pymol_command + args + [self._pipe.name]
        return subprocess.call(command, env=self._getEnviron())  
[docs]class VisRecord:
    """
    Represents a surface entry in a "vis_list" file
    :vartype name_pymol: str
    :ivar name_pymol: PyMOL object name
    :vartype visfile: str
    :ivar visfile: filename of vis file
    """
    _prefix = 'surf'
[docs]    def __init__(self, row, idx):
        """
        :type row: RowProxy
        :param row: project table row
        :type idx: int
        :param idx: zero-based index in "m_surface" table
        """
        self.cmd = row.cmd
        self.id = mm.m2io_get_int_indexed(row.vis_file, idx + 1, ["i_m_id"])[0]
        self.name = mm.m2io_get_string_indexed(row.vis_file, idx + 1,
                                               ["s_m_name"])[0]
        self.visfile = os.path.join(row.additional_data_dir,
                                    "%s%d.vis" % (self._prefix, self.id))
        if self._prefix == 'surf':
            v = mm.m2io_get_string_indexed(row.vis_file, idx + 1,
                                           ["s_m_volume_name"])
            self.volume_name = v[0] if v else "" 
    def __getattr__(self, key):
        if key == 'name_pymol':
            self._load()
            return self.name_pymol
        raise AttributeError(key)
    def _load(self):
        """
        Assign *name_pymol*
        """
        self.name_pymol = self.cmd.get_unused_name(
            PREFIX_SURFACES + (self.name or 'surf'), 0) 
[docs]class VisRecordVol(VisRecord):
    """
    Represents a volume entry in a "vis_list" file
    Volume gets auto-loaded when accessing *name_pymol*.
    """
    _prefix = 'vol'
    def _load(self):
        """
        Assign *name_pymol* and load the map in PyMOL.
        """
        self.name_pymol = self.cmd.get_unused_name(
            PREFIX_MAPS + (self.name or 'map'), 0)
        self.cmd.load(self.visfile, self.name_pymol, mimic=0) 
[docs]class RowProxy(project.ProjectRow):
    """
    Proxy for project table row to attach additional data.
    """
[docs]    def __init__(self, row, cmd):
        """
        :type row: `schrodinger.project.ProjectRow`
        :param row: project table row
        :type cmd: `PymolInstance`
        :param cmd: PyMOL API proxy
        """
        super().__init__(row._pt, row.index)
        self.cmd = cmd
        self.volume_recs = {}
        self.surface_recs = {}
        self.rep_surface = False
        self.vis_file = None
        self.additional_data_dir = mmproj.mmproj_index_entry_get_additional_data_dir(
            row._project_handle, row.index)
        if self.additional_data_dir:
            # for surfaces and volumes
            filename = os.path.join(self.additional_data_dir, "vis_list")
            if os.path.exists(filename):
                self.vis_file = mm.m2io_open_file(filename, mm.M2IO_READ)
        self.group_pymol = cmd._get_unused_group_name(row.group)
        self.name_pymol = cmd.get_unused_name(
            row.title or ('entry_%s' % row.entry_id), 0)
        cmd.row_pymol_names[row.index] = self.name_pymol 
    def __del__(self):
        if self.vis_file is not None:
            mm.m2io_close_file(self.vis_file)
[docs]    def doGroup(self, name):
        """
        Put *name* in PyMOL group, if row is in a Maestro group.
        :type name: str
        :param name: PyMOL object name
        """
        if self.group_pymol:
            self.cmd.group(self.group_pymol, name)
            self.cmd.enable(self.group_pymol)  
[docs]def select_surf_asl(row, surf_handle, name=''):
    """
    Make a PyMOL selection for surface ASL.
    :return: PyMOL selection name
    """
    asl = mmsurf.mmsurf_get_asl(surf_handle)
    if not asl:
        return ""
    asl_limit = (mmsurf.mmsurf_get_use_view_by(surf_handle) and
                 mmsurf.mmsurf_get_view_by_asl(surf_handle))
    if asl_limit:
        distance = mmsurf.mmsurf_get_viewing_distance(surf_handle)
        asl = f"({asl}) and within {distance} ({asl_limit})"
    # FIXME: FATAL search_mol(): error getting entry name for atom: 1
    # 'evaluate_asl' operates on single structure and does not support
    # global 'entry.id' and 'entry.name' selectors
    asl = re.sub(r'\bentry\.id\s+%s\b' % row.entry_id, 'all', asl)
    asl = re.sub(r'\bentry\.name\s+"([^"\\]|\\.)*"', 'all', asl)
    atom_index_list = analyze.evaluate_asl(row.getStructure(False, False), asl)
    if not atom_index_list:
        return ""
    if not name:
        name = row.cmd.get_unused_name('_mae_asl')
    atom_index_list = [i - 1 for i in atom_index_list]  # rank is 0-indexed
    row.cmd.select_list(name, row.name_pymol, atom_index_list, mode='rank')
    return name 
[docs]class WorkspaceIdMapper:
    """
    Maps workspace atom indices to (row.index, ID)
    """
[docs]    def __init__(self, prj_handle):
        """
        :type prj_handle: `schrodinger.project.Project`
        :param prj_handle: project handle
        """
        self.entries = []
        N = M = 0
        row_totals = [(row.index, row.getStructure(False, False).atom_total)
                      for row in prj_handle.included_rows]
        for row_index, total in sorted(row_totals):
            M += total
            self.entries.append((N, M, row_index))
            N = M 
    def __getitem__(self, i):
        for (N, M, row_index) in self.entries:
            if N < i <= M:
                return (row_index, i - N)
        raise LookupError 
[docs]def get_measurement_items(key, mmprojadmin):
    """
    Get workspace atom ids from the measurements table. If not running from
    Maestro, read the .tab files from the .mmproj-admin directory.
    :type key: str
    :param key: one of distance, angle or dihedral
    :rtype: list(int)
    :return: List of lists of atom ids (workspace)
    """
    try:
        return [x.split() for x in maestro.get_command_items(key)]
    except schrodinger.MaestroNotAvailableError:
        pass
    try:
        natom = {"distance": 2, "angle": 3, "dihedral": 4}[key]
    except KeyError:
        raise ValueError(key)
    basename = key if key != "dihedral" else "torsion"
    filename = os.path.join(mmprojadmin, basename + ".tab")
    if not os.path.exists(filename):
        return []
    measure_file = mm.m2io_open_file(filename, mm.M2IO_READ)
    mm.m2io_goto_next_block(measure_file, "f_m_table")
    mm.m2io_goto_next_block(measure_file, "m_row")
    index_dim = mm.m2io_get_index_dimension(measure_file)
    func, fmt = mm.m2io_get_string_indexed, "s_m_%s_atom%d"
    props = [fmt % (key, i + 1) for i in range(natom)]
    items = [func(measure_file, i + 1, props) for i in range(index_dim)]
    mm.m2io_close_file(measure_file)
    # trim element prefix
    items = [[i.rsplit(':')[-1] for i in ids] for ids in items]
    return items 
[docs]def process_measurements(cmd, prj_handle):
    """
    Send workspace measurements to PyMOL
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type prj_handle: `schrodinger.project.Project`
    :param prj_handle: project handle
    """
    idmapper = WorkspaceIdMapper(prj_handle)
    mmprojadmin = os.path.join(prj_handle.fullname, ".mmproj-admin")
    for key, color in [
        ("distance", "magenta"),
        ("angle", "green"),
        ("dihedral", "red"),
    ]:
        items = get_measurement_items(key, mmprojadmin)
        if not items:
            continue
        func = getattr(cmd, key)
        measure_name = cmd.get_unused_name(PREFIX_MEASUREMENTS + key)
        for atom_ids in items:
            try:
                atoms = [idmapper[int(i)] for i in atom_ids]
                atoms = [(cmd.row_pymol_names[r], a) for (r, a) in atoms]
            except (LookupError, ValueError):
                logger.error('workspace mapping failed for %s' % key)
                continue
            selections = ['%s & id %d' % (name, i) for (name, i) in atoms]
            func(measure_name, *selections)
        cmd.color(color, measure_name) 
[docs]def get_font_id(font_name, font_style):
    """
    Get the PyMOL label_font_id which best matches the given font name and style.
    :type font_name: str
    :type font_style: int
    :rtype: int
    """
    if 'Serif' in font_name and 'Sans' not in font_name or 'Times' in font_name:
        return 10 if font_style == 2 else \
               
17 if font_style == 3 else \
               
18 if font_style == 4 else \
               
9
    elif 'Mono' in font_name or 'Courier' in font_name:
        return 13 if font_style == 2 else \
               
12 if font_style == 3 else \
               
14 if font_style == 4 else \
               
11
    else:
        return 7 if font_style == 2 else \
               
6 if font_style == 3 else \
               
8 if font_style == 4 else \
               
5 
[docs]def process_highlights(cmd, prj_handle):
    """
    Send "highlights" (label+arrow annotation) to PyMOL
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type prj_handle: `schrodinger.project.Project`
    :param prj_handle: project handle
    """
    filename = os.path.join(prj_handle.fullname, ".mmproj-admin", "highlights")
    if not os.path.exists(filename):
        return
    handle = mm.m2io_open_file(filename, mm.M2IO_READ)
    mm.m2io_goto_next_block(handle, "f_m_table")
    mm.m2io_goto_next_block(handle, "m_row")
    index_dim = mm.m2io_get_index_dimension(handle)
    get_s = lambda *a: mm.m2io_get_string_indexed(handle, idx, list(a))
    get_i = lambda *a: mm.m2io_get_int_indexed(handle, idx, list(a))
    get_r = lambda *a: mm.m2io_get_real_indexed(handle, idx, list(a))
    get_b = lambda *a: mm.m2io_get_boolean_indexed(handle, idx, list(a))
    for idx in range(1, index_dim + 1):
        name, method, text, arrow_asl, asl = get_s("s_mhigh_Name",
                                                   "s_mhigh_Method",
                                                   "s_mhigh_Text",
                                                   "s_mhigh_Arrow_ASL",
                                                   "s_mhigh_ASL")
        show, = get_b("b_mhigh_Show")
        xy_text = get_r("r_mhigh_X_Text", "r_mhigh_Y_Text")
        xyz_head = get_r("r_mhigh_X_Head", "r_mhigh_Y_Head", "r_mhigh_Z_Head")
        rgb_text = get_r("r_mhigh_Text_Red", "r_mhigh_Text_Green",
                         "r_mhigh_Text_Blue")
        rgb_arrow = get_r("r_mhigh_Arrow_Red", "r_mhigh_Arrow_Green",
                          "r_mhigh_Arrow_Blue")
        bg_type, = get_i("i_mhigh_Text_Background_Type")
        font_size, = get_r("r_mhigh_Text_Font_Size")
        font_name, = get_s('s_mhigh_Text_Font_Name')
        font_style, = get_i('i_mhigh_Text_Font_Style')
        if not (show and text and arrow_asl):
            continue
        name = cmd.get_unused_name(PREFIX_CALLOUTS + name, 0)
        if math.isnan(xyz_head[0]):
            xyz_head = None
        font_id = get_font_id(font_name, font_style)
        # TODO: replace with proper callout object once implemented in PyMOL
        cmd.do(r'''
_ /if not hasattr(cmd, 'callout'):\
_    def callout(name, label, pos, *a, **kw):\
_        cmd.pseudoatom(name, label=label, pos=pos)\
_    cmd.callout = callout
''')
        cmd.callout(name,
                    text,
                    xyz_head, [i * 2 - 1 for i in xy_text],
                    color='0x%02x%02x%02x' %
                    tuple(int(255 * i) for i in rgb_arrow))
        sele = '(last ' + name + ')'
        cmd.set('label_connector_width', 3, name)
        cmd.set('label_color',
                '0x%02x%02x%02x' % tuple(int(255 * i) for i in rgb_text), sele)
        cmd.set('label_size', font_size, name)
        cmd.set('label_font_id', font_id, name)
        if bg_type == 0:
            cmd.set('label_bg_transparency', 0.3, name)
            cmd.set('label_bg_color', 'back', name)
        else:
            cmd.set('label_bg_transparency', 1.0, name)
        # global setting
        cmd.set('float_labels')
    mm.m2io_close_file(handle) 
[docs]def process_prj(cmd, prj_handle, limit="all", with_surf=True, mimic=True):
    """
    Send maestro project to PyMOL. By default send everything, optional
    filters may apply.
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type prj_handle: `schrodinger.project.Project`
    :param prj_handle: project handle
    :type limit: str
    :param limit: all, included or selected. The latter will not send
        workspace items like measurements and text highlights.
    :type with_surf: bool
    :param with_surf: send surfaces and maps (volumes)
    :type mimic: bool
    :param mimic: use PyMOL settings to match style as close as possible
    """
    if limit not in ["all", "included", "selected"]:
        raise ValueError(limit)
    # make sure files in additional_data_dir are up to date
    mmproj.mmproj_save(prj_handle.handle)
    cmd.wizard('message', 'processing...')
    cmd.set('suspend_updates')
    try:
        for row in prj_handle.all_rows:
            if (limit == "included" and not row.in_workspace or
                    limit == "selected" and not row.is_selected):
                continue
            process_row(cmd, row, limit == "included", with_surf, mimic)
        if limit != "selected":
            process_measurements(cmd, prj_handle)
            process_highlights(cmd, prj_handle)
    finally:
        cmd.set('suspend_updates', 0)
        cmd.wizard()
        cmd.refresh_wizard() 
[docs]def process_hypothesis(cmd, row, mae):
    """
    Import Phase pharmacophores as CGOs into PyMOL.
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type row: `schrodinger.project.ProjectRow`
    :param row: project table row
    :type mae: str
    :param mae: reference mae filename
    """
    if not pt_hypothesis.is_hypothesis_entry(row.entry_id):
        return
    tmpdir = tempfile.mkdtemp()
    try:
        hypo = pt_hypothesis.get_hypothesis_from_project(row.entry_id)
        sites = hypo.getHypoSites(True)
        # create old-style xyz files
        for with_Q, xyzfile in [
            (False, os.path.join(tmpdir, 'hypothesis.xyz')),
            (True, os.path.join(tmpdir, 'ALL.xyz')),
        ]:
            with open(xyzfile, 'w') as handle:
                for site in sites:
                    if not with_Q and site.getSiteTypeChar() == 'Q':
                        continue
                    handle.write(
                        '%d %s %f %f %f\n' %
                        (site.getSiteNumber(), site.getSiteTypeChar(),
                         site.getXCoord(), site.getYCoord(), site.getZCoord()))
        # load into PyMOL
        name = cmd.get_unused_name(row.name_pymol + '_hyp', 0)
        cmd.do('_ /import epymol.ph4')
        cmd.do('_ /epymol.ph4.load_hypothesis_xyz(%s, %s, %s)' %
               (repr(os.path.join(tmpdir, 'hypothesis.xyz')), repr(name),
                repr(os.path.join(tmpdir, 'ALL.xyz'))))
        # manage group
        row.doGroup(name)
        # load excluded volumes into PyMOL
        xvol_file = os.path.join(tmpdir, 'hypothesis.xvol')
        if hypo.visibleXvol():
            xvol = hypo.getXvol()
            xvol.exportToMMTableFile(xvol_file)
            name = cmd.get_unused_name(row.name_pymol + '_xvol', 0)
            cmd.do('_ /epymol.ph4.load_hypothesis_xvol(%s, %s)' %
                   (repr(xvol_file), repr(name)))
            # excluded volume not shown by default
            cmd.disable(name)
            row.doGroup(name)
    finally:
        cmd.do('_ /import shutil')
        cmd.do('_ /shutil.rmtree({})'.format(repr(tmpdir))) 
[docs]def process_row(cmd, row, limit_included=False, with_surf=True, mimic=True):
    """
    Send a row from the project table to PyMOL.
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type row: `schrodinger.project.ProjectRow`
    :param row: project table row
    :type limit_included: bool
    :param limit_included: limit surface export to workspace
    :type with_surf: bool
    :param with_surf: send surfaces
    :type mimic: bool
    :param mimic: use PyMOL settings to match style as close as possible
    """
    row = RowProxy(row, cmd)
    cmd.do("_ /import os")
    # load content
    tmp_mae = tempfile.mktemp(".mae")
    row.getStructure(props=True, copy=False).write(tmp_mae)
    cmd.load(tmp_mae, row.name_pymol, mimic=mimic)
    # set "ignore" flag for proper surface display (don't include
    # solvent and ligands when surfacing)
    cmd.flag('ignore', 'model %s & !polymer' % (row.name_pymol), 'set')
    # manage group
    row.doGroup(row.name_pymol)
    # pharmacophores
    process_hypothesis(cmd, row, tmp_mae)
    # done with mae file
    cmd.do('_ /os.unlink(%s)' % (repr(tmp_mae)))
    # load surfaces and volumes from vis files
    if not with_surf or not row.vis_file:
        return
    vis_file = row.vis_file
    mm.m2io_goto_next_block(vis_file, "f_m_vis_list")
    mm.m2io_goto_next_block(vis_file, "m_volume")
    index_dim = mm.m2io_get_index_dimension(vis_file)
    # Remember all volumes by name to load them on demand
    for idx in range(index_dim):
        vol = VisRecordVol(row, idx)
        row.volume_recs[vol.name] = vol
    mm.m2io_leave_block(vis_file)
    mm.m2io_goto_next_block(vis_file, "m_surface")
    index_dim = mm.m2io_get_index_dimension(vis_file)
    # map surface names to row indices
    for idx in range(index_dim):
        surf = VisRecord(row, idx)
        row.surface_recs[surf.name] = surf
    # loop over surfaces
    for e_surf in row.surfaces:
        if not limit_included or e_surf.included:
            process_surface(cmd, row, e_surf) 
[docs]def process_surface(cmd, row, e_surf):
    """
    Send a surface to PyMOL
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :type row: `RowProxy`
    :param row: project table row
    :type e_surf: `schrodinger.project.EntrySurface`
    :param e_surf: surface
    """
    import numpy
    try:
        surf = row.surface_recs[e_surf.name]
    except KeyError:
        logger.error('No VisRecord for surface "%s"' % e_surf.name)
        return
    cmd_color = cmd.color
    surf_handle = e_surf.surface_handle
    filename = surf.visfile
    color_scheme = mmsurf.mmsurf_get_color_scheme(surf_handle)
    styleint = mmsurf.mmsurf_get_style(surf_handle)
    style = Mapping.surface_cmd[styleint]
    isovalue = mmsurf.mmsurf_get_isovalue(surf_handle)
    strategy = "cgo"
    if surf.volume_name:
        strategy = None
        tmp_sele = select_surf_asl(row, surf_handle)
        isobuffer = 0.0
        isocarve = None
        if tmp_sele:
            isobuffer = mmsurf.mmsurf_get_viewing_distance(surf_handle)
            isocarve = isobuffer
        elif mmsurf.mmsurf_get_use_active_grid(surf_handle):
            tmp_sele = cmd.get_unused_name('_grid_center')
            cmd.pseudoatom(
                tmp_sele, pos=mmsurf.mmsurf_get_active_grid_center(surf_handle))
            isobuffer = mmsurf.mmsurf_get_active_grid_size(surf_handle) * 0.5
        # load into PyMOL as isosurface/mesh/dot
        func = getattr(cmd, style)
        vol = row.volume_recs[surf.volume_name]
        vol_name_pymol = vol.name_pymol
        func(surf.name_pymol, vol_name_pymol, isovalue, tmp_sele, isobuffer, 1,
             isocarve)
        if tmp_sele:
            cmd.delete(tmp_sele)
        row.doGroup(surf.name_pymol)
        row.doGroup(vol_name_pymol)
    elif mmsurf.mmsurf_get_surface_type(surf_handle) == "molecular surface" and \
            
not row.rep_surface:
        tmp_sele = select_surf_asl(row, surf_handle)
        if tmp_sele:
            strategy = "rep_surface"
    if strategy == "rep_surface":
        # first molecular surface as repr surface
        row.rep_surface = True
        if styleint:
            cmd.set('surface_type', 3 - styleint, row.name_pymol)
        cmd.flag('ignore', tmp_sele, 'clear')
        cmd.show('surface', tmp_sele)
        cmd.delete(tmp_sele)
        surf.name_pymol = row.name_pymol
        cmd_color = lambda *a: cmd.set('surface_color', *a)
    elif strategy == "cgo":
        # load into PyMOL as CGO
        cmd.load(filename, surf.name_pymol)
        row.doGroup(surf.name_pymol)
    # surface transparency
    cmd.set('transparency',
            mmsurf.mmsurf_get_transparency(surf_handle) * 0.01, surf.name_pymol)
    colorramp_name = None
    if color_scheme == "Color":
        color_rgb = mmsurf.mmsurf_get_rgb_color(surf_handle)
        cmd_color("0x%02x%02x%02x" % color_rgb, surf.name_pymol)
    elif color_scheme == "Grid Property":
        # Color by volume (grid property)
        vol_name = mmsurf.mmsurf_get_scheme_volume_name(surf_handle)
        vol = row.volume_recs[vol_name]
        vol_name_pymol = vol.name_pymol
        colorramp_name = mmsurf.mmsurf_get_colorramp_name(surf_handle)
        ramp_min = mmsurf.mmsurf_get_map_min(surf_handle)
        ramp_max = mmsurf.mmsurf_get_map_max(surf_handle)
    elif color_scheme == "Electrostatic Potential":
        # Color by ESP
        vol_name_pymol = cmd.get_unused_name(surf.name + '_esp', 0)
        colorramp_name = mmsurf.mmsurf_get_esp_colorramp_name(surf_handle)
        ramp_min = mmsurf.mmsurf_get_esp_min(surf_handle)
        ramp_max = mmsurf.mmsurf_get_esp_max(surf_handle)
        cmd.set('coulomb_cutoff', 4.0)
        cmd.set('coulomb_units_factor', 1.0)
        cmd.set('surface_ramp_above_mode', 0, surf.name_pymol)
        cmd.map_new(vol_name_pymol, "coulomb_local", 2.0, row.name_pymol, 2.0)
        row.doGroup(vol_name_pymol)
    else:
        logger.info("color scheme unknown: %s" % color_scheme)
    if colorramp_name:
        # color with color ramp in PyMOL
        ramp_name = cmd.get_unused_name(vol_name_pymol + '_ramp', 0)
        ramp_color = Mapping.ramp_colors[colorramp_name]
        ramp_range = numpy.linspace(ramp_min, ramp_max,
                                    ramp_color.count(',') + 1).tolist()
        cmd.ramp_new(ramp_name, vol_name_pymol, ramp_range, ramp_color)
        cmd.disable(ramp_name)
        cmd_color(ramp_name, surf.name_pymol)
        row.doGroup(vol_name_pymol)
        row.doGroup(ramp_name) 
[docs]def send_maestro_settings(cmd):
    """
    Map Maestro settings to closest matching PyMOL settings.
    :type cmd: `PymolInstance`
    :param cmd: PyMOL API proxy
    :rtype: bool
    :return: True on success and False if Maestro is not available
    """
    fopt = {
        ("repall", "tuberadius"): 0.16,
        ("ribbon", "ribbonwidth"): 1.61,
        ("ribbon", "ribbonthick"): 0.15,
        ("ribbon", "thintubewidth"): 0.25,
    }
    try:
        for key in fopt:
            fopt[key] = float(maestro.get_command_option(*key))
    except schrodinger.MaestroNotAvailableError:
        logger.info("maestro not available, using default settings")
    else:
        stereopymol = 'off'
        stereomethod = ''
        if maestro.get_command_option('displayopt', 'stereo') == 'True':
            stereomethod = maestro.get_command_option('displayopt',
                                                      'stereomethod')
            try:
                stereopymol = Mapping.stereomethods[stereomethod]
            except KeyError:
                logger.info("unknown stereo method: %s" % stereomethod)
        cmd.stereo(stereopymol)
        # chromadepth
        cmd.set('chromadepth', 2 if stereomethod == 'chromadepth' else 0)
        # cartoon highlight color
        cmd.set(
            'cartoon_highlight_color', 'gray50' if maestro.get_command_option(
                "ribbon", "helixcolor") == 'twocolors' else 'default')
        # background color
        color_idx = int(maestro.get_command_option("displayopt", "bgcindex"))
        color_rgb = get_rgb_from_color_index(color_idx)
        cmd.bg_color("0x%02x%02x%02x" % color_rgb)
        # angle dependent transparency
        if maestro.get_command_option('displayopt',
                                      'angledependenttransparency') == 'True':
            cmd.set('ray_transparency_oblique', 1)
            cmd.set('ray_transparency_oblique_power', 2.5)
        else:
            cmd.set('ray_transparency_oblique', 0)
    ribbonwidth_half = fopt["ribbon", "ribbonwidth"] * 0.5
    thintubewidth_half = fopt["ribbon", "thintubewidth"] * 0.5
    cmd.set('stick_radius', fopt["repall", "tuberadius"])
    cmd.set('stick_h_scale', 1.0)
    cmd.set('cartoon_oval_length', ribbonwidth_half)
    cmd.set('cartoon_rect_length', ribbonwidth_half)
    cmd.set('cartoon_oval_width', fopt["ribbon", "ribbonthick"])
    cmd.set('cartoon_rect_width', fopt["ribbon", "ribbonthick"])
    cmd.set('cartoon_loop_radius', thintubewidth_half) 
if __name__ == '__main__':
    # standalone test
    usage = """Usage: %s <project.prjzip>""" % sys.argv[0]
    try:
        prj_zipfile = sys.argv[1]
    except:
        print(usage)
        sys.exit(1)
    cmd = PymolInstance(['pymol', '-xK'])
    mm.mmzip_initialize(mm.MMERR_DEFAULT_HANDLER)
    mmsurf.mmsurf_initialize(mm.MMERR_DEFAULT_HANDLER)
    mm.mmct_initialize(mm.MMERR_DEFAULT_HANDLER)
    mmproj.mmproj_initialize(mm.MMERR_DEFAULT_HANDLER)
    mm.m2io_initialize(mm.MMERR_DEFAULT_HANDLER)
    mmsurf.mmvisio_initialize(mm.MMERR_DEFAULT_HANDLER)
    mmsurf.mmvol_initialize(mm.MMERR_DEFAULT_HANDLER)
    # open the project zip file
    prj_handle, prj_path, prj_temp_path = project.open_project(prj_zipfile)
    # Undefined project handle can indicate that the version is not compatible
    # with currently used mmshare version
    if prj_handle is None:
        print("incompatible version")
        sys.exit(1)
    process_prj(cmd, prj_handle)