GeoNeighIndex <- function(coordx = NULL, coordy = NULL, coordz = NULL, coordt = NULL,
                          coordx_dyn = NULL,
                          distance = "Eucl", neighb = 4, maxdist = NULL,
                          maxtime = 1, radius = 1, bivariate = FALSE,
                          p_neighb = 1, thin_method = "bernoulli", check.duplicates = FALSE)
{
  ##############################################################
  ## Checks
  ##############################################################

  if (!is.numeric(neighb) ||
      !(length(neighb) %in% c(1L, 3L)) ||
      any(!is.finite(neighb)) ||
      any(neighb <= 0) ||
      any(neighb %% 1 != 0)) {
    stop("neighb must be a positive integer, or a length-3 positive integer vector in the bivariate case.")
  }

  if (!bivariate && length(neighb) != 1L) {
    stop("neighb can have length 3 only when bivariate = TRUE.")
  }

  K <- as.integer(neighb)

  .GeoCheckDuplicateLocations(
    coordx = coordx, coordy = coordy, coordz = coordz, coordt = coordt,
    coordx_dyn = coordx_dyn,
    spacetime = !isTRUE(bivariate) && !is.null(coordt),
    bivariate = isTRUE(bivariate), grid = FALSE,
    context = "GeoNeighIndex", check.duplicates = check.duplicates
  )

  ## Canonical thinning-method names.
  ## Accepted inputs include "bernoulli", "Bernoulli",
  ## "fixedbudget", and "FixedBudget". The old aliases
  ## "targetbalanced" and "TargetBalanced" are also accepted
  ## for backward compatibility and mapped to "FixedBudget".
  ## Internally we use only "Bernoulli" and "FixedBudget".
  if (!is.character(thin_method) || length(thin_method) != 1L || is.na(thin_method)) {
    stop("thin_method must be a character scalar.")
  }

  thin_method <- tolower(thin_method)

  thin_method <- match.arg(
    thin_method,
    c("bernoulli", "fixedbudget", "targetbalanced")
  )

  thin_method <- switch(
    thin_method,
    bernoulli = "Bernoulli",
    fixedbudget = "FixedBudget",
    targetbalanced = "FixedBudget"
  )
 
  indices <- function(nn_idx, nn_dists, drop_first = FALSE, x_equals_y = FALSE) {

    if (drop_first) {
      n <- nrow(nn_idx)
      kout <- max(0L, ncol(nn_idx) - 1L)

      ## Fast normal case: nabor returns the query point itself first.
      ## With duplicated coordinates/ties this ordering is not guaranteed,
      ## so remove the actual target index rather than blindly dropping col 1.
      if (isTRUE(x_equals_y) &&
          length(nn_idx) &&
          all(nn_idx[, 1L] == seq_len(n))) {
        nn_idx   <- nn_idx[, -1L, drop = FALSE]
        nn_dists <- nn_dists[, -1L, drop = FALSE]
      } else if (isTRUE(x_equals_y)) {
        out_idx <- matrix(0L, nrow = n, ncol = kout)
        out_dst <- matrix(NA_real_, nrow = n, ncol = kout)
        if (kout > 0L) {
          for (ii in seq_len(n)) {
            keep <- which(nn_idx[ii, ] > 0L & nn_idx[ii, ] != ii)
            if (length(keep)) {
              keep <- keep[seq_len(min(kout, length(keep)))]
              out_idx[ii, seq_along(keep)] <- nn_idx[ii, keep]
              out_dst[ii, seq_along(keep)] <- nn_dists[ii, keep]
            }
          }
        }
        nn_idx <- out_idx
        nn_dists <- out_dst
      } else {
        nn_idx   <- nn_idx[, -1L, drop = FALSE]
        nn_dists <- nn_dists[, -1L, drop = FALSE]
      }
    }

    k <- ncol(nn_idx)

    rowidx <- rep(seq_len(nrow(nn_idx)), times = k)
    colidx <- as.vector(nn_idx)
    lags   <- as.vector(nn_dists)

    list(
      rowidx = rowidx,
      colidx = colidx,
      lags = lags
    )
  }

  ##############################################################
  ## Spatial nearest-neighbor helper
  ##############################################################

  nn2Geo <- function(x, y, K = 1, distance = 0, maxdist = NULL, radius = 1,
                     drop_first = FALSE, x_equals_y = FALSE) {

    if (!is.matrix(x) || !is.matrix(y)) {
      stop("x and y must be matrices")
    }

    n_x <- nrow(x)
    n_y <- nrow(y)

    K <- as.integer(K)
    K <- min(K, n_x, n_y)

    sphere_xyz <- function(coords) {
      lon <- coords[, 1L] * pi / 180
      lat <- coords[, 2L] * pi / 180
      clat <- cos(lat)
      cbind(clat * cos(lon), clat * sin(lon), sin(lat))
    }

    if (distance %in% c(1, 2)) {
      if (!requireNamespace("fields", quietly = TRUE)) {
        stop("Package 'fields' is required by GeoNeighIndex() for geodesic or chordal distance.",
             call. = FALSE)
      }
      ## Select spherical neighbours in 3D unit-sphere coordinates.  Chordal
      ## distance is monotone in great-circle distance, so this produces the
      ## correct NN ordering globally (including dateline/polar cases), unlike
      ## longitude/latitude or a planar projection.
      x_nn <- sphere_xyz(x)
      y_nn <- sphere_xyz(y)

      if (is.null(maxdist)) {
        nearest <- nabor::knn(x_nn, y_nn, k = K)
      } else {
        if (!is.finite(radius) || radius <= 0) stop("radius must be positive")
        radius_nn <- if (distance == 2) {
          2 * sin(0.5 * pmin(pi, maxdist / radius))
        } else {
          pmin(2, maxdist / radius)
        }
        nearest <- nabor::knn(x_nn, y_nn, k = K, radius = radius_nn)
      }
    } else if (is.null(maxdist)) {
      nearest <- nabor::knn(x, y, k = K)
    } else {
      nearest <- nabor::knn(x, y, k = K, radius = maxdist)
    }

    if (distance %in% c(1, 2)) {

      dists_mat <- matrix(NA_real_, nrow = n_y, ncol = K)

      for (i in seq_len(n_y)) {

        idx <- nearest$nn.idx[i, ]
        valid <- idx > 0

        if (!any(valid)) {
          next
        }

        di <- fields::rdist.earth.vec(
          x1 = matrix(y[i, ], nrow = 1),
          x2 = x[idx[valid], , drop = FALSE],
          miles = FALSE,
          R = 1
        )

        dists_mat[i, seq_len(length(di))] <- di
      }

      if (isTRUE(x_equals_y)) {
        dists_mat[, 1] <- 0
      }

      if (distance == 2) {
        nearest$nn.dists <- radius * dists_mat
      } else {
        nearest$nn.dists <- 2 * radius * sin(0.5 * dists_mat)
      }
    }

    sol <- indices(
      nearest$nn.idx,
      nearest$nn.dists,
      drop_first = drop_first,
      x_equals_y = x_equals_y
    )

    if (!is.null(maxdist)) {

      sel <- sol$colidx > 0

      return(list(
        lags = sol$lags[sel],
        rowidx = sol$rowidx[sel],
        colidx = sol$colidx[sel]
      ))

    } else {

      return(list(
        lags = sol$lags,
        rowidx = sol$rowidx,
        colidx = sol$colidx
      ))
    }
  }

  ##############################################################
  ## Space-time index
  ##############################################################

  spacetime_index <- function(coords, coordx_dyn = NULL, N, K = 4,
                              coordt = NULL, numtime,
                              maxtime = 1, maxdist = NULL,
                              distance = "Eucl", radius = 1) {

    if (is.null(coordt)) {
      stop("coordt is required for spacetime_index().")
    }

    if (numtime <= 1L) {
      stop("numtime must be >= 2.")
    }

    if (!is.numeric(K) || length(K) != 1L || !is.finite(K) || K < 1) {
      stop("K must be a positive integer.")
    }

    K <- as.integer(K)

    if (!is.numeric(maxtime) || length(maxtime) != 1L ||
        is.na(maxtime) || maxtime < 0) {
      stop("maxtime must be a non-negative numeric scalar.")
    }

    if (is.null(coordx_dyn)) {

      ns <- rep.int(as.integer(N), numtime)
      offset <- as.integer(N) * (0:(numtime - 1L))

      getX <- function(t) coords

    } else {

      if (!is.list(coordx_dyn) || length(coordx_dyn) != numtime) {
        stop("coordx_dyn must be a list of length = numtime.")
      }

      ns <- vapply(coordx_dyn, nrow, integer(1))

      if (any(ns <= 0L)) {
        stop("Each coordx_dyn[[t]] must have at least 1 row.")
      }

      offset <- c(0L, cumsum(ns))[seq_len(numtime)]
      offset <- as.integer(offset)

      getX <- function(t) coordx_dyn[[t]]
    }

    ct <- as.numeric(coordt)

    if (length(ct) != numtime) {
      stop("length(coordt) must equal numtime.")
    }

    if (any(!is.finite(ct))) {
      stop("coordt must contain finite values.")
    }

    glob <- function(idx_local, t) {
      as.integer(idx_local + offset[t])
    }

    ##############################################################
    ## Temporal-pair helper
    ##
    ## maxtime is interpreted as a temporal-distance threshold in
    ## the scale of coordt. Thus all pairs with temporal lag
    ## 0 < |coordt[to] - coordt[from]| <= maxtime are retained.
    ## For regularly spaced time points coordt = 1, ..., n and
    ## maxtime = q, this coincides with the old behavior based on
    ## index lags 1, ..., q.
    ##############################################################

    make_time_pairs <- function(ct, maxtime) {

      nt <- length(ct)

      if (nt <= 1L || maxtime <= 0) {
        return(data.frame(
          from = integer(0),
          to = integer(0),
          lag_time = numeric(0)
        ))
      }

      out <- vector("list", nt - 1L)
      out_i <- 1L

      for (from in seq_len(nt - 1L)) {

        to <- (from + 1L):nt
        lag_time <- abs(ct[to] - ct[from])
        ttol <- 100 * .Machine$double.eps *
          pmax(1, abs(ct[to]), abs(ct[from]))

        keep <- is.finite(lag_time) & lag_time > 0 & lag_time <= maxtime + ttol

        if (!any(keep)) {
          next
        }

        out[[out_i]] <- data.frame(
          from = rep.int(from, sum(keep)),
          to = to[keep],
          lag_time = lag_time[keep]
        )

        out_i <- out_i + 1L
      }

      if (out_i == 1L) {
        return(data.frame(
          from = integer(0),
          to = integer(0),
          lag_time = numeric(0)
        ))
      }

      do.call(rbind, out[seq_len(out_i - 1L)])
    }

    knn_dir <- function(ref, query, K,
                        drop_self = FALSE,
                        x_equals_y = FALSE) {

      inf <- nn2Geo(
        ref,
        query,
        if (drop_self) K + 1L else K,
        distance,
        maxdist,
        radius,
        drop_first = drop_self,
        x_equals_y = x_equals_y
      )

      if (length(inf$rowidx) == 0L) {
        return(list(
          tar = integer(0),
          nei = integer(0),
          d = numeric(0)
        ))
      }

      tar <- inf$rowidx
      nei <- inf$colidx
      d   <- inf$lags

      ok <- is.finite(nei) &
        is.finite(tar) &
        nei > 0 &
        tar > 0 &
        is.finite(d)

      list(
        tar = tar[ok],
        nei = nei[ok],
        d = d[ok]
      )
    }

    same_location_pairs <- function(ref, query) {
      if (ncol(ref) != ncol(query) || !nrow(ref) || !nrow(query)) {
        return(list(tar = integer(0L), nei = integer(0L)))
      }

      scale <- suppressWarnings(max(abs(c(ref, query)), na.rm = TRUE))
      if (!is.finite(scale) || scale < 1) scale <- 1
      tol <- 100 * .Machine$double.eps * scale

      available <- seq_len(nrow(ref))
      tar <- integer(0L)
      nei <- integer(0L)

      for (i in seq_len(nrow(query))) {
        if (!length(available)) break
        delta <- abs(sweep(ref[available, , drop = FALSE], 2L,
                           query[i, ], FUN = "-"))
        dmax <- apply(delta, 1L, max)
        jpos <- which.min(dmax)
        if (length(jpos) && is.finite(dmax[jpos]) && dmax[jpos] <= tol) {
          tar <- c(tar, i)
          nei <- c(nei, available[jpos])
          available <- available[-jpos]
        }
      }

      list(tar = tar, nei = nei)
    }

    out <- list()
    out_i <- 1L

    ##############################################################
    ## 1) Spatial within-time
    ##############################################################

    for (t in seq_len(numtime)) {

      Xt <- getX(t)

      inf <- knn_dir(
        ref = Xt,
        query = Xt,
        K = K,
        drop_self = TRUE,
        x_equals_y = TRUE
      )

      if (length(inf$tar) > 0L) {

        out[[out_i]] <- cbind(
          glob(inf$tar, t),
          glob(inf$nei, t),
          rep.int(0, length(inf$tar)),
          as.numeric(inf$d)
        )

        out_i <- out_i + 1L
      }
    }

    time_pairs <- make_time_pairs(ct, maxtime)

    if (nrow(time_pairs) == 0L) {

      if (!length(out)) return(matrix(numeric(0L), nrow = 0L, ncol = 4L))
      sol <- do.call(rbind, out)

      ok <- is.finite(sol[, 1]) &
        is.finite(sol[, 2]) &
        sol[, 1] > 0 &
        sol[, 2] > 0

      return(sol[ok, , drop = FALSE])
    }

    ##############################################################
    ## 2) Pure temporal same-site
    ##############################################################

    for (rr in seq_len(nrow(time_pairs))) {

      t_from <- time_pairs$from[rr]
      t_to   <- time_pairs$to[rr]

      if (is.null(coordx_dyn)) {
        bb <- min(ns[t_from], ns[t_to])
        if (bb <= 0L) next
        same <- list(tar = seq_len(bb), nei = seq_len(bb))
      } else {
        same <- same_location_pairs(getX(t_from), getX(t_to))
        bb <- length(same$tar)
        if (bb <= 0L) next
      }

      out[[out_i]] <- cbind(
        glob(same$tar, t_to),
        glob(same$nei, t_from),
        rep.int(time_pairs$lag_time[rr], bb),
        rep.int(0, bb)
      )

      out_i <- out_i + 1L
    }

    ##############################################################
    ## 3) Spatio-temporal cross-time
    ##############################################################

    for (rr in seq_len(nrow(time_pairs))) {

      t_from <- time_pairs$from[rr]
      t_to   <- time_pairs$to[rr]

      Xn <- getX(t_from)
      Xt <- getX(t_to)

      inf <- knn_dir(
        ref = Xn,
        query = Xt,
        K = K,
        drop_self = FALSE,
        x_equals_y = FALSE
      )

      if (length(inf$tar) > 0L) {

        sel <- is.finite(inf$d) & (inf$d > 0)

        if (any(sel)) {

          out[[out_i]] <- cbind(
            glob(inf$tar[sel], t_to),
            glob(inf$nei[sel], t_from),
            rep.int(time_pairs$lag_time[rr], sum(sel)),
            as.numeric(inf$d[sel])
          )

          out_i <- out_i + 1L
        }
      }
    }

    if (!length(out)) return(matrix(numeric(0L), nrow = 0L, ncol = 4L))
    sol <- do.call(rbind, out)

    ok <- is.finite(sol[, 1]) &
      is.finite(sol[, 2]) &
      sol[, 1] > 0 &
      sol[, 2] > 0

    sol[ok, , drop = FALSE]
  }

  ##############################################################
  ## Bivariate index
  ##############################################################

  bivariate_index <- function(coords, coordx_dyn = NULL, N, K = 4,
                              maxdist, distance, radius) {

    ## Build every bivariate pair block with the same five-column schema.
    ## In particular, an empty block must be a 0 x 5 matrix: base::cbind()
    ## would otherwise drop zero-length vectors and retain only scalar
    ## columns, causing rbind() failures when no positive-lag cross pair
    ## exists (for example, coincident grids with neighb = 1).
    pair_block <- function(rowidx, colidx, first, second, lags) {

      n <- length(rowidx)

      if (length(colidx) != n || length(lags) != n) {
        stop("Internal error: inconsistent bivariate pair-block lengths.")
      }

      if (n == 0L) {
        out <- matrix(numeric(0L), nrow = 0L, ncol = 5L)
        colnames(out) <- c("rowidx", "colidx", "first", "second", "lag")
        return(out)
      }

      out <- cbind(
        as.integer(rowidx),
        as.integer(colidx),
        rep.int(as.integer(first), n),
        rep.int(as.integer(second), n),
        as.numeric(lags)
      )

      colnames(out) <- c("rowidx", "colidx", "first", "second", "lag")
      out
    }

    if (length(K) == 3L) {
      K1 <- K[1]
      K2 <- K[2]
      K3 <- K[3]
    } else {
      K1 <- K2 <- K3 <- K
    }

    if (length(maxdist) == 3L) {
      maxdist1 <- maxdist[1]
      maxdist2 <- maxdist[2]
      maxdist3 <- maxdist[3]
    } else {
      maxdist1 <- maxdist2 <- maxdist3 <- maxdist
    }

    auto_tol <- function(A, B) {

      scaleA <- suppressWarnings(max(abs(range(A, finite = TRUE))))
      scaleB <- suppressWarnings(max(abs(range(B, finite = TRUE))))

      scale <- max(scaleA, scaleB)

      if (!is.finite(scale) || scale <= 0) {
        scale <- 1
      }

      100 * .Machine$double.eps * scale
    }

    ##############################################################
    ## CASE 1: coordx_dyn NULL
    ##############################################################

    if (is.null(coordx_dyn)) {

      n_half <- as.integer(N / 2)
      cm <- coords[1:n_half, , drop = FALSE]

      inf1 <- nn2Geo(
        cm,
        cm,
        K1 + 1,
        distance,
        maxdist1,
        radius,
        drop_first = TRUE,
        x_equals_y = TRUE
      )

      aa1 <- pair_block(
        inf1$rowidx,
        inf1$colidx,
        0L,
        0L,
        inf1$lags
      )

      inf2 <- nn2Geo(
        cm,
        cm,
        K3 + 1,
        distance,
        maxdist3,
        radius,
        drop_first = TRUE,
        x_equals_y = TRUE
      )

      aa2 <- pair_block(
        inf2$rowidx + n_half,
        inf2$colidx + n_half,
        1L,
        1L,
        inf2$lags
      )

      inf21 <- nn2Geo(
        cm,
        cm,
        K2 + 1,
        distance,
        maxdist2,
        radius,
        drop_first = TRUE,
        x_equals_y = TRUE
      )

      sel <- is.finite(inf21$lags) &
        (inf21$lags > 0) &
        (inf21$colidx > 0) &
        (inf21$rowidx > 0)

      aa_cross <- pair_block(
        inf21$rowidx[sel],
        inf21$colidx[sel] + n_half,
        0L,
        1L,
        inf21$lags[sel]
      )

      ii <- seq_len(n_half)

      aa_coll <- pair_block(
        ii,
        ii + n_half,
        0L,
        1L,
        rep.int(0, length(ii))
      )

      return(rbind(
        aa1,
        aa2,
        aa_cross,
        aa_coll
      ))
    }

    ##############################################################
    ## CASE 2: coordx_dyn list length 2
    ##############################################################

    if (!is.list(coordx_dyn) || length(coordx_dyn) != 2L) {
      stop("For bivariate case with coordx_dyn, it must be a list of length 2: list(var1_coords, var2_coords).")
    }

    X1 <- coordx_dyn[[1]]
    X2 <- coordx_dyn[[2]]

    n1 <- nrow(X1)

    inf1 <- nn2Geo(
      X1,
      X1,
      K1 + 1,
      distance,
      maxdist1,
      radius,
      drop_first = TRUE,
      x_equals_y = TRUE
    )

    aa1 <- pair_block(
      inf1$rowidx,
      inf1$colidx,
      0L,
      0L,
      inf1$lags
    )

    inf2 <- nn2Geo(
      X2,
      X2,
      K3 + 1,
      distance,
      maxdist3,
      radius,
      drop_first = TRUE,
      x_equals_y = TRUE
    )

    aa2 <- pair_block(
      inf2$rowidx + n1,
      inf2$colidx + n1,
      1L,
      1L,
      inf2$lags
    )

    inf_cross <- nn2Geo(
      X2,
      X1,
      K2,
      distance,
      maxdist2,
      radius,
      drop_first = FALSE,
      x_equals_y = FALSE
    )

    sel <- is.finite(inf_cross$lags) &
      (inf_cross$lags > 0) &
      (inf_cross$colidx > 0) &
      (inf_cross$rowidx > 0)

    aa_cross <- pair_block(
      inf_cross$rowidx[sel],
      inf_cross$colidx[sel] + n1,
      0L,
      1L,
      inf_cross$lags[sel]
    )

    tol <- auto_tol(X1, X2)

    same_grid <- (nrow(X1) == nrow(X2)) &&
      isTRUE(all(abs(X1 - X2) <= tol))

    if (same_grid) {

      ii <- seq_len(n1)

      aa_coll <- pair_block(
        ii,
        ii + n1,
        0L,
        1L,
        rep.int(0, length(ii))
      )

    } else {

      nn12 <- nn2Geo(
        X2,
        X1,
        K = 1,
        distance,
        maxdist2,
        radius,
        drop_first = FALSE,
        x_equals_y = FALSE
      )

      l12 <- nn12$lags

      l12[is.finite(l12) & (l12 <= tol)] <- 0

      selc <- is.finite(nn12$rowidx) &
        is.finite(nn12$colidx) &
        (nn12$rowidx > 0) &
        (nn12$colidx > 0) &
        is.finite(l12) &
        (l12 == 0)

      aa_coll <- pair_block(
        nn12$rowidx[selc],
        nn12$colidx[selc] + n1,
        0L,
        1L,
        l12[selc]
      )
    }

    rbind(
      aa1,
      aa2,
      aa_cross,
      aa_coll
    )
  }

  ##############################################################
  ## Calibrated Bernoulli probabilities
  ##############################################################

  calibrated_bernoulli_prob <- function(w, p) {

    n <- length(w)

    if (n == 0L) {
      return(numeric(0))
    }

    if (!is.numeric(p) ||
        length(p) != 1L ||
        !is.finite(p) ||
        p <= 0 ||
        p > 1) {
      stop("p_neighb must be in (0,1].")
    }

    if (p >= 1) {
      return(rep.int(1, n))
    }

    w <- as.numeric(w)

    good <- is.finite(w) & w > 0

    if (!any(good)) {
      w[] <- 1
    } else {
      w[!good] <- min(w[good])
    }

    target <- p * n

    f <- function(cst) {
      sum(pmin(1, cst * w)) - target
    }

    hi <- 1 / max(w)

    while (f(hi) < 0) {
      hi <- hi * 2
    }

    cstar <- stats::uniroot(
      f,
      interval = c(0, hi),
      tol = 1e-12
    )$root

    pi <- pmin(1, cstar * w)
    pi[!is.finite(pi)] <- 0

    pmin(1, pmax(0, pi))
  }

  ##############################################################
  ## Subset helper
  ##############################################################

  subset_gb_keep <- function(gb, keep_idx) {

    gb$colidx <- gb$colidx[keep_idx]
    gb$rowidx <- gb$rowidx[keep_idx]
    gb$lags   <- gb$lags[keep_idx]

    if (!is.null(gb$lagt)) {
      gb$lagt <- gb$lagt[keep_idx]
    }

    if (!is.null(gb$first)) {
      gb$first <- gb$first[keep_idx]
    }

    if (!is.null(gb$second)) {
      gb$second <- gb$second[keep_idx]
    }

    gb
  }

  ##############################################################
  ## Method II helpers: FixedBudget thinning
  ##
  ## The option name remains "FixedBudget" for compatibility with
  ## existing scripts. The implementation is target-wise/stratum-wise:
  ## the global budget is exact, but the allocation is made across
  ## natural strata. In the purely spatial case the stratum is the
  ## target location. In the space-time and bivariate cases the stratum
  ## also includes the available pair type descriptors, so that distinct
  ## temporal lags or variable-pair types are not mixed.
  ##
  ## This version avoids split(), which can be memory-expensive when the
  ## number of target locations is very large. It sorts edges by stratum
  ## and by a random tie-breaker, then keeps the first k_g randomized
  ## ranks within each stratum.
  ##############################################################

  make_stratum_id <- function(gb) {

    n_edges <- length(gb$rowidx)

    if (n_edges == 0L) {
      return(numeric(0))
    }

    if (is.null(gb$rowidx) || length(gb$rowidx) != n_edges) {
      stop("gb$rowidx must exist and have one entry per candidate edge.")
    }

    stratum_id <- as.double(gb$rowidx)

    if (any(!is.finite(stratum_id))) {
      stop("rowidx must be finite when constructing FixedBudget strata.")
    }

    ## Fold an additional discrete descriptor into the stratum id.
    ## factor() is used so that numeric descriptors such as temporal lags
    ## are converted to compact integer codes.
    fold_in <- function(base, extra, name) {

      if (length(extra) != length(base)) {
        stop(name, " must have one entry per candidate edge.")
      }

      if (any(!is.finite(extra))) {
        stop(name, " must be finite when constructing FixedBudget strata.")
      }

      code <- as.integer(factor(extra))
      L <- max(code, na.rm = TRUE)

      base * (L + 1L) + code
    }

    if (!is.null(gb$lagt)) {
      stratum_id <- fold_in(stratum_id, gb$lagt, "lagt")
    }

    if (!is.null(gb$first)) {
      stratum_id <- fold_in(stratum_id, gb$first, "first")
    }

    if (!is.null(gb$second)) {
      stratum_id <- fold_in(stratum_id, gb$second, "second")
    }

    stratum_id
  }

  fixed_budget_keep <- function(stratum_id, p, K_target) {

    n_edges <- length(stratum_id)

    if (n_edges <= 0L || K_target <= 0L) {
      return(integer(0))
    }

    if (!is.numeric(p) || length(p) != 1L || !is.finite(p) ||
        p <= 0 || p > 1) {
      stop("p must be in (0,1].")
    }

    K_target <- as.integer(round(K_target))
    K_target <- max(0L, min(K_target, as.integer(n_edges)))

    if (K_target <= 0L) {
      return(integer(0))
    }

    if (K_target >= n_edges || p >= 1) {
      return(seq_len(n_edges))
    }

    if (length(stratum_id) != n_edges || any(!is.finite(stratum_id))) {
      stop("stratum_id must be finite and have one entry per candidate edge.")
    }

    ## Random tie-breaker: sorting by (stratum_id, u) gives a uniformly
    ## random ordering of the candidate edges within each stratum.
    u <- stats::runif(n_edges)

    ord <- order(stratum_id, u)
    str_sorted <- stratum_id[ord]

    new_stratum <- c(TRUE, str_sorted[-1L] != str_sorted[-n_edges])
    stratum_index <- cumsum(new_stratum)
    stratum_start <- which(new_stratum)

    n_strata <- length(stratum_start)

    if (n_strata == 0L) {
      return(integer(0))
    }

    stratum_size <- diff(c(stratum_start, n_edges + 1L))

    rank_in_stratum <- seq_len(n_edges) -
      stratum_start[stratum_index] + 1L

    quota_raw <- p * stratum_size
    quota_base <- floor(quota_raw + 1e-12)
    quota_base <- pmin(quota_base, stratum_size)

    ## Exact global budget by randomized residual allocation.
    residual <- K_target - sum(quota_base)

    if (residual > 0L) {

      eligible <- which(quota_base < stratum_size)

      if (length(eligible) > 0L) {

        residual <- min(as.integer(residual), length(eligible))

        frac <- quota_raw[eligible] -
          floor(quota_raw[eligible] + 1e-12)

        frac[!is.finite(frac) | frac < 0] <- 0

        if (sum(frac) <= 0) {
          prob <- rep.int(1 / length(eligible), length(eligible))
        } else {
          prob <- frac / sum(frac)
        }

        add_strata <- sample(
          eligible,
          size = residual,
          replace = FALSE,
          prob = prob
        )

        quota_base[add_strata] <- quota_base[add_strata] + 1L
      }
    }

    ## Numerical safeguards. These should almost never be active, but they
    ## keep the retained size exact in irregular cases.
    deficit <- K_target - sum(quota_base)

    if (deficit > 0L) {

      eligible <- which(quota_base < stratum_size)

      if (length(eligible) > 0L) {

        add_strata <- sample(
          eligible,
          size = min(as.integer(deficit), length(eligible)),
          replace = FALSE
        )

        quota_base[add_strata] <- quota_base[add_strata] + 1L
      }
    }

    surplus <- sum(quota_base) - K_target

    if (surplus > 0L) {

      eligible <- which(quota_base > 0L)

      if (length(eligible) > 0L) {

        drop_strata <- sample(
          eligible,
          size = min(as.integer(surplus), length(eligible)),
          replace = FALSE
        )

        quota_base[drop_strata] <- quota_base[drop_strata] - 1L
      }
    }

    keep_sorted <- rank_in_stratum <= quota_base[stratum_index]
    keep_idx <- sort(ord[keep_sorted])

    if (length(keep_idx) != K_target) {
      stop(
        "FixedBudget retained ", length(keep_idx),
        " pairs instead of K_target = ", K_target, "."
      )
    }

    keep_idx
  }

  ##############################################################
  ## Start
  ##############################################################

  spatial <- TRUE
  spacetime <- FALSE

  if (!is.null(coordx_dyn)) {
    if (!is.list(coordx_dyn)) {
      stop("coordx_dyn must be a list")
    }
  }

  if (!is.null(coordt)) {
    if (is.numeric(coordt) && is.numeric(maxtime)) {
      if (length(coordt) > 1L && length(maxtime) >= 1L) {
        spacetime <- TRUE
      }
    }
  }

  distance <- CheckDistance(distance)
  spatial <- !spacetime && !bivariate

  ##############################################################
  ## Coordinate parsing
  ##############################################################

  if (is.null(coordx_dyn)) {

    if (is.null(coordx)) {
      stop("coordx must be provided when coordx_dyn is NULL.")
    }

    if (is.null(coordy)) {
      coords <- as.matrix(coordx)
    } else {
      coords <- if (is.null(coordz)) {
        cbind(coordx, coordy)
      } else {
        cbind(coordx, coordy, coordz)
      }
      coords <- as.matrix(coords)
    }

    storage.mode(coords) <- "double"

    if (!(ncol(coords) %in% c(2L, 3L))) {
      stop("coordx must be a coordinate matrix with 2 or 3 columns, or coordx/coordy/(coordz) must be coordinate vectors.")
    }

    if (bivariate && (nrow(coords) %% 2L != 0L)) {
      stop("For bivariate = TRUE with common coordinates, coordx must have an even number of rows.")
    }

    .GeoValidate3D(
      coordx = coords,
      distance = c("Eucl", "Chor", "Geod")[distance + 1L],
      context = "GeoNeighIndex"
    )

    numcoord <- nrow(coords)

  } else {

    coordx_dyn <- .GeoNormalizeDynamicCoordinates(
      coordx_dyn,
      coordt = if (!bivariate) coordt else NULL,
      context = "GeoNeighIndex"
    )

    if (!bivariate) {
      .GeoValidate3D(
        coordx_dyn = coordx_dyn,
        distance = c("Eucl", "Chor", "Geod")[distance + 1L],
        context = "GeoNeighIndex"
      )
      coords <- NULL
      numcoord <- nrow(coordx_dyn[[1L]])
    } else {
      coords <- 1
      numcoord <- 1
    }
  }

  ##############################################################
  ## Build pair indices
  ##############################################################

  if (spatial) {

    K_sp <- as.integer(K[1])

    sol <- nn2Geo(
      coords,
      coords,
      K_sp + 1L,
      distance,
      maxdist,
      radius,
      drop_first = TRUE,
      x_equals_y = TRUE
    )

    gb <- list(
      colidx = sol$colidx,
      rowidx = sol$rowidx,
      lags = sol$lags,
      maxdist = maxdist,
      neighb = neighb
    )
  }

  if (spacetime) {

    K_st <- as.integer(K[1])
    numtime <- length(coordt)

    sol <- spacetime_index(
      coords,
      coordx_dyn,
      numcoord,
      K_st,
      coordt,
      numtime,
      maxtime,
      maxdist,
      distance,
      radius
    )

    gb <- list(
      colidx = sol[, 2],
      rowidx = sol[, 1],
      lags = sol[, 4],
      lagt = sol[, 3]
    )
  }

  if (bivariate) {

    sol <- bivariate_index(
      coords,
      coordx_dyn,
      numcoord,
      K,
      maxdist,
      distance,
      radius
    )

    gb <- list(
      colidx = sol[, 2],
      rowidx = sol[, 1],
      lags = sol[, 5],
      first = sol[, 3],
      second = sol[, 4],
      maxdist = maxdist,
      neighb = neighb
    )
  }

  ##############################################################
  ## p_neighb checks
  ##############################################################

  if (!is.numeric(p_neighb) ||
      length(p_neighb) != 1L ||
      !is.finite(p_neighb)) {
    stop("p_neighb must be a finite numeric scalar.")
  }

  if (p_neighb <= 0 || p_neighb > 1) {
    stop("p_neighb must be in (0,1].")
  }

  thinning_on <- p_neighb < 1

  ##############################################################
  ## Defensive cleaning
  ##############################################################

  if (length(gb$rowidx) > 0L) {

    valid <- (gb$rowidx > 0L) &
      (gb$colidx > 0L) &
      is.finite(gb$rowidx) &
      is.finite(gb$colidx)

    if (!all(valid)) {

      gb$rowidx <- gb$rowidx[valid]
      gb$colidx <- gb$colidx[valid]
      gb$lags   <- gb$lags[valid]

      if (!is.null(gb$lagt)) {
        gb$lagt <- gb$lagt[valid]
      }

      if (!is.null(gb$first)) {
        gb$first <- gb$first[valid]
      }

      if (!is.null(gb$second)) {
        gb$second <- gb$second[valid]
      }
    }
  }

  ##############################################################
  ## No thinning if p_neighb = 1
  ##############################################################

  if (!thinning_on) {
    return(gb)
  }

  ##############################################################
  ## Method I: calibrated Bernoulli thinning
  ## thin_method = "bernoulli"
  ##############################################################

  if (thin_method == "Bernoulli") {

    n_all <- length(gb$rowidx)

    if (n_all == 0L) {
      return(gb)
    }

    w <- rep.int(1, n_all)

    p_ij <- calibrated_bernoulli_prob(
      w = w,
      p = p_neighb
    )

    B <- stats::rbinom(
      n_all,
      size = 1L,
      prob = p_ij
    )

    keep_idx <- which(B == 1L)

    gb <- subset_gb_keep(gb, keep_idx)

    return(gb)
  }

  ##############################################################
  ## Method II: FixedBudget thinning
  ##
  ## The option name is unchanged, but the retained set is now built
  ## target-wise. In the regular directed m-NN case and when
  ## p_neighb*m is an integer, this retains exactly p_neighb*m
  ## candidate edges for every target location. The total retained
  ## size is always exactly K_target = round(p_neighb * n_all).
  ##############################################################

  if (thin_method == "FixedBudget") {

    n_all <- length(gb$rowidx)

    if (n_all == 0L) {
      return(gb)
    }

    K_target <- as.integer(round(p_neighb * n_all))
    K_target <- max(0L, min(K_target, n_all))

    stratum_id <- make_stratum_id(gb)

    keep_idx <- fixed_budget_keep(
      stratum_id = stratum_id,
      p = p_neighb,
      K_target = K_target
    )

    gb <- subset_gb_keep(gb, keep_idx)

    return(gb)
  }

  return(gb)
}