#!/usr/bin/env Rscript
# -----------------------------------------------------------------------------
# Script 01: Fisherian Randomization Test
# Atlas: Causality Atlas — Chapter 03 (Fisher, Randomization, RCT)
# -----------------------------------------------------------------------------
# Demonstrates Fisher's exact test via the Lady Tasting Tea example and a
# general randomization test for a continuous outcome.
# -----------------------------------------------------------------------------

# ---- Part 1: Lady Tasting Tea (Fisher 1935) ----
cat("=== Fisher's Exact Test: Lady Tasting Tea ===\n")

# 8 cups: 4 milk-first, 4 tea-first. Lady identifies by taste.
# Under H0 (sharp null: no ability), all 4 correct guesses out of 8.

# Number of ways to choose 4 cups out of 8 (the "correct" ones):
n_choose_k <- choose(8, 4)
p_all_4_correct <- 1 / n_choose_k

cat(sprintf("P(correctly identify all 8 cups) = 1/%d = %.4f\n",
            n_choose_k, p_all_4_correct))
cat(sprintf("At any conventional α (0.05), we %s reject H0.\n\n",
            ifelse(p_all_4_correct < 0.05, "CAN", "CANNOT")))

# ---- Part 2: Randomization Test for Continuous Outcome ----
cat("=== Randomization Test: Treatment Effect ===\n")

set.seed(42)
N <- 40
control <- rnorm(N/2, mean = 5, sd = 2)
treated <- rnorm(N/2, mean = 6.5, sd = 2)
Y <- c(control, treated)
A <- c(rep(0, N/2), rep(1, N/2))

# Observed statistic
obs_diff <- mean(Y[A == 1]) - mean(Y[A == 0])
cat(sprintf("Observed ATE = %.3f\n", obs_diff))

# Randomization test under sharp null
# Under H0: all outcomes are fixed; only assignment is random
n_perm <- 5000
null_dists <- replicate(n_perm, {
  A_star <- sample(A)  # permute treatment labels
  mean(Y[A_star == 1]) - mean(Y[A_star == 0])
})

# Two-sided p-value
p_val <- mean(abs(null_dists) >= abs(obs_diff))
cat(sprintf("Randomization test p-value = %.4f\n", p_val))

# Inference
if (p_val < 0.05) {
  cat("Reject sharp null at α = 0.05: treatment has some effect.\n")
} else {
  cat("Fail to reject sharp null at α = 0.05.\n")
}

# ---- Part 3: Comparison with t-test ----
cat("\n=== Comparison: t-test ===\n")
tt <- t.test(Y ~ A)
cat(sprintf("Two-sample t-test p-value = %.4f\n", tt$p.value))

# ---- Part 4: Visualize ----
pdf("../outputs/figures/randomization_test.pdf", width = 8, height = 5)
hist(null_dists, breaks = 30, col = "lightblue",
     main = "Randomization Distribution (H0: no effect)",
     xlab = "Difference in means", ylab = "Frequency")
abline(v = obs_diff, col = "red", lwd = 2, lty = 2)
abline(v = -obs_diff, col = "red", lwd = 2, lty = 2)
legend("topright", legend = sprintf("Observed ATE = %.3f", obs_diff),
       col = "red", lty = 2, lwd = 2)
dev.off()
cat("Saved: randomization_test.pdf\n")

# ---- Part 5: Blocked randomization example ----
cat("\n=== Blocked (Stratified) Randomization ===\n")
# Blocks by sex (covariate)
blocks <- rep(c("M", "F"), each = N/2)

# Randomize within each block
A_blocked <- unlist(lapply(split(seq_len(N), blocks), function(idx) {
  sample(rep(0:1, each = length(idx)/2))
}))

# Analysis stratified by block
ate_blocked <- mean(Y[A_blocked == 1]) - mean(Y[A_blocked == 0])
cat(sprintf("Blocked ATE estimate = %.3f\n", ate_blocked))

# Variance reduction from blocking
var_simple <- var(Y[A == 1])/sum(A) + var(Y[A == 0])/sum(1-A)
var_blocked <- sum(sapply(split(data.frame(Y, A = A_blocked, B = blocks), blocks),
  function(d) {
    n <- nrow(d)
    v <- var(d$Y[d$A == 1])/sum(d$A) + var(d$Y[d$A == 0])/sum(1 - d$A)
    v * (n / N)^2
  }))
cat(sprintf("Variance (simple): %.3f\n", var_simple))
cat(sprintf("Variance (blocked): %.3f\n", var_blocked))
cat(sprintf("Efficiency gain: %.1f%%\n", 100 * (1 - var_blocked/var_simple)))

cat("\n=== Script complete ===\n")