Source code for schrodinger.application.matsci.jaguar_restart

"""
Utilities for Jaguar restarts.

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

import glob
import os
import zipfile

from schrodinger.application.jaguar.output import JaguarOutput
from schrodinger.application.matsci import jobutils
from schrodinger.utils import fileutils


def needs_restart(base_name):
    """
    Return True if the Jaguar job corresponding to the given base name
    needs to be restarted.

    :type base_name: str
    :param base_name: the base name

    :rtype: bool
    :return: True if the job needs to be restarted
    """

    out_file = f'{base_name}.out'
    if not os.path.exists(out_file):
        return True

    jout = JaguarOutput(out_file)
    if jout.status != JaguarOutput.OK or jout.fatal_error:
        return True

    files = [f'{base_name}.01.in', f'{base_name}.01.mae']
    return not all(map(os.path.exists, files))


def get_existing_files(files=None,
                       patterns=None,
                       base_names=None,
                       base_name_exts=None):
    """
    Return existing files.

    :type files: iterable or None
    :param files: file names

    :type patterns: iterable or None
    :param patterns: shell file patterns to glob for

    :type base_names: iterable or None
    :param base_names: base names of files

    :type base_name_exts: iterable or None
    :param base_name_exts: base name extensions of files, if None
        then Jaguar `*.out`, `*.01.in`, and `*.01.mae` are used

    :rtype: set
    :return: existing files
    """

    assert files or patterns or base_names

    files = files or []
    patterns = patterns or []
    base_names = base_names or []
    base_name_exts = base_name_exts or ['.out', '.01.in', '.01.mae']

    files = set(files)
    for base_name in base_names:
        for base_name_ext in base_name_exts:
            files.add(f'{base_name}{base_name_ext}')

    existing_files = set()
    for afile in files:
        if os.path.exists(afile):
            existing_files.add(afile)
    for pattern in patterns:
        existing_files.update(glob.glob(pattern))

    return existing_files


def set_restart_input_files(job_builder,
                            files=None,
                            patterns=None,
                            base_names=None):
    """
    Among the given files add those that exist to the given
    job builder.

    :type job_builder: `launchapi.JobSpecificationArgsBuilder`
    :param job_builder: the job specification builder object

    :type files: iterable or None
    :param files: the restart input file names

    :type patterns: iterable or None
    :param patterns: shell file patterns to glob for

    :type base_names: iterable or None
    :param base_names: base names of Jaguar `*.out`, `*.01.in`, and `*.01.mae`
    """

    for afile in get_existing_files(
            files=files,
            patterns=patterns,
            base_names=base_names,
            base_name_exts=['.out', '.01.in', '.01.mae']):
        job_builder.setInputFile(afile)


def get_restart_zip_files(base_names=None, base_name_patterns=None):
    """
    Return the restart zip files.

    :type base_names: iterable or None
    :param base_names: base names of `*zip` files

    :type base_name_patterns: iterable or None
    :param base_name_patterns: shell base name patterns to glob for

    :rtype: list
    :return: restart zip files
    """

    base_name_patterns = base_name_patterns or []
    base_name_patterns = [f'{pattern}.zip' for pattern in base_name_patterns]
    files = get_existing_files(patterns=base_name_patterns,
                               base_names=base_names,
                               base_name_exts=['.zip'])
    return list(files)


def prepare_restart_dirs(zip_files):
    """
    Prepare the restart directories.

    :type zip_files: iterable
    :param zip_files: the zip files
    """

    for afile in zip_files:
        obj = zipfile.ZipFile(afile, 'r')
        obj.extractall(fileutils.get_basename(afile))
        obj.close()
        fileutils.force_remove(afile)


def remove_driver_log_file(default_job_name):
    """
    Remove the driver log file.

    :type default_job_name: str
    :param default_job_name: the default job name
    """

    # for restarted jobs if a *-driver.log file using the same job name exists
    # it will not be overwritten, appended to, etc. (perhaps related to
    # MATSCI-6430) so force remove it here
    log_file_name = jobutils.get_jobname(default_job_name) + '-driver.log'
    fileutils.force_remove(log_file_name)