#!/usr/bin/env nextflow

metadata_ch = Channel
    .fromPath(params.metadata)
    .splitCsv(header:true)
    .map { row ->
        def id         = row.sample_id
        def input_dir  = row.input_dir
        def output_dir = row.output_dir
        def r1         = file("${input_dir}/${row.r1}")
        def r2         = row.r2 ? file("${input_dir}/${row.r2}") : null
        def read_type  = row.read_type
        def sequencing = row.sequencing
        [id, output_dir, r1, r2, read_type, sequencing]
    }

short_metaG_ch = metadata_ch.filter { it[4] == 'short' && it[5] == 'metagenomic' }.map { it[0..3] }
long_metaG_ch = metadata_ch.filter { it[4] == 'long' && it[5] == 'metagenomic' }.map { [it[0], it[1], it[2]] }
short_metaT_ch = metadata_ch.filter { it[4] == 'short' && it[5] == 'metatranscriptomic' }.map { it[0..3] }

process fastqc {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/FastQC_raw", mode: 'copy'
    container 'quay.io/biocontainers/fastqc:0.11.9--0'

    input:
    tuple val(id), val(output_dir), path(r1), path(r2)

    output:
    path("*")

    script:
    """
    fastqc -o . ${r1} ${r2}
    """
}

process bbtools_qc {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy'
    container 'quay.io/staphb/bbtools:39.01'

    input:
    tuple val(id), val(output_dir), path(r1), path(r2)

    output:
    tuple val(id), val(output_dir), 
          path("${id}_r1_qc.fastq.gz"),
          path("${id}_r2_qc.fastq.gz"),
          path("${id}_qc_stats.txt")   

    script:
    """
    mkdir -p cleaned_reads
    bbduk.sh in1=${r1} in2=${r2} \
             out1=${id}_r1_qc.fastq.gz \
             out2=${id}_r2_qc.fastq.gz \
             stats=${id}_qc_stats.txt \
             qtrim=rl trimq=10 minlength=50
    """
}

//process dorado {
//    tag "${id}"
//    cpus params.threads
//    container 'quay.io/nanoporetech/dorado:0.5.3'
//
//    input:
//    tuple val(id), val(output_dir), path(r1)
//
//    output:
//    tuple val(id), val(output_dir), path("${id}.bam"), path("${id}.fastq.gz")
//
//    publishDir "${output_dir}/${id}/base_calling", mode: 'copy'
//
//    script:
//    """
//    dorado basecaller \\
//        ${params.dorado_model_cache}/${params.dorado_model} \\
//        ${r1} \\
//        --device cpu \\
//        > ${id}.bam
//
//    samtools fastq ${id}.bam | gzip > ${id}.fastq.gz
//    """
//}

process chopper {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy'
    container 'quay.io/biocontainers/chopper:0.10.0--hcdda2d0_0'

    input:
    tuple val(id), val(output_dir), path(reads)   

    output:
    tuple val(id), val(output_dir), path("${id}_qc.fastq.gz")

    script:
    """
    chopper -q 10 -l 500 -i ${reads} -c ${params.human_db} | gzip > ${id}_qc.fastq.gz
    """
}

process nanoplot {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/", mode: 'copy'
    container 'quay.io/nanozoo/nanoplot:1.41.0--9bd2843'

    input:
    tuple val(id), val(output_dir), path(fastq)

    output:
    tuple val(id), val(output_dir), path("NanoPlot_raw")

    script:
    """
    mkdir -p NanoPlot_raw
    NanoPlot \\
        --fastq ${fastq} \\
        --plots hex dot \\
        --threads ${task.cpus} \\
        --outdir NanoPlot_raw
    """
}

process bbtools_rrna_removal_short_read {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy'
    container 'quay.io/staphb/bbtools:39.01'

    input:
    tuple val(id), val(output_dir), path(r1), path(r2)

    output:
    tuple val(id),
          val(output_dir),
          path("${id}_r1_norrna.fastq.gz"),
          path("${id}_r2_norrna.fastq.gz"),
          path("${id}_rrna_removal_stats.txt")

    script:
    """
    bbduk.sh in1=${r1} in2=${r2} \
             out1=${id}_r1_norrna.fastq.gz \
             out2=${id}_r2_norrna.fastq.gz \
             ref=$params.rrna_db \
             stats=${id}_rrna_removal_stats.txt \
             t=${task.cpus}
    """
}

process bbtools_rrna_removal_long_read {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy'
    container 'quay.io/staphb/bbtools:39.01'

    input:
    tuple val(id), val(output_dir), path(clean_reads)

    output:
    tuple val(id), val(output_dir),
          path("${id}_norrna.fastq.gz"),
          path("${id}_rrna_removal_stats.txt")

    script:
    """
    bbduk.sh in=${clean_reads} \
             out=${id}_norrna.fastq.gz \
             ref=$params.rrna_db \
             stats=${id}_rrna_removal_stats.txt \
             t=${task.cpus}
    """
}

