# =======1.calculate========

#' Calculate correlation for one or two t(otutab), or distance for one t(otutab).

#'

#' @param totu t(otutab), row are samples, column are features.

#' @param totu2 t(otutab2) or NULL, row are samples, column are features.

#' @param method "spearman" (default), "pearson", "sparcc", or distance index from \code{\link[vegan]{vegdist}}.

#' @param filename the prefix of saved .corr file or FALSE.

#' @param p.adjust.method see \code{\link[stats]{p.adjust}}

#' @param p.adjust.mode see \code{\link{p.adjust.table}}

#' @param threads threads, default: 1.

#' @param verbose verbose, default: TRUE.

#'

#' @return a corr object with 3 elements:

#' \item{r}{default: spearman correlation}

#' \item{p.value}{default: p-value of spearman correlation}

#' \item{p.adjust}{default p.adjust.method = NULL}

#' @family calculate

#' @export

#' @aliases c_net_cal

#' @examples

#' data("otutab", package = "pcutils")

#' t(otutab) -> totu

#' c_net_calculate(totu) -> corr

#' metadata[, 3:10] -> env

#' c_net_calculate(totu, env) -> corr2

c_net_calculate <- function(totu, totu2 = NULL, method = "spearman", filename = FALSE,

                            p.adjust.method = NULL, p.adjust.mode = "all", threads = 1, verbose = TRUE) {

  if (!is.null(totu2)) {

    tls <- check_tabs(totu, totu2)

    totu <- tls[[1]]

    totu2 <- tls[[2]]

  } else {

    totu <- check_tabs(totu)

  }

  if (method %in% c("spearman", "pearson")) {

    corr <- fast_cor(totu, totu2, method = method)

  } else if (method == "sparcc") {

    if (!is.null(totu2)) warning("sparcc only take the totu, ignore the totu2.")

    message("sparcc is not supported on CRAN version. Please install `zdk123/SpiecEasi` from github and use the function `sparcc`.")

    # corr <- par_sparcc(totu, threads = threads, verbose = verbose)

    corr <- NULL

  } else if (method %in% c(

    "manhattan", "euclidean", "canberra", "bray",

    "kulczynski", "gower", "morisita", "horn", "mountford",

    "jaccard", "raup", "binomial", "chao", "altGower", "cao",

    "mahalanobis", "clark", "chisq", "chord", "hellinger",

    "aitchison", "robust.aitchison"

  )) {

    corr <- cal_sim(totu, totu2, method = method)

  } else {

    stop("method should be one of 'spearman', 'pearson', 'sparcc', or distance index from vegan.")

  }

  if (!is.null(p.adjust.method)) {

    p.adjust <- p.adjust.table(corr$p.value, p.adjust.method)

    res <- list(r = corr$r, p.value = corr$p.value, p.adjust = p.adjust)

  } else {

    res <- list(r = corr$r, p.value = corr$p.value)

  }

  class(res) <- "corr"

  attributes(res)$method <- method

  # save the correlation result

  if (is.logical(filename)) {

    if (filename) filename <- paste0("c_net_", date())

  }

  if (is.character(filename)) {

    save_corr(res, filename)

  }

  return(res)

}

#' Check tables and extract common samples

#'

#' @param ... tables

#' @return formatted tables

#' @export

#' @examples

#' data("otutab", package = "pcutils")

#' check_tabs(otutab)

check_tabs <- function(...) {

  tables <- list(...)

  if (all(class(tables[[1]]) == "list")) tables <- tables[[1]]

  names(tables) <- NULL

  if (length(tables) > 1) {

    comm <- Reduce(intersect, lapply(tables, rownames))

    if (length(comm) < 2) stop("There are ", length(comm), " common sample! Can not calculate correlation.")

    if (all(lapply(tables, \(i)identical(rownames(i), comm)) %>% unlist())) {

      message("All samples matched.")

    } else {

      message("Extract ", length(comm), " commmon samples.")

    }

    dup <- lapply(tables, colnames) %>% do.call(c, .)

    dup <- dup[duplicated(dup)]

    if (length(dup) > 0) {

      stop("Duplicated colnames found: ", paste0(dup, collapse = ", "), "\nPlease check colnames of input tables.")

    } else {

      message("All features are OK.")

    }

    tables <- lapply(tables, \(i)i[comm, ])

  }

  if (length(tables) == 1) tables <- tables[[1]]

  return(tables)

}

#' Save a corr object

