## -----------------------------------------------------------------------------
library(randomizr)
library(ggplot2)
library(estimatr)

## ----echo=FALSE---------------------------------------------------------------
set.seed(20260822)
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 6,
  fig.height = 4
)
options(digits = 3)

## ----echo = FALSE-------------------------------------------------------------
knitr::kable(
  data.frame(
    Term = c(
      "Tight count",
      "Fair bet",
      "Direction",
      "Constraint",
      "Flight",
      "Landing",
      "Balancing matrix $X$",
      "First-order probability"
    ),
    Meaning = c(
      "A realized count sitting at the floor or the ceiling of its target. If the target is 1.5 the count is 1 or 2, never 0 and never 3.",
      "A random choice between two moves, with the odds set so that the average weight does not change. This is what keeps each unit's probability exactly as supplied.",
      "A recipe for a move: how much weight to add to each unit, and how much to take away. Written as a vector $u$, one number per unit.",
      "A quantity the design promises not to disturb, such as the total of all the weights (which is the expected number treated).",
      "The stage in which every move respects every constraint. Weight is only ever shifted between units, never created or destroyed.",
      "The stage reached when no move respects every constraint any more. Something then has to give: a constraint is set aside, or a last unit is settled by a coin.",
      "The table of covariates whose treated totals the design tries to hold near their targets. It is the model matrix of the `formula` you pass.",
      "The probability that a given unit ends up in a given condition. This is exact here."
    ),
    check.names = FALSE
  ),
  caption = "Terms used in this vignette."
)

## -----------------------------------------------------------------------------
blocks <- rep(1:2, each = 3)
balanced_ra(blocks = blocks)

## -----------------------------------------------------------------------------
reps <- replicate(5000, balanced_ra(blocks = blocks))

# individual assignment probabilities
cbind(target = .5, average = rowMeans(reps))

# block totals
table(colSums(reps[blocks == 1, ]), colSums(reps[blocks == 2, ]))

## -----------------------------------------------------------------------------
reps_blk <- replicate(5000, block_ra(blocks = blocks))
table(district_1 = colSums(reps_blk[blocks == 1, ]),
      district_2 = colSums(reps_blk[blocks == 2, ]))
table(total_treated = colSums(reps_blk))

## -----------------------------------------------------------------------------
set.seed(1)
N <- 100
x <- rnorm(N)
p <- runif(N)

n_draw <- 1000
Z_simple <- replicate(n_draw, simple_ra(N = 100, prob_unit = p))
Z_balanced <- replicate(n_draw, balanced_ra(formula = ~ x, prob_unit = p))

## -----------------------------------------------------------------------------
c(max_gap_simple   = max(abs(rowMeans(Z_simple) - p)),
  max_gap_balanced = max(abs(rowMeans(Z_balanced) - p)))

c(sum_p = sum(p), floor = floor(sum(p)), ceiling = ceiling(sum(p)))
table(balanced_count = colSums(Z_balanced))
range(colSums(Z_simple))

## -----------------------------------------------------------------------------
sx_simple <- colSums(x * Z_simple)
sx_balanced <- colSums(x * Z_balanced)
rbind(
  simple = c(mean = mean(sx_simple), var = var(sx_simple)),
  balanced = c(mean = mean(sx_balanced), var = var(sx_balanced))
)

## ----echo=FALSE, fig.width=6, fig.height=5.2, fig.cap="Treated $x$-total under `simple_ra` and cube-on-X with heterogeneous $p_i$. Cube-on-X is tighter in these draws."----
sx <- data.frame(
  design = factor(
    rep(c("simple_ra", "balanced_ra(formula = ~ x, prob_unit = p)"),
        each = n_draw),
    levels = c("simple_ra", "balanced_ra(formula = ~ x, prob_unit = p)")
  ),
  sum_x = c(sx_simple, sx_balanced)
)
x_lim <- range(sx$sum_x)
ggplot(sx, aes(sum_x)) +
  geom_histogram() +
  facet_wrap(~ design, ncol = 1, scales = "free_y") +
  labs(x = "treated total of x", y = "draws") +
  theme_bw(base_size = 11) 

## -----------------------------------------------------------------------------
set.seed(2)
N2 <- 100
x2 <- sort(rnorm(N2))
p2 <- seq(0.1, 0.9, length.out = N2)   # probability rises with x

gap <- function(Z) mean(x2[Z == 1]) - mean(x2[Z == 0])
g_bal <- replicate(1000, gap(balanced_ra(formula = ~ x2, prob_unit = p2)))
g_sim <- replicate(1000, gap(simple_ra(N = N2, prob_unit = p2)))

rbind(balanced = c(mean_gap = mean(g_bal), sd_gap = sd(g_bal)),
      simple   = c(mean_gap = mean(g_sim), sd_gap = sd(g_sim)))