process bbtools_human_removal {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy', pattern: "${id}_*"
    container 'quay.io/staphb/bbtools:39.01'

    input:
    tuple val(id), val(output_dir), path(r1_norrna), path(r2_norrna)

    output:
    tuple val(id), val(output_dir),
          path("${id}_r1_nohuman.fastq.gz"),
          path("${id}_r2_nohuman.fastq.gz"),
          path("${id}_human_removal_stats.txt")

    script:
    """
    HUMAN_REF="${params.ref_data_dir}/bbtools/human_genome.fa.gz"
    
    if [ ! -f "\${HUMAN_REF}" ]; then
        echo "Downloading human genome reference..."
        mkdir -p ${params.ref_data_dir}/bbtools
        wget -O "\${HUMAN_REF}" 'https://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/000/001/405/GCF_000001405.40_GRCh38.p14/GCF_000001405.40_GRCh38.p14_genomic.fna.gz'
        
        if [ ! -s "\${HUMAN_REF}" ]; then
            echo "ERROR: Download failed"
            exit 1
        fi
        echo "Download complete"
    else
        echo "Using cached human genome reference"
    fi
    
    bbduk.sh in1=${r1_norrna} in2=${r2_norrna} \
             out1=${id}_r1_nohuman.fastq.gz \
             out2=${id}_r2_nohuman.fastq.gz \
             ref=\${HUMAN_REF} \
             stats=${id}_human_removal_stats.txt \
             t=${task.cpus}
    """
}

process bbtools_phix_removal {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/cleaned_reads", mode: 'copy'
    container 'quay.io/staphb/bbtools:39.01'

    input:
    tuple val(id), val(output_dir), path(r1_norrna), path(r2_norrna)

    output:
    tuple val(id), val(output_dir),
          path("${id}_r1_nophix.fastq.gz"),
          path("${id}_r2_nophix.fastq.gz"),
          path("${id}_phix_removal_stats.txt")

    script:
    """
    bbduk.sh in1=${r1_norrna} in2=${r2_norrna} \
             out1=${id}_r1_nophix.fastq.gz \
             out2=${id}_r2_nophix.fastq.gz \
             ref=phix \
             stats=${id}_phix_removal_stats.txt \
             t=${task.cpus} 
    """
}

process fastqc_cleaned {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/FastQC_cleaned", mode: 'copy'
    container 'quay.io/biocontainers/fastqc:0.11.9--0'

    input:
    tuple val(id), val(output_dir), path(r1_nophix), path(r2_nophix)

    output:
    path("*")

    script:
    """
    fastqc -o . ${r1_nophix} ${r2_nophix}
    """
}

//process gottcha2 {
//    time params.time
//    tag "${id}"
//    cpus params.threads
//    publishDir "${output_dir}/${id}/gottcha2", mode : 'move'
//    container 'docker://poeli/gottcha2:2.1.11-amd64'
//
//    input:
//    tuple val(id), val(output_dir), path(r1_nophix), path(r2_nophix)
//
//    output:
//    path("*")
//
//    script:
//    """
//    mkdir -p ${output_dir}
//    mkdir -p ${output_dir}/${id}
//    gottcha2.py -d $params.gottcha2_db -i ${r1_nophix} ${r2_nophix} -t ${task.cpus}
//    """
//}
//
//process gottcha2_long_read {
//    time params.time
//    tag "${id}"
//    cpus params.threads
//    publishDir "${output_dir}/${id}/gottcha2", mode : 'move'
//    container 'quay.io/poeli/gottcha2:2.1.8'
//
//    input:
//    tuple val(id), val(output_dir), path(no_rrna_reads)
//
//    output:
//    path("*")
//
//    script:
//    """
//    mkdir -p ${output_dir}
//    mkdir -p ${output_dir}/${id}
//    gottcha2.py -d $params.gottcha2_db -i ${no_rrna_reads} -t ${task.cpus}
//    """
//}

process nanoplot_cleaned {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/", mode: 'move'
    container 'quay.io/nanozoo/nanoplot:1.41.0--9bd2843'

    input:
    tuple val(id), val(output_dir), path(clean_reads)

    output:
    tuple val(id), val(output_dir), path("NanoPlot_cleaned/**")

    script:
    """
    mkdir -p NanoPlot_cleaned
    NanoPlot \
        --fastq ${clean_reads} \
        --plots hex dot \
        --outdir NanoPlot_cleaned
    """
}

process metaspades {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/assembly", mode: 'copy'
    container 'quay.io/staphb/spades:4.0.0'

    input:
    tuple val(id), val(output_dir), path(r1_nophix), path(r2_nophix)

    output:
    tuple val(id), val(output_dir), path("contigs.fasta")  

    script:
    """
    metaspades.py -1 ${r1_nophix} -2 ${r2_nophix} \
                  -t ${task.cpus} \
                  -o .
    """
}

process metaquast {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/metaquast_output", mode: 'move'
    container 'quay.io/staphb/quast:5.2.0'

    input:
    tuple val(id), val(output_dir), path(contigs)

    output:
    tuple val(id), val(output_dir), path("*")

    script:
    """
    metaquast.py -m 100 -o . ${contigs}
    """
}

