shell: 'cd-hit -i {input} -o {output} -c {params.sim} -M 0'

shell:
    """
    assembly-stats -t {input.neff_v2} {input.neff_v1} > {params.assembly_tbl}
    Rscript scripts/01_assembly_stats.R {params.assembly_tbl} {output}
    """

shell: 'quast -t {threads} -e -o {output} {input.neff_v2} {input.ref_fa}'

shell: "seqkit grep -r -f {params.hgt_ids} {input} | seqkit translate > {output}"

run:
    blast = pd.read_csv(input['blast'], sep='\t')
    annot = pd.read_csv( input['annot'], sep='\t')
    strain = wildcards['strain']
    annot = annot.loc[annot.ID.str.startswith(strain), :]
    annot.ID = annot.ID.str.replace(f"{strain}"+r"_([a-zA-Z0-9_]*).*", r"\1")
    annot.chrom= annot.chrom.str.replace(f"{strain}_", "")
    blast['hgt'] = 1
    blast.target = blast.target.str.replace(r'([a-zA-Z0-9_]*).*', r'\1')
    annot = blast.merge(annot, how='right', left_on='target', right_on='ID')
    annot.hgt = annot.hgt.fillna(0).astype(int)
    annot = annot.loc[:, ['chrom', 'start', 'end', 'ID', 'hgt', 'n_exon']]
    annot.to_csv(output[0], sep='\t', header=None, index=False)

shell:
    """
    echo -e "chrom\tstart\tend\tgeneID\tHGT\tNEXON\tGC\tGCSKEW\tATSKEW\tENTRO\t2MER\t3MER\t4MER" > {output}
    bedtools intersect -a {input.win} -b {input.hgt} -wb \
        | sort --parallel={threads} -k11,11 -k12,12n \
        | bedtools groupby -g 11,12,13,14,15,16 -c 4,5,6,7,8,9,10 -o mean \
        | sort --parallel={threads} -k1,1 -k2,2n >> {output}
    """

shell:
    """
    bundlelinks -strict -min_bundle_size 1000 -max_gap 50000 \
                -links {input.mcsx} | sed 's/lgrey=$//' > {params.circos_dir}/bundles.txt
    circos -conf {input.cfg} -outputfile {output}
    """

shell:
    """
    minimap2 -c {input.neff} {input.c3} -x map-ont \
    > {output.paf}
    python ./scripts/04_compute_seq_divergence.py \
        {output.paf} \
        {output.div}
    """

shell: "genomepy install -g '{params.genomedir}/{wildcards.virus}' {params.assembly}"

run:
    import matplotlib.pyplot as plt
    import seaborn as sns
    import matplotlib as mpl
    mpl.use("Agg")
    vir = pd.read_csv(
        input[0], sep="\t", names=["chrom", "start", "end", "virus", "ident"]
        )
    vir["seqlen"] = vir.end - vir.start
    # Amount of each virus inserted
    sns.scatterplot(data=vir, x="seqlen", y="ident", hue="virus")
    plt.savefig(output[0])

shell:
    """
    ulimit -n 10000
    # Move all organisms proteomes to orthofinder workdir
    mkdir -p {params.orthofinder_dir}
    rm -f {params.orthofinder_dir}/*
    cp {input.org_proteomes} {params.orthofinder_dir}
    # Move Ac proteomes as well, but trim proteome from filename
    for strain in {input.ac_proteomes}; do
        fname=$(basename $strain)
        new_fname="${{fname/_proteome_filtered/}}"
        cp $strain "{params.orthofinder_dir}/$new_fname"
    done
    #cp "{{input.hgt_neff_v1}}" "{params.orthofinder_dir}/"
    orthofinder -f {params.orthofinder_dir} \
                -o {output} \
                -S diamond \
                -t {threads} \
                -n "amoeba"
    """

import time
import urllib
from Bio import Entrez


