def word_count(infilename, starter, stopper):
    with open(infilename, "r") as infile:
        words = []
        line = next(infile)
        while starter != line:
            line = next(infile)
        line = next(infile)  # make sure the first parsed line is not the starter
        while stopper != line:
            words += [word for word in line.strip().split()]
            line = next(infile)
    return len(words)


def check_wc(section_name, num_words, word_limit):
    if num_words > word_limit:
        raise ValueError(
            f"The {section_name} section is {num_words} words. You need to cut {num_words - word_limit} words."
        )


def main(src_filename, log_filename):
    with open(log_filename, "w") as outfile:
        outfile.write("section\tword_count\n")
        for section, word_limit, starter, stopper in zip(
            ["abstract", "importance"],
            [250, 150],
            ["## Abstract\n", "### Importance\n"],
            ["### Importance\n", "\\newpage\n"],
        ):
            wc = word_count(src_filename, starter, stopper)
            check_wc(section, wc, word_limit)
            outfile.write(f"{section}\t{wc}\n")


if __name__ == "__main__":
    if "snakemake" in locals() or "snakemake" in globals():
        main(snakemake.input.src, snakemake.output.txt)
    else:
        main("paper/paper.Rmd", "log/count_words_abstract.log")

library(here)
library(tidyverse)

rel_diff <- function(final, init, percent = TRUE) {
  mult <- if (isTRUE(percent)) 100 else 1
  return((final - init) / init * mult)
}
coeff_var <- function(x) {
  return(sd(x) / mean(x))
}


dat <- read_tsv(here("results", "summarized.tsv"))
agg <- read_tsv(here("results", "aggregated.tsv"))
################################################################################
# de novo datasets
opticlust_mcc <- agg %>%
  filter(
    method == "de_novo",
    tool == "mothur"
  ) %>%
  pull(mcc) %>%
  median()
opticlust_sec <- agg %>%
  filter(
    method == "de_novo",
    tool == "mothur"
  ) %>%
  pull(sec) %>%
  median()
opticlust_mem <- agg %>%
  filter(
    strategy == "de_novo",
    tool == "mothur"
  ) %>%
  pull(mem_gb) %>%
  median()
dn_vsearch_mcc <- agg %>%
  filter(strategy == "de_novo", tool == "vsearch") %>%
  pull(mcc) %>%
  median()
dn_vsearch_sec <- agg %>%
  filter(strategy == "de_novo", tool == "vsearch") %>%
  pull(sec) %>%
  median()
mcc_opticlust_vs_vsearch <- rel_diff(opticlust_mcc, dn_vsearch_mcc)
sec_opticlust_vs_vsearch <- abs(rel_diff(dn_vsearch_sec, opticlust_sec))

################################################################################
# de novo ref dbs

dn_dbs <- read_tsv("subworkflows/2_fit_reference_db/results/denovo_dbs.tsv") %>%
  group_by(ref) %>%
  summarize(med_mcc = median(mcc)) %>%
  full_join(read_tsv(here(
    "subworkflows", "0_prep_db", "data",
    "seq_counts.tsv"
  )),
  by = "ref"
  ) %>%
  mutate(refname = case_when(
    ref == "gg" ~ "Greengenes",
    TRUE ~ toupper(ref)
  ))

################################################################################
# ref db open
open_fit_db_mcc <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur"
  ) %>%
  pull(mcc) %>%
  median()

mcc_open_fit_db_vs_clust <- rel_diff(opticlust_mcc, open_fit_db_mcc)

open_fit_gg_mcc <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur",
    ref == "gg"
  ) %>%
  pull(mcc) %>%
  median()

open_fit_silva_mcc <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur",
    ref == "silva"
  ) %>%
  pull(mcc) %>%
  median()

open_fit_rdp_mcc <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur",
    ref == "rdp"
  ) %>%
  pull(mcc) %>%
  median()

open_vsearch_mcc <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "vsearch",
    ref == "gg"
  ) %>%
  pull(mcc) %>%
  median()
mcc_open_fit_db_vs_vsearch <- rel_diff(open_fit_gg_mcc, open_vsearch_mcc)