//process chimeracutter {
//    tag "${id}"
//    cpus params.threads
//    publishDir "${output_dir}/${id}", mode: 'move'
//
//    input:
//    tuple val(id), val(output_dir), path(contigs)
//
//    output:
//    tuple val(id), val(output_dir), path("chimeracutter_output")
//
//    script:
//    """
//    mkdir -p chimeracutter_output
//    source /g/g16/ruth6/miniforge3/etc/profile.d/conda.sh
//    conda activate /g/g16/ruth6/miniforge3/envs/chimeracutter
//    python /p/vast1/mlbiomon/analysis/securebio/test/other/wawa_wkflw/scripts/chimeracutter.py blast -f ${contigs} -d $params.chimeracutter_db -m someone@lanl.gov -o chimeracutter_output
//    python /p/vast1/mlbiomon/analysis/securebio/test/other/wawa_wkflw/scripts/chimeracutter.py bed -i chimeracutter_output/results.tsv -o chimeracutter_output/bed.tsv
//    """
//}

process bowtie2_build {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/bowtie_files", mode: 'copy'
    container 'quay.io/biocontainers/bowtie2:2.5.1--py310h8d7afc0_1'

    input:
    tuple val(id), val(output_dir), path(contigs)

    output:
    tuple val(id), val(output_dir), path("${id}_index.*.bt2")

    script:
    """
    bowtie2-build --threads ${task.cpus} ${contigs} ${id}_index
    """
}

process bowtie2_align_short {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/bowtie_files", mode: 'copy'
//    container 'quay.io/biocontainers/bowtie2:2.5.1--py310h8d7afc0_1'
    container 'quay.io/biocontainers/mulled-v2-ac74a7f02cebcfcc07d8e8d1d750af9c83b4d45a:a0ffedb52808e102887f6ce600d092675bf3528a-0'

    input:
    tuple val(id), val(output_dir), path(r1), path(r2), val(index_prefix)

    output:
    tuple val(id), val(output_dir), path("${id}.sorted.bam")

    script:
    """
    bowtie2 --very-sensitive-local -x ${index_prefix} \
            -1 ${r1} -2 ${r2} -p ${task.cpus} \
        | samtools view -b - \
        | samtools sort -@ ${task.cpus} -o ${id}.sorted.bam

    samtools index ${id}.sorted.bam
    """
}

process meta_mdbg {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/assembly", mode: 'copy'
    container 'quay.io/biocontainers/metamdbg:1.2--h077b44d_0'

    input:
    tuple val(id), val(output_dir), path(clean_reads)

    output:
    tuple val(id), val(output_dir), path("contigs.fasta"), path("metaMDBG.log")

    script:
    """
    metaMDBG asm --out-dir . --in-ont ${clean_reads}
    gunzip -c contigs.fasta.gz > contigs.fasta
    """
}

process minimap2_map {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/minimap2_align", mode: 'copy'
    container "${params.cache_dir}/biocontainers-minimap2-samtools.sif"
//container 'quay.io/biocontainers/minimap2:2.28--he4a0461_2'

    input:
    tuple val(id), val(output_dir), path(contigs), path(reads)

    output:
    tuple val(id), val(output_dir), path("${id}.sorted.bam"), path("${id}.sorted.bam.bai"), emit: bam
    tuple val(id), val(output_dir), path("${id}_mapping_stats.txt"), emit: stats

    script:
    """
    minimap2 -ax map-ont -t ${task.cpus} ${contigs} ${reads} \
        | samtools view -b -F 4 - \
        | samtools sort -@ ${task.cpus} -o ${id}.sorted.bam
    
    samtools index ${id}.sorted.bam
    
    samtools flagstat ${id}.sorted.bam > ${id}_mapping_stats.txt
    samtools coverage ${id}.sorted.bam >> ${id}_mapping_stats.txt
    """
}

process metabat2 {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/metabat2_bins", mode: 'copy'
    container 'quay.io/biocontainers/metabat2:2.15--h4da6f23_2'

    input:
    tuple val(id), val(output_dir), path(contigs), path(bam), path(bai)

    output:
    tuple val(id), val(output_dir), path("bins/*.fa"), emit: bins, optional: true
    tuple val(id), val(output_dir), path("${id}_depth.txt"), emit: depth
    tuple val(id), val(output_dir), path("${id}_metabat2.log"), emit: log

    script:
    """
    export OMP_NUM_THREADS=${task.cpus}
    mkdir -p bins

    jgi_summarize_bam_contig_depths --outputDepth ${id}_depth.txt ${bam}
    
    metabat2 -i ${contigs} -m 1500 -a ${id}_depth.txt -o bins/${id}_bin -t ${task.cpus} 2>&1 | tee ${id}_metabat2.log
    
    BIN_COUNT=\$(ls -1 bins/${id}_bin.*.fa 2>/dev/null | wc -l)
    echo "Total bins created: \$BIN_COUNT" | tee -a ${id}_metabat2.log
    
    if [ \$BIN_COUNT -eq 0 ]; then
        echo "No bins created" > bins/no_bins_created.txt
    fi
    """
}