def safe_request(fun):
    """
    Wraps function requesting data to allow safe errors and retry.

    Parameters
    ----------
    fun : python function
        The python function that queries a server

    Returns
    -------
    wrapped_f : python function
        A wrapped version of the input function which will call itself recursively
        every 5 seconds if the server is overloaded and will return None if the
        query record does not exist
    """

    def wrapped_f(*args, **kwargs):

        try:
            a = fun(*args, **kwargs)
            return a
        except urllib.error.HTTPError as e:
            if e.code in (429, 502):
                time.sleep(5)
                print(
                    "Sending too many requests, sleeping 5sec and retrying..."
                )
                wrapped_f(*args, **kwargs)
            else:
                print("No sequence found for %s" % args[0])
                return None

    return wrapped_f


@safe_request
def name_to_proteins(
    name, db="protein", email="someone@email.com", filters=""
):
    """
    Given an organism name, fetch all available protein sequences.

    Parameters
    ----------
    name : str
        Name of an organism or group of organism (e.g. Legionella).
    db : str
        Entrez db name to use for the search.
    email : str
        User email for identification.
    filters : str
        Additional filters to use for the query. E.g. " AND refseq[filter]" to
        filter results from refseq

    Returns
    -------
    seq_record : TextIOWrapper
        iterable object containing the query Fasta records.

    """
    Entrez.email = email
    query_string = name + "[Orgn]"
    time.sleep(0.1)
    # First search to see the number of hits (returns values for 10 by default)
    query = Entrez.read(
        Entrez.esearch(term=query_string, db=db, email=email, retmax=10 ** 9)
    )
    seqs = ""
    s = 0
    chunk_size = 1000
    # Fetch proteins by batch of 100 to avoid maximum number of queries
    for e in range(chunk_size, len(query["IdList"]), chunk_size):
        print(f"fetching entries {s} to {e}")
        fetched = False
        while not fetched:
            try:
                seqs += Entrez.efetch(
                    id=query["IdList"][s:e], rettype="fasta", db=db
                ).read()
                fetched = True
            except:
                print("Fetching failed, trying again.")
                time.sleep(5)
        s = e
        time.sleep(0.1)
    e = len(query["IdList"])
    print(f"fetching entries {s} to {e}")
    seqs += Entrez.efetch(
        id=query["IdList"][s:e], rettype="fasta", db=db
    ).read()
    return seqs


org_df = snakemake.params["org"]
organism = org_df.loc[
    org_df["clean_name"] == f"{snakemake.wildcards.organism}", "name"
]
time.sleep(0.1)  # Do not spam NCBI :)
proteome = name_to_proteins(
    organism, email="demo@email.com", filters=" AND refseq[filter]"
)
try:
    print(f"Writing {proteome.count('>')} proteins for {organism}.")
    with open(snakemake.output[0], "w") as outf:
        outf.write(proteome)
except TypeError:
    print(f"No proteome found for {organism[0]}")
    pass

library(readr)
library(stringr)
library(ggplot2)
library(dplyr)
library(tidyr)
library(tibble)
library(gridExtra)
args <- commandArgs(trailingOnly=TRUE)

### GET ARGS ###
v1_gff <- args[1]
neff_gff <- args[2]
c3_gff <- args[3]
out_tbl <- args[4]
out_plot <- args[5]

### LOAD DATA ###
load_gff <- function(in_gff, name){
    gff <- read_tsv(in_gff, comment="#", col_names=F)
    colnames(gff) <- c("chrom", "source", "type", "start", "end", "score", "strand", "phase", "attributes")
    gff <- gff %>% filter(type %in% c("mRNA", "CDS", "gene", "exon"))

    ### Clean data
    gff <- gff %>% mutate(attributes=gsub("Parent=", "ID=", attributes)) %>%
        mutate(ID=str_extract(attributes, 'ID=[^;\\.]*')) %>%
        mutate(ID=gsub("=[a-zA-Z]{1,}:", "=", ID)) %>%
        mutate(ID=sapply(ID, function(x){str_split(x,regex('[=-]'))[[1]][2]})) %>%
        mutate(name=name)
    return(gff)
}
gff <- load_gff(v1_gff, "NEFF_v1") %>%
    bind_rows(load_gff(neff_gff, "Neff")) %>%
    bind_rows(load_gff(c3_gff, "C3"))