open_vsearch_sec <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "vsearch",
    ref == "gg"
  ) %>%
  pull(sec) %>%
  median()
open_fit_db_sec <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur"
  ) %>%
  pull(sec) %>%
  median()
sec_vsearch_vs_open_fit_db <- rel_diff(open_vsearch_sec, open_fit_db_sec)
sec_opticlust_vs_open_fit_db <- rel_diff(opticlust_sec, open_fit_db_sec) %>% abs()

# human dataset to silva
open_fit_silva_human_sec <- agg %>%
  filter(
    method == "open",
    strategy == "database",
    tool == "mothur",
    ref == "silva",
    dataset == "human"
  ) %>%
  pull(sec) %>%
  median()
closed_fit_silva_human_sec <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "silva",
    dataset == "human"
  ) %>%
  pull(sec) %>%
  median()
opticlust_human_sec <- agg %>%
  filter(
    method == "de_novo",
    tool == "mothur",
    dataset == "human"
  ) %>%
  pull(sec) %>%
  median()

################################################################################
# ref db closed
closed_fit_db_mcc <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur"
  ) %>%
  pull(mcc) %>%
  median()
mcc_closed_fit_db_vs_clust <- rel_diff(closed_fit_db_mcc, opticlust_mcc) %>% abs()

closed_fit_gg_mcc <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "gg"
  ) %>%
  pull(mcc) %>%
  median()
closed_fit_silva_mcc <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "silva"
  ) %>%
  pull(mcc) %>%
  median()
closed_fit_rdp_mcc <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "rdp"
  ) %>%
  pull(mcc) %>%
  median()

frac_fit_db <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur"
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100
frac_fit_gg <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "gg"
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100
frac_fit_silva <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "silva"
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100
frac_fit_rdp <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur",
    ref == "rdp"
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100
frac_vsearch <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "vsearch",
    ref == "gg"
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100
frac_vsearch_vs_fit <- rel_diff(frac_vsearch, frac_fit_gg)

closed_fit_db_sec <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "mothur"
  ) %>%
  pull(sec) %>%
  median()
sec_closed_fit_db_vs_clust <- rel_diff(closed_fit_db_sec, opticlust_sec) %>% abs()

closed_vsearch_sec <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "vsearch"
  ) %>%
  pull(sec) %>%
  median()
sec_closed_fit_db_vs_vsearch <- rel_diff(closed_fit_db_sec, closed_vsearch_sec) %>% abs()

closed_vsearch_mcc <- agg %>%
  filter(
    method == "closed",
    strategy == "database",
    tool == "vsearch"
  ) %>%
  pull(mcc) %>%
  median()
mcc_closed_fit_db_vs_vsearch <- rel_diff(closed_fit_db_mcc, closed_vsearch_mcc) %>% abs()

################################################################################
# fit split
cv_fit_split_mcc <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(mcc) %>%
  coeff_var()

frac_fit_split <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    method == "closed",
    ref_frac == 0.5
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100

closed_fit_split_sec <- agg %>%
  filter(
    method == "closed",
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(sec) %>%
  median()
open_fit_split_sec <- agg %>%
  filter(
    method == "open",
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(sec) %>%
  median()

sec_closed_fit_split_vs_clust <- rel_diff(opticlust_sec, closed_fit_split_sec) %>% abs()
sec_open_fit_split_vs_clust <- rel_diff(opticlust_sec, open_fit_split_sec) %>% abs()
sec_open_fit_split_vs_db <- rel_diff(open_fit_db_sec, open_fit_split_sec) %>% abs()
sec_closed_fit_split_vs_db <- rel_diff(closed_fit_db_sec, closed_fit_split_sec) %>% abs()

closed_fit_split_mem <- agg %>%
  filter(
    method == "closed",
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(mem_gb) %>%
  median()
open_fit_split_mem <- agg %>%
  filter(
    method == "open",
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(mem_gb) %>%
  median()
mem_closed_fit_split_vs_clust <- rel_diff(closed_fit_split_mem, opticlust_mem)
mem_open_fit_split_vs_clust <- rel_diff(open_fit_split_mem, opticlust_mem)

cv_fit_split_mcc_human_simple <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple"
  ) %>%
  pull(mcc) %>%
  coeff_var()

cv_fit_split_mem_human_simple <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple"
  ) %>%
  pull(mem_gb) %>%
  coeff_var()

sec_fit_split_human_simple_1 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple",
    ref_frac == 0.1
  ) %>%
  pull(sec) %>%
  median()

sec_fit_split_human_simple_9 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple",
    ref_frac == 0.9
  ) %>%
  pull(sec) %>%
  median()