process checkm2_predict {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/checkm2_output", mode: 'copy'
    container 'quay.io/biocontainers/checkm2:1.1.0--pyh7e72e81_1'
 
    input:
    tuple val(id), val(output_dir), path(bins)
    path(checkm2_db)
    
    output:
    tuple val(id), val(output_dir), path("checkm2_output/quality_report.tsv"), emit: quality_report
    tuple val(id), val(output_dir), path("checkm2_output/*"), emit: all_results
    tuple val(id), val(output_dir), path(bins), emit: bins_passthrough
    
    script:
    """
    mkdir -p bins_input checkm2_output
    
    if [ -f "${bins}/no_bins_created.txt" ]; then
        echo "No bins available for CheckM2 quality assessment" > checkm2_output/no_bins.log
        touch checkm2_output/quality_report.tsv
        echo -e "Name\tCompleteness\tContamination\tCompleteness_Model_Used\tTranslation_Table_Used\tCoding_Density\tContig_N50\tAverage_Gene_Length\tGenome_Size\tGC_Content\tTotal_Coding_Sequences\tAdditional_Notes" > checkm2_output/quality_report.tsv
    else
        if [ -d "${bins}" ]; then
            cp ${bins}/*.fa bins_input/ 2>/dev/null || true
        else
            cp ${bins} bins_input/ 2>/dev/null || true
        fi
        
        BIN_COUNT=\$(ls -1 bins_input/*.fa 2>/dev/null | wc -l)
        
        if [ \$BIN_COUNT -eq 0 ]; then
            echo "No valid bins found for quality assessment" > checkm2_output/no_bins.log
            touch checkm2_output/quality_report.tsv
            echo -e "Name\tCompleteness\tContamination\tCompleteness_Model_Used\tTranslation_Table_Used\tCoding_Density\tContig_N50\tAverage_Gene_Length\tGenome_Size\tGC_Content\tTotal_Coding_Sequences\tAdditional_Notes" > checkm2_output/quality_report.tsv
        else
            echo "Running CheckM2 on \$BIN_COUNT bins"
            
            checkm2 predict \
                --threads ${task.cpus} \
                --input bins_input \
                --output-directory checkm2_output \
                --database_path ${checkm2_db} \
                --extension fa \
                --force
        fi
    fi
    
    echo "CheckM2 process complete" > checkm2_output/process_complete.txt
    """
}


//process checkm2_predict {
//    tag "$id"
//    cpus params.threads
//    publishDir "${output_dir}/${id}/checkm2_output", mode: 'copy'
//    container 'quay.io/biocontainers/checkm2:1.1.0--pyh7e72e81_1'
// 
//    input:
//    tuple val(id), val(output_dir), path(bins)
//    
//    output:
//    tuple val(id), val(output_dir), path("checkm2_output/quality_report.tsv"), emit: quality_report
//    tuple val(id), val(output_dir), path("checkm2_output/*"), emit: all_results
//    tuple val(id), val(output_dir), path(bins), emit: bins_passthrough
//    
//    script:
//    """
//    mkdir -p bins_input checkm2_output
//    
//    # Check if we have actual bins or just a placeholder
//    if [ -f "${bins}/no_bins_created.txt" ]; then
//        echo "No bins available for CheckM2 quality assessment" > checkm2_output/no_bins.log
//        touch checkm2_output/quality_report.tsv
//        echo -e "Name\\tCompleteness\\tContamination\\tCompleteness_Model_Used\\tTranslation_Table_Used\\tCoding_Density\\tContig_N50\\tAverage_Gene_Length\\tGenome_Size\\tGC_Content\\tTotal_Coding_Sequences\\tAdditional_Notes" > checkm2_output/quality_report.tsv
//    else
//        # Handle both single file and directory inputs
//        if [ -d "${bins}" ]; then
//            cp ${bins}/*.fa bins_input/ 2>/dev/null || true
//        else
//            cp ${bins} bins_input/ 2>/dev/null || true
//        fi
//        
//        # Count bins
//        BIN_COUNT=\$(ls -1 bins_input/*.fa 2>/dev/null | wc -l)
//        
//        if [ \$BIN_COUNT -eq 0 ]; then
//            echo "No valid bins found for quality assessment" > checkm2_output/no_bins.log
//            touch checkm2_output/quality_report.tsv
//            echo -e "Name\\tCompleteness\\tContamination\\tCompleteness_Model_Used\\tTranslation_Table_Used\\tCoding_Density\\tContig_N50\\tAverage_Gene_Length\\tGenome_Size\\tGC_Content\\tTotal_Coding_Sequences\\tAdditional_Notes" > checkm2_output/quality_report.tsv
//        else
//            echo "Running CheckM2 on \$BIN_COUNT bins"
//            
//            checkm2 predict \\
//                --threads ${task.cpus} \\
//                --input bins_input \\
//                --output-directory checkm2_output \\
//                --database_path ${params.checkm2_db} \\
//                --extension fa \\
//                --force
//        fi
//    fi
//    """
//}
process filter_quality_bins {
    tag "$id"
    publishDir "${output_dir}/${id}/filtered_bins", mode: 'copy'

    input:
    tuple val(id), val(output_dir), path(quality_report), path(bins)

    output:
    tuple val(id), val(output_dir), path("high_quality_bins/*.fa"), optional: true, emit: bins  // Add emit name
    path("high_quality_bins/filter_summary.txt"), emit: summary

    script:
    """
    mkdir -p high_quality_bins
    
    echo "=== DEBUG INFO ===" > debug.log
    echo "Quality report:" >> debug.log
    cat ${quality_report} >> debug.log
    echo "" >> debug.log
    echo "Current directory contents:" >> debug.log
    ls -lah >> debug.log
    echo "" >> debug.log
    
    awk -F'\\t' 'NR>1 && \$2>=30 && \$3<=10 {print \$1}' ${quality_report} > high_quality_bin_names.txt
    
    echo "High quality bin names:" >> debug.log
    cat high_quality_bin_names.txt >> debug.log
    echo "" >> debug.log
    
    BIN_COUNT=0
    while IFS= read -r bin_name; do
        if [ -f "\${bin_name}.fa" ]; then
            cp "\${bin_name}.fa" high_quality_bins/
            BIN_COUNT=\$((BIN_COUNT + 1))
            echo "Copied: \${bin_name}.fa" >> debug.log
        elif [ -f "\${bin_name}" ]; then
            cp "\${bin_name}" high_quality_bins/
            BIN_COUNT=\$((BIN_COUNT + 1))
            echo "Copied: \${bin_name}" >> debug.log
        else
            echo "NOT FOUND: \${bin_name}.fa or \${bin_name}" >> debug.log
            echo "Looking for files matching *\${bin_name}*:" >> debug.log
            ls -lah *\${bin_name}* 2>&1 >> debug.log || echo "No matches found" >> debug.log
        fi
    done < high_quality_bin_names.txt
    
    echo "Total bins copied: \$BIN_COUNT" | tee high_quality_bins/filter_summary.txt
    
    if [ \$BIN_COUNT -eq 0 ]; then
        echo "No high-quality bins found (Completeness >= 50%, Contamination <= 10%)" >> high_quality_bins/filter_summary.txt
        touch high_quality_bins/no_bins.txt
    fi
    
    cp debug.log high_quality_bins/
    """
}

