#!/usr/bin/env python3
"""
Script 05: Causal Forest for CATE (Python)
Atlas: Causality Atlas — Ch 10 (Causal Machine Learning)
"""
import numpy as np
import matplotlib.pyplot as plt
from econml.grf import CausalForest
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

print("=== Causal Forest for CATE Estimation ===\n")
np.random.seed(42)
N, K = 2000, 5
X = np.random.normal(0, 1, (N, K))

# Confounded treatment
p = 1 / (1 + np.exp(-(0.5*X[:,0] - 0.3*X[:,1])))
A = np.random.binomial(1, p)

# Heterogeneous effect: τ(X) = 1 + 2*X1
tau = 1 + 2 * X[:, 0]
Y = tau * A + np.sin(X[:, 1]) + 0.5*X[:, 2] + np.random.normal(0, 1, N)

# Causal Forest with asymptotic inference
cf = CausalForest(
    model_y=RandomForestRegressor(),
    model_t=RandomForestClassifier(),
    n_estimators=2000,
    min_samples_leaf=10,
    max_depth=20
)
cf.fit(X, A, Y)

# ATE
ate = cf.ate(X_test=None)
ate_ci = cf.ate_interval(X_test=None, alpha=0.05)
print(f"ATE = {ate:.3f} [95% CI: ({ate_ci[0]:.3f}, {ate_ci[1]:.3f})]")

# CATE
cate = cf.effect(X)
cate_cor = np.corrcoef(tau, cate)[0, 1]
print(f"CATE correlation with true: {cate_cor:.3f}")

if cate_cor > 0:
    plt.figure(figsize=(8, 6))
    plt.scatter(tau, cate, alpha=0.4, s=5)
    plt.plot([tau.min(), tau.max()], [tau.min(), tau.max()], 'r-', lw=2)
    plt.xlabel("True CATE")
    plt.ylabel("Estimated CATE")
    plt.title(f"Causal Forest: True vs Estimated CATE (r={cate_cor:.3f})")
    plt.tight_layout()
    plt.savefig("../outputs/figures/causal_forest_cate_py.png", dpi=150)

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