frac_fit_split_human_simple_1 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple",
    method == "closed",
    ref_frac == 0.1
  ) %>%
  pull(fraction_mapped) %>%
  median()

frac_fit_split_human_simple_9 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    dataset == "human",
    ref_weight == "simple",
    method == "closed",
    ref_frac == 0.9
  ) %>%
  pull(fraction_mapped) %>%
  median()

mcc_fit_split_simple <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(mcc) %>%
  median()
mcc_opticlust_vs_fit_split_simple <- rel_diff(opticlust_mcc, mcc_fit_split_simple)

#####
# fit split at ref_frac 0.5

mcc_fit_split_abun <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "abundance",
    ref_frac == 0.5
  ) %>%
  pull(mcc) %>%
  median()

mcc_fit_split_dist <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "distance",
    ref_frac == 0.5
  ) %>%
  pull(mcc) %>%
  median()

frac_fit_split_simple <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    method == "closed",
    ref_frac == 0.5
  ) %>%
  pull(fraction_mapped) %>%
  median()

frac_fit_split_abun <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "abundance",
    method == "closed",
    ref_frac == 0.5
  ) %>%
  pull(fraction_mapped) %>%
  median()

frac_fit_split_dist <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "distance",
    method == "closed",
    ref_frac == 0.5
  ) %>%
  pull(fraction_mapped) %>%
  median()

sec_fit_split_simple <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "simple",
    ref_frac == 0.5
  ) %>%
  pull(sec) %>%
  median()

sec_fit_split_abun <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "abundance",
    ref_frac == 0.5
  ) %>%
  pull(sec) %>%
  median()

sec_fit_split_dist <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    ref_weight == "distance",
    ref_frac == 0.5
  ) %>%
  pull(sec) %>%
  median()

##########

frac_fit_split_1 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    method == "closed",
    ref_frac == 1
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100

frac_fit_split_9 <- agg %>%
  filter(
    strategy == "self-split",
    tool == "mothur",
    method == "closed",
    ref_frac == 0.9
  ) %>%
  pull(fraction_mapped) %>%
  median() * 100


################################################################################
# save results
save.image(file = here("results", "stats.RData"))

devtools::load_all("../ggraph")
library(cowplot)
library(glue)
library(gridExtra)
library(ggtext)
library(here)
library(patchwork)
library(reticulate)
library(tidygraph)
library(tidyverse)
set.seed(20200308)
# use_python('/usr/local/bin/python3')
source_python(here("code", "py", "algorithm_diagram.py"))
optifit <- create_optifit()
optifit_iters <- optifit$iterate %>%
  lapply(function(x) {
    return(list(
      nodes = x[["nodes"]] %>% py_to_r(),
      edges = x[["edges"]] %>% py_to_r() %>%
        bind_rows(data.frame(from = 1, to = 1, mcc = NA)) %>%
        mutate(
          is_loop = from == to,
          loop_dir = ifelse(from == 1, 270, 90)
        )
    ))
  })