## ----echo = FALSE-------------------------------------------------------------
knitr::kable(
  data.frame(
    Call = c(
      "`balanced_ra(prob_unit = p)` or `balanced_ra(blocks = b)`",
      "`balanced_ra(prob_unit_each = P)`",
      "`balanced_ra(formula = ~ x)`"
    ),
    `C++` = c("`cube_two_arm_cpp`", "`cube_multi_cpp`", "`cube_on_x_cpp`"),
    Does = c(
      "Two-arm counts; leftover pairing if `blocks`",
      "Three or more arms",
      "Linear targets on a model matrix $X$"
    ),
    Paper = c(
      "Deville and Tillé (1998), pivotal method",
      "Deville and Tillé (2004), cube; Chauvet and Tillé (2006), window",
      "Deville and Tillé (2004), cube; Chauvet and Tillé (2006), window"
    ),
    check.names = FALSE
  ),
  caption = "Which C++ implementation a call uses."
)

## ----echo=FALSE---------------------------------------------------------------
along <- function(x1, y1, x2, y2, t) {
  data.frame(x = x1 + t * (x2 - x1), y = y1 + t * (y2 - y1))
}
fmt_u <- function(u) {
  lab <- paste0("u = ", ifelse(u > 0, "+", ""), u)
  lab <- gsub("u = \\+0", "u = 0", lab)
  gsub("u = -", "u = \u2212", lab)
}
fmt_abs <- function(x) {
  # Fractions the vignette already uses; leave tenths (0.2, 0.6, 0.4) as decimals.
  pairs <- list(
    "1/20" = 1 / 20, "1/12" = 1 / 12, "1/6" = 1 / 6,
    "1/4" = 1 / 4, "1/3" = 1 / 3, "5/12" = 5 / 12,
    "1/2" = 1 / 2, "2/3" = 2 / 3
  )
  for (nm in names(pairs)) {
    if (abs(x - pairs[[nm]]) < 1e-8) return(nm)
  }
  sub("\\.?0+$", "", sprintf("%.2f", x))
}
fmt_signed <- function(x) {
  paste0(if (x > 0) "+" else "\u2212", fmt_abs(abs(x)))
}
fmt_tuple <- function(x, signed = FALSE) {
  inner <- vapply(x, function(v) {
    if (signed) {
      if (abs(v) < 1e-12) "0" else paste0(if (v > 0) "+" else "\u2212",
                                          fmt_abs(abs(v)))
    } else {
      fmt_abs(v)
    }
  }, character(1))
  paste0("(", paste(inner, collapse = ", "), ")")
}
fmt_delta_note <- function(dplus, dminus) {
  paste0("\u03B4+ = ", fmt_abs(dplus),
         "      \u03B4\u2212 = ", fmt_abs(dminus))
}
# Largest plus/minus steps with z <- z + δ u, hitting 0 or 1.
cube_step_sizes <- function(z, u) {
  dplus <- Inf
  dminus <- Inf
  for (i in seq_along(z)) {
    if (abs(u[i]) < 1e-12) next
    if (u[i] > 0) {
      dplus <- min(dplus, (1 - z[i]) / u[i])
      dminus <- min(dminus, z[i] / u[i])
    } else {
      dplus <- min(dplus, z[i] / (-u[i]))
      dminus <- min(dminus, (1 - z[i]) / (-u[i]))
    }
  }
  list(dplus = dplus, dminus = dminus)
}
walk_step_sizes <- function(P, cu, ca) {
  dplus <- Inf
  dminus <- Inf
  for (e in seq_along(cu)) {
    z <- P[cu[e], ca[e]]
    if ((e - 1) %% 2 == 0) {
      dplus <- min(dplus, 1 - z); dminus <- min(dminus, z)
    } else {
      dplus <- min(dplus, z); dminus <- min(dminus, 1 - z)
    }
  }
  list(dplus = dplus, dminus = dminus)
}
# mass[i, j] is the plus-step transfer on that cell (δ+ u). Arrows follow its sign.
cube_net <- function(P, mass = NULL, xval = NULL, zlab = NULL, ulab = NULL,
                     arm_names = NULL, dplus = NULL, dminus = NULL) {
  n <- nrow(P)
  k <- ncol(P)
  if (is.null(mass)) mass <- matrix(0, n, k)
  if (is.null(arm_names)) {
    arm_names <- if (k == 2L) c("Control", "Treat") else as.character(seq_len(k))
  }
  frac <- P > 1e-8 & P < 1 - 1e-8
  open <- rowSums(frac) > 0
  if (is.null(zlab) && k == 2L) {
    z <- P[, 2L]
    zlab <- ifelse(open, sprintf("z = %.2f", z),
                   ifelse(z >= 1 - 1e-8, "z = 1", "z = 0"))
  }
  if (is.null(zlab) && k >= 3L) {
    zlab <- vapply(seq_len(n), function(i)
      paste0("z = ", fmt_tuple(P[i, ])), character(1))
  }
  vec_labs <- !is.null(zlab) && grepl("(", zlab[[1]], fixed = TRUE)
  z_size <- if (vec_labs) 2.45 else 3.0
  u_size <- if (vec_labs) 2.45 else 3.1
  xpad <- if (vec_labs) 1.05 else 0.55
  y_gap <- if (vec_labs) 0.26 else 0.20
  units <- data.frame(
    i = seq_len(n), x = seq_len(n), y = 0,
    open = open, kind = ifelse(open, "open", "settled"),
    zlab = if (is.null(zlab)) "" else zlab,
    ulab = if (is.null(ulab)) "" else ulab,
    xlab = if (is.null(xval)) "" else paste0("x = ", xval),
    stringsAsFactors = FALSE
  )
  pad <- if (n <= 2L) 0 else 0.45
  arms <- data.frame(
    name = arm_names,
    x = if (k == 1L) mean(units$x) else seq(min(units$x) + pad,
                                            max(units$x) - pad,
                                            length.out = k),
    y = 2.22
  )
  grey <- data.frame(x = numeric(0), y = numeric(0),
                     xend = numeric(0), yend = numeric(0))
  arr <- data.frame(x = numeric(0), y = numeric(0),
                    xend = numeric(0), yend = numeric(0),
                    lab = character(0), lx = numeric(0), ly = numeric(0))
  for (i in seq_len(n)) {
    for (j in seq_len(k)) {
      if (!frac[i, j]) next
      ux <- units$x[i]; uy <- 0
      ax <- arms$x[j]; ay <- arms$y[1]
      mval <- mass[i, j]
      if (abs(mval) < 1e-12) {
        grey <- rbind(grey, data.frame(x = ux, y = uy, xend = ax, yend = ay))
      } else {
        if (mval > 0) {
          a <- along(ux, uy, ax, ay, 0.12); b <- along(ux, uy, ax, ay, 0.88)
        } else {
          a <- along(ax, ay, ux, uy, 0.12); b <- along(ax, ay, ux, uy, 0.88)
        }
        mid <- along(a$x, a$y, b$x, b$y, if (mval > 0) 0.36 else 0.64)
        arr <- rbind(arr, data.frame(
          x = a$x, y = a$y, xend = b$x, yend = b$y,
          lab = fmt_signed(mval),
          lx = mid$x, ly = mid$y
        ))
      }
    }
  }
  has_u <- !is.null(ulab)
  has_z <- !is.null(zlab)
  has_x <- !is.null(xval)
  has_d <- !is.null(dplus) && !is.null(dminus)
  y_z <- -0.28
  y_x <- if (has_z) -0.28 - y_gap else -0.28
  y_u <- -0.28 - y_gap * (has_z + has_x)
  y_as <- y_u - 0.22
  y_lo <- y_as - if (vec_labs) 0.32 else 0.28
  y_hi <- if (has_d) 3.08 else 2.72
  g <- ggplot()
  if (nrow(grey) > 0) {
    g <- g + geom_segment(
      data = grey, aes(x = x, y = y, xend = xend, yend = yend),
      colour = "grey70", linewidth = 0.45
    )
  }
  if (nrow(arr) > 0) {
    g <- g + geom_segment(
      data = arr, aes(x = x, y = y, xend = xend, yend = yend),
      colour = "grey20", linewidth = 0.75,
      arrow = arrow(length = grid::unit(0.15, "cm"), type = "closed")
    ) + geom_label(
      data = arr, aes(x = lx, y = ly, label = lab),
      size = 2.9, fontface = "bold", linewidth = 0,
      fill = "white", colour = "grey15",
      label.padding = grid::unit(0.08, "lines")
    )
  }
  g <- g +
    geom_point(data = units, aes(x, y, fill = kind),
               shape = 21, size = 7.6, colour = "grey20", stroke = 0.7) +
    geom_text(data = units, aes(x, y, label = i),
              size = 3.3, fontface = "bold", colour = "grey20") +
    geom_point(data = arms, aes(x, y),
               shape = 21, size = 8.6, fill = "grey20", colour = "grey20") +
    geom_text(data = arms, aes(x, y = y + 0.30, label = name),
              size = 3.4, fontface = "bold", colour = "grey20", vjust = 0)
  if (has_d) {
    g <- g + annotate(
      "text", x = mean(units$x), y = 2.88,
      label = fmt_delta_note(dplus, dminus),
      size = 3.3, colour = "grey20", fontface = "bold"
    )
  }
  if (has_z) {
    g <- g + geom_text(
      data = units, aes(x, y = y_z, label = zlab, colour = kind),
      size = z_size, vjust = 1
    )
  }
  if (has_x) {
    g <- g + geom_text(
      data = units, aes(x, y = y_x, label = xlab, colour = kind),
      size = 3.0, vjust = 1
    )
  }
  if (has_u) {
    g <- g + geom_text(
      data = units, aes(x, y = y_u, label = ulab, colour = kind),
      size = u_size, fontface = "bold", vjust = 1
    )
  }
  settled <- subset(units, !open)
  if (nrow(settled) > 0) {
    if (k == 2L) {
      settled$slab <- "assigned"
    } else {
      hit <- max.col(P[settled$i, , drop = FALSE], ties.method = "first")
      settled$slab <- paste0("arm ", arm_names[hit])
    }
    g <- g + geom_text(
      data = settled, aes(x, y = y_as, label = slab),
      size = 3.0, colour = "grey45", vjust = 1
    )
  }
  g +
    scale_fill_manual(values = c(open = "white", settled = "grey78"),
                      guide = "none") +
    scale_colour_manual(values = c(open = "grey20", settled = "grey55"),
                        guide = "none") +
    coord_cartesian(xlim = c(min(units$x) - xpad, max(units$x) + xpad),
                    ylim = c(y_lo, y_hi), expand = FALSE, clip = "off") +
    theme_void(base_size = 11) +
    theme(plot.margin = margin(10, 10, 8, 8),
          plot.background = element_rect(fill = "white", colour = NA),
          panel.background = element_rect(fill = "white", colour = NA))
}
# Two-arm: plus step is z <- z + δ+ u. Control changes by −δ+ u.
cube_bip <- function(z, u, xval = NULL, zlab = NULL, ulab = NULL,
                     dplus = NULL, dminus = NULL) {
  n <- length(z)
  P <- cbind(1 - z, z)
  ss <- cube_step_sizes(z, u)
  if (is.null(dplus)) dplus <- ss$dplus
  if (is.null(dminus)) dminus <- ss$dminus
  mass <- cbind(-dplus * u, dplus * u)
  if (is.null(ulab)) ulab <- fmt_u(u)
  cube_net(P, mass, xval = xval, zlab = zlab, ulab = ulab,
           dplus = dplus, dminus = dminus)
}
# Multi-arm: alternating kernel on (cu, ca); arrows are the plus transfer.
# z and u under each unit are that unit's row of Z and of the kernel.
cube_walk <- function(P, cu, ca, arm_names = NULL,
                      dplus = NULL, dminus = NULL) {
  ss <- walk_step_sizes(P, cu, ca)
  if (is.null(dplus)) dplus <- ss$dplus
  if (is.null(dminus)) dminus <- ss$dminus
  n <- nrow(P)
  k <- ncol(P)
  U <- matrix(0, n, k)
  mass <- matrix(0, n, k)
  for (e in seq_along(cu)) {
    s <- if ((e - 1) %% 2 == 0) 1 else -1
    U[cu[e], ca[e]] <- s
    mass[cu[e], ca[e]] <- s * dplus
  }
  zlab <- vapply(seq_len(n), function(i)
    paste0("z = ", fmt_tuple(P[i, ])), character(1))
  ulab <- vapply(seq_len(n), function(i)
    paste0("u = ", fmt_tuple(U[i, ], signed = TRUE)), character(1))
  cube_net(P, mass, zlab = zlab, ulab = ulab,
           arm_names = arm_names, dplus = dplus, dminus = dminus)
}

