## ============================================================================
## Workshop CIPE | FGV EAESP | 19 de agosto de 2026
##
## Pedro Borrego
##
## Exercício R: dimensão da amostra para a estimar a comparação de uma média
## entre grupos de tratamento e controlo num estudo experimental aelatorizado.
##
## Caso de estudo fictício (inspirado em: Bruhn, M., Garber, G., Koyama, S., Zia, B. (2022).
## The Long-Term Impact of High School Financial Education: evidence from Brazil. WP 563.
## Banco Central do Brasil Working Papers)
## Um programa de educação financeira vai ser testado em N = 45.000 alunos da
## 3.ª série, distribuídos por M = 300 escolas (~150 alunos/escola). O resultado
## de interesse é a pontuação num teste de literacia financeira (escala 0-100,
## Pretende-se comparar a média do grupo de tratamento (recebe o programa)
## com a média do grupo de controlo (não recebe), i.e. testar H0: mu_T - mu_C = 0.
## 
##
## Alguns dados do artigo, usados para calibrar os parâmetros abaixo:
##  - "the financial education program caused a quarter of a standard
##     deviation improvement in student test scores on an SAT-like financial
##     proficiency test" -> delta = 0,25 * sigma.
##  - Tabela A1 do artigo (SD do grupo de controlo, escala 0-100): 14,808 a
##     15,374 consoante a amostra/vaga. Tabela 1 (proficiência na linha de
##     base, 2010): média=50,73, SD=15,02 no grupo de controlo. Vamos usar
##     sigma=15 (arredondado, dentro deste intervalo).
##  - Dimensão real do "cluster": uma turma do 11.º ano por escola
##     ("One eleventh grade class in each school participated"), com
##     ~25 000 alunos em 892 escolas -> m ~= 28 alunos/escola. É bastante
##     menor do que o m=150 usado no nosso caso fictício (número total das turmas 3.ª série), 
##     por isso a Parte B, abaixo, calcula também o cenário com m=28,
##     para mostrar o efeito de um cluster mais pequeno (mais realista se só uma turma participar).
##  - Importante: a randomização real do estudo foi ao nível da escola, por
##     pares emparelhados dentro de cada estado ("matched-pairs"): "schools
##     ... matched into pairs based on their pre-existing school- and
##     community-level characteristics. Within each pair, one school was
##     then randomly assigned as treatment" (p.9). O desenho amostral deste exercício
##     na Parte A é nível do aluno (como cenário hipotético/pedagógico para contraste)
##     e o da Parte B é por clusters mas sem emparelhamento
##
## Duas partes:
##   Parte A — Randomização ao nível do aluno ("unit-level randomisation")
##   Parte B — Randomização ao nível da escola ("cluster-level randomisation")
## Em cada parte: (1) cálculo analítico da dimensão da amostra;
##                (2) proposta de validação por simulação de Monte Carlo;
##                (3) gráfico de potência vs. n.
## ============================================================================


## -----------------------------------------------------------------------
## Carregar pacotes necessários
## -----------------------------------------------------------------------

library(ggplot2)   # para gráficos
# se não tiver instalado, usar install.packages("ggplot2") na primeira vez que 
# correr este script, mas depois já náão vair ser necessário instalar novamente

## ----------------------------------------------------------------------
## 0. Parâmetros do caso de estudo
## ----------------------------------------------------------------------

# Valor de inicialização para procedimentos aleatórios para reprodutibilidade das simulações
set.seed(2026)


alpha   <- 0.05          # nível de significância
power_target <- 0.80     # potência-alvo (1 - beta)
sigma   <- 15            # desvio-padrão do teste (0-100 pontos) 
delta   <- 0.25 * sigma  # efeito mínimo detetável (MDE)
rho     <- 0.15          # correlação intraclasse (ICC) para desempenho escolar.
                         # Não vem do artigo do Bruhn (que não reporta ICC para
                         # a proficiência financeira); Ajustar se tiver uma
                         # estimativa melhor.
