library(dplyr)
library(tidyr)
library(readr)
library(ggplot2)
library(tibble)
library(gridExtra)
library(ggtext)

knitr::opts_chunk$set(echo = TRUE,message = FALSE,warning = FALSE,fig.align = "center")

transparent_theme <- theme(
  panel.grid = element_blank(),
  panel.background = element_blank(),
  plot.background = element_rect(fill = "transparent", color = NA),
  legend.position = "none")

fst_df <- read_csv(fst_file, show_col_types = FALSE)
pi_df  <- read_csv(pi_file, show_col_types = FALSE)

required_fst_columns <- c("pop1", "pop2", "chromosome", "window_pos_1", "window_pos_2", "avg_wc_fst")
required_pi_columns <- c("pop", "chromosome", "window_pos_1", "window_pos_2", "avg_pi")

if (!all(required_fst_columns %in% names(fst_df))) {
  stop("The Fst file is missing: ",paste(setdiff(required_fst_columns, names(fst_df)), collapse = ", "))}

if (!all(required_pi_columns %in% names(pi_df))) {
  stop("The pi file is missing: ",paste(setdiff(required_pi_columns, names(pi_df)), collapse = ", "))}

pair_key <- function(pop_a, pop_b) {paste(pmin(pop_a, pop_b), pmax(pop_a, pop_b), sep = "__")}

comparison_map <- comparison_map %>% mutate(pair_key = pair_key(population_1, population_2))

fst_df_prepared <- fst_df %>%
  mutate(chromo.num = match(chromosome, chromosome_accessions),pair_key = pair_key(pop1, pop2)) %>%
  filter(!is.na(chromo.num)) %>%
  inner_join(comparison_map %>% select(comp, pair_key),by = "pair_key") %>%
  select(chromosome, chromo.num, window_pos_1, window_pos_2, comp,avg_wc_fst)

missing_fst_comparisons <- setdiff(comps, unique(fst_df_prepared$comp))
if (length(missing_fst_comparisons) > 0) {
  stop("No Fst rows were found for: ", paste(missing_fst_comparisons, collapse = ", "))}

# check to make sure this df looks good so far
# see knitted document version for tutorial example output here
glimpse(fst_df_prepared)

pi_df_wide <- pi_df %>%
  select(pop, chromosome, window_pos_1, window_pos_2, avg_pi) %>%
  pivot_wider(names_from = pop, values_from = avg_pi)

expected_populations <- unique(c(comparison_map$population_1,comparison_map$population_2))

missing_pi_populations <- setdiff(expected_populations,names(pi_df_wide))

if (length(missing_pi_populations) > 0) {
  stop("The pi file contains no data for: ",
    paste(missing_pi_populations, collapse = ", "))}

# Calculate the pi ratio for every comparison in comparison_map
pi_df_ratios <- lapply(seq_len(nrow(comparison_map)), function(i) {
  
  comparison_name <- comparison_map$comp[[i]]
  population_1 <- comparison_map$population_1[[i]]
  population_2 <- comparison_map$population_2[[i]]
  
  pi_df_wide %>%
    transmute(chromosome,window_pos_1,window_pos_2,
      comp = comparison_name,
      population_1 = population_1,population_2 = population_2,
      pi_ratio = .data[[population_1]] / .data[[population_2]]
    )
}) %>%
  bind_rows() %>%
  left_join(chromosome_map, by = "chromosome") %>%
  mutate(log2_pi.ratio = log2(pi_ratio)) %>%
  filter(!is.na(chromo.num))

# check to make sure this df looks good so far 
# see knitted document version for tutorial example output here
glimpse(pi_df_ratios)

pi_fst_df_merge_clean <- pi_df_ratios %>%
  inner_join(fst_df_prepared,
    by = c("chromosome", "chromo.num", "window_pos_1", "window_pos_2", "comp")) %>%
  filter(!is.na(avg_wc_fst), !is.na(log2_pi.ratio)) %>%
  mutate(chr_pos1_pos2_comp = paste(chromo.num, window_pos_1, window_pos_2, comp, sep = "_")) %>%
  select(chr_pos1_pos2_comp, chromo.num, window_pos_1, window_pos_2, comp, avg_wc_fst, pi_ratio, log2_pi.ratio)

