library(dplyr)
library(readr)
library(purrr)
library(ggplot2)
library(tidyr)
library(aplot)

# read in metadata and results --------------------------------------------

destfile <- "inputs/gtdb-rs202.taxonomy.v2.csv"
url <- "https://osf.io/p6z3w/download"
if (!file.exists(destfile)) {
  download.file(url, destfile, method="auto") 
}
gtdb_lineages <- read_csv(destfile)

destfile <- "inputs/hmp2_metadata.csv"
url <- "https://ibdmdb.org/tunnel/products/HMP2/Metadata/hmp2_metadata.csv"
if (!file.exists(destfile)) {
  download.file(url, destfile, method="auto") 
}
hmp_metadata <- read_csv(destfile)
h4017 <- hmp_metadata %>%
  select(participant_id = "Participant ID", data_id = "External ID", data_type, 
         week_num, diagnosis, antibiotics = "Antibiotics") %>%
  filter(data_type == "metagenomics") %>%
  filter(data_id %in% c('HSM67VF9', 'HSM67VFD', 'HSM67VFJ', 'HSM6XRQB',
                        'HSM6XRQI', 'HSM6XRQK', 'HSM6XRQM', 'HSM6XRQO',
                        'HSM7CYY7', 'HSM7CYY9', 'HSM7CYYB', 'HSM7CYYD'))

#gather_results <- unlist(snakemake@input[['gather']]) %>%
gather_results <- Sys.glob("outputs/sample_gather/*genomic.csv") %>%
  set_names() %>%
  map_dfr(read_csv, col_types = c("dddddlllcccddddcccd"), .id = "sample") %>%
  mutate(sample = gsub("outputs/sample_gather/", "", sample)) %>%
  mutate(sample = gsub("_gather_gtdb-rs202-genomic.csv", "", sample)) %>%
  separate(col = name, into = c("accession"), remove = F, sep = " ", extra = "drop")

gather_results <- left_join(gather_results, gtdb_lineages, by = c("accession" = "ident")) %>%
  left_join(h4017, by = c("sample" = "data_id"))

# plot --------------------------------------------------------------------
h4017$antibiotic <- c("ciprofloxacin", "ciprofloxacin", "ciprofloxacin",
                      "none", "none", "metronidazole", "metronidazole",
                      "metronidazole", "none", "none", "none", "unknown")
h4017$antibiotic <- factor(h4017$antibiotic, levels =  c("ciprofloxacin", "metronidazole", "unknown", "none"))
abx_plt <- ggplot(h4017, aes(x = week_num, y = diagnosis, shape = antibiotic)) +
  geom_point(size = 2.5) +
  theme_minimal() +
  theme(axis.title.x = element_blank(),
        axis.title.y = element_text(size = 8),
        axis.text = element_blank(),
        legend.title = element_text(size = 8),
        legend.text = element_text(size = 7),
        panel.grid = element_blank())+
  labs(x = "week number", y = "antibiotic") 
abx_plt

# decide which genome to query with ---------------------------------------

common_species <- gather_results %>%
  select(sample, species) %>%
  distinct() %>%
  group_by(species) %>%
  tally() %>%
  filter(n == 12)

# common_species_plt <- ggplot(gather_results %>%
#                                filter(species %in% common_species$species),
#                              aes(x = week_num, y = f_unique_to_query)) +
#   geom_col() +
#   facet_wrap(~species) +
#   theme_minimal() 

sgc_species <- gather_results %>%
  filter(species %in% common_species$species) %>%
  group_by(species) %>%
  summarize(sum_f_unique_to_query = sum(f_unique_to_query)) %>%
  filter(sum_f_unique_to_query > 0.2)

gather_results2 <- gather_results %>%
  mutate(species2 = ifelse(species %in% sgc_species$species, species, "other")) %>%
  mutate(species2 = gsub("s__", "", species2))
gather_results2$species2 <- factor(gather_results2$species2, 
                                   levels =c('other', 
                                             'Bacteroides fragilis',
                                             'Bacteroides uniformis',
                                             'Enterocloster bolteae',
                                             'Parabacteroides distasonis',
                                             'Parabacteroides merdae',
                                             'Phocaeicola vulgatus'))
gather_results2 <- left_join(gather_results2, h4017)
sgc_plt <- ggplot(gather_results2,
                     aes(x = week_num, y = f_unique_to_query, fill = species2)) +
  geom_col() +
  geom_point(aes(x = week_num, y = 1, shape = antibiotic)) +
  labs(x = "week number", y = "fraction of metagenome", fill = "species")+
  theme_minimal() +
  theme(legend.title = element_text(size = 8),
        legend.text = element_text(size = 7, face = "italic"),
        axis.title = element_text(size = 8),
        axis.text = element_text(size = 7)) +
  ylim(0, 1) + 
  scale_fill_manual(values = c("#999999", "#E69F00", "#56B4E9", "#F0E442", "#009E73",
                               "#CC79A7", "#0072B2")) +
  guides(fill = guide_legend(override.aes = list(shape = NA)))
