Snakefile

# May the data passing through this pipeline
# somehow help to bring just a little more peace 
# in this troubled world.

__author__ = "Carl M. Kobel"

__version__ = "2.10.1" # Haven't yet found a neat way of passing the version from the comparem2 executable, so for now it has to be defined here as well.

# https://semver.org/
version_major = ".".join(__version__.split(".")[:1])
version_minor = ".".join(__version__.split(".")[:2])
version_patch = ".".join(__version__.split(".")[:3])


import os
from os import listdir
from os.path import isfile, join
import pandas as pd
import numpy as np

# TODO: Test disabling these:
from shutil import copyfile
import csv # for quoting # Can be disabled?
import subprocess # For void_report
import datetime # For void_report


# Docker image is only updated for each minor version. This reduces overhead related to storage, transfer, and debugging.
containerized: f"docker://cmkobel/comparem2:v{version_minor}"

# When executing, Snakemake will fail with a reasonable error message if the variables below are undefined.
envvars:
    "COMPAREM2_BASE",
    "COMPAREM2_PROFILE",
    "COMPAREM2_DATABASES",


# --- Functions and routines ----------------------------------------

def get_fofn(path):
    """ Reads a newline separated list of input_genomes files to process.
        Uses the pandas csv reader and reads only the first value in each row.
    """
    df = pd.read_csv(path, header = None, delimiter = "\t", comment = "#")
    as_list = df[0].to_list()

    return as_list
    


def get_input_genomes(pattern):
    """ Runs ls on the pattern and parses what is returned. Uses module subprocess.
    """

    # Construct full command.
    # Using ls is the safest option, as the user will most likely be comfortable with how it globs.
    command = "ls -1 " + pattern

    # Run ls as a subprocess.
    ls = subprocess.run(command, shell = True, capture_output = True) # Apparently, shell = True is necessary when using advanced globbing symbols.

    # Parse and return
    decoded = ls.stdout.decode('utf-8').split("\n")
    decoded_nonempty = [i for i in decoded if i != ''] # If there are no matches, we must remove the empty result. Also, there is always an empty result in the end because ls returns a final newline.
    return decoded_nonempty
        
    
def interpret_true(text):
    return str(text).strip().lower() == "true"
    

def passthrough_parameter_unpack(rule_name):
    """ This function is responsible for the parameter passthrough 
    functionality. Given a rule name, it filters for the wanted keys and 
    formats the key-value pairs as a "shell" string. This shell string can
    then be pasted directly into the shell command for each program. 
    """
    
    # Add constant "set" prefix to differentiate the config settings and to avoid key name conflicts.
    prefix = "set_" + rule_name
    
    # Create a new dict that only contains the keys that start with the wanted prefix.
    rule_parameters = {k: v for k, v in config.items() if k.startswith(prefix)}
    #print(rule_parameters) # DEBUG
    
    
    # Create parameter list with spaces between key-value.
    formatted_list = [str(k[len(prefix):]) + " " + str(v) for k,v in rule_parameters.items()]    
    rv = " ".join(formatted_list)

    # Pretty print
    print("passthrough parameters to", rule_name, end = ": \"")
    print(rv, end = "")
    print("\"")
    
    # Format as "shell" and return
    return rv
    
    
# --- Field variables -----------------------------------------------

# User can manually set the title from the config.
if config["title"] != "CHANGE_THIS_TO_ACTIVATE": 
    batch_title = str(config["title"]).strip()
else: # Otherwise the default is to use the name of the current directory.
    batch_title = os.getcwd().split("/")[-1]

# Read environment variables.
base_variable = os.environ['COMPAREM2_BASE'] # rename to COMPAREM2_BASE
DATABASES = os.environ['COMPAREM2_DATABASES'] # Defines where the databases are stored. One for all. when snakemake issue 262 is solved I'll make this more flexible for each rule.

