cgmlst_main.nf

nextflow.enable.dsl=2

// ============================================================
// cgMLST / wgMLST pipeline using pyMLST
// Author: Mikra
// ============================================================

// -----------------------------
// Parameters and Help Section
// -----------------------------
params.species = ""
params.input   = ""
params.output  = "results_cgmlst"
params.help    = false
params.subgraph_threshold = 5

if (params.help) {
    println """
🧬 cgMLST / wgMLST pipeline — perform whole-genome MLST typing with pyMLST

Usage:
  nextflow run cgmlst_main.nf --species "Salmonella enterica" --input /path/to/fasta_dir --output results_cgmlst

Required arguments:
  --species   Name of the bacterial species (exactly as listed below)
  --input     Directory containing .fa or .fasta assemblies
  --output    Directory where results will be stored

Example:
  nextflow run cgmlst_main.nf --species "Salmonella enterica" --input ./assemblies --output ./cgmlst_out

✅ Supported species list (see full list on cgmlst.org)
------------------------------------------------------------
Acinetobacter baumannii
Bacillus anthracis
Bordetella pertussis
Brucella melitensis
Burkholderia pseudomallei
Campylobacter jejuni/coli
Clostridioides difficile
Corynebacterium diphtheriae
Enterococcus faecalis
Enterococcus faecium
Escherichia coli
Klebsiella pneumoniae
Legionella pneumophila
Listeria monocytogenes
Pseudomonas aeruginosa
Salmonella enterica
Staphylococcus aureus
Streptococcus pyogenes
Yersinia enterocolitica
------------------------------------------------------------
🔗 Visit https://cgmlst.org/ for complete scheme details.
"""
    System.exit(0)
}

// -----------------------------
// Validation
// -----------------------------
if (!params.species) exit 1, "❌ Missing required parameter: --species"
if (!params.input)   exit 1, "❌ Missing required parameter: --input"
if (!file(params.input).exists()) exit 1, "❌ Input folder not found: ${params.input}"

// -----------------------------
// Process: Run cgMLST with pyMLST
// -----------------------------
process cgmlst_run {
    tag "${params.species}"
    label 'cgmlst'
    conda "envs/pymlst.yml"

    publishDir "${params.output}", mode: 'copy'
    
    output:
    path "${params.output}/msa.fasta", emit: msa
    path "${params.output}/subgraph.tsv"
    path "${params.output}/matrix_inf.tsv"

    script:
    def DB_NAME = "DATABASE"
    def DB_PATH = "${params.output}/${DB_NAME}"
    def INPUT_DIR = file(params.input)

    """
    set -euo pipefail

    echo "▶ Starting cgMLST pipeline for species: ${params.species}"
    echo "📂 Input directory: ${INPUT_DIR}"
    echo "📁 Output directory: ${params.output}"
    echo "🗃️  Database: ${DB_PATH}"

    mkdir -p "${params.output}/tmp_fastas"

    echo "▶ Step 1: Importing cgMLST scheme for ${params.species}"
    (cd "${params.output}" && wgMLST import "${DB_NAME}" "${params.species}")

    echo "▶ Step 2: Renaming fasta files (replace '-' with '_') and copying"
    shopt -s nullglob
    for file in "${INPUT_DIR}"/*.fa "${INPUT_DIR}"/*.fasta; do
        [ -e "\$file" ] || continue
        base=\$(basename "\$file")
        clean="\${base//-/_}"
        cp "\$file" "${params.output}/tmp_fastas/\$clean"
    done
    shopt -u nullglob

    echo "▶ Step 3: Adding assemblies to database"
    for i in "${params.output}/tmp_fastas"/*.fa "${params.output}/tmp_fastas"/*.fasta; do
        [ -e "\$i" ] || continue
        base=\$(basename "\$i")
        base="\${base%.fa}"
        base="\${base%.fasta}"
        wgMLST add -s "\$base" "${DB_PATH}" "\$i"
    done

    echo "▶ Step 4: Calculating distance matrix"
    wgMLST distance "${DB_PATH}" -o "${params.output}/matrix_inf.tsv"

    echo "▶ Step 5: Creating subgraph (threshold ${params.subgraph_threshold})"
    wgMLST subgraph -e group -t ${params.subgraph_threshold} \\
        "${params.output}/matrix_inf.tsv" -o "${params.output}/subgraph.tsv"

    echo "▶ Step 6: Generating multiple sequence alignment (MSA)"
    wgMLST msa -o "${params.output}/msa.fasta" "${DB_PATH}"

    echo "✅ cgMLST pipeline completed successfully for ${params.species}"

    """
}

process run_grapetree {
    tag "grapetree"
    conda "envs/grapetree.yml"

    publishDir "${params.output}", mode: 'copy'


    input:
    path msa

    output:
    path "MSTree.tsv"

    script:
    """
    echo "▶ Running GrapeTree on msa.fasta"
    grapetree -p ${msa} > MSTree.tsv
    echo "✅ GrapeTree completed: MSTree.tsv generated"
    """
}


// -----------------------------
// Workflow definition
// -----------------------------
workflow {
    cgmlst_ch=cgmlst_run()
    run_grapetree(cgmlst_ch.msa)
}

// -----------------------------
// Completion message
// -----------------------------
workflow.onComplete {
    println ""
    println "🎉 cgMLST analysis completed successfully!"
    println "📁 Results saved to: ${params.output}"
    println "🔗 Visit https://cgmlst.org/ for more scheme details."
}