plot_optifit_graph <- function(graph, title = "",
                               hide_loops = FALSE) {
  loop_dir <- 90
  loop_color <- ifelse(hide_loops, "white", "black")
  create_layout(graph, "linear", sort.by = id) %>%
    ggraph() +
    geom_edge_arc(aes(
      label = mcc,
      start_cap = label_rect(node1.name),
      end_cap = label_rect(node2.name,
        padding = margin(1, 1, 2.8, 1, "mm")
      )
    ),
    arrow = arrow(
      length = unit(3, "mm"),
      angle = 35,
      type = "closed"
    ),
    edge_colour = "gray",
    angle_calc = "along",
    label_dodge = unit(-2, "mm")
    ) +
    geom_edge_loop(aes(
      span = 1,
      direction = loop_dir,
      strength = 0.5,
      color = is_loop
    )) +
    geom_node_label(aes(label = name)) +
    scale_edge_color_manual(values = c(loop_color)) +
    labs(title = title) +
    theme_void() +
    theme(
      plot.margin = unit(x = c(0, 0, 0, 0), units = "pt"),
      legend.position = "none"
    )
}

i <- 0
optifit_graphs <- lapply(optifit_iters, function(x) {
  i <<- i + 1
  tbl_graph(nodes = x$nodes, edges = x$edges) %>%
    plot_optifit_graph(
      title = glue("{i}. MCC = {x$edges %>% filter(is_loop) %>% pull(mcc)}"),
      hide_loops = TRUE
    )
})


base_color <- "#000000"
ref_color <- "#D95F02"
query_color <- "#1B9E77"
ref_seqs <- LETTERS[1:17]
query_seqs <- LETTERS[23:26]

dist_dat <- get_dists() %>%
  arrange(seq1, seq2) %>%
  mutate(
    color1 = ifelse(seq1 %in% ref_seqs, ref_color, query_color),
    color2 = ifelse(seq2 %in% ref_seqs, ref_color, query_color),
  )
dist_dat[["color3"]] <- rep.int("black", nrow(dist_dat))
dist_dat[["dist"]] <- runif(nrow(dist_dat), 1.0, 2.9) %>%
  format(digits = 2) %>%
  as.character()
table_colors <- dist_dat %>%
  select(color1, color2, color3) %>%
  as.matrix() %>%
  t()

table_plot <- plot_grid(ggdraw() +
  draw_label("0. List of sequence pairs within the distance threshold",
    x = 0,
    hjust = 0
  ) +
  theme(plot.margin = margin(5, 0, 5, 0)),
tableGrob(dist_dat %>%
  select(seq1, seq2, dist) %>%
  rename(
    `% distance` = dist,
    ` ` = seq1,
    `  ` = seq2
  ) %>%
  t(),
theme = ttheme_default(
  base_size = 10,
  padding = unit(c(4, 4), "pt"),
  core = list(
    fg_params = list(col = table_colors),
    bg_params = list(col = "white")
  ),
  rowhead = list(bg_params = list(col = NA)),
  colhead = list(bg_params = list(col = NA))
)
),
ncol = 1, rel_heights = c(0.1, 1)
)

plot_diagram <- table_plot /
  optifit_graphs +
  plot_layout(heights = c(0.75, 1, 1.5, 1, 0.3))


dims <- eval(parse(text = snakemake@params[["dim"]]))
ggsave(snakemake@output[["tiff"]],
  device = "tiff", dpi = 300,
  width = dims[1], height = dims[2], units = "in"
)

set.seed(20210509)
library(cowplot)
library(ggtext)
library(glue)
library(here)
library(knitr)
library(tidyverse)

mutate_perf <- function(dat) {
  dat %>%
    mutate(
      mem_mb = max_rss,
      mem_gb = mem_mb / 1024
    ) %>%
    rename(sec = s)
}
select_cols <- function(dat) {
  dat %>%
    select(
      dataset, strategy, method, tool, mcc, sec, mem_gb, fraction_mapped,
      ref_frac, ref_weight
    )
}

opticlust <- read_tsv(here("subworkflows/1_prep_samples/results/opticlust_results.tsv")) %>%
  full_join(read_tsv(here("subworkflows/1_prep_samples/results/dataset_sizes.tsv"))) %>%
  mutate_perf() %>%
  mutate(
    strategy = method,
    fraction_mapped = NA,
    ref_frac = 0,
    ref_weight = "NA"
  )
optifit_split <- read_tsv(here("subworkflows/3_fit_sample_split/results/optifit_split_results.tsv")) %>%
  mutate_perf() %>%
  mutate(strategy = "self-split")