## -----------------------------------------------------------------------------
p_piv <- c(0.2, 0.6, 0.7, 0.8)
sum(p_piv)
balanced_ra(prob_unit = p_piv)

## ----echo=FALSE, fig.width=6.4, fig.height=4.3, fig.cap="Start. Kernel on units 1 and 2. Arrows are $\\delta_+ u$. $\\delta_+ = 0.6$ assigns unit 2 to control; $\\delta_- = 0.2$ assigns unit 1 to control."----
cube_bip(
  z = p_piv, u = c(1, -1, 0, 0),
  ulab = c("u = +1", "u = \u22121", "u = 0", "u = 0")
)

## ----echo=FALSE, fig.width=6.4, fig.height=4.3, fig.cap="After the first plus step. Unit 2 is assigned. Next pair: units 1 and 3. Arrows are $\\delta_+ u$."----
cube_bip(
  z = c(0.8, 0, 0.7, 0.8), u = c(1, 0, -1, 0),
  ulab = c("u = +1", "u = 0", "u = \u22121", "u = 0")
)

## ----echo=FALSE, fig.width=6.4, fig.height=4.3, fig.cap="After the second plus step. Units 3 and 4 remain. $\\delta_+ = 0.5$ assigns unit 3 (leftover $z_4 = 0.3$); $\\delta_- = 0.2$ assigns unit 4."----
cube_bip(
  z = c(1, 0, 0.5, 0.8), u = c(0, 0, 1, -1),
  ulab = c("u = 0", "u = 0", "u = +1", "u = \u22121")
)

