shell:
    """
    spades.py -t {threads} --careful \
        -1 {input[0]} -2 {input[1]} \
        -o results/assembly/{wildcards.isolate} &> {log}
    """

shell:
    """
    trimmomatic PE -threads {threads} -phred33 \
    {input.r1} {input.r2} \
    {output.r1} {output.r1_unpaired} \
    {output.r2} {output.r2_unpaired} \
    ILLUMINACLIP:{params.illumina_clip}:2:30:10 \
    {params.extra} &> {log}
    """

shell:
    """
    seqkit seq --remove-gaps --min-len {params.min_contig_length} --threads {threads} {input} 1> {output} 2> {log}
    """

shell:
    """
    checkm lineage_wf {params.flag} \
    --extension fa \
    --tab_table --file {output} \
    --threads {threads} \
    results/quality_check/{wildcards.isolate}/ \
    results/quality_check/{wildcards.isolate}/checkm/ &> {log}
    """

shell:
    """
    bakta_db download --output {output}
    """

shell:
    """
    metaxa2_x -i {input} \
        -o SSU-{wildcards.isolate} -f fasta \
        -g ssu --mode genome \
        --table F --fasta T --graphical F --summary F \
        --cpu {threads} &> {log}
    """

shell:
    """
    quast --threads {threads} --labels "{wildcards.isolate}.raw,{wildcards.isolate}.final" \
    --no-icarus --min-contig 0 --output-dir {output.quast_dir} {input.raw_assembly} {input.final_assembly} > {log}
    """

import sys
import pandas as pd
from functools import reduce

sys.stderr = open(snakemake.log[0], "w")

# Read the summary csv files of each of the isolate processed and transposed the long table into a wide table
loaded_df = [pd.read_csv(csv, index_col=0).T for csv in snakemake.input]

# Concatenate all the genomes
aggregate_df = pd.concat(loaded_df)

# Make sure the isolate column does not remain in the index
aggregate_df = aggregate_df.rename_axis('isolate').reset_index()

# Write the aggregated table to file
aggregate_df.to_csv(snakemake.output[0], index=False)

import sys
import os
import pandas as pd
from functools import reduce
from datetime import date
import subprocess

sys.stderr = open(snakemake.log[0], "w")

# Read the csv files of basepairs metrics and SSU/LSU extraction
dict_csv = { key: pd.read_csv(snakemake.input[key], index_col=0) for key in ['metrics', 'ssu_lsu', 'trnas_5s'] }
# Read the tsv files of sample table provided in the configuration file and outputs from genome quality assessment tools
dict_table = { key: pd.read_table(snakemake.input[key], sep='\t') for key in ['samples', 'checkm', 'quast'] }
# Read the checksums of the raw fastq files and the genome fasta file
dict_md5 = { key: pd.read_table(snakemake.input[key], sep='\s+', names=['md5', 'file']) for key in ['fastq_md5', 'archive_md5', 'genome_md5'] }


# Keep the csv as a list
list_csv = list(dict_csv.values())
# Append the samples table and add to the csv list of genome metrics, SSU/LSU and tRNAs
dict_table['samples'].set_index('isolate', inplace=True)
list_csv.append(dict_table['samples'])

# Merge the table based on the index that is the isolate name
merged = reduce(lambda  left,right: pd.merge(left,right, left_index=True, right_index=True), list_csv)

# Add quality metrics from the dictionaries containing the tables
merged['compl_score'] = float(dict_table['checkm'].loc[ dict_table['checkm']['Bin Id']==snakemake.wildcards.isolate+'.genome', 'Completeness'])
merged['compl_software'] = 'checkm'
merged['contam_score'] = float(dict_table['checkm'].loc[ dict_table['checkm']['Bin Id']==snakemake.wildcards.isolate+'.genome', 'Contamination'])
merged['contam_software'] = 'checkm'

# Set the index for the QUAST data
dict_table['quast'].set_index('Assembly', inplace=True)

# Add assembly metrics
merged['N50'] = dict_table['quast'].loc[snakemake.wildcards.isolate+'.final', 'N50']
merged['genome_length'] = dict_table['quast'].loc[snakemake.wildcards.isolate+'.final', 'Total length']
merged['number_contig'] = dict_table['quast'].loc[snakemake.wildcards.isolate+'.final', '# contigs (>= 0 bp)']
merged['number_contig_below_1kb'] = dict_table['quast'].loc[snakemake.wildcards.isolate+'.raw', '# contigs (>= 0 bp)'] - dict_table['quast'].loc[snakemake.wildcards.isolate+'.raw', '# contigs (>= 1000 bp)']
merged['max_contig_length'] =  dict_table['quast'].loc[snakemake.wildcards.isolate+'.final','Largest contig']

# Add software information
merged['SSU_recover_software'] = merged['LSU_recover_software'] = 'metaxa2'
merged['trna_ext_software'] = 'tRNAscan-SE'
merged['assembly_software'] = 'SPAdes'