m       <- 150           # alunos por escola no caso fictício (N/M = 45000/300,
                         # coorte inteira da 3.ª série
m_real  <- 28            # alunos por escola no desenho de Bruhn et al (2022).
                          
N_pop   <- 45000
M_pop   <- 300

z_alpha <- qnorm(1 - alpha/2)
z_beta  <- qnorm(power_target)

cat(sprintf("Parametros: sigma=%.2f, delta=%.2f, delta/sigma=%.3f, rho=%.2f, m=%d\n",
            sigma, delta, delta/sigma, rho, m))

## ----------------------------------------------------------------------
## PARTE A — Randomização ao nível do aluno (unit-level)
## NOTA: este é um cenário fictício para fins pedagógicos. O desenho
## real de Bruhn et al. (2022) randomizou ao nível da escola,
## por pares emparelhados, e não ao nível do aluno. Serve de contraponto
## para mostrar quanto custaria (em alunos) obter a mesma potência SE a
## randomização pudesse ser feita individualmente.
## ----------------------------------------------------------------------

## A.1 Fórmula analítica, alocação
##     1:1 entre tratamento e controlo, variâncias iguais:
##       n* = 2 (t_{beta} + t_{alfa/2})^2 (sigma/delta)^2      por braço
## Resolvemos de forma iterativa porque os graus de liberdade da t dependem
## do próprio n.
n_unit_from_t <- function(sigma, delta, alpha, power_target, n_start = 300) {
  n <- n_start
  for (i in 1:50) {
    df <- 2 * n - 2
    t_a <- qt(1 - alpha/2, df)
    t_b <- qt(power_target, df)
    n_new <- 2 * (t_b + t_a)^2 * (sigma/delta)^2
    if (abs(n_new - n) < 0.01) break
    n <- n_new
  }
  ceiling(n)
}
n_unit <- n_unit_from_t(sigma, delta, alpha, power_target)
cat(sprintf("\n[Parte A] Dimensao da amostra por braco (randomizacao ao nivel do aluno): n* = %d\n", n_unit))
cat(sprintf("[Parte A] Amostra total (2 bracos): N* = %d alunos\n", 2 * n_unit))

## A.2 Função de potência analítica (para o gráfico potência vs. n)
power_unit <- function(n_per_arm, sigma, delta, alpha) {
  df <- 2 * n_per_arm - 2
  t_a <- qt(1 - alpha/2, df)
  ncp <- delta / (sigma * sqrt(2/n_per_arm))
  # aproximação normal (equivalente à usada no artigo)
  pnorm(ncp - t_a) # aproxima 1-beta ~ P(Z > z_alpha/2 - delta/SE)
}

## A.3 Validação por simulação de Monte Carlo (versão sem clusters).
##     Gera dados assumindo que o efeito verdadeiro é delta, 
##     e conta em quantas simulações o teste rejeita H0.
simulate_power_unit <- function(n_per_arm, sigma, delta, alpha, S = 1000) {
  rejects <- logical(S)
  for (s in 1:S) {
    y_control   <- rnorm(n_per_arm, mean = 50, sd = sigma)       # media base =50 (ponto medio da escala 0-100)
    y_treatment <- rnorm(n_per_arm, mean = 50 + delta, sd = sigma)
    treat <- c(rep(0, n_per_arm), rep(1, n_per_arm))
    y <- c(y_control, y_treatment)
    fit <- lm(y ~ treat)
    p_value <- summary(fit)$coefficients["treat", "Pr(>|t|)"]
    rejects[s] <- p_value < alpha
  }
  mean(rejects)
}

sim_power_at_n <- simulate_power_unit(n_unit, sigma, delta, alpha, S = 1000)
cat(sprintf("[Parte A] Potencia simulada em n*=%d (S=1000 simulacoes): %.3f (alvo: %.2f)\n",
            n_unit, sim_power_at_n, power_target))