dat <- list(optifit_split, opticlust) %>%
  lapply(select_cols) %>%
  reduce(bind_rows) %>%
  mutate(
    method = as.character(method),
    strategy = as.character(strategy)
  ) %>%
  mutate(fraction_mapped = case_when(
    method %>% as.character() != "closed" ~ NA_real_,
    TRUE ~ fraction_mapped
  )) %>%
  pivot_longer(c(mcc, fraction_mapped, sec),
    names_to = "metric"
  ) %>%
  mutate(
    metric = factor(
      case_when(
        metric == "mcc" ~ "MCC",
        metric == "fraction_mapped" ~ "Fraction Mapped",
        metric == "sec" ~ "Runtime (sec)",
        TRUE ~ metric
      ),
      levels = c("MCC", "Fraction Mapped", "Runtime (sec)")
    ),
    strategy = factor(
      case_when(
        strategy == "de_novo" ~ "_de novo_",
        strategy == "database_rdp" ~ "db: RDP",
        strategy == "database_silva" ~ "db: SILVA",
        strategy == "database_gg" ~ "db: Greengenes",
        TRUE ~ strategy
      ),
      levels = c(
        "db: RDP", "db: SILVA", "db: Greengenes",
        "self-split", "_de novo_"
      )
    ),
    method = factor(
      case_when(
        method == "de_novo" ~ "_de novo_",
        TRUE ~ method
      ),
      levels = c("open", "closed", "_de novo_")
    ),
    ref_weight = factor(
      case_when(
        ref_weight == "distance" ~ "similarity",
        TRUE ~ ref_weight
      ),
      levels = c("simple", "abundance", "similarity", "NA")
    )
  )

med_iqr <- function(x) {
  return(data.frame(
    y = median(x),
    ymin = quantile(x)[2],
    ymax = quantile(x)[4]
  ))
}
color_breaks <- list(
  simple = "#FF8C00",
  abundance = "#9932CC",
  similarity = "#008B8B"
)
color_labels <- lapply(
  names(color_breaks),
  function(name) {
    glue("<span style = 'color:{color_breaks[[name]]};'>{name}</span>")
  }
) %>% unlist()
color_values <- append(color_breaks, list(`NA` = "#000000"))
plot_results_split <- dat %>%
  filter(((ref_weight == "simple" | ref_frac == 0.5) | method == "_de novo_") & !is.na(value)) %>%
  ggplot(aes(ref_frac, value, color = ref_weight, shape = method)) +
  coord_flip() +
  stat_summary(
    geom = "point",
    fun = median,
    size = 2,
    position = position_dodge(width = 0.07)
  ) +
  facet_grid(dataset ~ metric, scales = "free", switch = "x") +
  scale_shape_manual(values = list(open = 1, closed = 19, `_de novo_` = 17)) +
  scale_color_manual(
    values = color_values,
    breaks = names(color_breaks),
    labels = color_labels
  ) +
  scale_x_continuous(
    breaks = seq(0, 1, 0.1),
    labels = c("NA", seq(0.1, 1, 0.1))
  ) +
  labs(x = "reference fraction", y = "") +
  theme_bw() +
  theme(
    legend.text = element_markdown(),
    legend.title = element_blank(),
    legend.position = "top",
    legend.margin = margin(t = 0, r = 0, b = 0, l = 0, unit = "pt"),
    legend.spacing.x = unit(0.5, "pt"),
    plot.margin = unit(x = c(0, 0, 0, 0), units = "pt"),
    panel.grid.minor.y = element_blank(),
    axis.title.x = element_blank(),
    strip.placement = "outside",
    strip.background = element_blank()
  ) +
  guides(
    shape = guide_legend(order = 1),
    colour = guide_legend(
      override.aes = list(size = -1)
    )
  )

dim <- eval(parse(text = snakemake@params[["dim"]]))
ggsave(snakemake@output[["tiff"]],
  device = "tiff", dpi = 300,
  width = dim[1], height = dim[2], units = "in"
)