## -----------------------------------------------------------------------------
P3 <- rbind(
  c(0.2, 0.4, 0.4),
  c(0.4, 0.3, 0.3),
  c(0.6, 0.2, 0.2),
  c(0.8, 0.1, 0.1)
)
P3
colSums(P3)
balanced_ra(prob_unit_each = P3, conditions = 1:3)

## ----echo=FALSE---------------------------------------------------------------
cube_move <- function(Z, cu, ca, u) {
  m <- length(cu)
  dplus <- Inf
  dminus <- Inf
  for (e in seq_len(m)) {
    z <- Z[cu[e], ca[e]]
    if ((e - 1) %% 2 == 0) {
      dplus <- min(dplus, 1 - z); dminus <- min(dminus, z)
    } else {
      dplus <- min(dplus, z); dminus <- min(dminus, 1 - z)
    }
  }
  up <- u < dminus / (dplus + dminus)
  for (e in seq_len(m)) {
    s <- if ((e - 1) %% 2 == 0) 1 else -1
    Z[cu[e], ca[e]] <- Z[cu[e], ca[e]] + if (up) s * dplus else -s * dminus
    if (Z[cu[e], ca[e]] < 1e-12) Z[cu[e], ca[e]] <- 0
    if (Z[cu[e], ca[e]] > 1 - 1e-12) Z[cu[e], ca[e]] <- 1
  }
  Z
}
Z3 <- P3
S3 <- list(start = Z3)
Z3 <- cube_move(Z3, c(1, 2, 2, 1), c(1, 1, 2, 2), 0.20)
S3$after_1 <- Z3
Z3 <- cube_move(Z3, c(1, 3, 3, 2, 2, 1), c(1, 1, 2, 2, 3, 3), 0.20)
S3$after_2 <- Z3
Z3 <- cube_move(Z3, c(3, 3, 4, 4), c(2, 1, 1, 2), 0.50)
S3$after_3 <- Z3
Z3 <- cube_move(Z3, c(2, 3, 3, 4, 4, 2), c(2, 2, 1, 1, 3, 3), 0.40)
S3$after_4 <- Z3
Z3 <- cube_move(Z3, c(2, 3, 3, 2), c(2, 2, 3, 3), 0.10)
S3$end <- Z3