# Other constants.
output_directory = config["output_directory"] # Should be renamed to "output_directory".
void_report = f"date -Iseconds >> {output_directory}/.comparem2_void_report.flag" # The modification time of this file tells the report subpipeline whether it needs to run. Thus, void_report is called in the end of every successful rule.



# --- Header --------------------------------------------------------

with open(f"{base_variable}/resources/logo.txt") as logo_file:
    for line in logo_file:
        print(line, end = "")

print()
print("                      Formerly known as Assemblycomparator2              ")
print("                       github.com/cmkobel/comparem2/issues               ")
print("                             comparem2.readthedocs.io                    ")
print(f"                                     v{__version__}")
print("                                                                         ")
print("  Variables                                                              ")
print("  ---------                                                              ")
print(f"    title         : '{batch_title}'")
print(f"    base          : '{base_variable}'")
print(f"    databases     : '{DATABASES}'")
print("                                                                         ")
print("  Available rules                                                        ")
print("  ---------------                                                        ")
print("    Q.C.          : copy assembly_stats sequence_lengths checkm2 busco   ")
print("    Annotate      : prokka bakta                                         ")
print("    Adv. annotate : eggnog interproscan dbcan kegg_pathway abricate mlst ")
print("    Core-pan      : panaroo                                              ")
print("    Phylogenetic  : mashtree fasttree gtdbtk iqtree treecluster snp_dists")
print("    Pseudo        : meta isolate downloads fast report                   ")
print("  (Use ´--until <rule> [<rule2>..]´ to activate only one or more rules)  ")
    

# --- Parse input files -------------------------------------------------------


# Uses snakemakes built in config system to set default and custom parameters. 

if config['fofn'] != "CHANGE_THIS_TO_ACTIVATE": # Read in the fofn, only if the default value is changed (in the config/config.yaml file or command line --config fofn=path/to/fofn.txt)
    input_genomes_parsed = get_fofn(config['fofn'])
else: # Not using, fofn, use input_genomes which has a default value of "*.fna *.fa *.fasta *.fas", but can also be changed.
    input_genomes_parsed = get_input_genomes(config['input_genomes']) # Run with 'comparem2 --config input_genomes="my_files*.fna' to customize.

if len(input_genomes_parsed) < 1:
    raise Exception(f"Could not find {config['input_genomes']}. Quitting ...")



# --- Construct sample table ----------------------------------------

df = pd.DataFrame(data = {'input_file': input_genomes_parsed})

# Check that the directory is not empty.
N = df.shape[0]
if N == 0:
    raise Exception("Error: No fasta files in the current directory. Quitting ...")
    #raise Exception("Zero genomic files present.")
    sys.exit(1)

# Check if there are missing files.
missing_files = [file for file in df["input_file"].tolist() if not os.path.isfile(file)]
if len(missing_files) > 0: # If any files do not exist, inform and exit.
    print("Pipeline: Error, all input files must exist.")
    raise FileNotFoundError(f"The following files are not found: {missing_files}")


df['basename'] = [os.path.basename(i) for i in df['input_file'].tolist()]

# Since mashtree doesn't like spaces in filenames, we must convert those.
df['basename_clean'] = df['basename'].str.replace(' ','_').str.replace(',','_').str.replace('"','_').str.replace('\'','_') # Convert punctuation marks to underscores: Makes everything easier.
df['sample'] = [".".join((i.split(".")[:-1])) for i in df['basename_clean']] # Remove extension by splitting, removing last, and joining.
df['extension'] =  [i.split(".")[-1] for i in df['input_file'].tolist()] # Extract extension
df['input_file_fasta'] = output_directory + "/samples/" + df['sample'] + "/" + df['sample'] + ".fna" # This is where the input file is copied to in the first snakemake rule "copy".
df['1-index'] = [i+1 for i in range(len(df))] # Python usually uses 0-index, which may not be intuitive for lay person users. To make sure no one is confused, we name the index as a "one"-index.

