# R/identify_critical_links.R
# Fast, mathematically exact implementation.

#' Identify Critical Edges That Explain Population Differences
#'
#' Iteratively identifies the edges that drive the observed difference between
#' two or more populations of brain networks, using the permutation-based
#' L1-distance ANOVA test of Fraiman and Fraiman (2018). Edges are ranked by
#' their marginal p-value, then removed in batches until the global test is no
#' longer significant. The permutation null is re-evaluated incrementally via
#' a prefix-sum decomposition of T, which avoids recomputing the statistic
#' from scratch at every step.
#'
#' @param populations A named \code{list} with at least two elements. Each
#'   element is itself a list of square, binary, symmetric adjacency matrices
#'   (zero diagonal, no self-loops). All graphs must share the same number of
#'   nodes.
#' @param alpha Numeric significance level used both for the global test and
#'   for the iterative edge-removal stopping rule. Default \code{0.05}.
#' @param method Character string selecting the marginal edge test used to
#'   rank candidate edges. One of \code{"fisher"} (default),
#'   \code{"chi.squared"}, or \code{"prop"}. See
#'   \code{\link{compute_edge_pvalues}}.
#' @param adjust_method Character string passed to \code{\link[stats]{p.adjust}}
#'   for multiplicity correction of the marginal p-values. Default
#'   \code{"none"}.
#' @param batch_size Positive integer giving the number of edges to drop per
#'   iteration. Larger values speed up the procedure at the cost of coarser
#'   resolution. Default \code{1}.
#' @param n_permutations Positive integer giving the number of permutation
#'   replicates used to approximate the null distribution of T. Default
#'   \code{1000}.
#' @param a Positive numeric normalisation constant for the test statistic T
#'   (see \code{\link{compute_test_statistic}}). Default \code{1}.
#' @param seed Optional integer seed for reproducibility of the permutation
#'   replicates. Default \code{NULL}.
#'
#' @return A \code{list} with three components: \code{critical_edges} (a
#'   \code{data.frame} with columns \code{node1}, \code{node2}, \code{p_value},
#'   listing the edges removed until the test became non-significant, ordered
#'   from most to least significant; \code{NULL} if the global test was not
#'   significant), \code{edges_removed} (a list of length-2 integer vectors
#'   giving the \code{(i, j)} indices of the removed edges, in removal order),
#'   and \code{modified_populations} (the input \code{populations} with the
#'   critical edges zeroed out in every graph).
#'
#' @details The implementation exploits the fact that T decomposes as a sum of
#'   per-edge contributions \eqn{\Delta_e}. Removing an edge e is equivalent
#'   to subtracting \eqn{\Delta_e} from T, so prefix sums of the sorted
#'   \eqn{\Delta_e} values give the test statistic after any number of
#'   removals in O(1). The permutation null is built once via a single BLAS
#'   matrix multiplication and reused for every iteration, reducing the
#'   overall complexity from O(K * B * |E| * m) (naive) to O(B * |E| * m).
#'
#'   Permutation replicates are generated by uniformly permuting the group
#'   labels across the pooled sample of all graphs, while preserving the
#'   original group sizes (sampling without replacement).
#'
#' @references
#' Fraiman, D. and Fraiman, R. (2018) An ANOVA approach for statistical
#' comparisons of brain networks. \emph{Scientific Reports}, 8, 4746.
#' \doi{10.1038/s41598-018-21688-0}.
#'
#' @seealso \code{\link{compute_test_statistic}},
#'   \code{\link{compute_edge_pvalues}}, \code{\link{get_critical_nodes}}.
#'
#' @export
#' @examples
#' \donttest{
#' set.seed(1)
#' control <- generate_category_graphs(
#'   n_graphs = 15, n_nodes = 10, n_communities = 2,
#'   base_intra_prob = 0.8, base_inter_prob = 0.2, seed = 1)
#' patient <- generate_category_graphs(
#'   n_graphs = 15, n_nodes = 10, n_communities = 2,
#'   base_intra_prob = 0.5, base_inter_prob = 0.5, seed = 2)
#' populations <- list(Control = control, Patient = patient)
#' result <- identify_critical_links(
#'   populations, alpha = 0.05, method = "fisher",
#'   n_permutations = 200, seed = 42)
#' head(result$critical_edges)
#' }
identify_critical_links <- function(populations,
                                    alpha          = 0.05,
                                    method         = "fisher",
                                    adjust_method  = "none",
                                    batch_size     = 1,
                                    n_permutations = 1000,
                                    a              = 1,
                                    seed           = NULL) {

  if (!is.null(seed)) set.seed(seed)
  if (!is.list(populations) || length(populations) < 2)
    stop("`populations` must be a list with at least 2 groups.")
  if (batch_size  < 1L) stop("`batch_size` must be >= 1.")
  if (n_permutations < 1L) stop("`n_permutations` must be >= 1.")

  Npop <- sapply(populations, length)
  m    <- length(populations)
  n    <- sum(Npop)

  make_key <- function(i, j) paste(i, j, sep = "-")

  freq       <- compute_edge_frequencies(populations)
  edge_pvals <- compute_edge_pvalues(freq$edge_counts, Npop,
                                     method        = method,
                                     adjust_method = adjust_method)
  edge_df    <- rank_edges(edge_pvals)
  n_edges    <- nrow(edge_df)
  if (n_edges == 0L)
    return(list(critical_edges       = NULL,
                edges_removed        = list(),
                modified_populations = populations))

  .edge_deltas <- function(edge_counts) {
    n_nodes <- dim(edge_counts)[1]
    idx     <- which(upper.tri(matrix(0, n_nodes, n_nodes)), arr.ind = TRUE)
    if (nrow(idx) == 0L)
      return(list(indices = idx, deltas = numeric(0)))

    counts <- do.call(cbind, lapply(seq_len(m),
                     function(k) edge_counts[ , , k][idx]))
    if (!is.matrix(counts))
      counts <- matrix(counts, nrow = 1L)

    p_mat <- sweep(counts, 2, Npop, "/")
    p_tot <- rowSums(counts) / n
    Ptot  <- matrix(p_tot, nrow = nrow(counts), ncol = m)

    d_mat <- 2 * p_mat * (1 - p_mat)
    D_mat <- p_mat + Ptot - 2 * p_mat * Ptot

    coef_d <- sqrt(Npop) * (Npop / (Npop - 1))
    coef_D <- sqrt(Npop) * (n    / (n    - 1))

    delta  <- (sqrt(m) / a) *
              (d_mat %*% coef_d - D_mat %*% coef_D)[, 1]

    list(indices = idx, deltas = delta)
  }

  .edge_deltas_matrix <- function(counts_list) {
    E   <- nrow(counts_list[[1]])
    B   <- ncol(counts_list[[1]])

    p_list  <- lapply(seq_len(m), function(k) counts_list[[k]] / Npop[k])
    C_total <- Reduce("+", counts_list)
    p_tot   <- C_total / n

    coef_d <- sqrt(Npop) * (Npop / (Npop - 1))
    coef_D <- sqrt(Npop) * (n    / (n    - 1))

    delta_mat <- matrix(0, nrow = E, ncol = B)
    for (k in seq_len(m)) {
      pk  <- p_list[[k]]
      d_k <- 2 * pk * (1 - pk)
      D_k <- pk + p_tot - 2 * pk * p_tot
      delta_mat <- delta_mat + coef_d[k] * d_k - coef_D[k] * D_k
    }
    (sqrt(m) / a) * delta_mat
  }

  obs_delta_info <- .edge_deltas(freq$edge_counts)

  key_all   <- make_key(obs_delta_info$indices[, 1], obs_delta_info$indices[, 2])
  map_idx   <- match(make_key(edge_df$node1, edge_df$node2), key_all)
  delta_ord <- obs_delta_info$deltas[map_idx]

  prefix_obs <- cumsum(delta_ord)
  T0         <- sum(obs_delta_info$deltas)

  all_graphs <- unlist(populations, recursive = FALSE)
  ut_idx     <- obs_delta_info$indices
  X <- do.call(rbind, lapply(all_graphs, function(A) A[ut_idx]))
  storage.mode(X) <- "double"

  perm_mat <- replicate(n_permutations, sample.int(n))
  cumNpop  <- c(0L, cumsum(Npop))
  W_list <- lapply(seq_len(m), function(k) {
    rows    <- (cumNpop[k] + 1L):cumNpop[k + 1L]
    idx_all <- as.vector(perm_mat[rows, , drop = FALSE])
    rep_id  <- rep(seq_len(n_permutations), each = Npop[k])
    lin_idx <- (rep_id - 1L) * n + idx_all
    W <- matrix(tabulate(lin_idx, nbins = n * n_permutations),
                nrow = n, ncol = n_permutations)
    storage.mode(W) <- "double"
    W
  })

  C_list      <- lapply(W_list, function(W) crossprod(X, W))
  delta_all   <- .edge_deltas_matrix(C_list)
  perm_deltas <- t(delta_all[map_idx, , drop = FALSE])

  T_perm0     <- rowSums(perm_deltas)
  prefix_perm <- t(apply(perm_deltas, 1, cumsum))

  initial_p <- mean(T_perm0 < T0)
  if (initial_p > alpha) {
    warning("Initial test is not significant (p = ",
            round(initial_p, 4),
            "). Populations may be identical or differences too small to detect.")
    return(list(critical_edges       = NULL,
                edges_removed        = list(),
                modified_populations = populations))
  }

  batch_steps <- seq(batch_size, n_edges, by = batch_size)
  if (length(batch_steps) == 0L || batch_steps[length(batch_steps)] != n_edges)
    batch_steps <- c(batch_steps, n_edges)

  delta_vec  <- T_perm0 - T0
  R_at_steps <- sweep(prefix_perm[, batch_steps, drop = FALSE],
                      2, prefix_obs[batch_steps])
  p_at_steps <- colMeans(R_at_steps > delta_vec)

  nonsig_idx <- which(p_at_steps > alpha)[1]
  k_removed  <- if (is.na(nonsig_idx)) n_edges else batch_steps[nonsig_idx]

  edges_removed <- vector("list", k_removed)
  for (idx in seq_len(k_removed)) {
    i <- edge_df$node1[idx]; j <- edge_df$node2[idx]
    for (pop in seq_along(populations))
      for (g in seq_len(Npop[pop])) {
        populations[[pop]][[g]][i, j] <- 0
        populations[[pop]][[g]][j, i] <- 0
      }
    edges_removed[[idx]] <- c(i, j)
  }

  critical_edges <- if (k_removed > 0) edge_df[seq_len(k_removed), ] else NULL

  list(critical_edges       = critical_edges,
       edges_removed        = edges_removed,
       modified_populations = populations)
}
