library(ggplot2)

list_to_array <- function(lst) {
  array(
    do.call(c, lst),
    dim = c(dim(lst[[1]]), length(lst)),
    dimnames = c(dimnames(lst[[1]]), list(names(lst)))
  )
}

#' Apply a filter to a trait table
#'
#' Removes rows corresponding to species that have been discarded in a
#' given filter
#'
#' @param filter The name of the filter as a string.
#' @param trait_table A table containing trait information, with a
#'   column "iucn2020_binomial" with species names that have to match
#'   the ones provided in `filter_table`.
#' @param filter_table A table indicating if a species has been
#'   filtered (0/1).  One of its columns has to be named as indicated
#'   in the `filter` parameter.
#' @return A table derived from `trait_table` without the species
#'   discarded by the filter
filter_trait_table <- function(filter, trait_table, filter_table) {
  sel_filt <- as.logical(filter_table[[filter]])
  sp_list_filt <- filter_table$Spp_names[sel_filt]
  trait_sp_names <- trait_table$iucn2020_binomial
  trait_filt_sel <- trait_sp_names %in% sp_list_filt
  trait_table[trait_filt_sel, ]
}

category_count_table <- function(trait_data, cat1, cat2) {
  table(trait_data[[cat1]], trait_data[[cat2]])
}

#' Nested list of category counts for each filter (seed, percentage,
#' whole family/taxonomy weighted)
#'
#' rapply() does not work here because data frames are also lists.  But
#' these are too many nested lapplys
category_count_tables <- function(filtered_trait_data, cat1, cat2) {
  lapply(filtered_trait_data, function(filter_seed) {
    lapply(filter_seed, function(filter_perc) {
      lapply(filter_perc, function(filtered_table) {
        category_count_table(filtered_table, cat1, cat2)
      })
    })
  })
}

#' A function to collect tables from all seeds into 3D arrays for each
#' filtering percentage and mode of taxonomic filtering
collect_tables <- function(tbl_list) {
  # Create empty list with the same nested structure as each of the
  # seeds
  collected_tbl <- rapply(
    tbl_list[[1]], function(x) NULL, how = "replace"
  )

  for (perc in names(collected_tbl)) {
    for (tax_method in names(collected_tbl[[perc]])) {
      coll_tbl_list <- lapply(tbl_list, function(seed) {
        seed[[perc]][[tax_method]]
      })
      coll_arr <- list_to_array(coll_tbl_list)
      collected_tbl[[perc]][[tax_method]] <- coll_arr
    }
  }

  collected_tbl
}

plot_heatmap <- function(mat, font_size = 11, x_lab_rot = 0) {
  mat_df <- as.data.frame(as.table(mat))
  names(mat_df) <- c("size_cat", "diet_cat", "quantity")
  ggplot(mat_df) +
    geom_tile(aes(diet_cat, size_cat, fill = quantity)) +
    geom_label(
      data = subset(mat_df, !is.nan(quantity)),
      aes(diet_cat, size_cat, label = as.character(round(quantity, 2))),
      size = font_size, size.unit = "pt"
    ) +
    xlab("Diet category") +
    ylab("Size category") +
    guides(fill = "none") +
    theme_gray(base_size = font_size) +
    theme(axis.text.x = element_text(angle = x_lab_rot))
}

#' Generate table with information on taxonomic composition for each filter
#'
#' @param preserv_rate Preservation rate as string ("25", "50", or
#'   "75").
#' @param filters An object with a table containing information on
#'   which species have been kept after each filter (multiple
#'   replicates, and a separate table for each retention percentage
#'   (25, 50, 75)).
#' @param trait_table A table with taxonomic information on orders.
#' @return A data frame with species counts, percentages of retention,
#'   and 5% and 95% quantiles for these, for each order and filter, as
#'   well as the proportion of the total species represented by each
#'   order after each filter.
taxonomy_by_filter <- function(preserv_rate, filters, trait_table) {
  sequential_filter <- lapply(filters, function(filter_tables) {
    filter_table <- filter_tables[[preserv_rate]]
    filter_steps <- names(filter_table)[4:10]
    filtered <- lapply(filter_steps, function(step) {
      filter_trait_table(step, trait_table, filter_table)
    })
    names(filtered) <- filter_steps
    filtered
  })

  seq_filter_count <- lapply(sequential_filter, function(seed) {
    sapply(seed, function(x) table(x$order_cat))
  })
  seq_filter_perc <- lapply(seq_filter_count, function(tbl) tbl / tbl[, 1])
  seq_filter_count_arr <- list_to_array(seq_filter_count)
  seq_filter_perc_arr <- list_to_array(seq_filter_perc)
  seq_filter_count_mean <- apply(seq_filter_count_arr, c(1, 2), mean)
  seq_filter_count_95 <- apply(seq_filter_count_arr, c(1, 2), quantile, 0.95)
  seq_filter_count_05 <- apply(seq_filter_count_arr, c(1, 2), quantile, 0.05)
  seq_filter_perc_mean <- apply(seq_filter_perc_arr, c(1, 2), mean)
  seq_filter_perc_95 <- apply(seq_filter_perc_arr, c(1, 2), quantile, 0.95)
  seq_filter_perc_05 <- apply(seq_filter_perc_arr, c(1, 2), quantile, 0.05)

  seq_filter_df <- data.frame(
    expand.grid(dimnames(seq_filter_perc_mean)),
    count_mean = c(seq_filter_count_mean),
    count_q05 = c(seq_filter_count_05),
    count_q95 = c(seq_filter_count_95),
    perc_mean = c(seq_filter_perc_mean),
    perc_q05 = c(seq_filter_perc_05),
    perc_q95 = c(seq_filter_perc_95)
  )
  names(seq_filter_df)[1:2] <- c("order", "filter")
  seq_filter_df$filt_total_perc <- tapply(
    seq_filter_df$count_mean,
    seq_filter_df$filter,
    function(x) {
      x / sum(x)
    }
  ) |> unlist()
  seq_filter_df
}