Source code for schrodinger.application.livedesign.tasks

"""
LiveDesign protocol tasks.

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

import os
import shlex

from schrodinger.tasks import tasks
from schrodinger.tasks import jobtasks

from schrodinger.models import parameters

from schrodinger.application.glide import glide


def run_task(task, task_type):
    """
    Check the task to see if it completed successfully. Will raise a
    RuntimeError if the job failed or a FileNotFoundError is the output file
    does not exist

    :param task: job_utils task object
    :type task: 'schrodinger.tasks.jobtasks.CmdJobTask'

    :param task_type: Calculation performed by the task
    :type task_type: str
    """

    task.start()
    task.wait()

    if task.status == tasks.Status.FAILED:
        raise RuntimeError(f'ERROR: {task_type} job failed')

    if not os.path.exists(task.output.out_file):
        raise RuntimeError(
            f'ERROR: {task_type} failed to generate {task.output.out_file}')


class LigPrepTask(jobtasks.CmdJobTask):
    """
    Create a LigPrep task using in the input, output and job settings
    """

    class JobConfig(jobtasks.JobConfig):
        driver_host_settings: jobtasks.HostSettings = jobtasks.HostSettings(
            allowed_host_types=jobtasks.AllowedHostTypes.CPU_ONLY)

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        ligprep_infile: jobtasks.TaskFile
        ligprep_args: str

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile
        dropped_file: jobtasks.TaskFile

    def makeCmd(self):
        cmd = ['ligprep']

        args = ['-isd', self.input.ligands_file, '-omae', self.output.out_file]

        if self.input.ligprep_infile:
            args.extend(['-inp', self.input.ligprep_infile])

        if self.input.ligprep_args:
            args.extend(shlex.split(self.input.ligprep_args))

        cmd.extend(args)
        cmd.extend(
            ['-DRIVERHOST',
             self.job_config.driver_host_settings.toCmdArg()])

        return cmd

    @tasks.preprocessor()
    def processOutput(self):
        self.output.dropped_file = f'{self.name}-dropped.sdf'


class LigFilterTask(jobtasks.CmdJobTask):
    """
    Create a LigFilter task using in the input, output and job settings
    """

    class JobConfig(jobtasks.JobConfig):
        driver_host_settings: jobtasks.HostSettings = jobtasks.HostSettings(
            allowed_host_types=jobtasks.AllowedHostTypes.CPU_ONLY)

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        filter_file: jobtasks.TaskFile

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile

    def makeCmd(self):
        cmd = [
            'ligfilter', self.input.ligands_file, '-f', self.input.filter_file,
            '-j', self.name
        ]

        if self.job_config.host_settings.num_subjobs:
            cmd.extend([
                '-NJOBS',
                str(self.job_config.host_settings.num_subjobs), '-DRIVERHOST',
                self.job_config.driver_host_settings.toCmdArg()
            ])

        return cmd


class GlideTask(jobtasks.CmdJobTask):
    """
    Create a Glide task using in the input, output and job settings
    """

    class JobConfig(jobtasks.JobConfig):
        driver_host_settings: jobtasks.HostSettings = jobtasks.HostSettings(
            allowed_host_types=jobtasks.AllowedHostTypes.CPU_ONLY)

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        grid_file: jobtasks.TaskFile
        glide_infile: jobtasks.TaskFile
        num_poses: int
        ref_file: jobtasks.TaskFile
        report_strain_energy: bool

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile
        csv_file: jobtasks.TaskFile
        parsed_inp_file: jobtasks.TaskFile
        keyword_dict: dict

    @tasks.preprocessor()
    def format_inp_file(self):
        """
        Parse the Glide input file and combine command line arguments and Glide
        defaults into a final validated Glide input file
        """

        # Read the Glide input file and combine it with all default Glide
        # settings
        try:
            glide_keywords = glide.get_glide_job(self.input.glide_infile)
        except (ValueError, RuntimeError) as err:
            return False, err

        # Overwrite keywords for ligand infile, grid infile and the command line
        # number of poses per ligand
        glide_keywords['LIGANDFILE'] = self.input.ligands_file
        glide_keywords['GRIDFILE'] = self.input.grid_file
        glide_keywords['POSES_PER_LIG'] = self.input.num_poses

        # If a reference ligand file is given, enable its use
        if self.input.ref_file:
            # Check if the input file was setup to use core constraints
            if not glide_keywords['CORE_RESTRAIN']:
                return False, f'CORE_RESTRAIN keyword is not True in {self.input.glide_inp_file}'

            glide_keywords['REF_LIGAND_FILE'] = self.input.ref_file
            glide_keywords['USE_REF_LIGAND'] = True

        else:
            # Disable core constraints if no reference file is given but the
            # input file contains constraint information
            glide_keywords['CORE_DEFINITION'] = 'allheavy'
            glide_keywords['CORE_RESTRAIN'] = False
            glide_keywords['CORECONS_FALLBACK'] = False
            glide_keywords['USE_REF_LIGAND'] = False
            glide_keywords['REF_LIGAND_FILE'] = ''

        if self.input.report_strain_energy:
            glide_keywords['POSTDOCKSTRAIN'] = True

        self.output.keyword_dict = glide_keywords
        self.output.parsed_inp_file = f'{self.name}_parsed.inp'
        glide_keywords.writeSimplified(self.output.parsed_inp_file)

    def makeCmd(self):
        cmd = ['glide', self.output.parsed_inp_file]

        cmd.extend(
            ['-DRIVERHOST',
             self.job_config.driver_host_settings.toCmdArg()])

        return cmd

    @tasks.postprocessor()
    def processOutput(self):
        self.output.csv_file = self.name + '.csv'


class AlignLigandsTask(jobtasks.CmdJobTask):
    """
    Create a align_ligands task using in the input, output and job settings. The
    alignment_args is a list of arguments that align_ligands takes
    """

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        alignment_args: str

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile

    def makeCmd(self):
        cmd = [
            'align_ligands', self.input.ligands_file, '-o',
            self.output.out_file, '-JOBNAME', self.name
        ] + self.input.alignment_args

        return cmd


class FlexAlignTask(jobtasks.CmdJobTask):
    """
    Create a flex_align task using in the input, output and job settings. The
    alignment_args is a list of arguments that align_ligands takes
    """

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        refligs_file: jobtasks.TaskFile
        flex_align_args: str

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile

    def makeCmd(self):
        cmd = [
            'flex_align', '-screen', self.input.ligands_file, '-shape',
            self.input.refligs_file, '-JOB', self.name
        ]
        if self.input.flex_align_args:
            cmd.extend(shlex.split(self.input.flex_align_args))

        return cmd


class EpikTask(jobtasks.CmdJobTask):
    """
    Create a Epik task using in the input, output and job settings
    """

    class JobConfig(jobtasks.JobConfig):
        driver_host_settings: jobtasks.HostSettings = jobtasks.HostSettings(
            allowed_host_types=jobtasks.AllowedHostTypes.CPU_ONLY)

    class Input(parameters.CompoundParam):
        ligands_file: jobtasks.TaskFile
        epik_infile: jobtasks.TaskFile
        epik_args: str

    class Output(jobtasks.CmdJobTask.Output):
        out_file: jobtasks.TaskFile

    def makeCmd(self):
        cmd = ['epik']

        args = ['-imae', self.input.ligands_file, '-omae', self.output.out_file]

        if self.input.epik_infile:
            args.extend(['-inp', self.input.epik_infile])

        if self.input.epik_args:
            args.extend(shlex.split(self.input.epik_args))

        cmd.extend(args)
        cmd.extend(
            ['-DRIVERHOST',
             self.job_config.driver_host_settings.toCmdArg()])

        return cmd