# Add the location and checksums of the genome
merged['archive_file'] = dict_md5['archive_md5'].at[0, 'file']
merged['archive_file_md5'] = dict_md5['archive_md5'].at[0, 'md5']
merged['genome_md5'] = dict_md5['genome_md5'].at[0, 'md5']

# Add information about the raw sequences: location, checksums, metrics
merged['forward_file_md5'] = dict_md5['fastq_md5'].at[0, 'md5']
merged['reverse_file_md5'] = dict_md5['fastq_md5'].at[1, 'md5']


# Rename columns
merged.rename(columns={'SSU_bp':'16S_SSU_rRNA_length', 'LSU_bp':'23S_LSU_rRNA_length', 'tRNAs':'trnas'}, inplace=True)


# Compute criteria of the MIMAG and SeqCode
hq_criteria = {
    'is_compl_grtr_90': merged.at[snakemake.wildcards.isolate, 'compl_score'] > 90,
    'is_contam_less_5': merged.at[snakemake.wildcards.isolate, 'contam_score'] < 5,
    'is_coverage_grtr_10': merged.at[snakemake.wildcards.isolate, 'coverage'] > 10,
    'are_contigs_less_100': merged.at[snakemake.wildcards.isolate, 'number_contig'] < 100,
    'is_N50_grtr_25kb': merged.at[snakemake.wildcards.isolate, 'N50'] > 25000,
    'is_max_contig_grtr_100kb': merged.at[snakemake.wildcards.isolate, 'max_contig_length'] > 100000,
    'is_trnas_grtr_18': merged.at[snakemake.wildcards.isolate, 'trnas'] > 18,
    'is_SSU_grtr_0': merged.at[snakemake.wildcards.isolate, '16S_SSU_rRNA_length'] != "0",
    'is_LSU_grtr_0': merged.at[snakemake.wildcards.isolate, '23S_LSU_rRNA_length'] != "0",
    'is_5S_grtr_0': merged.at[snakemake.wildcards.isolate, '5S_rRNA_length'] != "0",
}

# Assess which criteria of the MIMAG and SeqCode did not pass
to_check = [x for x in hq_criteria.keys() if not hq_criteria[x]]
criteria_false = ', '.join(to_check)
# Indicate the assembly quality if compliance to the MIMAG and SeqCode criteria
merged['assembly_qual'] = 'High-quality draft' if all(hq_criteria.values()) else 'Manual review: ' + criteria_false

# Add the flags of the criteria to the main table
merged = pd.merge(merged, pd.DataFrame(hq_criteria, index = [snakemake.wildcards.isolate]),
                  left_index=True, right_index=True)

# Add the plasmids lengths if a non-empty file was provided
if os.path.getsize(snakemake.input['plasmids']) > 0:
    # Read the report
    plasmids = pd.read_table(snakemake.input['plasmids'], names = ['file', 'length'])
    # Extract the lengths and sort in descending order
    plasmids = plasmids.loc[:,'length'].sort_values(ascending = False).tolist()
    # Format the lengths as string and concatenate. No concatenation if only one element
    merged['plasmid_length'] = ';'.join([ str(x) for x in plasmids ])
else: # if the report is empty
    merged['plasmid_length'] = '0'


# Add the adapters file used
merged['adapters_file'] = snakemake.params.adapters

# Reorder columns
genome_csv = merged.reindex(columns=['genome_md5',
                'assembly_qual','genome_length',
                'number_contig', 'N50',
                'number_contig_below_1kb', 'max_contig_length',
                'coverage', 'assembly_software',
                'compl_score', 'compl_software',
                'contam_score', 'contam_software',
                '16S_SSU_rRNA_length', 'SSU_recover_software',
                '23S_LSU_rRNA_length', 'LSU_recover_software',
                'trnas', 'trna_ext_software', '5S_rRNA_length',
                'archive_file', 'archive_file_md5',
                'forward_file', 'forward_file_md5',
                'reverse_file', 'reverse_file_md5',
                'sequence_count', 'basepairs_count', 'average_length',
                'sequence_count_qual', 'basepairs_count_qual', 'adapters_file', 'plasmid_length',
                'is_compl_grtr_90', 'is_contam_less_5',
                'is_coverage_grtr_10', 'are_contigs_less_100',
                'is_N50_grtr_25kb','is_max_contig_grtr_100kb',
                'is_trnas_grtr_18',
                'is_SSU_grtr_0', 'is_LSU_grtr_0', 'is_5S_grtr_0'])

# Add the date when the assembly was done
genome_csv['assembly_date'] = date.today()
# Get the version of the workflow (and the last commit if no tags are present) and add to the table
try:
    get_tag = subprocess.check_output("git describe --always", shell=True, text=True)
    genome_csv['workflow_version'] = get_tag.strip()
except subprocess.CalledProcessError:
    genome_csv['workflow_version'] = "NA"
    print("Either git is not installed or the workflow is not run in the git repository")

# Convert the index to a column 'isolate' for clarity
# src: https://stackoverflow.com/a/63926255
genome_csv = genome_csv.rename_axis('isolate').reset_index()