process gtdbtk_classify {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}", mode: 'copy'
    container 'quay.io/biocontainers/gtdbtk:2.4.0--pyhdfd78af_1'
 
    input:
    tuple val(id), val(output_dir), path(bins)
    
    output:
    tuple val(id), val(output_dir), path("gtdbtk_output/*"), emit: all_results
    tuple val(id), val(output_dir), path("gtdbtk_output/*.summary.tsv"), emit: summary, optional: true
    path("gtdbtk_output/gtdbtk.log"), emit: log, optional: true
    
    script:
    """
    mkdir -p gtdbtk_output bins_input

    if [ -f "no_bins.txt" ] || [ ! \$(ls *.fa 2>/dev/null | wc -l) -gt 0 ]; then
        echo "No high-quality bins available for GTDB-Tk classification" > gtdbtk_output/no_bins.log
        touch gtdbtk_output/gtdbtk.log
    else
        cp *.fa bins_input/
        
        BIN_COUNT=\$(ls -1 bins_input/*.fa 2>/dev/null | wc -l)
        echo "Running GTDB-Tk on \$BIN_COUNT bins"
        
        export GTDBTK_DATA_PATH=${params.gtdb_tk_db}
        
        gtdbtk classify_wf \\
            --genome_dir bins_input \\
            --out_dir gtdbtk_output \\
            --extension fa \\
            --cpus ${task.cpus} \\
            --pplacer_cpus ${task.cpus} \\
            --skip_ani_screen \\
            2>&1 | tee gtdbtk_output/gtdbtk.log
        
        if [ -f "gtdbtk_output/gtdbtk.bac120.summary.tsv" ] || [ -f "gtdbtk_output/gtdbtk.ar53.summary.tsv" ]; then
            echo "Classification complete. Results:"
            wc -l gtdbtk_output/*.summary.tsv 2>/dev/null || echo "No summary files generated"
        fi
    fi
    """
}