### COMPUTE STATS ###

gff_gene_len <- gff %>%
    filter(type == 'gene') %>%
    mutate(gene_len = end - start)

# N exon / gene (drop duplicated exons)
gff_exons <- gff %>% 
    filter(type != 'gene') %>%
    group_by(ID) %>%
    distinct(chrom, start, end, type, .keep_all=T) %>%
    mutate(n_exon=sum(type == 'exon')) %>%
    filter(type == 'mRNA') %>% 
    mutate(mrna_len = end - start)

# Combine informations and reorder columns
gff_stats <- gff_gene_len %>%
    select(ID, gene_len) %>%
    inner_join(gff_exons, by='ID') %>%
    select(ID, chrom, start, end, type, strand, n_exon, gene_len, mrna_len, attributes)


### VISUALISE ###

# Label generator to get number of samples
n_obs <- function(x){return(data.frame(y=0, label=paste0("n=", length(x))))}

feature_len <- ggplot(data=gff %>% filter(type %in% c("CDS", "gene")), 
                      aes(x=name, y=log10(end - start), fill=name)) +
    geom_violin() +
    stat_summary(fun.data = n_obs, geom = "text") +
    theme_minimal() +
    xlab("") + 
    ylab("log10 length") + 
    ggtitle("Gene length") +
    facet_wrap(.~type) +
    scale_y_continuous(limits = c(0, 6)) + 
    geom_boxplot(width=0.1)

exon_quantiles <- gff_exons %>% 
  group_by(name) %>% 
  summarise(x=list(enframe(quantile(n_exon, probs=c(0.25,0.5,0.75)), "quantiles", "exons"))) %>% 
  unnest(x) %>%
  spread(key=quantiles, value=exons)

# Assign a quantile label to each exon number based on its strain's exon number distribution
exons_tbl <- gff_exons %>%
    inner_join(exon_quantiles, by="name") %>%
    mutate(
        n_exon_quant = case_when(
            n_exon <= `25%` ~ "25% quantile",
            n_exon >= `75%` ~ "75% quantile",
            TRUE            ~ "Median"
        )
    ) %>%
    select(name, n_exon, n_exon_quant)

exons_tbl <- exons_tbl %>% 
    mutate(
        n_exon_quant = factor(
            n_exon_quant,
            ordered=T, 
            levels=c("25% quantile", "Median", "75% quantile")
        )
    )

exon_per_gene <- ggplot(data=exons_tbl, aes(x=n_exon, fill=n_exon_quant)) +
    scale_fill_manual(values=c("#6666ee", "#999999", "#ee6666")) +
    geom_histogram(binwidth=1) +
    theme_minimal() +
    stat_summary(aes(x = 0.1, y = n_exon, xintercept = stat(y), group = name), 
               fun.y = median, geom = "vline") +
    ylab("Genes") +
    xlab("Exons per genes") +
    facet_grid(name~.) +
    scale_x_continuous(limits = c(-0, 100))

### SAVE OUTPUT ###
svg(out_plot, height=10, width=12)
grid.arrange(feature_len, exon_per_gene)
dev.off()
write_tsv(gff_stats, out_tbl)

library(ggplot2)
library(patchwork)
library(readr)
library(tidyr)
library(dplyr)
library(tools)
library(stringr)


args <- commandArgs(trailingOnly=T)

# No need to scale, now
#library(scales)
# minscale <- function(x){s = scale(x); return(s + abs(min(s)))}
#  mutate_at(vars(-assembly), minscale ) %>%

asb <- read_tsv(args[1]) %>%
  rename(n_contigs=number, assembly=filename) %>%
  mutate(assembly=file_path_sans_ext(basename(assembly))) %>%
  select(assembly, total_length, n_contigs, N50, N70, N90, N50n, N70n, N90n) %>%
  pivot_longer(-assembly, 'metric') %>%
  mutate(assembly = gsub("_genome", "", assembly))

svg(args[2])

