Source code for schrodinger.application.phase.packages.shape_ligprep

'''
Thin LigPrep wrapper and related subroutines (for PHASE-2070).
'''

import json

from voluptuous import Invalid
from voluptuous import Schema

from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.infra import phase
from schrodinger.utils import fileutils
from schrodinger.utils import log

logger = log.get_output_logger('shape.shape_ligprep')

try:
    from schrodinger.application.macromodel.packages.ligprep3 import arguments
    from schrodinger.application.macromodel.packages.ligprep3 import driver
    from schrodinger.application.macromodel.packages.ligprep3 import ligprepworkflow

    split_structures = driver.split_structures
except ImportError:
    logger.error("LigPrep is not available")

# -----------------------------------------------------------------------------

LIGPREP_INPUT_OPTION = {
    fileutils.SD: '-isd',
    fileutils.MAESTRO: '-imae',
    fileutils.SMILES: '-ismi',
    fileutils.SMILESCSV: '-ismi'
}

# -----------------------------------------------------------------------------


def get_structure_file_format(filename):
    '''
    Much like fileutils.get_structure_file_format() except also
    allows for compressed SMILES and SMILESCSV.
    '''

    format = fileutils.get_structure_file_format(filename)
    if format is not None:
        return format

    for (end, format) in [(".smi.gz", fileutils.SMILES),
                          (".csv.gz", fileutils.SMILESCSV)]:
        if filename.lower().endswith(end):
            return format
    else:
        return None


# =============================================================================
# Business logic
# =============================================================================


class LigPrep(object):

    def __init__(self, input_file, options='[]'):
        '''
        :param input_file: Input file name.
        :type input_file: str

        :param options: Ligprep command line options (in JSON format).
        :type options: str

        :raises RuntimeError: If something goes wrong.
        '''

        input_format = get_structure_file_format(input_file)
        if input_format not in LIGPREP_INPUT_OPTION:
            raise RuntimeError("ligprep does not support '%s' input format" %
                               input_format)

        argv = json.loads(options)

        argv += [
            LIGPREP_INPUT_OPTION[input_format], input_file, '-omae', 'void.mae',
            '-px',
            mm.mmpipeline_initialize()
        ]

        self._lp = ligprepworkflow.LigPrepWorkflow(logger, argv)

    def close(self):

        self._lp.close()

    def prepareStructure(self, ct):
        '''
        Applies ligprep to `ct`.

        :return: (outcomes, failures, summary) tuple
        :rtype: (list(structure.Structure), list(structure.Structure), dict)
        '''

        return self._lp.run(ct)

    def yieldStructures(self):
        '''
        Reads consecutive entries from the input file, applies LigPrep
        and yields resulting structures.
        '''

        with self._lp.makeInputReader() as reader:
            while True:
                try:
                    st = next(reader)
                except StopIteration:
                    break
                except Exception as e:
                    logger.warning('entry #%d: %s', reader.getOrigPos(), e)
                    continue

                (outcomes, _, _) = self.prepareStructure(st)
                if outcomes:
                    yield outcomes


# =============================================================================
# Command line
# =============================================================================


def add_ligprep_arguments(parser):
    '''
    Registers ligprep-specific options.

    :param parser: Command line argument parser.
    :type parser: `argparse.ArgumentParser`
    '''

    parser.add_argument(
        '-ligprep',
        action='store_true',
        help='Process input structures using LigPrep. Required '
        'for SMILES and 2D input or if structures lack Hydrogens. '
        'Recommended for 3D input structures not in a low-energy '
        'ionization/tautomeric state. Not compatible with '
        '-distinct, -connect, -stereo, or -title.')

    # custom options
    parser.add_argument(
        '-ligprep_options',
        metavar='<text>',
        default='["-epik", "-pht", "1.0", "-s", "16"]',
        help='Options to pass to the ligprep work flow '
        '(JSON-formatted list of strings. Default: \'%(default)s\').')


# -----------------------------------------------------------------------------


def validate_ligprep_arguments(args):
    '''
    Validates ligprep-specific command line options.

    :param args: Populated `argparse.Namespace` instance.
    :type args: `argparse.Namespace`

    :return: Tuple of validation success, and error message.
    :rtype: bool, str
    '''

    if args.ligprep:
        ok, msg = validate_ligprep_options(args.ligprep_options)
        if not ok:
            return False, msg

    return True, ''


# -----------------------------------------------------------------------------


def validate_ligprep_options(text):
    '''
    Validates ligprep-specific command-line arguments.

    :param text: Ligprep options as list of strings in JSON format.
    :type text: str

    :return: Tuple of validation success, and error message.
    :rtype: bool, str
    '''

    validate = Schema([str])
    try:
        obj = validate(json.loads(text))
    except json.JSONDecodeError as e:
        return False, f'malformed JSON: {e}'
    except Invalid as e:
        return False, f'unexpected type: {e}'

    for bad in ('-imae', '-isd', '-ismi', '-icsv'):
        if bad in obj:
            return False, f'must not include {bad}'

    lp_parser = arguments.make_parser()

    def raise_error(msg):
        raise RuntimeError(msg)

    lp_parser.error = lambda msg: raise_error(msg)

    try:
        lp_parser.parse_args(arguments.preprocess_args(obj))
    except RuntimeError as e:
        return False, f'{e}'

    return True, ''


# =============================================================================