#!/usr/bin/env python3
"""
Script 01: Double/Debiased Machine Learning (DML) for ATE (Python)
Atlas: Causality Atlas — Ch 10 (Causal Machine Learning)
Demonstrates DML with Neyman orthogonality and cross-fitting.
"""
import numpy as np
from sklearn.linear_model import RidgeClassifier, Ridge
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.model_selection import KFold

print("=== Double/Debiased Machine Learning for ATE ===\n")
np.random.seed(42)
N, K_dim = 2000, 10

# High-dimensional confounders
X = np.random.normal(0, 1, (N, K_dim))

# Complex nuisance functions
def pscore(x):
    lp = -1 + 0.3*x[:,0] + 0.5*x[:,1] - 0.4*x[:,2]**2 + 0.2*np.sin(x[:,3])
    return 1 / (1 + np.exp(-lp))

ps = pscore(X)
A = np.random.binomial(1, ps)
mu = 2 + 0.3*X[:,0] - 0.5*X[:,1] - 0.2*np.sin(X[:,2]) + 0.1*X[:,3]**2
Y = 1.5 * A + mu + 0.5 * np.random.randn(N)

# Naive estimate
print(f"Naive ATE = {Y[A==1].mean() - Y[A==0].mean():.3f} (confounded)")

# ---- DML with Cross-Fitting ----
print("\n--- DML with Cross-Fitting ---")
K = 5
kf = KFold(K, shuffle=True, random_state=42)
ate_folds = []

for train_idx, test_idx in kf.split(np.arange(N)):
    X_tr, X_te = X[train_idx], X[test_idx]
    A_tr, A_te = A[train_idx], A[test_idx]
    Y_tr, Y_te = Y[train_idx], Y[test_idx]

    # Nuisance: propensity score
    ml_m = GradientBoostingClassifier(n_estimators=100, max_depth=3)
    ml_m.fit(X_tr, A_tr)
    e_hat = ml_m.predict_proba(X_te)[:, 1]

    # Nuisance: outcome regression (with treatment as feature)
    ml_g = GradientBoostingRegressor(n_estimators=100, max_depth=3)
    ml_g.fit(np.column_stack([X_tr, A_tr]), Y_tr)
    mu_hat = ml_g.predict(np.column_stack([X_te, A_te]))

    # Neyman-orthogonal score
    psi = (A_te * (Y_te - mu_hat) / e_hat -
           (1 - A_te) * (Y_te - mu_hat) / (1 - e_hat))
    ate_folds.append(psi.mean())

dml_ate = np.mean(ate_folds)
print(f"DML ATE = {dml_ate:.3f}")
print(f"True ATE = 1.500")

print("\n--- Neyman Orthogonality Note ---")
print("ψ(Y, A, X; τ, η) = A(Y - μ₁(X))/e(X) - (1-A)(Y - μ₀(X))/(1-e(X)) + μ₁ - μ₀")
print("∂E[ψ]/∂η|η₀ = 0 → first-order errors in η do not bias τ̂")

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