## ----echo=FALSE, fig.width=6.4, fig.height=4.8, fig.cap="Start. Cycle on units 1--2, arms 1--2. Under each unit, $z$ is the row of $Z$ and $u$ is the kernel row. Arrows are $\\delta_+ u$. $\\delta_+ = 0.4$ sends unit 2's arm-1 cell to 0; $\\delta_- = 0.2$ sends unit 1's arm-1 cell to 0."----
cube_walk(P3, c(1, 2, 2, 1), c(1, 1, 2, 2), arm_names = c("1", "2", "3"))

## ----echo=FALSE, fig.width=6.4, fig.height=4.8, fig.cap="After the first plus step. Next cycle on units 1, 3, 2. $\\delta_+ = 0.4$ assigns unit 1 to arm 1; $\\delta_- = 0.2$ sends unit 3's arm-2 cell to 0."----
cube_walk(S3$after_1, c(1, 3, 3, 2, 2, 1), c(1, 1, 2, 2, 3, 3),
          arm_names = c("1", "2", "3"))

## ----echo=FALSE, fig.width=6.4, fig.height=4.8, fig.cap="After the second plus step. Unit 1 is assigned. Cycle brings in unit 4. $\\delta_+ = 0.1$ sends unit 4's arm-2 cell to 0; $\\delta_- = 0.6$ sends unit 3's arm-2 cell to 0."----
cube_walk(S3$after_2, c(3, 3, 4, 4), c(2, 1, 1, 2),
          arm_names = c("1", "2", "3"))

## ----echo=FALSE, fig.width=6.4, fig.height=4.8, fig.cap="After the third plus step. Cycle on units 2, 3, 4. $\\delta_+ = 0.7$ assigns unit 2 to arm 2; $\\delta_- = 0.1$ assigns unit 4 to arm 1."----
cube_walk(S3$after_3, c(2, 3, 3, 4, 4, 2), c(2, 2, 1, 1, 3, 3),
          arm_names = c("1", "2", "3"))

## ----echo=FALSE, fig.width=6.4, fig.height=4.7, fig.cap="Start. Window of three units. Kernel $u = (1, -2, 1, 0)$. Arrows are $\\delta_+ u$; unit 2 moves twice as far. $\\delta_\\pm = 1/4$."----
cube_bip(
  z = c(0.5, 0.5, 0.5, 0.5),
  xval = c(1, 2, 3, 6),
  u = c(1, -2, 1, 0),
  ulab = c("u = +1", "u = \u22122", "u = +1", "u = 0")
)

## ----echo=FALSE, fig.width=6.4, fig.height=4.7, fig.cap="After the first plus step. Unit 2 is assigned. Arrows are $\\delta_+ u$ with $u = (3, 0, -5, 2)$. $\\delta_+ = 1/12$; $\\delta_- = 1/20$."----
cube_bip(
  z = c(0.75, 0, 0.75, 0.5),
  xval = c(1, 2, 3, 6),
  u = c(3, 0, -5, 2),
  ulab = c("u = +3", "u = 0", "u = \u22125", "u = +2")
)

## ----echo=FALSE, fig.width=6.4, fig.height=4.7, fig.cap="After $\\delta_+ = 1/12$. No kernel remains on both columns, so $x$ is dropped. Arrows are the count-only landing move $\\delta_+ u$. $\\delta_+ = 2/3$; $\\delta_- = 1/3$."----
cube_bip(
  z = c(1, 0, 1 / 3, 2 / 3),
  xval = c(1, 2, 3, 6),
  u = c(0, 0, 1, -1),
  zlab = c("z = 1", "z = 0", "z = 1/3", "z = 2/3"),
  ulab = c("u = 0", "u = 0", "u = +1", "u = \u22121")
)

