main.py

import logging
import os.path

import gradio as gr

import run_utils


def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> str:
    if other_args_file is not None:
        kwargs["other_arg_file"] = other_args_file.name
    return run_utils.run_cli_command(
        protein_file.name, ligand_file.name, *args, **kwargs
    )


def run():
    iface = gr.Interface(
        fn=run_wrapper,
        inputs=[
            gr.File(label="Protein PDB", file_types=[".pdb"]),
            gr.File(label="Ligand SDF", file_types=[".sdf"]),
            gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None),
            gr.Number(
                label="Samples Per Complex",
                value=1,
                minimum=1,
                maximum=100,
                precision=0,
            ),
            gr.Checkbox(label="Keep Local Structures", value=True),
            gr.Checkbox(label="Save Visualisation", value=True),
        ],
        outputs=gr.File(label="Result"),
        title="DiffDock-Pocket",
        description="""
        Run [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket) for a single protein and ligand.
        We have provided the most important inputs as UI elements.  
        Additional values can be included in "Other arguments", and will be passed 
        to [inference.py](https://github.com/plainerman/DiffDock-Pocket/blob/main/inference.py). 
        Must be a YAML file without any nesting.  
        For example, to specify a pocket of (0.0, 1.0, 2.0), the YAML file should contain:  
        ```
        pocket_center_x: 0.0  
        pocket_center_y: 1.0  
        pocket_center_z: 2.0  
        ```
        """
    )

    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)


if __name__ == "__main__":
    run_utils.set_env_variables()
    run_utils.configure_logging()

    run()