library(knitr)
library(kableExtra)
library(dplyr)
library(tidyr)
library(rlang)
library(stringr)

ellipse_table <- function(df, nr = 5, nc = 5, format = "latex") {
    nrows <- nrow(df)
    ncols <- ncol(df)

    # Ellipsis symbols depending on format
    if (format == "latex") {
        h_ellipsis <- "$\\dots$" # horizontal (for columns)
        v_ellipsis <- "$\\vdots$" # vertical (for rows)
    } else {
        h_ellipsis <- "..."
        v_ellipsis <- "..."
    }

    # Select top/bottom rows
    row_idx <- c(
        seq_len(ceiling(nr / 2)),
        nrows - floor(nr / 2) + 1:nr - nr
    )
    row_idx <- unique(pmax(pmin(row_idx, nrows), 1)) # clamp

    # Select left/right cols
    col_idx <- c(
        seq_len(ceiling(nc / 2)),
        ncols - floor(nc / 2) + 1:nc - nc
    )
    col_idx <- unique(pmax(pmin(col_idx, ncols), 1))

    # Subset
    sub_df <- df[row_idx, col_idx, drop = FALSE]

    # Insert ellipsis row if needed
    if (nrows > nr) {
        ellipsis_row <- as.list(rep(h_ellipsis, length(col_idx)))
        sub_df <- rbind(
            sub_df[1:ceiling(nr / 2), , drop = FALSE],
            ellipsis_row,
            sub_df[(ceiling(nr / 2) + 1):nrow(sub_df), , drop = FALSE]
        )
        rownames(sub_df)[ceiling(nr / 2) + 1] <- v_ellipsis
    }

    # Insert ellipsis col if needed
    if (ncols > nc) {
        ellipsis_col <- rep(v_ellipsis, nrow(sub_df))
        sub_df <- cbind(
            sub_df[, 1:ceiling(nc / 2), drop = FALSE],
            "..." = ellipsis_col,
            sub_df[, (ceiling(nc / 2) + 1):ncol(sub_df), drop = FALSE]
        )
        names(sub_df)[ceiling(nc / 2) + 1] <- h_ellipsis
    }

    # Print with kable
    return(sub_df)
}

collapse_otu_at_taxo <- function(phylo_data, NArm = TRUE, renameOTUs = TRUE) {
    tax_ranks <- colnames(tax_table(phylo_data))
    tax_df <- tax_table(phylo_data)
    otu_tables_by_rank <- lapply(tax_ranks, function(rank) {
        ps_glom <- tax_glom(the_data, taxrank = rank, NArm = NArm)
        otu <- as.data.frame(otu_table(ps_glom))

        # Ensure taxa are rows
        if (!taxa_are_rows(ps_glom)) {
            otu <- t(otu)
        }
        otu_matrix <- as.matrix(otu)
        if (renameOTUs) {
            id_in_taxo <- match(rownames(otu_matrix), rownames(tax_df))
            taxo_names <- tax_df[id_in_taxo, seq(which(tax_ranks == rank))] %>%
                as.data.frame() %>%
                mutate_if(.predicate = rlang::is_character, .funs = purrr::partial(str_replace_all, pattern = "-", replacement = ".")) %>%
                tidyr::unite(col = "tax_full", sep = ";_;") %>%
                unique()
            rownames(otu_matrix) <- as.vector(taxo_names[, 1])
        }
        return(otu_matrix)
    })
    names(otu_tables_by_rank) <- tax_ranks
    return(otu_tables_by_rank)
}


propagate_taus <- function(
  tau_matrix,
  physeq,
  taxrank = phyloseq::rank_names(physeq)[1],
  NArm = TRUE,
  bad_empty = c(NA, "", " ", "\t")
) {
    if (is.null(phyloseq::access(physeq, "tax_table"))) {
        stop("The tax_glom() function requires that physeq contain a taxonomyTable")
    }
    if (!taxrank[1] %in% phyloseq::rank_names(physeq)) {
        stop("Bad taxrank argument. Must be among the values of rank_names(physeq)")
    }
    CN <- which(phyloseq::rank_names(physeq) %in% taxrank[1])
    tax <- as(phyloseq::access(physeq, "tax_table"), "matrix")[, CN]
    if (NArm) {
        keep_species <- names(tax)[!(tax %in% bad_empty)]
        physeq <- phyloseq::prune_taxa(keep_species, physeq)
    }
    tax <- as(phyloseq::access(physeq, "tax_table"), "matrix")[, 1:CN,
        drop = FALSE
    ]
    tax <- apply(tax, 1, function(i) {
        paste(i, sep = ";_;", collapse = ";_;")
    })
    tax_names <- unique(unname(tax))
    spCliques <- tapply(names(tax), factor(tax), list)

    old_tax_names <- rownames(tau_matrix)

    nb_desc <- sapply(old_tax_names, function(name) length(grep(name, tax_names)))
    cumsum_nb_desc <- cumsum(nb_desc)

    nb_desc_tot <- sum(nb_desc)
    new_tau <- matrix(0, nrow = nb_desc_tot, ncol = ncol(tau_matrix))

    old_and_new <- lapply(old_tax_names, function(name) tax_names[grep(name, tax_names)])
    names(old_and_new) <- old_tax_names
    for (idx in seq_along(old_and_new)) {
        old_taxa <- names(old_and_new)[[idx]]
        if (idx == 1) {
            old_idx <- 1
        } else {
            old_idx <- idx - 1
        }
        new_tau[cumsum_nb_desc[old_idx]:cumsum_nb_desc[idx], ] <- tau_matrix[old_taxa, ]
    }
    rownames(new_tau) <- tax_names
    return(new_tau)
}