#!/usr/bin/env Rscript
# -----------------------------------------------------------------------------
# Script 02: Propensity Score — Estimation, Matching, IPW
# Atlas: Causality Atlas — Ch 05 (Neyman-Rubin Potential Outcomes)
# -----------------------------------------------------------------------------
# Demonstrates PS estimation, nearest-neighbor matching, IPW, and balance checks.
# -----------------------------------------------------------------------------

library(MatchIt); library(WeightIt); library(cobalt); library(MASS)

cat("=== Propensity Score Analysis ===\n")
set.seed(42)
N <- 1000

# Simulate: confounders affect both treatment and outcome
C1 <- rnorm(N); C2 <- rnorm(N)
C3 <- sample(0:1, N, replace = TRUE)

# Treatment assignment depends on confounders
logit_ps <- -1 + 0.5*C1 - 0.8*C2 + 0.6*C3
ps <- plogis(logit_ps)
A <- rbinom(N, 1, ps)

# Outcome depends on treatment + confounders
Y <- 2 + 1.5*A + 0.3*C1 - 0.5*C2 + 0.4*C3 + rnorm(N)

dat <- data.frame(Y, A, C1, C2, C3)

# Naive (unadjusted) estimate (biased due to confounding)
naive <- lm(Y ~ A, data = dat)
cat(sprintf("Naive ATE = %.3f (confounded)\n", coef(naive)["A"]))

# ---- 1. Propensity Score Estimation ----
ps_model <- glm(A ~ C1 + C2 + C3, data = dat, family = binomial)
dat$ps <- predict(ps_model, type = "response")
cat(sprintf("PS range: [%.3f, %.3f]\n", min(dat$ps), max(dat$ps)))

# ---- 2. Nearest-Neighbor Matching ----
cat("\n--- Nearest-Neighbor Matching (ATT) ---\n")
m_out <- matchit(A ~ C1 + C2 + C3, data = dat, method = "nearest",
                  ratio = 1, caliper = 0.2)
m_data <- match.data(m_out)
m_att <- lm(Y ~ A, data = m_data, weights = weights)
cat(sprintf("ATT (1:1 matching) = %.3f\n", coef(m_att)["A"]))

# ---- 3. Inverse Probability Weighting ----
cat("\n--- IPW (ATE) ---\n")
w_out <- weightit(A ~ C1 + C2 + C3, data = dat, method = "ps")
w_ipw <- lm(Y ~ A, data = dat, weights = w_out$weights)
cat(sprintf("ATE (IPW) = %.3f\n", coef(w_ipw)["A"]))

# Stabilized weights
dat$sw <- ifelse(dat$A == 1, mean(dat$A)/dat$ps, (1-mean(dat$A))/(1-dat$ps))
w_sw <- lm(Y ~ A, data = dat, weights = dat$sw)
library(sandwich); library(lmtest)
coef_sandwich <- coeftest(w_sw, vcov = vcovHC(w_sw, type = "HC3"))
cat(sprintf("ATE (Stabilized IPW) = %.3f (SE=%.3f)\n",
            coef(w_sw)["A"], coef_sandwich["A", "Std. Error"]))

# ---- 4. Doubly Robust (AIPW) ----
cat("\n--- Doubly Robust (AIPW) ---\n")
outcome_model <- glm(Y ~ A + C1 + C2 + C3, data = dat)
ps2 <- glm(A ~ C1 + C2 + C3, data = dat, family = binomial)
dat$ps2 <- predict(ps2, type = "response")
dat$mu1 <- predict(update(outcome_model, .~.), newdata = transform(dat, A=1), type = "response")
dat$mu0 <- predict(update(outcome_model, .~.), newdata = transform(dat, A=0), type = "response")

dr <- mean((dat$A * dat$Y - (dat$A - dat$ps2) * dat$mu1) / dat$ps2 -
           ((1 - dat$A) * dat$Y + (dat$A - dat$ps2) * dat$mu0) / (1 - dat$ps2))
cat(sprintf("AIPW ATE = %.3f\n", dr))

# ---- 5. Balance Assessment ----
cat("\n--- Covariate Balance ---\n")
bal_tab <- bal.tab(m_out, un = TRUE, thresholds = c(m = 0.1))
print(bal_tab$Balance)
cat("Mean absolute standardized difference:\n")
cat(sprintf("  Unadjusted: %.3f\n", mean(abs(bal_tab$Diff.Unadjusted))))
cat(sprintf("  Adjusted:   %.3f\n", mean(abs(bal_tab$Diff.Adjusted))))

# ---- 6. PS Distribution Plot ----
pdf("../outputs/figures/propensity_score.pdf", width=8, height=5)
par(mfrow = c(1, 2))
# Overlap
hist(dat$ps[dat$A==1], col=rgb(1,0,0,0.4), breaks=30,
     xlab="Propensity Score", main="PS Overlap", ylim=c(0, 80))
hist(dat$ps[dat$A==0], col=rgb(0,0,1,0.4), breaks=30, add=TRUE)
legend("topright", c("Treated","Control"), fill=c(rgb(1,0,0,0.4), rgb(0,0,1,0.4)))
# Love plot
love.plot(m_out, binary = "std", thresholds = c(m = 0.1),
          title = "Covariate Balance (Love Plot)")
dev.off()
cat("Saved: propensity_score.pdf\n")

cat("\n=== Script complete. Truth: ATE = 1.5 ===\n")