shell:
    """
    prodigal -i {input} -d {output.genes} -a {output.faa} -o {output.gff} \
        -f gff -p meta 2>{log}
    """

shell:
    """
    tRNAscan-SE -G -b {output.bed} -o {output.file} -a {output.fasta} \
        --thread {threads} {input} >{log} 2>&1
    """

shell:
    """
    cmscan --cpu {threads} --oskip --rfam --cut_ga --nohmmonly \
        --tblout {output} --fmt 2 --clanin {input.cl} {params.db} \
        {input.fastafile} > /dev/null 2>{log}
    """

shell:
    """
    mkdir -p {params.tmpdir}
    emapper.py {params.flags} --cpu {threads} --temp_dir {params.tmpdir} \
    -i {input[0]} -o {params.out} --output_dir {params.tmpdir} \
        --data_dir {params.resource_dir} >{log} 2>&1
    mv {params.tmp_out}.emapper.seed_orthologs {output[0]}
    rm -rf {params.tmpdir}
    """

shell:
    """
    mkdir -p {params.tmpdir}
    rgi load -i {input.db} --local > {log} 2>&1
    sed 's/*//g' {input.faa} > {params.faa}
    rgi main -i {params.faa} -o {params.out} \
        -n {threads} {params.settings} >>{log} 2>>{log}
    rm -r {params.tmpdir}
    """

shell:
    """
    # Create directories
    mkdir -p {params.tmp}
    # Clean output dir
    #rm -rf {params.output_dir}/*
    # Clean temp dir
    rm -rf {params.tmp}/*
    # Only use single-end if present
    if [ -s {input.se} ]; then
        single="-s {input.se}"
    else
        single=""
    fi
    metaspades.py \
        -t {threads} -1 {input.R1} -2 {input.R2} $single \
        -o {params.tmp} > {log} 2>&1

    # If set to keep intermediate contigs, move to intermediate folder before deleting
    if [ "{config[metaspades][keep_intermediate]}" == "True" ]; then
        mkdir -p {params.intermediate_contigs}
        cp -r {params.tmp}/K* {params.intermediate_contigs}
    fi
    if [ "{config[metaspades][keep_corrected]}" == "True" ]; then
        mkdir -p {params.corrected}
        cp -r {params.tmp}/corrected {params.corrected}
    fi

    # Clear intermediate contigs
    rm -rf {params.tmp}/K*
    # Clear corrected reads dir
    rm -rf {params.tmp}/corrected
    # Sync tmp output to outdir before removing
    cp -r {params.tmp}/* {params.output_dir}
    rm -rf {params.tmp}
    mv {params.output_dir}/scaffolds.fasta {params.output_dir}/final_contigs.fa
    """

shell:
    """
    mkdir -p {config[paths][temp]}
    rm -rf {params.tmp}
    # Only use paired-end if present
    if [ -s {input.R1} ]; then
        R1=$(cat {input.R1})
        R2=$(cat {input.R2})
        paired="-1 $R1 -2 $R2"
    else
        paired=""
    fi
    # Only use single-end if present
    se=$(cat {input.se})
    if [ -s {input.se} ]; then
        single="-r $se"
    else
        single=""
    fi

    # Run Megahit
    megahit -t {threads} $paired $single -o {params.tmp} \
        {params.additional_settings} >{log} 2>&1

    # Sync intermediate contigs if asked for
    if [ "{config[megahit][keep_intermediate]}" == "True" ]; then
        mkdir -p {params.intermediate_contigs}
        cp -r {params.tmp}/intermediate_contigs/* {params.intermediate_contigs}
    fi

    # Cleanup intermediate
    rm -rf {params.tmp}/intermediate_contigs

    # Sync tmp output to outdir before removing
    cp -r {params.tmp}/* {params.output_dir}
    rm -rf {params.tmp}
    mv {params.output_dir}/final.contigs.fa {params.output_dir}/final_contigs.fa
    """

shell:
    """
    bowtie2-build \
        --large-index \
        --threads {threads} \
        {params.prefix} \
        {params.prefix} > /dev/null 2>&1
    """

shell:
    """
    bowtie2 {params.setting} -p {threads} -x {params.prefix} -1 {input.R1} -2 {input.R2} 2> {log} \
        | samtools view -bh - | samtools sort - -o {params.temp_bam}
    samtools index {params.temp_bam}
    mv {params.temp_bam} {output.bam}
    mv {params.temp_bam}.bai {output.bai}
    """

shell:
    """
    metabat2 -i {input.fa} -a {input.depth} -m {wildcards.l} -t {threads} \
        -o {params.n} > {log} 2>&1
    """