# check to make sure this df looks good so far
# see knitted document version for tutorial example output here
glimpse(pi_fst_df_merge_clean)

main_df_clean <- pi_fst_merge_no.neg.fst_merged_individ.comp %>%
  mutate(
    log2_pi.ratio_inf2NA = if_else(
      is.finite(log2_pi.ratio), log2_pi.ratio, NA_real_
    ),
    tri.coded_inf.or.no = case_when(
      log2_pi.ratio == -Inf ~ "inf_bc.pop1.pi.0", 
      log2_pi.ratio == Inf  ~ "inf_bc.pop2.pi.0",
      TRUE                  ~ "not.inf"
    ),
    outlier.labels_inf.labels = case_when(
      outlier.labels_individ.comp == "outlier" & log2_pi.ratio == -Inf ~ "outlier_pop1.pi.0",
      outlier.labels_individ.comp == "outlier" & log2_pi.ratio == Inf  ~ "outlier_pop2.pi.0",
      outlier.labels_individ.comp == "outlier" ~ "outlier_not.inf",
      log2_pi.ratio == -Inf ~ "not.out_pop1.pi.0",
      log2_pi.ratio == Inf  ~ "not.out_pop2.pi.0",
      TRUE                  ~ "not.out_not.inf"
    ),
    alpha.for.asterisk = case_when(
      outlier.labels_inf.labels %in%
        c("outlier_pop1.pi.0", "outlier_pop2.pi.0") ~ 1,
      outlier.labels_inf.labels %in%
        c("not.out_pop1.pi.0", "not.out_pop2.pi.0") ~ 0.6,
      TRUE ~ 0
    ),
    alpha.for.main.pts = case_when(
      outlier.labels_inf.labels == "outlier_not.inf" ~ 1,
      outlier.labels_inf.labels == "not.out_not.inf" ~ 0.6,
      TRUE ~ 0
    ),
    tri.coded_inf.or.no = factor(
      tri.coded_inf.or.no,
      levels = c("not.inf", "inf_bc.pop1.pi.0", "inf_bc.pop2.pi.0")
    ),
    outlier.labels_individ.comp = factor(
      outlier.labels_individ.comp,
      levels = c("not.outlier", "outlier")
    )
  )

outlier_counts <- main_df_clean %>%
  filter(outlier.labels_individ.comp == "outlier") %>%
  mutate(
    outlier_population = if_else(
      log2_pi.ratio < 0, "population_1", "population_2"
    )
  ) %>%
  count(comp, outlier_population, name = "n_outlier_windows") %>%
  pivot_wider(
    names_from = outlier_population,
    values_from = n_outlier_windows,
    values_fill = 0,
    names_prefix = "n_outliers_"
  ) %>%
  rename(comparison = comp)

outlier_counts

write_csv(
  outlier_counts,
  file.path(wrangled_dir, "outlier_counts_by_comparison.csv")
)

remove_chromosome_strips <- function(plot_object) {
  plot_object + theme(strip.background = element_blank(),strip.text = element_blank())}

add_combined_panel_border <- function(plot) {
  grob <- ggplotGrob(plot)
  panels <- grob$layout[grepl("^panel", grob$layout$name), ]

  gtable::gtable_add_grob(
    grob,
    grid::rectGrob(gp = grid::gpar(
      color = "black", lwd = 1)), # edit panel border color here
    t = min(panels$t), b = max(panels$b),
    l = min(panels$l), r = max(panels$r),
    z = Inf, clip = "off"
  )
}

main_point_panels <- lapply(seq_along(plot_order), function(i) {
  plot <- meadow_plots[[plot_order[i]]]$main_points + labs(x = "Chromosome")

  if (i < length(plot_order)) {
    plot <- plot + theme(strip.background = element_blank(),
                         strip.text.x = element_text(color = "transparent", size = 1),
                         axis.title.x = element_text(color = "transparent", size = 1))
  } else {
    plot <- plot + theme(strip.background = element_blank(),
                         strip.text.x = element_text(),
                         axis.title.x = element_text(size = 9, color = "black")) # edit shared x-axis label here
  }

  add_combined_panel_border(plot)
})