## A.4 Gráfico: potência vs. dimensão da amostra por braço
n_grid_A <- seq(20, 400, by = 20)
power_grid_A <- sapply(n_grid_A, power_unit, sigma = sigma, delta = delta, alpha = alpha)
df_A <- data.frame(n = n_grid_A, potencia = power_grid_A)

p_A <- ggplot(df_A, aes(x = n, y = potencia)) +
  geom_line(color = "#2C6E91", linewidth = 1) +
  geom_hline(yintercept = power_target, linetype = "dashed", color = "grey40") +
  geom_vline(xintercept = n_unit, linetype = "dashed", color = "grey40") +
  annotate("point", x = n_unit, y = power_target, color = "#B5443A", size = 3) +
  labs(title = "Parte A - Potencia vs. dimensao da amostra por braco (nivel do aluno)",
       subtitle = sprintf("sigma=%.2f, delta=%.2f, alfa=%.2f | n*=%d por braco", sigma, delta, alpha, n_unit),
       x = "Alunos por braco (n)", y = "Potencia (1-beta)") +
  theme_minimal(base_size = 12)
# Visualiza o gráfico
p_A


## ----------------------------------------------------------------------
## PARTE B — Randomização ao nível da escola (cluster-level)
## ----------------------------------------------------------------------

## B.1 Fórmula analítica com efeito de desenho n* = m*k* = 2 (t_{alfa/2}+t_{beta})^2 (sigma/delta)^2
##     (1 + (m-1)*rho), com m fixo (alunos/escola) e k* = número de escolas/braço.
deff <- 1 + (m - 1) * rho
n_cluster_total_from_t <- function(sigma, delta, alpha, power_target, m, rho, k_start = 30) {
  k <- k_start
  for (i in 1:50) {
    df <- 2 * k - 2   # graus de liberdade em funcao do numero de clusters
    t_a <- qt(1 - alpha/2, df)
    t_b <- qt(power_target, df)
    deff_local <- 1 + (m - 1) * rho
    n_total <- 2 * (t_b + t_a)^2 * (sigma/delta)^2 * deff_local
    k_new <- ceiling(n_total / m)
    if (k_new == k) break
    k <- k_new
  }
  k
}
k_cluster <- n_cluster_total_from_t(sigma, delta, alpha, power_target, m, rho)
n_cluster_per_arm <- k_cluster * m

cat(sprintf("\n[Parte B] Efeito de desenho (design effect) = 1+(m-1)*rho = %.2f\n", deff))
cat(sprintf("[Parte B] Escolas por braco: k* = %d  (=> %d escolas no total)\n", k_cluster, 2*k_cluster))
cat(sprintf("[Parte B] Alunos por braco: n* = %d  (=> %d alunos no total)\n", n_cluster_per_arm, 2*n_cluster_per_arm))
cat(sprintf("[Parte B] Comparacao direta com a Parte A: %d alunos/braco (aluno) vs. %d alunos/braco (escola) -- racio = %.1fx\n",
            n_unit, n_cluster_per_arm, n_cluster_per_arm / n_unit))

## B.1bis Cenario alternativo: m = 28 alunos/escola, o tamanho de cluster
## REAL do estudo de Bruhn et al. (uma turma do 11.º ano por escola, em vez
## da coorte inteira da 3.ª série). Um cluster mais pequeno reduz o efeito de
## desenho, mas exige mais escolas para o mesmo numero de alunos.
deff_real <- 1 + (m_real - 1) * rho
k_cluster_real <- n_cluster_total_from_t(sigma, delta, alpha, power_target, m_real, rho)
n_cluster_real_per_arm <- k_cluster_real * m_real
cat(sprintf("\n[Parte B - cenario real, m=%d alunos/escola, como em Bruhn et al.(2022)] Efeito de desenho = %.2f\n", m_real, deff_real))
cat(sprintf("[Parte B - cenario real] Escolas por braco: k* = %d  (=> %d escolas no total)\n", k_cluster_real, 2*k_cluster_real))
cat(sprintf("[Parte B - cenario real] Alunos por braco: n* = %d  (=> %d alunos no total)\n", n_cluster_real_per_arm, 2*n_cluster_real_per_arm))
cat(sprintf("[Parte B - cenario real] Nota: com clusters mais pequenos (uma turma), precisamos de mais escolas (%d vs. %d por braco)\n",
            k_cluster_real, k_cluster))