len_asb <- ggplot(
             data=asb %>% filter(metric == 'total_length'),
             aes(x=assembly, y=value, fill=assembly)
           ) +
        geom_bar(stat='identity') +
        theme_minimal() +
        ylab("Length") +
        xlab("") +
        ggtitle("Assembly length") +
        theme(legend.position="none")

num_tigs <- ggplot(
             data=asb %>% filter(metric == 'n_contigs'),
             aes(x=assembly, y=value, fill=assembly)
            ) +
        geom_bar(stat='identity') +
        theme_minimal() +
        ylab("# contigs") +
        xlab("") +
        ggtitle("Total number of contigs") +
        theme(legend.position="none")

Nxx <- ggplot(
         data=asb %>% filter(str_detect(metric, '0$')),
         aes(x=assembly, y=value, fill=assembly)
       ) +
       geom_bar(stat='identity') +
       facet_wrap(~metric) +
       theme_minimal() +
       xlab("") +
       ylab("Length") +
       theme(legend.position="none")

Nxxn <- ggplot(
          data=asb %>% filter(str_detect(metric, 'n$')),
          aes(x=assembly, y=value, fill=assembly)
        ) +
        geom_bar(stat='identity') +
        xlab("") +
        ylab("# contigs") +
        theme_minimal() +
        facet_wrap(~metric)

(len_asb + num_tigs + plot_spacer()) / (Nxx) / (Nxxn)

dev.off()

from collections import OrderedDict
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt


def read_busco(tbl: str, name: str) -> pd.DataFrame:
    """
    Read a busco summary table into a
    dataframe and create relevant columns
    """
    df = pd.read_csv(tbl, names=["number", "type"], sep="\t", header=None)
    df["strain"] = name
    # Single letter busco class as a column
    df["busco"] = df["type"].str.extract(r".*\(([A-Z])\).*")
    df.loc[pd.isnull(df.busco), "busco"] = "T"
    # Transform number of buscos to percentages
    df["perc"] = 100 * np.round(
        df["number"] / df.number[df.busco == "T"].values[0], decimals=2
    )
    return df


# Concatenate busco tables from all strains and set relevant
# variables as indices
busco = (
    pd.concat([read_busco(tbl, name) for name, tbl in snakemake.input.items()])
    .reset_index(drop=True)
    .sort_values("strain", ascending=False)
    .set_index(["busco", "strain"])
)
mpl.use("Agg")
# Keep track of the order in which strains will be plotted
str_to_num = OrderedDict({"v1": 0, "neff": 1, "c3": 2})
# Map each BUSCO class to a color
col_dict = OrderedDict(
    [("M", "#F04441"), ("F", "#F0E441"), ("D", "#3592C6"), ("S", "#58B4E8"),]
)
# We plot one busco class per iteration
for i, t in enumerate(col_dict.keys()):
    # Subset busco table for current strain
    busco_type = busco.loc[t]
    # Retrieve x index from strains (deterministic order)
    r = [str_to_num[s] for s in busco_type.index.values]
    # Stacked barplot, uses cumulative height of previous bars
    if i:
        plt.barh(r, busco_type.perc.values, color=col_dict[t], left=cum_height)
        cum_height += busco_type.perc.values
    else:
        cum_height = busco_type.perc.values
        plt.barh(r, busco_type.perc.values, color=col_dict[t])


def busco_string(strain):
    """Use the BUSCO summary table to construct the standard busco string"""
    mis = busco.loc[("M", strain), "number"]
    fra = busco.loc[("F", strain), "number"]
    com = busco.loc[("C", strain), "number"]
    sin = busco.loc[("S", strain), "number"]
    dup = busco.loc[("D", strain), "number"]
    tot = busco.loc[("T", strain), "number"]
    return f"{strain} - C:{com}[S:{sin}, D:{dup}], F:{fra}, M:{mis}, n={tot}"


# Write the strain name and busco string besides each bar
plt.yticks(r, [busco_string(s) for s in str_to_num.keys()])
plt.xlabel("% BUSCOs")
plt.ylabel("Strain")
plt.savefig(str(snakemake.output))