shared_y_title <- grid::textGrob(
  expression(paste(
    "Average ", italic(F)[italic(ST)], "(15-kb windows, 5-kb step)")), # edit shared y-axis label text here
  rot = 90,
  gp = grid::gpar(fontsize = 9) # edit shared y-axis label text size here
)

meadow.plots_main.pts_combined <- arrangeGrob(
  grobs = main_point_panels, ncol = 1,
  left = shared_y_title)

ggsave(
  file.path(plot_dir, "meadow_plots.png"),
  meadow.plots_main.pts_combined,
  width = 20, # edit final figure width here
  height = min(8*nrow(comparison_map),60), # edit final figure height here
  units = "cm", dpi = 600,
  bg = "transparent"
)

max_intervening_nonoutliers <- 3

if (length(max_intervening_nonoutliers) != 1 ||
    is.na(max_intervening_nonoutliers) ||
    max_intervening_nonoutliers < 0 ||
    max_intervening_nonoutliers %% 1 != 0) {
  stop("max_intervening_nonoutliers must be one non-negative whole number.")
}

comparison_peak_summary <- candidate_peaks %>%
  group_by(comp) %>%
  summarise(
    n_candidate_peaks = n(),
    mean_outlier_windows_per_peak = mean(n_outlier_windows),
    max_outlier_windows_per_peak = max(n_outlier_windows),
    n_outlier_windows = sum(n_outlier_windows),
    .groups = "drop"
  ) %>%
  relocate(n_outlier_windows, .after = n_candidate_peaks) %>%
  arrange(match(comp, plot_order))

peak_size_summary <- candidate_peaks %>%
  count(comp, n_outlier_windows, name = "n_candidate_peaks") %>%
  arrange(match(comp, plot_order), n_outlier_windows)

comparison_peak_summary
peak_size_summary

write_csv(candidate_peaks, file.path(wrangled_dir, "candidate_peaks.csv"))
write_csv(comparison_peak_summary, file.path(wrangled_dir, "candidate_peak_summary_by_comparison.csv"))
write_csv(peak_size_summary, file.path(wrangled_dir, "candidate_peak_size_counts.csv"))

candidate_peaks_for_plot <- candidate_peaks %>%
  mutate(comp = factor(comp, levels = plot_order))

candidate_peak_size_plot <- ggplot(
  candidate_peaks_for_plot,
  aes(comp, n_outlier_windows, group = comp)) +
  
  geom_violin(width = 1, 
              color = "#797979", # violin plot outline color
              fill = "#797979",  # violin plot fill color
              alpha = 0.25) +  # violin plot fill transparency
  
  geom_jitter(aes(size = n_outlier_windows), 
              color = "black", # individual point outline color
              fill = "#797979",  # individual point fill color
              alpha = 0.9, # individual point fill transparency
              width = 0.25, # width of point jitter zone
              shape = 21) + # shape of point (21 = filled circle)
 
  scale_size(range = c(1.5, 3), # scale point size to reduce crowding at low values
             guide = "none") +
  scale_x_discrete(drop = FALSE, labels = function(x) gsub(".v.", "\nvs\n", x, fixed = TRUE)) +
  scale_y_continuous(breaks = scales::breaks_pretty(n = 4)) +
  labs(x = "Population comparison", y = "Outlier windows per candidate peak") +
  theme_classic() +
  transparent_theme +
  theme(axis.text.x = element_text(angle = 0, hjust = 0.5, lineheight = 0.9),
      panel.border = element_rect(color = "black", fill = NA, linewidth = 0.5))

ggsave(
  file.path(plot_dir, "candidate_peak_size_distributions.png"),
  candidate_peak_size_plot,
  width = min(5*nrow(comparison_map),60), # edit final figure width here
  height = 15, # edit final figure height here
  units = "cm", dpi = 600,
  bg = "transparent"
)

sessionInfo()