#!/usr/bin/env nextflow

nextflow.enable.dsl=2

params.pdb = 's3://omic/eureka/ligandmpnn/1BC8'
params.outdir = 's3://omic/eureka/ligandmpnn/output'
params.model_type = 'ligand_mpnn'
params.temperature = 0.1
params.seed = 111
params.batch_size = 1
params.number_of_batches = 1
params.chains_to_design = ''
params.fixed_residues = ''
params.pack_side_chains = 0

process LIGANDMPNN {
    container 'harbor.cluster.omic.ai/omic/ligandmpnn:latest'
    publishDir params.outdir, mode: 'copy'
    stageInMode 'copy'

    input:
        path pdb

    output:
        path "${pdb.simpleName}/seqs/*.fa"
        path "${pdb.simpleName}/backbones/*.pdb"
        path "${pdb.simpleName}/packed/*.pdb", optional: true
        path "run.log"

    script:
        // Set checkpoint path based on model type
        def checkpoint_arg = ''
        if (params.model_type == 'ligand_mpnn') {
            checkpoint_arg = '--checkpoint_ligand_mpnn /app/LigandMPNN/model_params/ligandmpnn_v_32_010_25.pt'
        } else if (params.model_type == 'protein_mpnn') {
            checkpoint_arg = '--checkpoint_protein_mpnn /app/LigandMPNN/model_params/proteinmpnn_v_48_020.pt'
        } else if (params.model_type == 'soluble_mpnn') {
            checkpoint_arg = '--checkpoint_soluble_mpnn /app/LigandMPNN/model_params/solublempnn_v_48_020.pt'
        } else if (params.model_type == 'global_label_membrane_mpnn') {
            checkpoint_arg = '--checkpoint_global_label_membrane_mpnn /app/LigandMPNN/model_params/global_label_membrane_mpnn_v_48_020.pt'
        } else if (params.model_type == 'per_residue_label_membrane_mpnn') {
            checkpoint_arg = '--checkpoint_per_residue_label_membrane_mpnn /app/LigandMPNN/model_params/per_residue_label_membrane_mpnn_v_48_020.pt'
        }

        """
        mkdir -p ${pdb.simpleName}/seqs ${pdb.simpleName}/backbones ${pdb.simpleName}/packed

        ligandmpnn \\
            --pdb_path \$PWD/${pdb} \\
            --out_folder \$PWD/${pdb.simpleName} \\
            --model_type ${params.model_type} \\
            ${checkpoint_arg} \\
            --temperature ${params.temperature} \\
            --seed ${params.seed} \\
            --batch_size ${params.batch_size} \\
            --number_of_batches ${params.number_of_batches} \\
            --pack_side_chains ${params.pack_side_chains} \\
            ${params.chains_to_design ? "--chains_to_design ${params.chains_to_design}" : ''} \\
            ${params.fixed_residues ? "--fixed_residues \"${params.fixed_residues}\"" : ''} \\
            2>&1 | tee run.log
        """
}

workflow {
    LIGANDMPNN(Channel.fromPath(params.pdb))
}