## ----echo=FALSE---------------------------------------------------------------
claim_tick <- function(ok, claim) {
  stopifnot(isTRUE(ok))
  cat(sprintf(
    '<p><span style="color:#2e7d32;font-weight:bold;">&#10003;</span> %s</p>\n',
    claim
  ))
}
claim_cross <- function(held, claim) {
  stopifnot(!isTRUE(held))
  cat(sprintf(
    '<p><span style="color:#c62828;font-weight:bold;">&#10007;</span> %s</p>\n',
    claim
  ))
}
assign_id <- function(z) paste(which(z == 1), collapse = ",")
vertex_share <- function(Z, n = 4) {
  pairs <- combn(n, 2)
  ids <- apply(pairs, 2, paste, collapse = ",")
  tab <- table(factor(apply(Z, 2, assign_id), levels = ids))
  as.numeric(tab) / ncol(Z)
}

## -----------------------------------------------------------------------------
chances <- c(0.5, 0.3, 0.15, 0.05)
n_race <- 100000
set.seed(1)
which(balanced_ra(prob_unit = chances) == 1)
Z_race <- replicate(n_race, balanced_ra(prob_unit = chances))
win_rate <- rowMeans(Z_race)
race <- rbind(chance = chances, win_rate = win_rate)
colnames(race) <- paste0("contestant ", seq_along(chances))
race

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(all(colSums(Z_race) == 1),
           "Every draw has exactly one winner.")
claim_tick(max(abs(win_rate - chances)) < 0.02,
           "Each contestant's win rate tracks the chance supplied (max absolute gap below 0.02).")

## -----------------------------------------------------------------------------
set.seed(12)
blocks10 <- rep(1:10, each = 3)
balanced_ra(blocks = blocks10)
r_blk <- replicate(2000, balanced_ra(blocks = blocks10))
table(colSums(r_blk))
block_range <- sapply(1:10, function(b)
  range(colSums(r_blk[blocks10 == b, , drop = FALSE])))
rownames(block_range) <- c("min", "max")
block_range

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(
  all(colSums(r_blk) == 15) &&
    all(block_range["min", ] == 1) &&
    all(block_range["max", ] == 2),
  "Every draw treats 15 units, and every block contributes 1 or 2."
)

## -----------------------------------------------------------------------------
P23 <- cbind(c(0.15, 0.47), c(0.65, 0.48), c(0.20, 0.05))
P23
colSums(P23)
set.seed(4)
balanced_ra(prob_unit_each = P23, conditions = 1:3)
Z23 <- replicate(2000, 
  balanced_ra(prob_unit_each = P23, conditions = 1:3))
table(colSums(Z23 == 2))

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(all(colSums(Z23 == 2) %in% 1:2),
           "Arm 2 receives 1 or 2 units on every draw.")

## -----------------------------------------------------------------------------
set.seed(3)
n <- 80
p <- runif(n)
blocks_h <- sample(1:5, n, replace = TRUE, prob = 1:5)
reps_h <- replicate(2000, balanced_ra(prob_unit = p, blocks = blocks_h))
share <- rowMeans(reps_h)
ggplot(data.frame(p, share), aes(p, share)) +
  geom_abline(slope = 1, intercept = 0, colour = "grey40") +
  geom_point(size = 1.5) +
  coord_equal(xlim = c(0, 1), ylim = c(0, 1), expand = FALSE) +
  labs(x = "supplied probability", y = "share treated") +
  theme_bw(base_size = 11) +
  theme(panel.grid.minor = element_blank())

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(max(abs(share - p)) < 0.08,
           "Unit-level shares track the supplied probabilities (max absolute gap below 0.08).")

## -----------------------------------------------------------------------------
set.seed(8)
clusters <- rep(1:6, times = c(3, 1, 4, 2, 5, 3))
p_cluster <- c(0.2, 0.4, 0.6, 0.8, 0.5, 0.5)
z_cl <- balanced_ra(prob_unit = p_cluster[clusters], clusters = clusters)
table(clusters, z_cl)

## ----echo=FALSE---------------------------------------------------------------
Z_cl <- replicate(2000, balanced_ra(prob_unit = p_cluster[clusters],
                                    clusters = clusters))
n_treated_cl <- apply(Z_cl, 2, function(z)
  sum(tapply(z, clusters, function(v) v[1])))
cluster_constant <- all(apply(Z_cl, 2, function(z)
  all(tapply(z, clusters, function(v) length(unique(v)) == 1L))))

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(cluster_constant && all(n_treated_cl == 3),
           "Units in a cluster share an assignment, and exactly three clusters are treated on every draw.")

## -----------------------------------------------------------------------------
set.seed(19)
clusters <- rep(1:6, times = c(3, 1, 4, 2, 5, 3))
x_cl <- c(-2, -1, 0, 1, 2, 3)[clusters]
Z_clx <- replicate(2000, balanced_ra(prob_unit = p_cluster[clusters],
                                     clusters = clusters, formula = ~ x_cl))