# Check that list of input files is non-empty.
if df.shape[0] == 0:
    raise Exception("Error: No fasta files in the current directory. Quitting ...(2)")
    #raise Exception("Zero genomic files present.")
    sys.exit(1)


  
# --- Display sample table ------------------------------------------

print() # Visual padding
print("  Sample overview")
print("  ---------------")
df = df.reset_index(drop = True)
#print(df[['input_file', 'sample', 'extension']])
#print(df[['input_file', 'extension', 'input_file_fasta']])
#print(df[['1-index', 'sample', 'extension']].to_string(index = False))
#print(df[['1-index', 'input_file', 'basename', 'sample', 'extension', 'input_file_fasta']].to_string(index = False))
print(df[['1-index', 'input_file', 'sample']].to_string(index = False))

print("//")
print()

# Warn the user if there exists spaces in the file names.
if any([" " in i for i in df['input_file'].tolist()]): # TODO test if it still works after new globbing system.
    print("Warning: One or more file names contain space(s). These have been replaced with underscores \" \" -> \"_\"")

# Check if the sample names are unique
duplicates = df[df.duplicated(['sample'])]
n_duplicates = duplicates.shape[0]
if n_duplicates > 0:
    raise Exception(f"Error: Sample names are not unique. The following ({n_duplicates}) input genome(s) are duplicated:\n{duplicates.to_string(index = False)}")


# The DATABASES directory must exist, otherwise apptainer gets confused and throws the following:
# WARNING: skipping mount of /home/thylakoid/comparem2/databaseas: stat /home/thylakoid/comparem2/databaseas: no such file or directory
if not os.path.isdir(DATABASES):
    os.mkdir(DATABASES)


# --- Make sure the output directory exists. ------------------------
if not os.path.isdir(output_directory):
    os.mkdir(output_directory) # If running with local profile, the directory won't be created. This is necessary in the edge case that the user _only_ runs "--until report".


# --- Localrules and ruleorders -------------------------------------

# The localrules are routine tasks that take up very little cpu. In some cases there won't be internet access on compute nodes, why it is smart that downloads are run on the local (aka. frontend) node. 
localrules: metadata, annotate, antismash_download, bakta_download, eggnog_download, checkm2_download, dbcan_download, busco_download, gtdb_download, report, install_report_environment_aot

#ruleorder: prokka > bakta > eggnog # I solved this by having an external output called ".annotation" that requests the appropriate annotator based on the config parameter "annotator". Could have otherwise been a nice solution but would be harder to keep track of which annotation was used for what.

#ruleorder: gapseq_find > gapseq # Most of the time, we just want the pathways completion fractions. Drafting and gapfilling a complete GEM is a bit overkill, but should of course be possible if the user wants it.


# --- Main rule to collect all targets ------------------------------
rule all:
    input: expand([\
        "{output_directory}/metadata.tsv", \
        "{output_directory}/.install_report_environment_aot.flag", \
        "{output_directory}/assembly-stats/assembly-stats.tsv", \
        "{output_directory}/samples/{sample}/sequence_lengths/{sample}_seqlen.tsv", \
        "{output_directory}/samples/{sample}/busco/short_summary_extract.tsv", \
        "{output_directory}/samples/{sample}/.annotation/{sample}.gff", \
        "{output_directory}/samples/{sample}/eggnog/{sample}.emapper.genepred.gff", \
        "{output_directory}/samples/{sample}/dbcan/overview.txt", \
        "{output_directory}/samples/{sample}/interproscan/{sample}_interproscan.tsv", \
        "{output_directory}/checkm2/quality_report.tsv", \
        "{output_directory}/kegg_pathway/kegg_pathway_enrichment_analysis.tsv", \
        "{output_directory}/gtdbtk/gtdbtk.summary.tsv", \
        "{output_directory}/abricate/card_detailed.tsv", \
        "{output_directory}/mlst/mlst.tsv", \
        "{output_directory}/panaroo/summary_statistics.txt", \
        "{output_directory}/snp-dists/snp-dists.tsv", \
        "{output_directory}/snp-dists/.done.flag", \
        "{output_directory}/mashtree/mashtree.newick", \
        "{output_directory}/treecluster/treecluster.tsv", \
        "{output_directory}/fasttree/fasttree.newick", \
        "{output_directory}/iqtree/core_genome_iqtree.treefile"], \
        output_directory = output_directory, sample = df["sample"]) 

        # "{output_directory}/snp-dists/done.flag", \
        # "{output_directory}/samples/{sample}/gapseq/gapseq_done.flag", \ # Still working on this one.
        # "{output_directory}/samples/{sample}/antismash/{sample}.json", \ # Issue when running with apptainer


