#!/usr/bin/env nextflow


params.fasta = 's3://omic/eureka/immunebuilder/antibody_test.fasta'
params.outdir = 's3://omic/eureka/immunebuilder/output'
params.mode = 'antibody'
params.verbose = true
params.original_weights = false

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

    input:
        path fasta

    output:
        path "${fasta.simpleName}.pdb", emit: structure
        path "run.log", emit: logfile

    script:
        """
        export HOME=/home/mambauser
        /opt/conda/bin/ABodyBuilder2 \\
            -f ${fasta} \\
            -o ${fasta.simpleName}.pdb \\
            ${params.verbose ? '-v' : ''} 2>&1 | tee run.log
        """
}

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

    input:
        path fasta

    output:
        path "${fasta.simpleName}.pdb", emit: structure
        path "run.log", emit: logfile

    script:
        """
        export HOME=/home/mambauser
        /opt/conda/bin/NanoBodyBuilder2 \\
            -f ${fasta} \\
            -o ${fasta.simpleName}.pdb \\
            ${params.verbose ? '-v' : ''} 2>&1 | tee run.log
        """
}

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

    input:
        path fasta

    output:
        path "${fasta.simpleName}.pdb", emit: structure
        path "run.log", emit: logfile

    script:
        """
        export HOME=/home/mambauser
        /opt/conda/bin/TCRBuilder2 \\
            -f ${fasta} \\
            -o ${fasta.simpleName}.pdb \\
            ${params.verbose ? '-v' : ''} \\
            ${params.original_weights ? '--original_weights' : ''} 2>&1 | tee run.log
        """
}

workflow PREDICT_ANTIBODY {
    take:
        fasta_ch

    main:
        ABODYBUILDER2(fasta_ch)

    emit:
        structure = ABODYBUILDER2.out.structure
        logfile = ABODYBUILDER2.out.logfile
}

workflow PREDICT_NANOBODY {
    take:
        fasta_ch

    main:
        NANOBODYBUILDER2(fasta_ch)

    emit:
        structure = NANOBODYBUILDER2.out.structure
        logfile = NANOBODYBUILDER2.out.logfile
}

workflow PREDICT_TCR {
    take:
        fasta_ch

    main:
        TCRBUILDER2(fasta_ch)

    emit:
        structure = TCRBUILDER2.out.structure
        logfile = TCRBUILDER2.out.logfile
}

workflow {
    if (!params.fasta) {
        error "Please provide input FASTA file(s) using --fasta parameter"
    }

    fasta_ch = Channel.fromPath(params.fasta, checkIfExists: true)

    if (params.mode == 'antibody') {
        PREDICT_ANTIBODY(fasta_ch)
    } else if (params.mode == 'nanobody') {
        PREDICT_NANOBODY(fasta_ch)
    } else if (params.mode == 'tcr') {
        PREDICT_TCR(fasta_ch)
    } else {
        error "Invalid mode: ${params.mode}. Please use 'antibody', 'nanobody', or 'tcr'"
    }
}