sgc_plt

pdf("figures/common_species_breakdown2.pdf", width = 6, height = 3.6)
#pdf(snakemake@output[['pdf']], width = 6, height = 3)
#abx_plt %>% insert_bottom(sgc_plt, height = 5)
sgc_plt
dev.off()

# output sgc queries as gather file ---------------------------------------

# get rep accession
sgc_genomes <- gather_results %>%
  filter(species %in% sgc_species$species) %>%
  group_by(name, species) %>%
  summarize(sum_f_unique_to_query = sum(f_unique_to_query)) %>%
  ungroup() %>%
  group_by(species) %>%
  arrange(desc(sum_f_unique_to_query)) %>%
  slice(1)

gather_sgc_queries <- unlist(snakemake@input[['gather']])[6] %>%
#gather_sgc_queries <- Sys.glob("outputs/sample_gather/*genomic.csv")[6] %>%
  map_dfr(read_csv, col_types = c("dddddlllcccddddcccd")) %>%
  filter(name %in% sgc_genomes$name)

write_csv(gather_sgc_queries, snakemake@output[['gather_grist']])

import sys
import argparse
import urllib.request
import csv

from lxml import etree


def url_for_accession(accession):
    db, acc = accession.strip().split("_")
    if '.' in acc:
        number, version = acc.split(".")
    else:
        number, version = acc, '1'
    number = "/".join([number[p : p + 3] for p in range(0, len(number), 3)])
    url = f"https://ftp.ncbi.nlm.nih.gov/genomes/all/{db}/{number}"
    print(f"opening directory: {url}", file=sys.stderr)

    with urllib.request.urlopen(url) as response:
        all_names = response.read()

    print("done!", file=sys.stderr)

    all_names = all_names.decode("utf-8")

    full_name = None
    for line in all_names.splitlines():
        if line.startswith(f'<a href='):
            name=line.split('"')[1][:-1]
            db_, acc_, *_ = name.split("_")
            if db_ == db and acc_.startswith(acc):
                full_name = name
                break

    if full_name is None:
        return None
    else:
        url = "htt" + url[3:]
        return (
            f"{url}/{full_name}/{full_name}_genomic.fna.gz",
            f"{url}/{full_name}/{full_name}_assembly_report.txt",
        )


def get_taxid_from_assembly_report(url):
    print(f"opening assembly report: {url}", file=sys.stderr)
    with urllib.request.urlopen(url) as response:
        content = response.read()
    print("done!", file=sys.stderr)

    content = content.decode("utf-8").splitlines()
    for line in content:
        if "Taxid:" in line:
            line = line.strip()
            pos = line.find("Taxid:")
            assert pos >= 0
            pos += len("Taxid:")
            taxid = line[pos:]
            taxid = taxid.strip()
            return taxid

    assert 0


def get_tax_name_for_taxid(taxid):
    tax_url = (
        f"https://www.ncbi.nlm.nih.gov/taxonomy/?term={taxid}&report=taxon&format=text"
    )
    print(f"opening tax url: {tax_url}", file=sys.stderr)
    with urllib.request.urlopen(tax_url) as response:
        content = response.read()

    print("done!", file=sys.stderr)

    root = etree.fromstring(content)
    notags = etree.tostring(root).decode("utf-8")
    if notags.startswith("<pre>"):
        notags = notags[5:]
    if notags.endswith("</pre>"):
        notags = notags[:-6]
    notags = notags.strip()

    return notags


def main():
    p = argparse.ArgumentParser()
    p.add_argument("accession")
    p.add_argument("-o", "--output")
    args = p.parse_args()

    fieldnames = ["acc", "genome_url", "assembly_report_url", "ncbi_tax_name"]
    fp = None
    if args.output:
        fp = open(args.output, "wt")
        w = csv.DictWriter(fp, fieldnames=fieldnames)
    else:
        w = csv.DictWriter(sys.stdout, fieldnames=fieldnames)
    w.writeheader()

    acc = args.accession

    genome_url, assembly_report_url = url_for_accession(acc)
    taxid = get_taxid_from_assembly_report(assembly_report_url)
    tax_name = get_tax_name_for_taxid(taxid)

    d = dict(
        acc=acc,
        genome_url=genome_url,
        assembly_report_url=assembly_report_url,
        ncbi_tax_name=tax_name,
    )

    w.writerow(d)
    print(f"retrieved for {acc} - {tax_name}", file=sys.stderr)

    if fp:
        fp.close()

    return 0


if __name__ == "__main__":
    sys.exit(main())

install.packages("pagoo", repos="http://cran.us.r-project.org")
library(pagoo)
file.create(snakemake@output[['pagoo']])

library(readr)
library(dplyr)
library(tidyr)
library(purrr)