shell:
    """
    set +e
    mkdir -p {params.dir}
    mkdir -p {params.tmp_dir}
    run_MaxBin.pl -markerset {params.markerset} -contig {input} \
        {params.reads} -min_contig_length {wildcards.l} -thread {threads} \
        -out {params.tmp_dir}/maxbin >{log} 2>{log}
    exitcode=$?
    if [ $exitcode -eq 255 ]; then
        exit 0
    else
        # Rename fasta files
        ls {params.tmp_dir} | grep ".fasta" | while read f;
        do
            mv {params.tmp_dir}/$f {params.dir}/${{f%.fasta}}.fa
        done
        # Move output from temporary dir
        ls {params.tmp_dir} | while read f;
        do
            mv {params.tmp_dir}/$f {params.dir}/
        done
    fi
    # Clean up
    rm -r {params.tmp_dir}
    """

shell:
    """
    for f in {input.bam} ;
        do
            n=$(basename $f);
            s=$(echo -e $n | sed 's/_[ps]e{params.p}.bam//g');
            echo $s;
        done > {params.samplenames}
    concoct_coverage_table.py \
        --samplenames {params.samplenames} \
        {input.bed} {input.bam} > {output.cov}
    rm {params.samplenames}
    """

shell:
    """
    # Download
    curl -L https://data.ace.uq.edu.au/public/CheckM_databases/checkm_data_2015_01_16.tar.gz -o {params.tar} -s
    # Extract
    tar -C {params.dir} -xf {params.tar}
    # Set root
    checkm data setRoot {params.dir} > {log} 2>&1
    """

shell:
    """
    bins=$(wc -l {input.tsv} | cut -f1 -d ' ')
    if [ $bins == 0 ] ; then
        echo "NO BINS FOUND" > {output}
    else
        export PYTHONPATH=$(which python)
        export GTDBTK_DATA_PATH={params.dbdir}
        gtdbtk classify_wf -x {params.suff} --out_dir {params.outdir} \
            --cpus {threads} --pplacer_cpus {threads} \
            --genome_dir {params.indir} > {log} 2>&1
    fi
    """

shell:
    """
    bins=$(wc -l {input.tsv} | cut -f1 -d ' ')
    if [ $bins == 0 ] ; then
        touch {output}
    else
        cat {params.gtdbtk_dir}/gtdbtk.*.summary.tsv | cut -f1 -d ';' | grep -v "user_genome" | \
            while read line;
            do
                d=$(echo -e "$line" | cut -f2)
                g=$(echo -e "$line" | cut -f1)
                if [ "$d" == "d__Bacteria" ]; then
                    k="bac"
                else
                    k="arc"
                fi
                barrnap --kingdom $k {params.indir}/$g.fa > {params.outdir}/out 2>>{log}
                lines=$(wc -l {params.outdir}/out | cut -f1 -d ' ')
                if [ $lines -gt 1 ]; then
                    egrep -v "^#" {params.outdir}/out | sed "s/$/;genome=$g/g" >> {output}
                else
                    touch {output}
                fi
            done
        rm {params.outdir}/out
    fi
    """

shell:
    """
    fastANI --rl {input[0]} --ql {input[1]} -k {params.k} -t {threads} \
        --fragLen {params.frag_len} --minFraction {params.fraction} \
        --matrix -o {output[0]} > {log} 2>&1
    """

shell:
    """
    kraken2-build --standard --db {params.dir} --threads {threads} > {log.build} 2>&1
    kraken2-build --clean {params.dir} > {log.clean} 2>&1
    """

shell:
    """
    mkdir -p {params.tmpdir}
    centrifuge -k {params.k} -x {params.prefix} -1 {input.R1} -2 {input.R2} \
        -S {params.tmp_out} -p {threads} --report-file {params.tmp_report} \
        > {log} 2>&1
    mv {params.tmp_out} {output[0]}
    mv {params.tmp_report} {output[1]}
    """

shell:
    """
    metaphlan --install --bowtie2db {params.dir} \
        --nproc {threads} -x {params.index} >{log} 2>&1
    """

shell:
    """
    ktImportTaxonomy -t 1 -m 2 -o {output} -tax {params.dbdir} \
        {params.input_string} > {log} 2>&1
    """

shell:
     """
     seqtk sample -s {wildcards.s} {input} \
        {params.example_dataset_size} | gzip -c > {output}
     """

