# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/af3/01_docker.ipynb.

# %% auto 0
__all__ = ['get_docker_command', 'docker_single_full', 'docker_multi_full', 'docker_multi_msa', 'docker_multi_infer']

# %% ../../nbs/af3/01_docker.ipynb 3
from fastcore.meta import delegates

# %% ../../nbs/af3/01_docker.ipynb 6
def get_docker_command(
    input_dir="af_input", # Directory for input data
    output_dir="af_output", # Directory for output results
    model_dir="af_model", # Directory containing models
    db_dir="af_database", # Directory for databases. If None, this option is ommitted
    cache_dir="af_cache", # Directory for JAX compilation cache. If None, this option is omitted
    gpus=0, # GPU devices to allocate (e.g., 0,1), if None, ommitted
    docker_name="sky1ove/alphafold3", # Docker image name
    embedding=False, # If True, includes the --save_embeddings=true flag
    skip_search=False, # if MSA is precalculated and present in json; If True, includes the --norun_data_pipeline flag
    search_only=False, # search MSA only; If True, sets skip_search to False and includes the --norun_inference flag
    json_path=None  # Path to JSON file. If not None, uses json_file instead of input_dir
):

    "Generate a Docker run command for Alphafold with customizable parameters."
    
    # Start building the Docker command
    cmd_parts = ["docker run --rm \\"]

    if json_path:
        parent_dir=json_path.split('/')[0]
        cmd_parts.append(f'    --volume "$HOME/{parent_dir}:/root/af_input" \\')
    else:
        parent_dir = input_dir.split('/')[0]
        cmd_parts.append(f'    --volume "$HOME/{parent_dir}:/root/af_input" \\')

    cmd_parts.extend([f'    --volume "$HOME/{output_dir}:/root/af_output" \\',
                      f'    --volume "$HOME/{model_dir}:/root/models" \\',
                     ])
    if db_dir:
        cmd_parts.append(f'    --volume "$HOME/{db_dir}:/root/public_databases" \\')

    if cache_dir:
        cmd_parts.append(f'    --volume "$HOME/{cache_dir}:/root/cache" \\')

    if gpus is not None:
        cmd_parts.append(f'    --gpus "device={gpus}" \\')

    cmd_parts.extend([f'    {docker_name} \\',
                      '    python run_alphafold.py \\'])
    
    if json_path:
        mount_path= '/'.join(json_path.split('/')[1:])
        cmd_parts.append(f'    --json_path=/root/af_input/{mount_path} \\')
    else:
        len_dir = len(input_dir.split('/')[1:])
        if len_dir:
            mount_dir= '/'.join(input_dir.split('/')[1:])
            cmd_parts.append(f'    --input_dir=/root/af_input/{mount_dir} \\')
        else:
            cmd_parts.append('    --input_dir=/root/af_input \\')

    cmd_parts.extend(['    --output_dir=/root/af_output \\',
                      '    --model_dir=/root/models \\'])
    
    if cache_dir:
        cmd_parts.append('    --jax_compilation_cache_dir=/root/cache \\')

    if embedding:
        cmd_parts.append('    --save_embeddings=true \\')
    
    if search_only:
        skip_search = False
        cmd_parts.append('    --norun_inference \\')
    elif skip_search:
        cmd_parts.append('    --norun_data_pipeline \\')

    if cmd_parts[-1].endswith('\\'):
        cmd_parts[-1] = cmd_parts[-1].rstrip(' \\')

    docker_command = "\n".join(cmd_parts)
    print(docker_command)

# %% ../../nbs/af3/01_docker.ipynb 12
@delegates(get_docker_command)
def docker_single_full(json_path,output_dir,cache_dir=False,**kwargs):
    "Single json task with full pipeline."
    return get_docker_command(
        json_path=json_path,
        output_dir=output_dir,
        cache_dir=cache_dir,
        **kwargs)

# %% ../../nbs/af3/01_docker.ipynb 16
@delegates(get_docker_command)
def docker_multi_full(input_dir,output_dir,**kwargs):
    "Folder of json as input with full pipeline."
    return get_docker_command(
        input_dir=input_dir,
        output_dir=output_dir,
        **kwargs)

# %% ../../nbs/af3/01_docker.ipynb 19
@delegates(get_docker_command)
def docker_multi_msa(input_dir,output_dir,search_only=True,**kwargs):
    "MSA search only, without structure inference; CPU only."
    return get_docker_command(
        input_dir=input_dir,
        output_dir=output_dir,
        search_only=search_only,
        **kwargs)

# %% ../../nbs/af3/01_docker.ipynb 22
@delegates(get_docker_command)
def docker_multi_infer(input_dir,output_dir,skip_search=True,**kwargs):
    "Infer only with pre-calculated MSA; GPU is needed."
    return get_docker_command(
        input_dir=input_dir,
        output_dir=output_dir,
        skip_search=skip_search,
        **kwargs)