gtdb_lineages <- read_csv(snakemake@input[['lineages']])

# this specified as a sysglob to not mess up the sample wildcard in the workflow
gather_bins <- Sys.glob("outputs/metabat2_gather/*csv") %>%
#gather_bins <- unlist(snakemake@input[['gather']]) %>%
  map_dfr(read_csv, col_types = c("dddddlllcccddddcccd")) %>%
  separate(col = name, into = c("accession"), remove = F, sep = " ", extra = "drop") %>%
  mutate(sample = basename(query_filename)) %>%
  mutate(sample = gsub("-.*", "", sample)) %>%
  left_join(gtdb_lineages, by = c("accession" = "ident"))

species_of_interest <- read_csv(unlist(snakemake@input[['acc_to_db']])) %>%
  #filter(!species %in% c("s__Ruminococcus_B_gnavus", "s__Clostridium_Q_symbiosum", "s__Roseburia_intestinalis")) %>%
  mutate(species = gsub("(s__.*?)_", "\\1 ", species))

acc_to_db <- read_csv(unlist(snakemake@input[['acc_to_db']])) %>%
  #filter(!species %in% c("s__Ruminococcus_B_gnavus", "s__Clostridium_Q_symbiosum", "s__Roseburia_intestinalis")) %>%
  mutate(acc_to_db = paste0(accession, "-", species)) %>%
  mutate(species = gsub("(s__.*?)_", "\\1 ", species))

# only consider best match for the bin
gather_bins <- gather_bins %>%
  filter(gather_result_rank == 0)

named_bins <- gather_bins %>% 
  left_join(acc_to_db, by = "species") %>%
  select(acc_to_db, query_filename) %>%
  mutate(sig_filename = gsub("metabat2", "metabat2_prokka_sigs", query_filename),
         sig_filename = gsub("/bin", "_bin", sig_filename),
         sig_filename = gsub(".fa", ".sig", sig_filename)) %>%
  select(acc_to_db, sig_filename)

# filter to single acc_db using wildcard
# and only write out sig path
tmp <- named_bins %>%
  filter(acc_to_db %in% snakemake@wildcards[["acc_db"]]) %>%
  select(sig_filename)

write_tsv(tmp, snakemake@output[['txt']], col_names = FALSE)

from sourmash import signature
import pandas as pd
import os
import sys
import argparse

def main():
    p = argparse.ArgumentParser()
    p.add_argument('signature')       # sourmash signature
    p.add_argument('output')          # output csv file name
    args = p.parse_args()

    # load the signature from disk
    sigfp = open(args.signature, 'rt')
    siglist = list(signature.load_signatures(sigfp))
    loaded_sig = siglist[0]

    # Get the minhashes
    mins = loaded_sig.minhash.hashes.keys()

    name = os.path.basename(args.signature)
    df = pd.DataFrame(mins, columns=[name])

    # write to a csv
    df.to_csv(args.output, index = False)

if __name__ == '__main__':
    sys.exit(main())

shell:'''
fastp --in1 {input.R1} \
  --in2 {input.R2} \
  --out1 {output.R1} \
  --out2 {output.R2} \
  --detect_adapter_for_pe \
  --qualified_quality_phred 4 \
  --length_required 31 --correction \
  --json {output.json} \
  --html {output.html}
'''

shell:'''
bbduk.sh -Xmx64g t=3 in={input.r1} in2={input.r2} out={output.r1} out2={output.r2} outm={output.human_r1} outm2={output.human_r2} k=31 ref={input.human}
'''

shell:'''
megahit -1 {input.r1} -2 {input.r2} --out-dir {output.tmp_dir} --out-prefix {wildcards.sample}
mv {output.tmp_dir}/{wildcards.sample}.contigs.fa {output.fa}
'''

shell:'''
bowtie2-build {input} {input}    
'''

shell:'''
bowtie2 -x {input.fa} -1 {input.r1} -2 {input.r2} |
    samtools view -bS -o {output}
'''

shell:'''
prokka {input} --outdir {params.outdir} --prefix {wildcards.sample}_bin.{wildcards.bin} \
    --metagenome --force --locustag {wildcards.sample}_bin.{wildcards.bin} --cpus {threads} --centre X --compliant
'''

run:
    # something weird is happening here where snakemake solves the inputs correctly, but 
    # "queries" ends up being the gather csv file. So...parse the CSV file for accession
    # numbers. Sigh.
    gather_csv = input[0]
    genome_accs = []
    with open(gather_csv, 'rt') as fp:
        r = csv.DictReader(fp)
        for row in r:
            acc = row['name'].split(' ')[0]
            acc = "outputs/genbank_genomes/" + acc + "_genomic.fna.gz"
            genome_accs.append(acc)

    query_list = "\n- ".join(genome_accs)
    #query_list = "\n- ".join(input.queries)
    with open(output.conf, 'wt') as fp:
       print(f"""\