shell:
    """
    trimmomatic PE \
        -threads {threads} \
        {input.R1} {input.R2} \
        {output.R1P} {output.R1U} \
        {output.R2P} {output.R2U} \
        {params.trim_string} \
        2>{log.R1log}
    sed \
        's/{wildcards.sample}_{wildcards.unit}_R1/{wildcards.sample}_{wildcards.unit}_R2/g' \
        {log.R1log} > {log.R2log}
    """

shell:
    """
    cutadapt \
        {params.extra_params} \
        -e {params.error_rate} \
        -a {params.adapter} \
        -A {params.rev_adapter} \
        -o {output.fastq1} \
        -p {output.fastq2} \
        -j {threads} \
        {input.R1} {input.R2} > {log.R1log} 2>{log.err}
    cp {log.R1log} {log.R2log}
    """

shell:
    """
    fastqc -q --noextract -o {params.dir} {input} >{log} 2>&1
    """

shell:
    """
    multiqc -f -c {params.config} -n samples_report.html \
        -o {params.output_dir} {input} >{log} 2>{log}
    """

shell:
    """
    mkdir -p {params.temp_dir}
    # Fix bam header
    samtools view -H {input} | egrep -v "^@PG" > {params.header}
    samtools reheader -P {params.header} {input} > {params.rehead_bam}
    # Set memory max
    mem="-Xmx$((6 * {threads}))g"
    java -Xms2g $mem -XX:ParallelGCThreads={threads} \
        -jar $CONDA_PREFIX/share/picard-*/picard.jar MarkDuplicates \
        I={params.rehead_bam} M={output[2]} O={params.temp_bam} REMOVE_DUPLICATES=TRUE \
        USE_JDK_DEFLATER=TRUE USE_JDK_INFLATER=TRUE ASSUME_SORT_ORDER=coordinate \
        PROGRAM_RECORD_ID=null ADD_PG_TAG_TO_READS=FALSE 2> {log}
    # Re sort the bam file using samtools
    samtools_threads="$(({threads} - 1))"
    samtools sort -@ $samtools_threads -o {params.temp_sort_bam} {params.temp_bam} > /dev/null 2>&1
    # Index the bam file
    samtools index {params.temp_sort_bam}
    mv {params.temp_sort_bam} {output[0]}
    mv {params.temp_sort_bam}.bai {output[1]}
    rm {params.temp_bam} {params.rehead_bam} {params.header}
    """

shell:
    """
    mkdir -p {params.tmpdir}
    featureCounts -a {input.gff} -o {output[0]} -t CDS -g gene_id -M \
        {params.setting} -T {threads} --tmpDir {params.tmpdir} {input.bam} > {log} 2>&1
    """

script:
    "../scripts/edger.R"

script:
    "../scripts/metagenomeseq.R"

shell:
    """
    contigtax download taxonomy -t {params.taxdir} >{log} 2>&1
    """

import pandas as pd


def orf2feat(d, val_name="vals", regex=""):
    import re
    keys = []
    vals = []
    for key, val in d.items():
        for item in val.split(","):
            if regex != "":
                m = re.search(regex, item)
                if m == None:
                    continue
                item = m.group()
            keys.append(key)
            vals.append(item)
    df = pd.DataFrame(data={"orf": keys, val_name: vals})
    df.set_index("orf", inplace=True)
    return df


def parse_emapper(sm):
    from pandas.errors import EmptyDataError
    db = sm.wildcards.db
    db_att = {'kos': {'val_name': 'ko', 'regex': "K\d{5}"},
              'pathways': {'val_name': 'pathway', 'regex': "map\d{5}"},
              'modules': {'val_name': 'module', 'regex': ""},
              'enzymes': {'val_name': 'enzyme', 'regex': ""}}
    df = pd.read_csv(sm.input.annotations, sep="\t", index_col=0)
    df.rename(columns={'KEGG_ko': 'kos', 'KEGG_Pathway': 'pathways',
                          'KEGG_Module': 'modules', 'EC': 'enzymes'},
               inplace=True)
    df.fillna("-", inplace=True)
    d = df.loc[df[db] != "-", db].to_dict()
    how = "inner"
    try:
        info_df = pd.read_csv(sm.input.info, sep="\t", index_col=0)
    except EmptyDataError:
        info_df = pd.DataFrame()
        how = "right"
    annot = orf2feat(d, val_name=db_att[db]["val_name"],
                     regex=db_att[db]["regex"])
    annot = pd.merge(info_df, annot, left_index=True,
                     right_on=db_att[db]["val_name"], how=how)
    annot.to_csv(sm.output[0], sep="\t", index=True, header=True)


