# ============================================================
# Algorithm 3 -- Subsampling comparison between unmatchable
#                networks with unequal edge densities.
# ============================================================


#' Two-sample test for unmatchable networks with unequal densities (Algorithm 3)
#'
#' Tests whether two unmatchable networks --- differing in node set, size, and
#' possibly edge density --- arise from the same underlying graphon model.
#' The procedure avoids the density-mismatch problem via a
#' \emph{split-and-sparsify} strategy: each network is randomly split into two
#' halves, one used to estimate the sparsification probability and the other to
#' compute density-normalised moments at a shared target density.  Node
#' subsampling from the larger network then generates a reference distribution
#' against which the observed statistic is compared.
#'
#' @section Algorithm:
#' \enumerate{
#'   \item Randomly partition the nodes of \code{G_prime} into two equal
#'     halves \eqn{G'_1} and \eqn{G'_2}; compute the observed statistic
#'     \eqn{\bar\Psi_{\rho^\dagger}(G'_1, G'_2)} via Algorithm 2.
#'   \item For \eqn{i = 1, \ldots, N_{\mathrm{sub}}}: independently draw
#'     two sets of \eqn{\lfloor b/2 \rfloor} nodes from \eqn{G} and compute
#'     \eqn{\bar\Psi_{\rho^\dagger}(G^*_{i1}, G^*_{i2})} via Algorithm 2.
#'   \item Compare the observed statistic against the empirical null
#'     distribution using the selected \code{test}.
#' }
#'
#' @section Choice of \code{rho_target}:
#' A recommended default is
#' \eqn{\rho^\dagger = \kappa \cdot \min(\hat\rho_G, \hat\rho_{G'})},
#' with \eqn{\kappa \in (0.5,\, 0.9)}.  Values close to 0 discard too many
#' edges and inflate variance; values too close to
#' \eqn{\min(\hat\rho_G, \hat\rho_{G'})} leave little room for edge removal
#' and can break the density-matching property.  A value of \eqn{\kappa = 0.7}
#' (corresponding to
#' \code{rho_target = 0.7 * min(rho_hat_G, rho_hat_G_prime)}) works well
#' across a range of settings.
#'
#' @section Tests:
#' \describe{
#'   \item{Mahalanobis (default)}{Computes a multivariate distance that fully
#'     exploits the joint distribution of all moments.  The p-value is the
#'     proportion of null distances exceeding the observed distance.  This test
#'     is more powerful when the signal lies in the \emph{joint} rather than
#'     the \emph{marginal} distribution of the moments.}
#'   \item{Cauchy combination}{Combines marginal p-values via the Cauchy
#'     combination test of Liu & Xie (2020).  Controls type I error under
#'     arbitrary dependence but ignores cross-moment correlation.}
#' }
#'
#' @param G A square symmetric binary adjacency matrix (dense or sparse
#'   \code{Matrix}) of size \eqn{n \times n}.  Typically the larger or
#'   reference network.
#' @param G_prime A square symmetric binary adjacency matrix of size
#'   \eqn{b \times b}.  The network to compare against \code{G}.
#' @param rho_target Numeric in \eqn{(0, 1)}; the common target edge density
#'   \eqn{\rho^\dagger}.  See the \emph{Choice of \code{rho_target}} section.
#' @param N_sub Integer; number of subsampling replicates.  Default 2000.
#' @param motifs A list of motif specifications; see
#'   \code{\link{subsample_moments}} for the format.  Defaults to
#'   \code{list("v_shape", "triangle", "star3")}.
#' @param test Character; the multivariate test to apply.  Either
#'   \code{"mahalanobis"} (default) or \code{"cauchy"}.
#' @param parallel Logical; whether to parallelise the subsampling loop.
#'   Default \code{FALSE}.  Requires the \pkg{netsubsamp} package to be
#'   \emph{installed} on all worker nodes when \code{TRUE}.
#' @param n_cores Integer or \code{NULL}; number of cores when
#'   \code{parallel = TRUE}.
#'
#' @return A list with:
#' \describe{
#'   \item{\code{p_value}}{The p-value of the test.}
#'   \item{\code{test_statistic}}{Observed test statistic: Mahalanobis
#'     distance \eqn{D_0} when \code{test = "mahalanobis"}, or Cauchy
#'     combination statistic \eqn{T} when \code{test = "cauchy"}.}
#'   \item{\code{observed_moments}}{Named numeric vector; the observed
#'     \eqn{\bar\Psi} statistic.}
#'   \item{\code{null_moments}}{\eqn{N_{\mathrm{sub}} \times m} matrix of
#'     null \eqn{\bar\Psi} statistics (including any \code{NA} rows from
#'     degenerate subsamples).}
#'   \item{\code{null_statistic}}{Numeric vector of per-replicate null
#'     Mahalanobis distances (or marginal p-values for \code{"cauchy"}).}
#'   \item{\code{marginal_p_values}}{(Only for \code{test = "cauchy"}.) Named
#'     numeric vector of per-motif marginal p-values.}
#'   \item{\code{method}}{The test method used.}
#'   \item{\code{motif_names}}{Character vector of moment names.}
#' }
#'
#' @references
#' Qi, Hua, Li and Zhou (2024). Multivariate Inference of Network Moments by
#' Subsampling. \emph{Biometrika}, \strong{111}(1), 1--18.
#' \doi{10.1093/biomet/asad056}.
#'
#' Liu, Y. and Xie, J. (2020). Cauchy combination test: a powerful test with
#' analytic p-value calculation under arbitrary dependency structures.
#' \emph{Journal of the American Statistical Association}, \strong{115},
#' 393--402. \doi{10.1080/01621459.2018.1554485}.
#'
#' @examples
#' \donttest{
#' library(randnet)
#' set.seed(3)
#' G       <- BlockModel.Gen(lambda = 20, n = 400, K = 3)$A
#' G_prime <- BlockModel.Gen(lambda = 15, n = 100, K = 3)$A
#'
#' # Choose rho_target as 0.7 * min(density(G), density(G_prime))
#' rho_G       <- sum(G)       / (400 * 399)
#' rho_G_prime <- sum(G_prime) / (100 * 99)
#' rho_target  <- 0.7 * min(rho_G, rho_G_prime)
#'
#' res <- two_sample_test(G, G_prime,
#'                        rho_target = rho_target,
#'                        N_sub      = 500,
#'                        test       = "mahalanobis")
#' res$p_value
#' }
#'
#' @importFrom foreach foreach "%dopar%"
#' @importFrom stats complete.cases cov setNames
#' @export
two_sample_test <- function(
    G,
    G_prime,
    rho_target,
    N_sub    = 2000L,
    motifs   = list("v_shape", "triangle", "star3"),
    test     = c("mahalanobis", "cauchy"),
    parallel = FALSE,
    n_cores  = NULL
) {
  test  <- match.arg(test)
  validate_adjacency(G,       "G")
  validate_adjacency(G_prime, "G_prime")

  n     <- nrow(G)
  b     <- nrow(G_prime)
  N_sub <- as.integer(N_sub)

  if (b >= n) {
    warning(
      "'G_prime' is not smaller than 'G'; results may be unreliable. ",
      "Algorithm 3 is designed for the case nrow(G_prime) < nrow(G)."
    )
  }

  b_half <- floor(b / 2L)
  if (b_half < 3L) {
    stop("nrow(G_prime) must be at least 6 to allow an equal split.")
  }
  if (N_sub < 10L) {
    stop("'N_sub' must be at least 10.")
  }
  if (!is.numeric(rho_target) || length(rho_target) != 1L ||
      rho_target <= 0 || rho_target >= 1) {
    stop("'rho_target' must be a single numeric value strictly in (0, 1).")
  }

  motif_specs <- init_motif_specs(motifs)
  m           <- length(motif_specs$names)

  # ------------------------------------------------------------------
  # Step 1: Split G_prime and compute the observed statistic.
  # ------------------------------------------------------------------
  G_prime_mat <- as.matrix(G_prime)
  all_nodes   <- seq_len(b)
  half1_idx   <- sample(all_nodes, b_half, replace = FALSE)
  half2_idx   <- setdiff(all_nodes, half1_idx)

  G_prime_1 <- G_prime_mat[half1_idx, half1_idx]
  G_prime_2 <- G_prime_mat[half2_idx, half2_idx]

  psi_obs <- .sparsify_moments_internal(
    G_prime_1, G_prime_2, rho_target, motif_specs
  )

  if (anyNA(psi_obs)) {
    warning(
      "The observed statistic contains NA (zero density after edge removal). ",
      "Consider using a smaller 'rho_target' or a larger 'G_prime'."
    )
  }

  # ------------------------------------------------------------------
  # Step 2: Subsampling loop -- generate null distribution.
  # ------------------------------------------------------------------
  par_setup <- setup_foreach(parallel, n_cores)
  on.exit(par_setup$cleanup(), add = TRUE)
  `%op%` <- par_setup$op

  G_local          <- G
  b_half_local     <- b_half
  rho_target_local <- rho_target
  specs_local      <- motif_specs

  pkg_arg <- if (parallel) "netsubsamp" else character(0L)

  null_list <- foreach::foreach(
    i         = seq_len(N_sub),
    .packages = pkg_arg,
    .export   = c(".sparsify_moments_internal", "sparsify_graph",
                  "compute_network_moments")
  ) %op% {
    n_w    <- nrow(G_local)
    idx1   <- sample(n_w, b_half_local, replace = FALSE)
    idx2   <- sample(n_w, b_half_local, replace = FALSE)
    G_sub1 <- as.matrix(G_local[idx1, idx1])
    G_sub2 <- as.matrix(G_local[idx2, idx2])
    .sparsify_moments_internal(G_sub1, G_sub2, rho_target_local, specs_local)
  }

  null_mat           <- do.call(rbind, null_list)
  colnames(null_mat) <- motif_specs$names

  # Rows with NA arise from degenerate (zero-density) subsamples.
  valid       <- stats::complete.cases(null_mat)
  null_valid  <- null_mat[valid, , drop = FALSE]
  n_valid     <- nrow(null_valid)

  if (n_valid < 10L) {
    stop(
      "Fewer than 10 valid null replicates obtained (", n_valid, " of ",
      N_sub, "). Consider a larger 'N_sub', a larger 'G', or a smaller ",
      "'rho_target'."
    )
  }

  # ------------------------------------------------------------------
  # Step 3: Compute test statistic and p-value.
  # ------------------------------------------------------------------
  eta_hat <- colMeans(null_valid)

  if (test == "mahalanobis") {

    Sigma_hat <- stats::cov(null_valid)

    Sigma_inv <- tryCatch(
      solve(Sigma_hat),
      error = function(e) {
        warning(
          "Covariance matrix is singular or near-singular; ",
          "applying ridge regularisation (lambda = 1e-6 * trace / m)."
        )
        lambda <- 1e-6 * sum(diag(Sigma_hat)) / m
        solve(Sigma_hat + diag(lambda, m))
      }
    )

    .mahal_dist <- function(x) {
      d <- x - eta_hat
      as.numeric(crossprod(d, Sigma_inv %*% d))
    }

    D_null <- apply(null_valid, 1, .mahal_dist)
    D_obs  <- .mahal_dist(psi_obs)

    p_value <- mean(D_null > D_obs, na.rm = TRUE)

    list(
      p_value          = p_value,
      test_statistic   = D_obs,
      observed_moments = psi_obs,
      null_moments     = null_mat,
      null_statistic   = D_null,
      method           = "mahalanobis",
      motif_names      = motif_specs$names
    )

  } else {
    # Cauchy combination test (Liu & Xie, 2020).
    obs_dev  <- abs(psi_obs  - eta_hat)

    null_dev <- abs(sweep(null_valid, 2L, eta_hat, FUN = "-"))

    # Per-motif marginal p-values: proportion of null |deviations| that
    # exceed the observed |deviation|.
    p_marginal <- vapply(seq_len(m), function(j) {
      mean(null_dev[, j] > obs_dev[j], na.rm = TRUE)
    }, numeric(1L))
    names(p_marginal) <- motif_specs$names

    # Cauchy statistic and combined p-value.
    T_cauchy <- mean(tan((0.5 - p_marginal) * pi))
    p_cauchy <- 0.5 - atan(T_cauchy) / pi

    list(
      p_value           = p_cauchy,
      test_statistic    = T_cauchy,
      marginal_p_values = p_marginal,
      observed_moments  = psi_obs,
      null_moments      = null_mat,
      null_statistic    = p_marginal,
      method            = "cauchy",
      motif_names       = motif_specs$names
    )
  }
}