n_treated_clx <- apply(Z_clx, 2, function(z)
  sum(tapply(z, clusters, function(v) v[1])))
table(treated_clusters = n_treated_clx)
table(treated_units = colSums(Z_clx))

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(
  all(n_treated_clx == 3) &&
    all(apply(Z_clx, 2, function(z)
      all(tapply(z, clusters, function(v) length(unique(v))) == 1L))),
  "With `formula` and `clusters` together, exactly three clusters are treated on every draw, and no cluster is split."
)

## ----error=TRUE---------------------------------------------------------------
try({
x_fb <- c(1, 2, 3, 6)
blocks_fb <- rep(1:2, each = 2)
balanced_ra(formula = ~ x_fb, blocks = blocks_fb)
})

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(
  inherits(try(balanced_ra(formula = ~ x_fb, blocks = blocks_fb),
               silent = TRUE),
           "try-error"),
  "`formula` plus `blocks` is refused."
)

## -----------------------------------------------------------------------------
set.seed(16)
P_mb <- matrix(1 / 3, 6, 3)
blocks_mb <- rep(1:3, each = 2)
balanced_ra(prob_unit_each = P_mb, blocks = blocks_mb, conditions = 1:3)
Z_mb <- replicate(2000, 
  balanced_ra(prob_unit_each = P_mb, blocks = blocks_mb, conditions = 1:3))
table(colSums(Z_mb == 1))

## ----echo=FALSE---------------------------------------------------------------
within_block_tight <- TRUE
for (b in 1:3) {
  for (j in 1:3) {
    tot <- colSums(Z_mb[blocks_mb == b, , drop = FALSE] == j)
    if (!all(tot %in% 0:1)) within_block_tight <- FALSE
  }
}
overall_always_two <- all(colSums(Z_mb == 1) == 2)

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(within_block_tight,
           "Within each block, each arm's count is 0 or 1.")
claim_cross(
  overall_always_two,
  "Overall arm counts are always the floor or the ceiling of the overall target (here, always 2)."
)

## -----------------------------------------------------------------------------
x <- c(1, 2, 3, 4)
pairs <- combn(4, 2)
n_draw <- 8000
set.seed(20260822)
Z_complete <- replicate(n_draw, complete_ra(N = 4, m = 2))
Z_balanced <- replicate(n_draw, balanced_ra(formula = ~ x))

tab <- data.frame(
  treated = apply(pairs, 2, paste, collapse = ","),
  `x treated` = apply(pairs, 2, function(j) paste(x[j], collapse = ",")),
  `sum x` = apply(pairs, 2, function(j) sum(x[j])),
  complete = 1 / 6,
  balanced = vertex_share(Z_balanced),
  check.names = FALSE
)
knitr::kable(tab, digits = 3,
             caption = "Shares of the six assignments of two treated units with $x = (1, 2, 3, 4)$. The target pair-sum is 5 and is attained by treating units 1 and 4 or units 2 and 3. `balanced_ra` does not concentrate on those two.")

## -----------------------------------------------------------------------------
rbind(
  complete = c(p = mean(Z_complete), treated = mean(colSums(Z_complete))),
  balanced = c(p = mean(Z_balanced), treated = mean(colSums(Z_balanced)))
)

## -----------------------------------------------------------------------------
set.seed(31)
Z_simple <- replicate(n_draw, simple_ra(N = 4, prob = 0.5))
sx_complete <- colSums(x * Z_complete)
sx_balanced <- colSums(x * Z_balanced)
sx_simple   <- colSums(x * Z_simple)
rbind(
  simple   = c(mean = mean(sx_simple),   var = var(sx_simple)),
  complete = c(mean = mean(sx_complete), var = var(sx_complete)),
  balanced = c(mean = mean(sx_balanced), var = var(sx_balanced))
)

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(
  all(colSums(Z_balanced) == 2) && max(abs(rowMeans(Z_balanced) - 0.5)) < 0.02,
  "First-order probabilities remain 1/2, and every draw treats exactly two units."
)
claim_cross(
  all(abs(sx_balanced - 5) < 1e-8),
  "The treated x-total is 5 on every draw."
)

## -----------------------------------------------------------------------------
x <- c(1, 2, 2, 3)
set.seed(20260822)
Z <- replicate(1000, balanced_ra(formula = ~ x))
sx <- colSums(x * Z)
table(sx)

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(
  all(colSums(Z) == 2) && max(abs(rowMeans(Z) - 0.5)) < 0.02,
  "First-order probabilities remain 1/2, and every draw treats exactly two units."
)

claim_cross(
  all(abs(sx - 4) < 1e-8),
  "The treated x-total is 4 on every draw."
)