#'

#' @param corr a corr object

#' @param filename filename without extension, default: "corr"

#'

#' @return a .corr file

#' @export

#'

save_corr <- function(corr, filename = "corr") {

  stopifnot(inherits(corr, "corr"))

  if (!grepl("\\.corr$", filename)) filename <- paste0(filename, ".corr")

  if (t_flag(corr$r)) {

    # 节约一半的储存空间

    corr$r[upper.tri(corr$r)] -> r

    corr$p.value[upper.tri(corr$p.value)] -> p.value

    if (!is.null(corr$p.adjust)) {

      corr$p.adjust[upper.tri(corr$p.adjust)] -> p.adj

      if (all(p.value == p.adj)) {

        p.adj <- FALSE

      }

    } else {

      p.adj <- NULL

    }

    corr_names <- rownames(corr$r)

    saveRDS(

      list(

        r = r, p.value = p.value, p.adjust = p.adj,

        corr_names = corr_names, corr_attr = attributes(corr)

      ),

      file = filename

    )

  } else {

    saveRDS(corr, file = filename)

  }

}

#' Read a corr object

#'

#' @param filename filename of .corr

#'

#' @return a corr object

#' @export

#' @family calculate

read_corr <- function(filename) {

  # r <- read.csv(paste0(filename, "_r.csv"), row.names = 1, check.names = FALSE)

  # p.value <- read.csv(paste0(filename, "_p.csv"), row.names = 1, check.names = FALSE)

  # p.adjust <- read.csv(paste0(filename, "_p_adj.csv"), row.names = 1, check.names = FALSE)

  if (!grepl("\\.corr$", filename)) filename <- paste0(filename, ".corr")

  in_corr <- readRDS(filename)

  if ("corr_names" %in% names(in_corr)) {

    corr_names <- in_corr$corr_names

    r <- in_corr$r

    p.value <- in_corr$p.value

    p.adj <- in_corr$p.adj

    corr_attr <- in_corr$corr_attr

    new_p.adjust <- new_p <- new_r <- matrix(0, nrow = length(corr_names), ncol = length(corr_names), dimnames = list(corr_names, corr_names))

    new_r[upper.tri(new_r)] <- r

    new_r <- new_r + t(new_r)

    diag(new_r) <- 1

    new_p[upper.tri(new_p)] <- p.value

    new_p <- new_p + t(new_p)

    if (is.null(p.adj)) {

      new_p.adjust <- NULL

    } else if (is.logical(p.adj)) {

      new_p.adjust <- new_p

    } else {

      new_p.adjust[upper.tri(new_p.adjust)] <- p.adj

      new_p.adjust <- new_p.adjust + t(new_p.adjust)

      diag(new_p.adjust) <- 1

    }

    in_corr <- list()

    in_corr$r <- new_r

    in_corr$p.value <- new_p

    if (!is.null(new_p.adjust)) in_corr$p.adjust <- new_p.adjust

    attributes(in_corr) <- corr_attr

  }

  return(in_corr)

}

#' Fast correlation calculation

#'

#' @param totu t(otutab), row are samples, column are features.

#' @param totu2 t(otutab) or NULL, row are samples, column are features.

#' @param method "spearman" or "pearson"

#'

#' @export

#'

#' @return a list with 2 elements:

#' \item{r}{default: spearman correlation}

#' \item{p.value}{default: p-value of spearman correlation}

#' @family calculate

#' @examples

#' data("otutab", package = "pcutils")

#' t(otutab[1:100, ]) -> totu

#' fast_cor(totu, method = "spearman") -> corr

fast_cor <- function(totu, totu2 = NULL, method = c("pearson", "spearman")) {

  method <- match.arg(method, c("pearson", "spearman"))

  totu <- as.matrix(totu)

  if (!is.null(totu2)) {

    totu2 <- as.matrix(totu2)

    r <- stats::cor(totu, totu2, method = method)

    df <- dim(totu)[1] - 2

    t <- r * sqrt(df / (1 - r^2))

    p <- -2 * expm1(stats::pt(abs(t), df, log.p = TRUE))

  } else {

    r <- stats::cor(totu, method = method)

    p <- r

    p[p != 0] <- 0

    r_tri <- r[upper.tri(r)]

    df <- dim(totu)[1] - 2

    t <- r_tri * sqrt(df / (1 - r_tri^2))

    p[upper.tri(p)] <- -2 * expm1(stats::pt(abs(t), df, log.p = TRUE))

    p <- p + t(p)

  }

  return(list(r = r, p.value = p))

}