# Write the sample table for later reference.
rule metadata:
    input: df['input_file_fasta']
    output: "{output_directory}/metadata.tsv"
    params: dataframe = df.to_csv(None, index_label = "index", sep = "\t")
    resources:
        runtime = "1h"
    shell: """
        echo '''{params.dataframe}''' > "{output:q}"
        {void_report}
    """


# --- Downloads -----------------------------------------------------

include: "rules/downloads.smk"


# --- Rules run per sample ------------------------------------------

# QC
include: "rules/sample_quality_control.smk"

# Annotation
include: "rules/sample_annotation.smk"

# Advanced annotation
include: "rules/sample_advanced_annotation.smk"



# --- Rules run per batch -------------------------------------------

# QC
include: "rules/batch_quality_control.smk"

# Advanced annotation
include: "rules/batch_advanced_annotation.smk"

# Clinical relevance
include: "rules/batch_clinical.smk"

# Core/pan 
include: "rules/batch_core_pan.smk"

# Phylogeny
include: "rules/batch_phylogeny.smk"



# --- Pro forma rules -----------------------------------------------

# This rule might seem silly, but it makes sure that the report environment is ready to rock when the report subpipeline eventually is run: This has two pros:
#    1) The vastly faster mamba configuration in the comparem2 pipeline is used
#    2) The conda/mamba debugging is taken care of, without having to wait for jobs to finish on fresh installations.
# Since all snakemake conda environments are installed in $SNAKEMAKE_CONDA_PREFIX set to ${COMPAREM2_BASE}/conda_base, reuse is guaranteed.
rule install_report_environment_aot:
    output: touch(f"{output_directory}/.install_report_environment_aot.flag")
    conda: "../dynamic_report/workflow/envs/r-markdown.yaml"
    shell: """

        echo "Report conda environment OK ..."

    """

# Just a dummy rule if you wanna force the report
# comparem2 --until report
# It isn't enough to just touch the file. The dynamic_report will not be triggered if the file is empty. Thus we add the date, and we have a nice debug log for seeing when the report was triggered.
# Will only but run if asked to. No need to use --forcerun, since snakemake states this in the output: "reason: Rules with neither input nor output files are always executed."
# Rule report does not depend on metadata, as the metadata is not interesting in itself.
rule report:
    shell: """
        
        {void_report}

    """



# --- Pseudo targets ------------------------------------------------

# Makes it easy to check that all databases are installed properly. Eventually for touching the database representatives in case of using prior installations.
# Max 6 databases. Can't take adding any more.
rule downloads:
    input:
        DATABASES + "/checkm2/comparem2_checkm2_database_representative.flag",
        DATABASES + "/busco/comparem2_busco_database_representative.flag",
        DATABASES + "/dbcan/comparem2_dbcan_database_representative.flag",
        DATABASES + "/gtdb/comparem2_gtdb_database_representative.flag",
        DATABASES + "/bakta/comparem2_bakta_database_representative.flag",
        DATABASES + "/eggnog/comparem2_eggnog_database_representative.flag",
        