process gunc_run {
    tag "${id}"
    cpus params.threads
    publishDir "${output_dir}/${id}/gunc_output", mode: 'copy'
    container 'quay.io/biocontainers/gunc:1.0.5--pyhdfd78af_0'

    input:
    tuple val(id), val(output_dir), path(quality_bins)

    output:
    tuple val(id), val(output_dir), path("final_bins/*.fa"), optional: true, emit: filtered_bins
    tuple val(id), val(output_dir), path("gunc_out/*"), emit: gunc_results
    path("gunc_summary.txt"), emit: summary

    script:
    """
    mkdir -p gunc_out final_bins bins_input
    
    # Check if we have bins or just a placeholder
    if [ -f "no_bins.txt" ] || [ ! \$(ls *.fa 2>/dev/null | wc -l) -gt 0 ]; then
        echo "No bins available for GUNC analysis" > gunc_out/no_bins.log
        echo "No bins to filter" > gunc_summary.txt
        touch final_bins/no_bins.txt
    else
        # Copy bins to input directory
        cp *.fa bins_input/ 2>/dev/null || true
        
        BIN_COUNT=\$(ls -1 bins_input/*.fa 2>/dev/null | wc -l)
        echo "Running GUNC on \$BIN_COUNT bins"
        
        if [ \$BIN_COUNT -eq 0 ]; then
            echo "No valid bins found for GUNC analysis" > gunc_out/no_bins.log
            echo "No bins to filter" > gunc_summary.txt
            touch final_bins/no_bins.txt
        else
            # Run GUNC
            gunc run --db_file ${params.gunc_db} \\
                     --input_dir bins_input \\
                     --threads ${task.cpus} \\
                     --out_dir gunc_out \\
                     --file_suffix .fa
            
            # Filter bins that pass GUNC (pass.GUNC column is True)
            if [ -f "gunc_out/GUNC.progenomes_2.1.maxCSS_level.tsv" ]; then
                PASSED=0
                FAILED=0
                
                # Column 13 is pass.GUNC field
                awk -F'\\t' 'NR>1 && \$13=="True" {print \$1}' gunc_out/GUNC.progenomes_2.1.maxCSS_level.tsv | while read bin; do
                    if [ -f "bins_input/\${bin}.fa" ]; then
                        cp "bins_input/\${bin}.fa" final_bins/
                        PASSED=\$((PASSED + 1))
                    fi
                done
                
                TOTAL=\$(awk 'NR>1' gunc_out/GUNC.progenomes_2.1.maxCSS_level.tsv | wc -l)
                PASSED=\$(ls -1 final_bins/*.fa 2>/dev/null | wc -l)
                FAILED=\$((TOTAL - PASSED))
                
                echo "GUNC Filtering Summary:" > gunc_summary.txt
                echo "Total bins analyzed: \$TOTAL" >> gunc_summary.txt
                echo "Bins passed: \$PASSED" >> gunc_summary.txt
                echo "Bins failed (potential contamination/chimerism): \$FAILED" >> gunc_summary.txt
                
                if [ \$PASSED -eq 0 ]; then
                    echo "No bins passed GUNC filtering" >> gunc_summary.txt
                    touch final_bins/no_bins.txt
                fi
            else
                echo "GUNC output file not found" > gunc_summary.txt
                touch final_bins/no_bins.txt
            fi
        fi
    fi
    """
}


process checkv_run {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/checkv_output", mode: 'move'
    container 'quay.io/staphb/checkv:1.0.3'

    input:
    tuple val(id), val(output_dir), path(contigs)

    output:
    tuple val(id), path("*")

    script:
    """
    checkv end_to_end "${contigs}" checkv_output -t ${task.cpus} -d ${params.checkv_db}
    """
}

process genomad_run {
    tag "$id"
    cpus params.threads
    publishDir "${output_dir}/${id}/genomad_output", mode: 'move'
    container 'quay.io/biocontainers/genomad:1.11.2--pyhdfd78af_0'

    input:
    tuple val(id), val(output_dir), path(contigs)

    output:
    tuple val(id), path("*")

    script:
    """
    genomad end-to-end --cleanup ${contigs} genomad_output ${params.genomad_db} -t ${task.cpus}
    """
}

process download_checkm2_db {
    tag "checkm2_database"
    cpus 1
    storeDir "${params.ref_data_dir}/checkM2_DB"
    
    output:
    path("*"), emit: db_dir
    
    script:
    """
    if [ -f "uniref100.KO.1.dmnd" ]; then
        echo "CheckM2 database already exists, skipping download"
    else
        echo "Downloading CheckM2 database..."
        
        MAX_RETRIES=5
        RETRY_COUNT=0
        SUCCESS=0
        
        while [ \$RETRY_COUNT -lt \$MAX_RETRIES ] && [ \$SUCCESS -eq 0 ]; do
            echo "Download attempt \$((RETRY_COUNT + 1)) of \$MAX_RETRIES..."
            
            if wget -T 60 -O checkm2_database.tar.gz \
                https://zenodo.org/records/14897628/files/checkm2_database.tar.gz; then
                SUCCESS=1
                echo "Download successful!"
            else
                RETRY_COUNT=\$((RETRY_COUNT + 1))
                if [ \$RETRY_COUNT -lt \$MAX_RETRIES ]; then
                    echo "Download failed. Waiting 10 seconds before retry..."
                    sleep 10
                    rm -f checkm2_database.tar.gz
                fi
            fi
        done
        
        if [ \$SUCCESS -eq 0 ]; then
            echo "ERROR: Failed to download database after \$MAX_RETRIES attempts"
            exit 1
        fi
        
        echo "Extracting database..."
        tar -xzf checkm2_database.tar.gz --strip-components=1
        rm checkm2_database.tar.gz
        echo "Database setup complete!"
    fi
    """
}

process download_gtdbtk_db {
    tag "gtdbtk_database"
    cpus 16
    memory '8 GB'
    container 'docker://alpine:latest'
    containerOptions '--writable-tmpfs'
    storeDir "${params.ref_data_dir}/gtdbtk_release226"
    
    output:
    path("release226"), emit: db_dir
    
    shell:
    '''
    if [ -d "release226" ] && [ -f "release226/metadata/metadata.txt" ]; then
        echo "GTDB-Tk database already exists, skipping download"
    else
        apk add --no-cache aria2
        
        aria2c \
            -x 16 \
            -s 16 \
            -k 1M \
            -c \
            --max-tries=0 \
            --retry-wait=10 \
            --max-connection-per-server=16 \
            --min-split-size=1M \
            --file-allocation=none \
            --auto-file-renaming=false \
            --allow-overwrite=true \
            --disk-cache=64M \
            -d . \
            -o gtdbtk_data.tar.gz \
            https://data.ace.uq.edu.au/public/gtdb/data/releases/latest/auxillary_files/gtdbtk_package/full_package/gtdbtk_data.tar.gz
        
        tar -xzf gtdbtk_data.tar.gz
        rm gtdbtk_data.tar.gz
    fi
    '''
}


