main.nf

#!/usr/bin/env nextflow

nextflow.enable.dsl=2

// Modules to get genomes and list of BUSCO genes
include { download_NCBI } from './modules/download.nf'
include { extract_longest_transcript } from './modules/extract_longest_transcript.nf'
include { runBusco; plotBusco } from './modules/run_busco.nf'
include { runOrthoFinder } from './modules/run_orthofinder.nf'
include { annotateGO } from './modules/annotate_og.nf'

// Make BUSCO gene trees
include { extractBuscoGenes; extractBusco; alignBusco; trimAlign; runBuscoTrees } from './modules/makeGeneTree.nf'

// Species tree preping
include { makeConsensusMCMC; prepMCMCtree; plotMCMCtree } from './modules/makeSpeciesTree.nf'

// Run multiple fossil-calibrated trees
include { MCMCTREE as MCMCTREE_1 } from './modules/makeSpeciesTree.nf'
include { MCMCTREE as MCMCTREE_2 } from './modules/makeSpeciesTree.nf'
include { MCMCTREE as MCMCTREE_3 } from './modules/makeSpeciesTree.nf'

// Run multiple CAFE runs
include { filterHogs } from './modules/cafe.nf'
include { runCAFE as runCAFE_1 } from './modules/cafe.nf'
include { runCAFE as runCAFE_2 } from './modules/cafe.nf'
include { runCAFE as runCAFE_3 } from './modules/cafe.nf'
include { runCAFE as runCAFE_4 } from './modules/cafe.nf'

// Downstream analyses
include { enrichmentGO } from './modules/go_enrichment.nf'

// Define the workflow
workflow {

    // Run download workflow
    def downloaded_files = download_NCBI(params.species_list)

    // Map downloaded file paths to tuple (file, species)
    def proteomes = downloaded_files.protein_faa
        .flatten()
        .map { file ->
            def filename = file.toString().trim().tokenize('/').last()
            def species = filename.tokenize('.')[0]
            tuple(file, species)
        }

    // Extract longest transcript
    def orf = extract_longest_transcript(proteomes)
   
    // Run BUSCO on proteomes
    def busco = runBusco(proteomes)

     // Annotate GO terms
    def annot_go = annotateGO(busco.species)

    // Collect unique BUSCO result directories
    def busco_results = busco.busco_dir.unique()
        .collect()
        .map { list -> list[0] }
    
    // Plot BUSCO results after all BUSCO processes are complete
    plotBusco(busco_results)

    // Collect all longest transcripts paths
    def longest_transcripts_dirs = orf.output_dir
        .unique()
        .collect()
        .map { list -> list[0] }

    // Run OrthoFinder on longest transcripts
    def orthofinder = runOrthoFinder(longest_transcripts_dirs)

    // Part II: Make Gene Trees & Run CAFE

    // Extract species names from the ncbi_genomes file
    Channel
        .fromPath(params.species_list)
        .splitCsv(header: false, sep: '\t')
        .map { columns -> columns[1] }
        .set { species }

    // 1) Extract BUSCO genes
    def busco_genes_ext = extractBuscoGenes(orthofinder.fasta_dir)

    // Directly use sco_list
    busco_genes_ext.sco_list
        .distinct()
        .set { sco_busco_results }

    // 2) Extract and align BUSCO genes
    def ext_busco = extractBusco(species, 
                                 sco_busco_results)

    // Collect and distinct SCO files, then merge them into a single file
    Channel
        ext_busco.sco
        .collect()
        .distinct()
        .map { list -> list[0] }
        .splitText() { it.trim() }
        .set { ext_busco_results }

    // 3) Align BUSCO genes
    def align_busco = alignBusco(ext_busco_results)
    
    // 4) Trim and realign BUSCO genes
    trimAlign(align_busco.flatten())

    // Make tree from each BUSCO
    trimAlign.out
        .map { tuple -> tuple[0] }
        .flatten()
        .set { trim_faa }

    // 5) Make Busco trees
    runBuscoTrees(trim_faa)

    // Collect all treefiles
    runBuscoTrees.out
        .collect()
        .set { busco_trees }

    // 6) Make consensus species-tree
    def makeConsensus = makeConsensusMCMC(busco_trees)

    // 7) Make mega-MSA
    prepMCMCtree(makeConsensus.nwk)

    // 8) Estimate divergence time
    MCMCTREE_1(prepMCMCtree.out.ali, makeConsensus.nwk, 1)
    MCMCTREE_2(prepMCMCtree.out.ali, makeConsensus.nwk, 2)
    MCMCTREE_3(prepMCMCtree.out.ali, makeConsensus.nwk, 3)

    // 9) Plot MCMCtree and make prep-file for CAFE
    MCMCTREE_1.out.mcmc
        .concat(MCMCTREE_2.out.mcmc,
                MCMCTREE_3.out.mcmc)
        .unique()
        .set { mcmc_files }

    plotMCMCtree(mcmc_files)

    // 10) Run CAFE on species-tree
    filterHogs(params.hog, 
               orthofinder.fasta_dir)

    // 11) Run CAFE with multiple settings
    runCAFE_1(filterHogs.out,
              plotMCMCtree.out.cafe_input_tree,
              "-o r1_hogs")

    runCAFE_2(filterHogs.out,
              plotMCMCtree.out.cafe_input_tree,
              "-k 3 -o r1_k3_hogs")

    runCAFE_3(filterHogs.out,
              plotMCMCtree.out.cafe_input_tree,
              "-k 3 -p -o r1_k3_p_hogs")

    runCAFE_4(filterHogs.out,
              plotMCMCtree.out.cafe_input_tree,
              "-e -o r1_e_hogs")

    // 12) Run GO term enrichment}
    runCAFE_1.out.cafe_out
        .concat(runCAFE_2.out.cafe_out,
                runCAFE_3.out.cafe_out,
                runCAFE_4.out.cafe_out)
        .unique()
        .set { cafe_ch }
    enrichmentGO(cafe_ch)

}