import sys
import re
from collections import defaultdict
import csv
from Bio import SeqIO
from Bio.Graphics import BasicChromosome
from Bio.SeqFeature import SeqFeature, FeatureLocation
from reportlab.lib.units import cm
from matplotlib import cm as mcm
from matplotlib.colors import to_hex

logh = open(snakemake.log[0], "w")
# {chrom: length}
entries = {}
for rec in SeqIO.parse(snakemake.params["genome"], "fasta"):
    if len(rec.seq) > 100000:
        entries[rec.id] = len(rec.seq)
# all available colors
# cols = list(colors.getAllNamedColors().values())[8:]
cols = [to_hex(mcm.tab10(c)) for c in range(len(entries))]
# {chrom: {type: [features]}} where a feature is [start, end, type, color]
features = defaultdict(list)
type_id = 0
ftype_to_col = {}
with open(snakemake.input["features"]) as bedh:
    bed = csv.DictReader(
        bedh, delimiter="\t", fieldnames=["chrom", "start", "end", "type"]
    )
    for feat in bed:
        start, end = int(feat["start"]), int(feat["end"])
        if start > end:
            start, end = end, start
        try:
            if end > entries[feat["chrom"]]:
                continue
        except KeyError:
            continue
        if feat["type"] not in ftype_to_col.keys():
            ftype_to_col[feat["type"]] = cols[type_id]
            logh.write(f"{feat['type']}: {cols[type_id]}\n")
            type_id += 1
        features[feat["chrom"]].append(
            SeqFeature(
                location=FeatureLocation(int(feat["start"]), int(feat["end"])),
                type=feat["type"],
                qualifiers={"color": [ftype_to_col[feat["type"]]]},
            )
        )

max_len = max(entries.values())
telomere_length = 10000  # For illustration

chr_diagram = BasicChromosome.Organism(output_format="svg")
chr_diagram.page_size = (29.7 * cm, 21 * cm)  # A4 landscape

for name, length in entries.items():
    # features = [f for f in record.features if f.type == "tRNA"]
    chrom_feat = features[name]

    chrom_num = re.sub(r"^.*_([0-9]+)$", r"\1", name)
    cur_chromosome = BasicChromosome.Chromosome(chrom_num)
    # Set the scale to the MAXIMUM length plus the two telomeres in bp,
    # want the same scale used on all five chromosomes so they can be
    # compared to each other
    cur_chromosome.scale_num = max_len + 2 * telomere_length

    # Add an opening telomere
    start = BasicChromosome.TelomereSegment()
    start.scale = telomere_length
    cur_chromosome.add(start)

    # Add a body - again using bp as the scale length here.
    body = BasicChromosome.AnnotatedChromosomeSegment(length, chrom_feat)
    body.scale = length
    cur_chromosome.add(body)

    # Add a closing telomere
    end = BasicChromosome.TelomereSegment(inverted=True)
    end.scale = telomere_length
    cur_chromosome.add(end)

    # This chromosome is done
    chr_diagram.add(cur_chromosome)
chr_diagram.draw(snakemake.output[0], snakemake.params["title"])
logh.close()

import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import patches as mpatches
import matplotlib as mpl
import matplotlib.gridspec as gridspec
import cooler

## LOAD CLI ARGS
cool_path = sys.argv[1]
rdna_path = sys.argv[2]
out_path = sys.argv[3]

# Discard contigs smaller than this size
# Note: assumes chromosomes are sorted by size in the cool file
MIN_SIZE = 100000

## LOAD INPUT FILES
c = cooler.Cooler(cool_path)
gff = pd.read_csv(rdna_path, sep="\t", comment="#", usecols=range(9))
gff.columns = [
    "seqname",
    "source",
    "feature",
    "start",
    "end",
    "score",
    "strand",
    "frame",
    "attribute",
]

