#!/usr/bin/env Rscript
# -----------------------------------------------------------------------------
# Script 03: Double/Debiased Machine Learning (DML) for ATE
# Atlas: Causality Atlas — Ch 10 (Causal Machine Learning)
# -----------------------------------------------------------------------------
# Demonstrates DML with Neyman orthogonality and cross-fitting.
# -----------------------------------------------------------------------------

library(ranger); library(xgboost)

cat("=== Double/Debiased Machine Learning for ATE ===\n\n")
set.seed(42)
N <- 2000
K_dim <- 10

# High-dimensional confounders
X <- matrix(rnorm(N * K_dim), N, K_dim)

# Nuisance functions (complex, nonlinear)
ps <- plogis(-1 + 0.3*X[,1] + 0.5*X[,2] - 0.4*X[,3]^2 + 0.2*sin(X[,4]))
A <- rbinom(N, 1, ps)

# Outcome: ATE = 1.5
mu <- 2 + 0.3*X[,1] - 0.5*X[,2] - 0.2*sin(X[,3]) + 0.1*X[,4]^2
Y <- 1.5 * A + mu + rnorm(N, 0, 0.5)

# Naive estimate
naive <- mean(Y[A == 1]) - mean(Y[A == 0])
cat(sprintf("Naive ATE = %.3f (confounded)\n", naive))

# ---- DML with Cross-Fitting ----
cat("\n--- DML with Cross-Fitting (ATE) ---\n")
K <- 5
folds <- sample(rep(1:K, length.out = N))
ate_folds <- numeric(K)

for (k in 1:K) {
  test_idx <- which(folds == k)
  train_idx <- which(folds != k)

  X_train <- X[train_idx, ]; X_test <- X[test_idx, ]
  A_train <- A[train_idx];   A_test <- A[test_idx]
  Y_train <- Y[train_idx];   Y_test <- Y[test_idx]

  # ML nuisance: propensity score
  ps_data <- data.frame(A = A_train, X_train)
  rf_ps <- ranger(as.factor(A) ~ ., data = ps_data, probability = TRUE)
  e_hat <- predict(rf_ps, data.frame(X_test))$predictions[, 2]

  # ML nuisance: outcome regression
  y_data <- data.frame(Y = Y_train, A = A_train, X_train)
  rf_mu <- ranger(Y ~ ., data = y_data)
  mu_hat <- predict(rf_mu, data.frame(Y = NA, A = A_test, X_test))$predictions

  # Compute separate mu_0 and mu_1 via one model with A as feature
  # Then the score function
  psi <- (A_test * (Y_test - mu_hat) / e_hat -
          (1 - A_test) * (Y_test - mu_hat) / (1 - e_hat))
  ate_folds[k] <- mean(psi) + mean(mu_hat)
}

dml_ate <- mean(ate_folds)
cat(sprintf("DML ATE = %.3f\n", dml_ate))
cat(sprintf("True ATE = 1.500\n"))

cat("\n--- Neyman Orthogonality ---\n")
cat("The DML score function is Neyman-orthogonal:\n")
cat("∂E[ψ(W; τ, η)]/∂η|_{η=η₀} = 0\n")
cat("Small errors in nuisance functions do NOT bias τ̂.\n")

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