Source code for schrodinger.application.mopac.mopac_startup

"""
Startup driver for running semiempirical calculations (i.e. MOPAC)
"""

# Contributors: Mark A. Watson

import os
import re

from schrodinger.application.mopac import mopac_backend
from schrodinger.application.mopac import mopac_parser
from schrodinger.application.mopac import utils
from schrodinger.application.mopac.exceptions import MopacIOError
from schrodinger.infra import mm
from schrodinger.job import jobcontrol
from schrodinger.job import launchapi
from schrodinger.job import launchparams

_external_re = re.compile(r"EXTERNAL\s*=\s*(\S+)", re.I)


def _get_EXTERNAL_files(infiles, keywords):
    """
    Look for EXTERNAL specification in MOPAC input files and cmdline flags to
    indicate files which need to be registered as additional input files.

    :type  infiles: list
    :param infiles: known input file paths

    :type  keywords: str
    :param keywords: string of additional MOPAC keywords from cmdline

    :return: set of input file names
    """

    def add_file(f, efiles):
        if not os.path.isfile(f):
            raise RuntimeError(f"EXTERNAL file '{f}' not found.")
        efiles.add(f)

    efiles = set()

    # Handle EXTERNAL specification in .mop files
    for infile in infiles:
        if utils.is_mopac_file(infile):
            with open(infile, 'r') as fh:
                last_line_continued = True
                for ix, line in enumerate(fh):
                    match = _external_re.search(line)
                    if match:
                        add_file(match.group(1), efiles)
                    if ix > 0 and not last_line_continued:
                        break
                    last_line_continued = re.search("[&+]", line)

    # Handle EXTERNAL=filename from cmdline flags
    if keywords:
        match = _external_re.search(keywords)
        if match:
            add_file(match.group(1), efiles)

    return efiles


def _get_input_files(parsed_args):
    """
    Return set of input files that need to be registered with jobcontrol
    and/or copied to the work directory.

    :type  parsed_args: ParsedArgs object
    :param parsed_args: parsed cmdline arguments

    :return: set of input file names
    """

    infiles = set(parsed_args.infiles)
    for infile in infiles:
        if not os.path.isfile(infile):
            msg = f'Cannot find input file {infile} in {os.getcwd()}'
            raise MopacIOError(msg)

    # Include infiles obtained from the EXTERNAL specification
    infiles.update(_get_EXTERNAL_files(infiles, parsed_args.keywords))

    return infiles


def _get_output_files(parsed_args):
    """
    Return set of output files expected from this workflow.

    Currently all output files are registered on-the-fly in the
    backend scripts.  Hopefully this can be improved in the future
    to make the API more clearly defined.

    :type  parsed_args: ParsedArgs object
    :param parsed_args: parsed cmdline arguments

    :return: set of output file names
    """

    outfiles = set()
    return outfiles


def _register_input_files(jsb, infiles):
    """
    Register input files with job specification.

    :type  jsb: `schrodinger.job.launchapi.JobSpecificationArgsBuilder`
    :param jsb: jobcontrol job specification builder instance

    :type  infiles: set of strings
    :param infiles: input file names
    """

    for infile in infiles:
        if os.path.exists(infile):
            jsb.setInputFile(infile)
        else:
            msg = f'Failed to find expected file {infile}'
            raise MopacIOError(msg)


def _register_output_files(jsb, outfiles):
    """
    Register output files with job specification.

    :type  jsb: `schrodinger.job.launchapi.JobSpecificationArgsBuilder`
    :param jsb: jobcontrol job specification builder instance

    :type  outfiles: list of strings
    :param outfiles: output file names
    """

    for outfile in outfiles:
        jsb.setOutputFile(outfile, stream=True)


def _get_job_params():
    """
    Return the jobcontrol launch parameters

    :type  parsed_args: ParsedArgs object
    :param parsed_args: parsed cmdline arguments
    """

    params = launchparams.LaunchParameters()

    # Set jobcontrol parallel resources
    params.setNumberOfProcessorsManyNodes(1)  # no MPI
    params.setNumberOfProcessorsOneNode(1)  # launch 1-CPU driver

    return params


def _get_job_spec(cmdline, infiles, outfiles, jobname):
    """
    Return the jobcontrol job specification

    :type  cmdline: str
    :param cmdline: cmdline to launch under jobcontrol

    :type  infiles: list
    :param infiles: list of input files to register

    :type  outfiles: list
    :param outfiles: list of (known) output files to register

    :type  jobname: str
    :param jobname: name of job

    :return `schrodinger.job.launchapi.JobSpecification`
    """

    # Create job_spec_builder instance
    jsb = launchapi.JobSpecificationArgsBuilder(cmdline,
                                                use_schrodinger_run=True,
                                                use_jobname_log=True)

    # Set jobname and .log filename
    jsb.setJobname(jobname)

    # Set Program name so Maestro knows how to incorporate
    jsb.setProgramName("SemiEmp")

    # Register input and output files
    _register_input_files(jsb, infiles)
    _register_output_files(jsb, outfiles)

    return jsb.getJobSpec()


def _run_backend(args):
    """
    Launch backend process under jobcontrol

    :type  args: list
    :param args: cmdline arguments
    """

    status = 0
    parsed_args = mopac_parser.parse_args(args)

    print("Launching MOPAC under jobcontrol.")
    # Get jobcontrol jobspec and launch parameters
    cmdline = ['semi_emp_backend.py'] + args
    infiles = _get_input_files(parsed_args)
    outfiles = _get_output_files(parsed_args)
    job_spec = _get_job_spec(cmdline, infiles, outfiles, parsed_args.jobname)
    # Get jobcontrol launch parameters
    job_params = _get_job_params()
    # Launch jobcontrol job
    job = jobcontrol.launch_from_job_spec(job_spec, job_params)
    print("Release:", mm.mmfile_get_release_name())
    print("Exec:", os.getenv('MMSHARE_EXEC'))
    print(f"JobId: {job.JobId}")

    if parsed_args.wait:
        # Barrier until jobcontrol job is finished
        job.wait()

    return status


def main(args):
    """
    Startup driver for running MOPAC

    :type  args: list
    :param args: cmdline arguments
    """

    parsed_args = mopac_parser.parse_args(args)

    if parsed_args.nojobid:
        # Run script directly
        status = mopac_backend.main(args)
    else:
        # Run script under jobcontrol
        status = _run_backend(args)

    return status