def parse_rgi(sm):
    annot = pd.read_csv(sm.input.txt, sep="\t", index_col=0)
    annot = annot.loc[:, ["Model_ID", "AMR Gene Family", "Resistance Mechanism"]]
    annot.loc[:, "Model_ID"] = ["RGI_{}".format(x) for x in annot.Model_ID]
    annot.rename(index=lambda x: x.split(" ")[0], inplace=True)
    annot.to_csv(sm.output.tsv, sep="\t", index=True)


def parse_pfam(sm):
    annot = pd.read_csv(sm.input[0], comment="#", header=None, sep=" +",
                        usecols=[0, 5, 7, 14], engine="python",
                        names=["orf", "pfam", "pfam_type", "pfam_clan"])
    clans = pd.read_csv(sm.input[1], header=None, names=["clan", "clan_name"],
                        usecols=[0, 3], sep="\t")
    info = pd.read_csv(sm.input[2], header=None,
                       names=["pfam", "clan", "pfam_name"], usecols=[0, 1, 4],
                       sep="\t")
    # Strip suffix for pfams
    annot.loc[:, "pfam"] = [x.split(".")[0] for x in annot.pfam]
    # Select unique orf->pfam mappings
    # TODO: This masks multiple occurrences of domains on the same orf. Figure out if this is wanted or not.
    # Merge with pfam info and clan info
    annot = annot.groupby(["orf", "pfam"]).first().reset_index()
    annot = pd.merge(annot, info, left_on="pfam", right_on="pfam")
    annot = pd.merge(annot, clans, left_on="clan", right_on="clan", how="left")
    annot.fillna("No_clan", inplace=True)
    annot = annot.loc[:, ["orf", "pfam", "pfam_name", "clan", "clan_name"]]
    annot.sort_values("orf", inplace=True)
    # Write to file
    annot.to_csv(sm.output[0], sep="\t", index=False)


def main(sm):
    toolbox = {"parse_pfam": parse_pfam,
               "parse_emapper": parse_emapper,
               "parse_rgi": parse_rgi}
    toolbox[sm.rule](sm)


if __name__ == "__main__":
    main(snakemake)

import pandas as pd
from Bio.SeqIO import parse
from pathlib import Path


# file io

def fasta2bed(sm):
    with open(sm.input[0], 'r') as fhin, open(sm.output[0], 'w') as fhout:
        for record in parse(fhin, "fasta"):
            fhout.write("{}\t{}\t{}\n".format(record.id, 0, len(record)))


# statistics

def store_lengths(f, minlen=False):
    """
    Reads lengths of contigs from fasta
    :param f: fasta file
    :param minlen: minimum length to store
    :return: pandas DataFrame of lengths
    """
    r = {}
    for record in parse(f, "fasta"):
        if minlen:
            if len(record.seq) < minlen:
                continue
        r[record.id] = len(record.seq)
    df = pd.DataFrame(r, index=["length"]).T
    return df


def size_distribute(df, lengths=None):
    """
    Calculates the distribution of an assembly in length bins

    For each <l> in <lengths> calculate for contigs >= <l>:
     n = the number of contigs
     s = the total length in bp
     p = the fraction of lengths / total assembly size
    :param df: pandas DataFrame of lengths
    :param lengths: intervals at which to calculate stats
    :return: pandas DataFrame
    """
    if lengths is None:
        lengths = [0, 100, 250, 500, 1000, 2500, 5000, 10000, 15000, 20000,
                   25000, 30000, 35000, 40000, 45000, 50000, 75000, 100000,
                   125000, 150000, 200000, 250000, 500000]
    size_dist = {}
    for i, l in enumerate(lengths):
        if len(df.loc[df.length >= l]) == 0:
            break
        n = len(df.loc[df.length >= l])
        s = int(df.loc[df.length >= l].sum())
        p = int(df.loc[df.length >= l].sum()) / float(df.sum()) * 100
        size_dist[i] = {"min_length": l, "num_contigs": n, "total_length": s,
                        "%": p}
    size_dist_df = pd.DataFrame(size_dist).T
    size_dist_df = size_dist_df[
        ["min_length", "num_contigs", "total_length", "%"]]
    return size_dist_df


def calculate_n_stats(df):
    """
    Calculates n50 and n90 statistics from a list of lengths
    :param df: pandas DataFrame of contig lengths
    :return:
    """
    df.sort_values("length", inplace=True, ascending=True)
    size = int(df.sum())
    N50_length = N90_length = 0
    cumulative = 0
    for contig in df.index:
        l = df.loc[contig, "length"]
        cumulative += l
        if float(cumulative) >= 0.5 * size and not N50_length:
            N50_length = l
        elif float(cumulative) >= 0.1 * size and not N90_length:
            N90_length = l
    return N50_length, N90_length