cat(sprintf("                          mas de MENOS alunos no total (%d vs. %d por braco) -- o efeito de desenho pesa menos.\n",
            n_cluster_real_per_arm, n_cluster_per_arm))

## B.2 Validação por simulação de Monte Carlo (com estrutura de cluster).
##       Erro-composto = componente ao nível da escola (rho) + componente ao nível do aluno
simulate_power_cluster <- function(k_per_arm, m, sigma, delta, rho, alpha, S = 500) {
  rejects <- logical(S)
  sigma_c <- sigma * sqrt(rho)        # desvio-padrao do componente "escola"
  sigma_u <- sigma * sqrt(1 - rho)    # desvio-padrao do componente "aluno"
  total_clusters <- 2 * k_per_arm
  for (s in 1:S) {
    school_id <- rep(1:total_clusters, each = m)
    treat_school <- rep(c(rep(0, k_per_arm), rep(1, k_per_arm)), each = m)
    school_error <- rep(rnorm(total_clusters, 0, sigma_c), each = m)
    student_error <- rnorm(total_clusters * m, 0, sigma_u)
    y <- 50 + delta * treat_school + school_error + student_error
    fit <- lm(y ~ treat_school)
    # erro-padrao clusterizado (aproximacao simples via agregacao ao nivel da escola,
    #para robustez ao ICC)
    school_means <- tapply(y, school_id, mean)
    school_treat <- tapply(treat_school, school_id, mean)
    fit_agg <- lm(school_means ~ school_treat)
    p_value <- summary(fit_agg)$coefficients["school_treat", "Pr(>|t|)"]
    rejects[s] <- p_value < alpha
  }
  mean(rejects)
}

sim_power_cluster_at_k <- simulate_power_cluster(k_cluster, m, sigma, delta, rho, alpha, S = 500)
cat(sprintf("[Parte B] Potencia simulada em k*=%d escolas/braco (S=500 simulacoes): %.3f (alvo: %.2f)\n",
            k_cluster, sim_power_cluster_at_k, power_target))

## B.3 Gráfico: potência vs. número de escolas por braço, para 3 valores de ICC
##     (ilustra a sensibilidade ao ICC)
power_cluster_analytic <- function(k_per_arm, m, sigma, delta, rho, alpha) {
  df <- 2 * k_per_arm - 2
  if (df < 1) return(NA)
  t_a <- qt(1 - alpha/2, df)
  deff_local <- 1 + (m - 1) * rho
  se <- sigma * sqrt(2 * deff_local / (k_per_arm * m))
  ncp <- delta / se
  pnorm(ncp - t_a)
}

k_grid <- seq(5, 100, by = 5)
rho_values <- c(0.05, 0.15, 0.30)
df_B <- do.call(rbind, lapply(rho_values, function(r) {
  data.frame(k = k_grid, rho = factor(r),
             potencia = sapply(k_grid, power_cluster_analytic, m = m, sigma = sigma,
                                delta = delta, rho = r, alpha = alpha))
}))

p_B <- ggplot(df_B, aes(x = k, y = potencia, color = rho)) +
  geom_line(linewidth = 1) +
  geom_hline(yintercept = power_target, linetype = "dashed", color = "grey40") +
  labs(title = "Parte B - Potencia vs. numero de escolas por braco (nivel da escola)",
       subtitle = sprintf("sigma=%.2f, delta=%.2f, m=%d alunos/escola, alfa=%.2f", sigma, delta, m, alpha),
       x = "Escolas por braco (k)", y = "Potencia (1-beta)", color = "ICC (rho)") +
  theme_minimal(base_size = 12)
# Visualiza o gráfico
p_B


## -----------------------------------------------------------------------
## Fim do script
## -----------------------------------------------------------------------