Hmisc_cor <- function(totu, totu2 = NULL, method = c("spearman", "pearson")[1]) {

  lib_ps("Hmisc", library = FALSE)

  totu <- as.matrix(totu)

  if (!is.null(totu2)) {

    totu2 <- as.matrix(totu2)

    tmp <- Hmisc::rcorr(totu, totu2, type = method)

    r <- tmp$r[1:ncol(totu), (ncol(totu) + 1):ncol(tmp$r)]

    p <- tmp$P[1:ncol(totu), (ncol(totu) + 1):ncol(tmp$r)]

    return(list(r = r, p.value = p))

  }

  tmp <- Hmisc::rcorr(totu, type = method)

  p <- tmp$P

  p[is.na(p)] <- 0

  return(list(r = tmp$r, p.value = p))

}

#' Calculate similarity for one t(otutab)

#'

#' @param totu t(otutab), row are samples, column are features.

#' @param method Dissimilarity index, see \code{\link[vegan]{vegdist}}.

#' @param totu2 t(otutab) or NULL, row are samples, column are features.

#'

#' @family calculate

#' @return similarity = 1-distance

#' @export

#' @seealso \code{\link[vegan]{vegdist}}

#' @examples

#' if (requireNamespace("vegan")) {

#'   data("otutab", package = "pcutils")

#'   t(otutab) -> totu

#'   cal_sim(totu) -> sim_corr

#' }

cal_sim <- function(totu, totu2 = NULL, method = "bray") {

  lib_ps("vegan", library = FALSE)

  if (is.null(totu2)) {

    vegan::vegdist(t(totu), method = method) %>% as.matrix() -> dist

  } else {

    n1 <- ncol(totu)

    n2 <- ncol(totu2)

    dist <- matrix(NA, nrow = n1, ncol = n2, dimnames = list(colnames(totu), colnames(totu2)))

    for (i in seq_len(n1)) {

      for (j in seq_len(n2)) {

        dist[i, j] <- vegan::vegdist(rbind(totu[, i], totu2[, j]), method = method)

      }

    }

  }

  sim <- 1 - dist

  p <- dist

  message("p-value is not supported for distance index, all set as 0.")

  p[p != 0] <- 0

  return(list(r = sim, p.value = p))

}

cal_KLD <- function(totu, totu2 = NULL, method = "KLD") {

  # Kullback-Leibler divergence

  # KLD = function(p, q) {

  #   sum(p * log(p / q))

  # }

  # p = c(0.1, 0.2, 0.3, 0.4)

  # q = c(0.2, 0.3, 0.2, 0.3)

  # KLD(p, q)

  lib_ps("philentropy", library = FALSE)

  dat <- t(totu) / rowSums(t(totu))

  philentropy::KL(dat, unit = "log")

}

#' p.adjust apply on a correlation table (matrix or data.frame)

#'

#' @param pp table of p-values

#' @param method see \code{\link[stats]{p.adjust}}, default: "BH".

#' @param mode "all" for all values; "rows" adjust each row one by one; "columns" adjust each column one by one. Default: "all".

#'

#' @return a table of adjusted p-values

#' @export

#'

#' @family calculate

#' @examples

#' matrix(abs(rnorm(100, 0.01, 0.1)), 10, 10) -> pp

#' p.adjust.table(pp, method = "BH", mode = "all") -> pp_adj

p.adjust.table <- \(pp, method = "BH", mode = "all"){

  mode <- match.arg(mode, c("all", "rows", "columns"))

  pp <- as.matrix(pp)

  if (mode == "all") {

    if (t_flag(pp)) {

      lp <- lower.tri(pp)

      pa <- pp[lp]

      pa <- p.adjust(pa, method)

      pp[lower.tri(pp, diag = FALSE)] <- pa

      pp[upper.tri(pp, diag = FALSE)] <- 0

      pp <- pp + t(pp)

    } else {

      pp1 <- p.adjust(pp, method)

      pp <- matrix(pp1, nrow(pp), ncol(pp), dimnames = list(rownames(pp), colnames(pp)))

    }

  } else if (mode == "rows") {

    pp <- t(apply(pp, 1, p.adjust, method = method))

  } else if (mode == "columns") {

    pp <- apply(t(pp), 1, p.adjust, method = method)

  }

  return(pp)

}