def calculate_length_stats(df):
    """
    Calculates length statistics from a dataframe
    :param df: pandas DataFrame with contig lengths
    :return:
    """
    contigs = len(df)
    total_size = int(df.sum())
    min_length = int(df["length"].min())
    max_length = int(df["length"].max())
    avg_length = float(df["length"].mean())
    median_length = float(df["length"].median())
    return contigs, total_size, min_length, max_length, avg_length, median_length


def generate_stat_df(contig_lengths):
    """
    Generates statistics from a dataframe of contig lengths
    :param contig_lengths: pandas DataFrame
    :return:
    """
    index = ["contigs", "total_size_bp", "min_length", "max_length",
             "avg_length", "median_length", "N50_length", "N90_length"]
    stat_items = calculate_length_stats(contig_lengths)
    n50_length, n90_length = calculate_n_stats(contig_lengths)
    stat_df = pd.DataFrame([stat_items[0], stat_items[1], stat_items[2],
                            stat_items[3], stat_items[4], stat_items[5],
                            n50_length, n90_length], index=index).T
    return stat_df


def stats(sm):
    """
    Reads a list of assembly fasta files and generates statistics

    :param sm: snakemake object
    :return:
    """
    stat_result = pd.DataFrame()
    sizedist_result = pd.DataFrame()
    for f in sm.input.fa:
        p = Path(f)
        name = p.parent.name
        contig_lengths = store_lengths(f)
        stat_df = generate_stat_df(contig_lengths)
        size_dist = size_distribute(contig_lengths)
        stat_df["assembly"] = [name]*len(stat_df)
        size_dist["assembly"] = [name]*len(size_dist)
        stat_result = pd.concat([stat_result, stat_df])
        sizedist_result = pd.concat([sizedist_result,size_dist])
    stat_result = stat_result[["assembly", "contigs", "total_size_bp",
                               "min_length", "max_length", "avg_length",
                               "median_length", "N50_length", "N90_length"]]
    stat_result.to_csv(sm.output[0], sep="\t", index=False)
    sizedist_result = sizedist_result[["assembly", "min_length",
                                       "num_contigs", "total_length", "%"]]
    sizedist_result.to_csv(sm.output[1], sep="\t", index=False)

# assembly input

def metaspades_input(sm):
    """
    Generates fastq files to use as input for metaspades assembler

    :param sm: snakemake object
    :return:
    """
    from common import rename_records
    files = {"R1": [], "R2": [], "se": []}
    assembly_dict = sm.params.assembly
    # Collect all files belonging to the assembly group
    for sample in assembly_dict.keys():
        for unit in assembly_dict[sample]:
            for pair in assembly_dict[sample][unit].keys():
                files[pair].append(
                    assembly_dict[sample][unit][pair][0])
    # Rename and concatenate reads (required for Metaspades)
    with open(sm.output.R1, 'w') as fh1, open(sm.output.R2, 'w') as fh2, open(
        sm.output.se, 'w') as fhse:
        i = 0
        for f in files["R1"]:
            f2 = files["R2"][i]
            fh1 = rename_records(f, fh1, i)
            fh2 = rename_records(f2, fh2, i)
            i += 1
        for i, f in enumerate(files["se"], start=i):
            fhse = rename_records(f, fhse, i)


def megahit_input(sm):
    """
    Genereate input lists for megahit assembler

    :param sm: snakemake object
    :return:
    """
    files = {"R1": [], "R2": [], "se": []}
    assembly_dict = sm.params.assembly
    for sample in assembly_dict.keys():
        for unit in assembly_dict[sample]:
            for pair in assembly_dict[sample][unit].keys():
                files[pair].append(assembly_dict[sample][unit][pair][0])
    with open(sm.output.R1, 'w') as fh1, \
        open(sm.output.R2, 'w') as fh2, \
        open(sm.output.se, 'w') as fhse:
        fh1.write(",".join(files["R1"]))
        fh2.write(",".join(files["R2"]))
        fhse.write(",".join(files["se"]))


def main(sm):
    toolbox = {"assembly_stats": stats,
               "generate_megahit_input": megahit_input,
               "generate_metaspades_input": metaspades_input,
               "fasta2bed": fasta2bed}
    toolbox[sm.rule](sm)


if __name__ == "__main__":
    main(snakemake)