## -----------------------------------------------------------------------------
# Draw a design many times, and compare the standard error an analyst would
# report against the standard deviation the estimator actually has. Passing x
# adjusts for it; leaving it NULL does not.
assess <- function(assign, p, y0, tau, nrep = 1000, x = NULL) {
  est <- se <- covered <- numeric(nrep)
  for (r in seq_len(nrep)) {
    Z <- as.numeric(as.character(assign()))
    Y <- y0 + tau * Z
    w <- Z / p + (1 - Z) / (1 - p)          # inverse-probability weights
    fit <- if (is.null(x)) {
      lm_robust(Y ~ Z, weights = w, se_type = "HC2")
    } else {
      lm_lin(Y ~ Z, covariates = ~ x, weights = w, se_type = "HC2")
    }
    est[r] <- fit$coefficients[["Z"]]
    se[r] <- fit$std.error[["Z"]]
    covered[r] <- fit$conf.low[["Z"]] <= tau && tau <= fit$conf.high[["Z"]]
  }
  c(true_sd = sd(est), mean_se = mean(se),
    ratio = mean(se) / sd(est), coverage = mean(covered))
}

## -----------------------------------------------------------------------------
set.seed(20260822)
N4 <- 200
p4 <- runif(N4, 0.2, 0.8)
y0 <- 3 * p4 + rnorm(N4)
tau <- 1

unweighted <- weighted <- numeric(1000)
for (r in 1:1000) {
  Z <- balanced_ra(prob_unit = p4, check_inputs = FALSE)
  Y <- y0 + tau * Z
  unweighted[r] <- mean(Y[Z == 1]) - mean(Y[Z == 0])
  weighted[r] <- lm_robust(Y ~ Z, weights = Z / p4 + (1 - Z) / (1 - p4),
                           se_type = "HC2")$coefficients[["Z"]]
}
rbind(unweighted = c(mean = mean(unweighted), bias = mean(unweighted) - tau),
      weighted   = c(mean = mean(weighted),   bias = mean(weighted)   - tau))

## ----echo=FALSE, results='asis'-----------------------------------------------
claim_tick(abs(mean(weighted) - tau) < 0.05 &&
             abs(mean(unweighted) - tau) > 0.2,
           "The weighted estimator recovers the true effect; the unweighted one does not.")

## -----------------------------------------------------------------------------
p6 <- c(0.2, 0.4, 0.6, 0.8, 0.5, 0.5)
Zb <- replicate(4000, balanced_ra(prob_unit = p6, check_inputs = FALSE))
Zs <- replicate(4000, simple_ra(N = 6, prob_unit = p6, check_inputs = FALSE))
mean_pair_cor <- function(S) { C <- cor(t(S)); mean(C[upper.tri(C)]) }
rbind(balanced = c(var_treated = var(colSums(Zb)), pair_cor = mean_pair_cor(Zb)),
      simple   = c(var_treated = var(colSums(Zs)), pair_cor = mean_pair_cor(Zs)))

## -----------------------------------------------------------------------------
rbind(
  balanced = assess(function() balanced_ra(prob_unit = p4, check_inputs = FALSE),
                    p4, y0, tau),
  simple   = assess(function() simple_ra(N = N4, prob_unit = p4,
                                         check_inputs = FALSE), p4, y0, tau)
)

## -----------------------------------------------------------------------------
x4 <- rnorm(N4)
p_half <- rep(0.5, N4)
y0_x <- 3 * x4 + rnorm(N4)

rbind(
  "balanced ~ x" = assess(function() balanced_ra(N = N4, formula = ~ x4,
                                                 check_inputs = FALSE),
                          p_half, y0_x, tau),
  "complete"     = assess(function() complete_ra(N = N4, check_inputs = FALSE),
                          p_half, y0_x, tau)
)

## -----------------------------------------------------------------------------
rbind(
  "balanced ~ x, adjusted" = assess(function() balanced_ra(N = N4, formula = ~ x4,
                                                           check_inputs = FALSE),
                                    p_half, y0_x, tau, x = x4),
  "complete, adjusted"     = assess(function() complete_ra(N = N4,
                                                           check_inputs = FALSE),
                                    p_half, y0_x, tau, x = x4)
)

## -----------------------------------------------------------------------------
y0_q <- 3 * x4^2 + rnorm(N4)
rbind(
  unadjusted = assess(function() balanced_ra(N = N4, formula = ~ x4,
                                             check_inputs = FALSE),
                      p_half, y0_q, tau),
  adjusted   = assess(function() balanced_ra(N = N4, formula = ~ x4,
                                             check_inputs = FALSE),
                      p_half, y0_q, tau, x = x4)
)

## -----------------------------------------------------------------------------
set.seed(5)
d_probs <- declare_ra(N = 6, prob_unit = c(0.2, 0.4, 0.6, 0.8, 0.5, 0.5),
                      ra_type = "balanced")
table(treated = replicate(500, sum(conduct_ra(d_probs))))

x5 <- rnorm(20)
d_formula <- declare_ra(N = 20, formula = ~ x5)
table(treated = replicate(500, sum(conduct_ra(d_formula))))

## -----------------------------------------------------------------------------
Z5 <- conduct_ra(d_probs)
cbind(Z = Z5, prob = obtain_condition_probabilities(d_probs, Z5))