workflow mt_short_read_workflow {
//    fastqc(short_metaT_ch)
    bbtools_qc_out = bbtools_qc(short_metaT_ch)
    bbtools_rrna_out = bbtools_rrna_removal_short_read(bbtools_qc_out.map { id, output_dir, r1, r2, stats -> tuple(id, output_dir, r1, r2) })
    bbtools_human_out = bbtools_human_removal(bbtools_rrna_out.map { id, output_dir, r1_norrna, r2_norrna, rrna_stats -> tuple(id, output_dir, r1_norrna, r2_norrna) })
//    bbtools_phix_out = bbtools_phix_removal(bbtools_human_out.map { id, output_dir, r1_nohuman, r2_nohuman, human_stats -> tuple(id, output_dir, r1_nohuman, r2_nohuman) })
//    gottcha_out = gottcha2(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) })
//    fastqc_cleaned(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) })
//    metaspades_out = metaspades(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) })
//    metaquast_out = metaquast(metaspades_out)
//    bowtie2_build_out = bowtie2_build(metaspades_out)
//    build_ch_flat = bowtie2_build_out.map { id, output_dir, files -> def f = files.find { it.name.endsWith('.1.bt2') }; tuple(id, output_dir, f.toString().replaceAll(/\.1\.bt2$/, '')) }
//    phix_ch_flat = bbtools_phix_out.map { id, output_dir, r1, r2, stats -> tuple(id, r1, r2) }
//    bowtie2_align_out = bowtie2_align_short(build_ch_flat.join(phix_ch_flat).map { id, output_dir, index_prefix, r1, r2 -> tuple(id, output_dir, r1, r2, index_prefix) })
//    chimeracutter_out = chimeracutter(metaspades_out)
//    checkv_out = checkv_run(metaspades_out)
//    genomad_output = genomad_run(metaspades_out)
}

workflow mg_short_read_workflow {
    checkm2_db = download_checkm2_db()
    gtdbtk_db = download_gtdbtk_db()
//    fastqc(short_metaG_ch)
//    bbtools_qc_out = bbtools_qc(short_metaG_ch)
//    bbtools_rrna_out = bbtools_rrna_removal_short_read(bbtools_qc_out.map { id, output_dir, r1, r2, stats -> tuple(id, output_dir, r1, r2) })
//    bbtools_human_out = bbtools_human_removal(bbtools_rrna_out.map { id, output_dir, r1_norrna, r2_norrna, rrna_stats -> tuple(id, output_dir, r1_norrna, r2_norrna) })
//    bbtools_phix_out = bbtools_phix_removal(bbtools_human_out.map { id, output_dir, r1_nohuman, r2_nohuman, human_stats -> tuple(id, output_dir, r1_nohuman, r2_nohuman) })
//    metaspades_out = metaspades(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) })
//    bowtie2_build_out = bowtie2_build(metaspades_out)
//    build_ch_flat = bowtie2_build_out.map { id, output_dir, files -> 
//        def f = files.find { it.name.endsWith('.1.bt2') }
//        tuple(id, output_dir, f.toString().replaceAll(/\.1\.bt2$/, '')) 
//    }
//    phix_ch_flat = bbtools_phix_out.map { id, output_dir, r1, r2, stats -> tuple(id, r1, r2) }
//    bowtie2_align_out = bowtie2_align_short(build_ch_flat.join(phix_ch_flat).map { id, output_dir, index_prefix, r1, r2 -> tuple(id, output_dir, r1, r2, index_prefix) })
//    
//    metabat2_input = metaspades_out.join(bowtie2_align_out, by: [0, 1]).map { id, output_dir, contigs, bam -> 
//        tuple(id, output_dir, contigs, bam, file("${bam}.bai"))
//    }
//    metabat2_out = metabat2(metabat2_input)
//   gottcha_out = gottcha2(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) }) 
//    checkm2_out = checkm2_predict(metabat2_out.bins, checkm2_db.db_dir)
//    checkm2_out = checkm2_predict(metabat2_out.bins)
//    filter_input = checkm2_out.quality_report.join(checkm2_out.bins_passthrough, by: [0, 1]).map { id, output_dir, quality_report, bins -> 
//        tuple(id, output_dir, quality_report, bins) 
//    }
//    filtered_bins = filter_quality_bins(filter_input)
//    gunc_out = gunc_run(filtered_bins.bins)
//    gtdbtk_out = gtdbtk_classify(filtered_bins.bins)
//    fastqc_cleaned(bbtools_phix_out.map { id, output_dir, r1_nophix, r2_nophix, phix_stats -> tuple(id, output_dir, r1_nophix, r2_nophix) })
//    metaquast_out = metaquast(metaspades_out)
//    chimeracutter_out = chimeracutter(metaspades_out)
//    checkv_out = checkv_run(metaspades_out)
//    genomad_output = genomad_run(metaspades_out)
}