## EXTRACT COOL DATA
chroms = c.chroms()[:]
# Discard all chromosomes after the first shorter than MIN_SIZE
chroms["include"] = chroms.length > MIN_SIZE
stop_bin = c.extent(chroms.loc[np.where(~chroms.include)[0][0], "name"])[1]
bins = c.bins()

# Remove chromosomes that are absent from cool or too small
gff = gff.loc[np.isin(gff.seqname, chroms.loc[chroms.include, "name"]), :]

## TRANSFORM

# Transform rDNA coords to bins
gff["minbin"] = gff.apply(
    lambda r: np.min(
        bins.fetch(f"{r.seqname}:{r.start}-{r.end}").index.values
    ),
    axis=1,
)

gff["maxbin"] = gff.apply(
    lambda r: np.max(
        bins.fetch(f"{r.seqname}:{r.start}-{r.end}").index.values
    ),
    axis=1,
)

gff = gff.drop_duplicates(subset=["attribute", "minbin", "maxbin"])

## VISUALIZE
mat = c.matrix(balance=False)[:stop_bin, :stop_bin]
# np.fill_diagonal(mat, 0)

# Define color saturation threshold
# sat_thr = np.percentile(mat[mat > 0], 85)

# Plot heatmap with rDNA overlay

plt.style.use("seaborn-white")
gs = gridspec.GridSpec(4, 1, height_ratios=[5, 1, 1, 1])
mpl.rcParams["figure.figsize"] = (5, 10)
fig = plt.figure()
ax1 = fig.add_subplot(gs[0])
ax2 = fig.add_subplot(gs[1])
ax3 = fig.add_subplot(gs[2])
ax4 = fig.add_subplot(gs[3])
fig.align_xlabels()
ax2.margins(x=0)
ax3.margins(x=0)
ax4.margins(x=0)

coldict = {
    "18s_rRNA": ("b", "-.", "18S", ax3),
    "28s_rRNA": ("g", ":", "28S", ax2),
    "8s_rRNA": ("r", "--", "5S", ax4),
}
legend_entries = [
    mpatches.Patch(color=v[0], label=v[2], ls=v[1]) for k, v in coldict.items()
]
plt.legend(handles=legend_entries, loc="upper right")
# mat[np.isnan(mat)] = 0
ax1.imshow(np.log(mat), cmap="Reds", interpolation="none", rasterized=True)

# Add vertical lines to line plots at rdna positions
for r in gff.iterrows():
    coldict[r[1]["attribute"]][3].axvline(
        x=r[1]["maxbin"],
        alpha=0.4,
        # linestyle=coldict[r[1]["attribute"]][1],
        lw=0.3,
        c=coldict[r[1]["attribute"]][0],
    )

# Add chromosome grid
for r in chroms.iterrows():
    if r[1].include:
        ax2.axvline(x=c.extent(r[1]["name"])[1], alpha=0.4, c="grey", lw=0.5)
        ax3.axvline(x=c.extent(r[1]["name"])[1], alpha=0.4, c="grey", lw=0.5)
        ax4.axvline(x=c.extent(r[1]["name"])[1], alpha=0.4, c="grey", lw=0.5)


# Subset bins containing rRNA
rdna_bins = {
    sub: np.unique(gff.maxbin[gff.attribute == sub])
    for sub in np.unique(gff.attribute)
}
# Compute contact sum per bin for each subunit
rdna_sig = {}
rdna_sig["28S"] = mat[rdna_bins["28s_rRNA"], :].mean(axis=0)
rdna_sig["18S"] = mat[rdna_bins["18s_rRNA"], :].mean(axis=0)
rdna_sig["5S"] = mat[rdna_bins["8s_rRNA"], :].mean(axis=0)

ax2.plot(
    range(mat.shape[0]), np.log10(rdna_sig["28S"]), label="28S", lw=0.5, c="g"
)
ax3.plot(
    range(mat.shape[0]), np.log10(rdna_sig["18S"]), label="18S", lw=0.5, c="b"
)
ax4.plot(
    range(mat.shape[0]), np.log10(rdna_sig["5S"]), label="5S", lw=0.5, c="r"
)

fig.savefig(out_path, dpi=2500)