# Blink-of-an-eye analysis
rule fast:
    input: expand(\
        ["{output_directory}/samples/{sample}/sequence_lengths/{sample}_seqlen.tsv", \
        "{output_directory}/assembly-stats/assembly-stats.tsv", \
        "{output_directory}/mashtree/mashtree.newick"], \
        output_directory = output_directory, \
        sample = df["sample"]) # TODO: define the expansion in each rule instead.


# Rules designed for bins of metagenomic origin
# Could be cool to find a way of setting command line arguments like panaroo thresholds and stuff.
rule meta:
    input: expand(\
        ["{output_directory}/metadata.tsv", \
        "{output_directory}/.install_report_environment_aot.flag", \
        "{output_directory}/assembly-stats/assembly-stats.tsv", \
        "{output_directory}/samples/{sample}/sequence_lengths/{sample}_seqlen.tsv", \
        "{output_directory}/samples/{sample}/busco/short_summary_extract.tsv", \
        "{output_directory}/checkm2/quality_report.tsv", \
        "{output_directory}/samples/{sample}/eggnog/{sample}.emapper.genepred.gff", \
        "{output_directory}/kegg_pathway/kegg_pathway_enrichment_analysis.tsv", \
        "{output_directory}/samples/{sample}/dbcan/overview.txt", \
        "{output_directory}/samples/{sample}/interproscan/{sample}_interproscan.tsv", \
        "{output_directory}/gtdbtk/gtdbtk.summary.tsv", \
        "{output_directory}/samples/{sample}/.annotation/{sample}.gff", \
        "{output_directory}/mashtree/mashtree.newick"], \
        output_directory = output_directory, \
        sample = df["sample"])


# Rules designed for cultured isolates
rule isolate:
    input: expand(\
        ["{output_directory}/metadata.tsv", \
        "{output_directory}/.install_report_environment_aot.flag", \
        "{output_directory}/assembly-stats/assembly-stats.tsv", \
        "{output_directory}/samples/{sample}/sequence_lengths/{sample}_seqlen.tsv", \
        "{output_directory}/samples/{sample}/.annotation/{sample}.gff", \
        "{output_directory}/samples/{sample}/eggnog/{sample}.emapper.genepred.gff", \
        "{output_directory}/kegg_pathway/kegg_pathway_enrichment_analysis.tsv", \
        "{output_directory}/gtdbtk/gtdbtk.summary.tsv", \
        "{output_directory}/mlst/mlst.tsv", \
        "{output_directory}/abricate/card_detailed.tsv", \
        "{output_directory}/panaroo/summary_statistics.txt", \
        "{output_directory}/fasttree/fasttree.newick", \
        "{output_directory}/snp-dists/snp-dists.tsv", \
        "{output_directory}/mashtree/mashtree.newick"], \
        output_directory = output_directory, sample = df["sample"])



# --- Dynamic report ------------------------------------------------
# For calling the report subpipeline we need some variables. The easiest way to communicate these from the main pipeline to the report pipeline, is to write a config.yaml.

# Why doesn't this happen on onsuccess and onerror? But does that really matter?
onstart:
    print("Writing config for dynamic report pipeline.")
    shell(f"""
    
        echo "# config for dynamic report pipeline" > .report_config.yaml
        echo batch_title: \"{batch_title}\" >> .report_config.yaml
        echo output_directory: \"{output_directory}\" >> .report_config.yaml
        echo base_variable: \"{base_variable}\" >> .report_config.yaml
        echo __version__: \"{__version__}\" >> .report_config.yaml
        
    """)
    
    
# TODO: A speedup could be to list the possible locations, and skipping the missing ones. By doing a full run and listing all possible files, these can be easily enlisted.
final = f"""echo "CompareM2 v{__version__}" > {output_directory}/.software_version.txt; find {output_directory} -name ".software_version.txt" | xargs cat | sort | uniq > {output_directory}/versions.txt"""
onsuccess:
    shell(final)
onerror:
    shell(final)