workflow long_read_workflow {
    long_metaG_ch.branch {
        pod5: it[2].name.endsWith('.pod5')
        fastq: true  
    }.set { branched_reads }
    
    dorado_out = dorado(branched_reads.pod5)
    fastq_from_dorado = dorado_out.map { id, output_dir, bam_file, fastq -> 
        tuple(id, output_dir, fastq) 
    }
    all_fastq = fastq_from_dorado.mix(branched_reads.fastq)
    chopper_out = chopper(all_fastq)
    bbtools_rrna_out = bbtools_rrna_removal_long_read(chopper_out)
//    gottcha_out = gottcha2_long_read(bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
//        tuple(id, output_dir, clean_reads) 
//    })
    meta_mdbg_out = meta_mdbg(bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
        tuple(id, output_dir, clean_reads) 
    })
    contigs_ch = meta_mdbg_out.map { id, output_dir, contigs, log -> 
        tuple(id, output_dir, contigs) 
    }
    reads_ch = bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
        tuple(id, output_dir, clean_reads) 
    }
    minimap2_input = contigs_ch.join(reads_ch, by: [0, 1]).map { id, output_dir, contigs, reads -> 
        tuple(id, output_dir, contigs, reads) 
    }
    minimap2_out = minimap2_map(minimap2_input)
    metabat2_input = contigs_ch.join(minimap2_out.bam, by: [0, 1]).map { id, output_dir, contigs, bam, bai -> 
        tuple(id, output_dir, contigs, bam, bai) 
    }
    metabat2_out = metabat2(metabat2_input)
    checkm2_out = checkm2_predict(metabat2_out.bins)
    filter_input = checkm2_out.quality_report.join(checkm2_out.bins_passthrough, by: [0, 1]).map { id, output_dir, quality_report, bins -> 
        tuple(id, output_dir, quality_report, bins) 
    }
    filtered_bins = filter_quality_bins(filter_input)
    
    // Run GUNC on CheckM2-filtered bins (for contamination detection)
    gunc_out = gunc_run(filtered_bins.bins)
    
    // GTDB-Tk runs on ALL CheckM2-filtered bins (not just GUNC-passed bins)
    gtdbtk_out = gtdbtk_classify(filtered_bins.bins)
    
    checkv_out = checkv_run(contigs_ch)
    genomad_out = genomad_run(contigs_ch)
}


//workflow long_read_workflow {
//    long_metaG_ch.branch {
//        pod5: it[2].name.endsWith('.pod5')
//        fastq: true  
//    }.set { branched_reads }
//    
//    dorado_out = dorado(branched_reads.pod5)
//    fastq_from_dorado = dorado_out.map { id, output_dir, bam_file, fastq -> 
//        tuple(id, output_dir, fastq) 
//    }
//    all_fastq = fastq_from_dorado.mix(branched_reads.fastq)
//    chopper_out = chopper(all_fastq)
//    bbtools_rrna_out = bbtools_rrna_removal_long_read(chopper_out)
//    gottcha_out = gottcha2_long_read(bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
//    tuple(id, output_dir, clean_reads) 
//    })
//    meta_mdbg_out = meta_mdbg(bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
//        tuple(id, output_dir, clean_reads) 
//    })
//    contigs_ch = meta_mdbg_out.map { id, output_dir, contigs, log -> 
//        tuple(id, output_dir, contigs) 
//    }
//    reads_ch = bbtools_rrna_out.map { id, output_dir, clean_reads, rrna_stats -> 
//        tuple(id, output_dir, clean_reads) 
//    }
//    minimap2_input = contigs_ch.join(reads_ch, by: [0, 1]).map { id, output_dir, contigs, reads -> 
//        tuple(id, output_dir, contigs, reads) 
//    }
//    minimap2_out = minimap2_map(minimap2_input)
//    metabat2_input = contigs_ch.join(minimap2_out.bam, by: [0, 1]).map { id, output_dir, contigs, bam, bai -> 
//        tuple(id, output_dir, contigs, bam, bai) 
//    }
//    metabat2_out = metabat2(metabat2_input)
//    checkm2_out = checkm2_predict(metabat2_out.bins)
//    filter_input = checkm2_out.quality_report.join(checkm2_out.bins_passthrough, by: [0, 1]).map { id, output_dir, quality_report, bins -> 
//        tuple(id, output_dir, quality_report, bins) 
//    }
//    filtered_bins = filter_quality_bins(filter_input)
//    gunc_out = gunc_run(filtered_bins.bins)
//
//    gtdbtk_out = gtdbtk_classify(filtered_bins.bins)
////    chimeracutter_out = chimeracutter(contigs_ch)
//    checkv_out = checkv_run(contigs_ch)
//    genomad_out = genomad_run(contigs_ch)
//}

workflow {
//    mt_short_read_workflow()
    mg_short_read_workflow()
//    long_read_workflow()
}
