#!/usr/bin/env Rscript
# -----------------------------------------------------------------------------
# Script 05: Causal Forest for CATE (R)
# Atlas: Causality Atlas — Ch 10 (Causal Machine Learning)
# -----------------------------------------------------------------------------

library(grf); library(ggplot2)

cat("=== Causal Forest for CATE Estimation ===\n\n")
set.seed(42)
N <- 2000; K_dim <- 5
X <- matrix(rnorm(N * K_dim), N, K_dim)
# Confounded treatment
p <- plogis(0.5 * X[,1] - 0.3 * X[,2])
A <- rbinom(N, 1, p)
# Heterogeneous treatment effect: τ(X) = 1 + 2*X[,1]
tau <- 1 + 2 * X[,1]
Y <- tau * A + sin(X[,2]) + 0.5 * X[,3] + rnorm(N)

# Causal forest (honest, asymptotic inference)
cf <- causal_forest(X, Y, A, num.trees = 2000)
ate <- average_treatment_effect(cf)
cat(sprintf("ATE estimate = %.3f (SE = %.3f)\n", ate[1], ate[2]))
cat(sprintf("95%% CI: [%.3f, %.3f]\n", ate[1] - 1.96*ate[2], ate[1] + 1.96*ate[2]))

# CATE predictions
cate_hat <- predict(cf, estimate.variance = TRUE)
tau_hat <- cate_hat$predictions
tau_se <- sqrt(cate_hat$variance.estimates)

# Evaluation: calibration
# Best Linear Predictor test
blp <- best_linear_projection(cf)
cat("\nBest Linear Projection (CATE heterogeneity test):\n")
print(blp)

# Variable importance
varimp <- variable_importance(cf)
cat("\nTop variables by importance:\n")
for (i in order(varimp, decreasing = TRUE)[1:3]) {
  cat(sprintf("  X[,%d]: %.3f\n", i, varimp[i]))
}

# Plot: true CATE vs estimated CATE
pdf("../outputs/figures/causal_forest_cate.pdf", width=8, height=6)
plot(tau, tau_hat, pch=16, cex=0.5, col=rgb(0,0,1,0.3),
     xlab="True CATE", ylab="Estimated CATE",
     main="Causal Forest: True vs Estimated CATE")
abline(0, 1, col="red", lwd=2)
legend("topleft", sprintf("Cor = %.3f", cor(tau, tau_hat)), bty="n")
dev.off()
cat("Saved: causal_forest_cate.pdf\n")

# Rank correlation of CATE estimates
cat(sprintf("Rank correlation: %.3f\n", cor(tau, tau_hat, method = "spearman")))
cat("\n=== Script complete. True ATE = %.2f ===\n", mean(tau))