Source code for schrodinger.application.glide.utils

"""
Glide utility functions.

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

import os
import zipfile

import schrodinger.structure as structure
from schrodinger.infra import mm


def is_valid_pv_file(file_name, require_poses=False):
    """
    Returns bool indicating if the file appears to be a valid PV file.

    This function with throw an exception is the file is not a valid Maestro
    structure file.

    The checks for this are simple, and may not be conclusive. The test is
    content, and not file-extension based, since there are valid PV or EPV files
    that lack any indication that they are PV files in the file extension. If
    you need to check the extension as well see fileutils.is_poseviewer_file.

    A PV file has exactly one structure with a true value for b_glide_receptor.

    An EPV file has one or more structures with a true value for
    b_glide_receptor at the beginning of the file. Single-receptor EPV files are
    therefore considered PV files, but not EPV files with two or more receptors.

    :param file_name: path to the structure file
    :type file_name: str

    :param require_poses: pv file must contain more than one structure
    :type require_poses: bool

    :return: bool indicating if the file appears to be a valid PV file
    :rtype: bool
    """
    reader = structure.StructureReader(file_name)

    try:
        first_st = next(reader)
    except StopIteration:
        # There are no structures
        return False

    if not first_st.property.get("b_glide_receptor"):
        # The first structure is not a receptor
        return False

    try:
        second_st = next(reader)
    except StopIteration:
        return not require_poses

    # If there is a second structure, we rule out the possibility that this is
    # an EPV file with multiple receptors.
    return not second_st.property.get("b_glide_receptor", False)


def is_valid_epv_file(file_name):
    """
    Returns bool indicating if the file appears to be a valid EPV file.

    :param file_name: path to the structure file
    :type file_name: str

    :return: bool indicating if the file appears to be a valid EPV file
    :rtype: bool
    """
    reader = structure.StructureReader(file_name)

    try:
        first_st = next(reader)
    except StopIteration:
        # There are no structures
        return False

    if not first_st.property.get("b_glide_receptor"):
        # The first structure is not a receptor
        return False
    if not first_st.property.get("i_epv_receptor"):
        # EPV format requires receptor to have i_epv_receptor set
        return False

    try:
        second_st = next(reader)
    except StopIteration:
        # There are no more structures
        return False

    # If the second structure is a receptor or epv marked ligand
    # then this is an EPV file.
    is_receptor = second_st.property.get("b_glide_receptor", False)
    is_epv_ligand = second_st.property.get("i_epv_best_receptor", False)
    return bool(is_receptor or is_epv_ligand)


def get_recep_structure_from_grid(gridfile):
    """
    Return a Structure object given a grid file (may be .grd or .zip)
    """

    if zipfile.is_zipfile(gridfile):
        zipobj = zipfile.ZipFile(gridfile, 'r')
        for name in zipobj.namelist():
            if name.endswith('_recep.mae'):
                recep_string = zipobj.read(name).decode('utf-8')
                reader = structure.StructureReader.fromString(recep_string)
                return next(reader)
        raise IOError("*_recep.mae not found in '%s'" % gridfile)
    else:
        root, ext = os.path.splitext(gridfile)
        recep_mae = "%s_recep.mae" % root
        return structure.Structure.read(recep_mae)


def is_grid_good_for_ligand(gridfile, lig_st) -> bool:
    """
    Check whether the ligand structure fits in the outer box of the grid.

    :param gridfile: Filename of gridfile
    :type gridfile: str

    :param lig_st: Ligand structure
    :type lig_st: schrodinger.structure.Structure
    """
    try:
        grid = get_recep_structure_from_grid(gridfile)
    except IOError:
        return False
    lig_xyz = lig_st.getXYZ()
    for axis_idx, axis in enumerate("xyz"):
        box_center = grid.property[f"r_glide_gridbox_{axis}cent"]
        box_size = grid.property[f"r_glide_gridbox_{axis}range"]
        half_size = box_size / 2
        box_min = box_center - half_size
        lig_axis = lig_xyz[:, axis_idx]
        if lig_axis.min() < box_min:
            return False
        box_max = box_center + half_size
        if lig_axis.max() > box_max:
            return False
    return True


def check_required_gridfiles(gridfile):
    """
    Check whether all required gridfiles exist for the given uncompressed gridfile

    :param gridfile: Path to the grid file (.grd)
    :type gridfile: str
    """
    # prevent glide-src dependency at import time
    from schrodinger.application.glide.packages import driver
    grid_base, _ = os.path.splitext(gridfile)
    # Check all of the required grid files
    grid_files = [grid_base + ext for ext in driver.REQUIRED_GRID_EXTS]
    return all(os.path.isfile(fn) for fn in grid_files)


def extract_file_from_grid(grid_file, ext):
    """
    Extract the file with requested extension form the given grid file archive.
    Will return None if the requested file is not present in the archive.
    If the specified grid file is uncompressed, will simply return the file
    with the same extension in the same location as the grid.

    :type grid_file: str
    :param grid_file: Path to the grid file (.zip or .grd)

    :type ext: str
    :param ext: Extension of the requested file.

    :rtype: str or None
    :return: File path or None
    """
    gridbase, grid_ext = os.path.splitext(grid_file)

    if grid_ext == ".zip":
        mm.mmim_initialize(mm.error_handler)
        try:
            return mm.mmim_glide_extract_gridfile(grid_file, ext)
        except:
            return None
        finally:
            mm.mmim_terminate()
    elif grid_ext == ".grd":
        filename = gridbase + ext
        if os.path.isfile(filename):
            return filename
        else:
            return None
    else:
        raise ValueError("Invalid grid file extension: %s" % grid_ext)


def parse_xvol_file(xvol_file):
    """
    Returns a list of excluded volumes in the specified excluded volumes file.

    :type xvol_file: str
    :param xvol_file: Path to the excluded volumes file

    :rtype: List of tuples
    :return: Each tuple defines an excluded volume. First item is the name (str),
            second item is a tuple of XYZ coordinates, third item is a radius (float).
    """

    if not xvol_file:
        return []

    mm.mmim_initialize(mm.error_handler)
    xs, ys, zs, radii, ignored, names = mm.mmim_read_exvol_spheres(xvol_file)
    mm.mmim_terminate()

    volumes = []
    for i, name in enumerate(names):
        center = (xs[i], ys[i], zs[i])
        radius = radii[i]
        volumes.append((name, center, radius))

    return volumes