#!/usr/bin/env nextflow

nextflow.enable.dsl=2

params.input_receptor = 'YNKIVHLLVAEPEKIYAMPDPTVPDSDIKALTTLCDLADRELVVIIGWAKHIPGFSTLSLADQMSLLQSAWMEILILGVVYRSLFEDELVYADDYIMDEDQSKLAGLLDLNNAILQLVKKYKSMKLEKEEFVTLKAIALANSDSMHIEDVEAVQKLQDVLHEALQDYEAGQHMEDPRRAGKMLMTLPLLRQTSTKAVQHFYNKLEGKVPMHKLFLEMLEAKV'
params.input_ligand = 'c1cc2c(cc1O)CCCC2'
params.input_template = ''
params.sample_id = 'flowdock_sample'
params.outdir = 's3://omic/eureka/flowdock/output/'
params.n_samples = 5

process FLOWDOCK {
    container 'harbor.cluster.omic.ai/omic/flowdock:latest'
    containerOptions '--rm --gpus all'
    publishDir params.outdir, mode: 'copy'
    stageInMode 'copy'

    input:
        val receptor
        val ligand

    output:
        path "${params.sample_id}/*.pdb", optional: true
        path "${params.sample_id}/*.sdf", optional: true
        path "${params.sample_id}/**/*.pdb", optional: true
        path "run.log"

    script:
        """
        set +u
        source /opt/conda/etc/profile.d/conda.sh
        conda activate FlowDock
        
        mkdir -p ${params.sample_id}

        python /software/flowdock/flowdock/sample.py \\
            ckpt_path=/software/flowdock/checkpoints/esmfold_prior_paper_weights-EMA.ckpt \\
            model.cfg.prior_type=esmfold \\
            sampling_task=batched_structure_sampling \\
            input_receptor='${receptor}' \\
            input_ligand='"${ligand}"' \\
            ${params.input_template ? "input_template=${params.input_template}" : "input_template=null"} \\
            sample_id='${params.sample_id}' \\
            out_path=\$PWD/${params.sample_id}/ \\
            paths.output_dir=\$PWD/${params.sample_id}/ \\
            hydra.run.dir=\$PWD/${params.sample_id}/ \\
            n_samples=${params.n_samples} \\
            chunk_size=5 \\
            num_steps=40 \\
            sampler=VDODE \\
            sampler_eta=1.0 \\
            start_time='1.0' \\
            use_template=true \\
            separate_pdb=true \\
            visualize_sample_trajectories=false \\
            auxiliary_estimation_only=false \\
            esmfold_chunk_size=null \\
            trainer=gpu 2>&1 | tee run.log
        
        # Copy any outputs from FlowDock's log directory if they exist there
        if ls /software/flowdock/logs/sample/runs/*/rank*.pdb 2>/dev/null; then
            cp /software/flowdock/logs/sample/runs/*/rank*.pdb ${params.sample_id}/ 2>/dev/null || true
            cp /software/flowdock/logs/sample/runs/*/rank*.sdf ${params.sample_id}/ 2>/dev/null || true
        fi
        
        # List what was generated
        echo "=== Generated files ===" >> run.log
        find ${params.sample_id}/ -type f >> run.log 2>&1 || true
        """
}

workflow {
    FLOWDOCK(Channel.of(params.input_receptor), Channel.of(params.input_ligand))
}