# Write the table to file
genome_csv.T.to_csv(snakemake.output[0], header=False)

__author__ = "Filipe G. Vieira"
__copyright__ = "Copyright 2021, Filipe G. Vieira"
__license__ = "MIT"

from snakemake.shell import shell
from snakemake_wrapper_utils.java import get_java_opts

java_opts = get_java_opts(snakemake)
extra = snakemake.params.get("extra", "")
adapters = snakemake.params.get("adapters", "")
log = snakemake.log_fmt_shell(stdout=True, stderr=True)


n = len(snakemake.input.sample)
assert (
    n == 1 or n == 2
), "input->sample must have 1 (single-end) or 2 (paired-end) elements."

if n == 1:
    reads = "in={}".format(snakemake.input.sample)
    trimmed = "out={}".format(snakemake.output.trimmed)
else:
    reads = "in={} in2={}".format(*snakemake.input.sample)
    trimmed = "out={} out2={}".format(*snakemake.output.trimmed)


singleton = snakemake.output.get("singleton", "")
if singleton:
    singleton = f"outs={singleton}"


discarded = snakemake.output.get("discarded", "")
if discarded:
    discarded = f"outm={discarded}"


stats = snakemake.output.get("stats", "")
if stats:
    stats = f"stats={stats}"


shell(
    "bbduk.sh {java_opts} t={snakemake.threads} "
    "{reads} "
    "{adapters} "
    "{extra} "
    "{trimmed} {singleton} {discarded} "
    "{stats} "
    "{log}"
)

__author__ = "Patrik Smeds"
__copyright__ = "Copyright 2016, Patrik Smeds"
__email__ = "patrik.smeds@gmail.com"
__license__ = "MIT"

from os.path import splitext

from snakemake.shell import shell

log = snakemake.log_fmt_shell(stdout=False, stderr=True)

# Check inputs/arguments.
if len(snakemake.input) == 0:
    raise ValueError("A reference genome has to be provided!")
elif len(snakemake.input) > 1:
    raise ValueError("Only one reference genome can be inputed!")

# Prefix that should be used for the database
prefix = snakemake.params.get("prefix", splitext(snakemake.output.idx[0])[0])

if len(prefix) > 0:
    prefix = "-p " + prefix

# Contrunction algorithm that will be used to build the database, default is bwtsw
construction_algorithm = snakemake.params.get("algorithm", "")

if len(construction_algorithm) != 0:
    construction_algorithm = "-a " + construction_algorithm

shell(
    "bwa index" " {prefix}" " {construction_algorithm}" " {snakemake.input[0]}" " {log}"
)

__author__ = "Johannes Köster, Julian de Ruiter"
__copyright__ = "Copyright 2016, Johannes Köster and Julian de Ruiter"
__email__ = "koester@jimmy.harvard.edu, julianderuiter@gmail.com"
__license__ = "MIT"


import tempfile
from os import path
from snakemake.shell import shell
from snakemake_wrapper_utils.java import get_java_opts
from snakemake_wrapper_utils.samtools import get_samtools_opts


# Extract arguments.
extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=False, stderr=True)
sort = snakemake.params.get("sorting", "none")
sort_order = snakemake.params.get("sort_order", "coordinate")
sort_extra = snakemake.params.get("sort_extra", "")
samtools_opts = get_samtools_opts(snakemake)
java_opts = get_java_opts(snakemake)


index = snakemake.input.idx
if isinstance(index, str):
    index = path.splitext(snakemake.input.idx)[0]
else:
    index = path.splitext(snakemake.input.idx[0])[0]


# Check inputs/arguments.
if not isinstance(snakemake.input.reads, str) and len(snakemake.input.reads) not in {
    1,
    2,
}:
    raise ValueError("input must have 1 (single-end) or 2 (paired-end) elements")


if sort_order not in {"coordinate", "queryname"}:
    raise ValueError("Unexpected value for sort_order ({})".format(sort_order))


# Determine which pipe command to use for converting to bam or sorting.
if sort == "none":
    # Simply convert to bam using samtools view.
    pipe_cmd = "samtools view {samtools_opts}"

elif sort == "samtools":
    # Add name flag if needed.
    if sort_order == "queryname":
        sort_extra += " -n"

    # Sort alignments using samtools sort.
    pipe_cmd = "samtools sort {samtools_opts} {sort_extra} -T {tmpdir}"

elif sort == "picard":
    # Sort alignments using picard SortSam.
    pipe_cmd = "picard SortSam {java_opts} {sort_extra} --INPUT /dev/stdin --TMP_DIR {tmpdir} --SORT_ORDER {sort_order} --OUTPUT {snakemake.output[0]}"

else:
    raise ValueError(f"Unexpected value for params.sort ({sort})")

with tempfile.TemporaryDirectory() as tmpdir:
    shell(
        "(bwa mem"
        " -t {snakemake.threads}"
        " {extra}"
        " {index}"
        " {snakemake.input.reads}"
        " | " + pipe_cmd + ") {log}"
    )