# If you are running this notebook in a fresh environment, uncomment the next line.
!pip install -q econml dowhy dice-ml lightgbm rdatasets

import warnings
warnings.filterwarnings("ignore")

import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from sklearn.linear_model import LogisticRegression
from lightgbm import LGBMRegressor, LGBMClassifier

# EconML estimators
from econml.metalearners import SLearner, TLearner, XLearner
from econml.dr import DRLearner
from econml.dml import LinearDML, CausalForestDML

np.random.seed(42)
pd.set_option("display.max_columns", 200)

url = "https://raw.githubusercontent.com/vincentarelbundock/Rdatasets/master/csv/MatchIt/lalonde.csv"
df = pd.read_csv(url)

# Keep only the useful columns
df = df.drop(columns=["rownames"], errors="ignore")

# Treatment and outcome naming consistent with the chapter
df = df.rename(columns={"treat": "treatment", "re78": "outcome"})

df.head()

df.describe(include="all").T

X = df.drop(columns=["treatment", "outcome"])
X = pd.get_dummies(X,drop_first=True)
T = df["treatment"].astype(int)
Y = df["outcome"].astype(float)

print("Rows:", len(df))
print("Treatment rate:", T.mean().round(4))
print("Average outcome:", Y.mean().round(2))
X.head()

# Check marginal treatment distribution
treatment_dist = T.value_counts(normalize=True).sort_index()
treatment_dist

X_train_eda, X_test_eda, T_train_eda, T_test_eda = train_test_split(
    X, T, test_size=0.5, random_state=42, stratify=T
)

clf_eda = LGBMClassifier(
    n_estimators=100,
    max_depth=4,
    learning_rate=0.05,
    verbosity=-1,
    random_state=42
)
clf_eda.fit(X_train_eda, T_train_eda)

T_pred_eda = clf_eda.predict(X_test_eda)
eda_accuracy = accuracy_score(T_test_eda, T_pred_eda)
eda_accuracy

p1 = T.mean()
n_test = len(T_test_eda)

random_scores = []
for _ in range(10000):
    random_pred = np.random.binomial(1, p1, size=n_test)
    random_scores.append((random_pred == T_test_eda.to_numpy()).mean())

ci_low, ci_high = np.quantile(random_scores, [0.025, 0.975])

print("Observed treatment-prediction accuracy :", round(eda_accuracy, 4))
print("95% empirical random-accuracy interval :", (round(ci_low, 4), round(ci_high, 4)))

plt.figure(figsize=(8, 4.5))
plt.hist(random_scores, bins=40, alpha=0.8)
plt.axvline(eda_accuracy, linestyle="--", linewidth=2, label="Model accuracy")
plt.title("Randomization check: model accuracy vs random baseline")
plt.xlabel("Accuracy")
plt.ylabel("Count")
plt.legend()
plt.show()

X_train, X_test, y_train, y_test, T_train, T_test = train_test_split(
    X, Y, T, test_size=0.5, random_state=42, stratify=T
)

print("Train rows:", len(X_train))
print("Test rows :", len(X_test))
print("Treatment rate train:", T_train.mean().round(4))
print("Treatment rate test :", T_test.mean().round(4))

def create_regressor():
    return LGBMRegressor(
        n_estimators=200,
        max_depth=4,
        learning_rate=0.05,
        subsample=0.9,
        colsample_bytree=0.9,
        verbosity=-1,
        random_state=42
    )

def create_classifier():
    return LGBMClassifier(
        n_estimators=200,
        max_depth=4,
        learning_rate=0.05,
        subsample=0.9,
        colsample_bytree=0.9,
        verbosity=-1,
        random_state=42
    )

s_learner = SLearner(overall_model=create_regressor())

t_learner = TLearner(models=[create_regressor(), create_regressor()])

x_learner = XLearner(
    models=[create_regressor(), create_regressor()],
    cate_models=[create_regressor(), create_regressor()],
)

dr_learner = DRLearner(
    model_propensity=LogisticRegression(max_iter=2000),
    model_regression=create_regressor(),
    model_final=create_regressor(),
    cv=5,
)

linear_dml = LinearDML(
    model_y=create_regressor(),
    model_t=create_classifier(),
    discrete_treatment=True,
    cv=5,
    random_state=42,
)

causal_forest = CausalForestDML(
    model_y=create_regressor(),
    model_t=create_classifier(),
    discrete_treatment=True,
    cv=5,
    random_state=42,
)

models = {
    "SLearner": s_learner,
    "TLearner": t_learner,
    "XLearner": x_learner,
    "DRLearner": dr_learner,
    "LinearDML": linear_dml,
    "CausalForestDML": causal_forest,
}

fit_times = {}

for model_name, model in models.items():
    start = time.time()
    if model_name in {"LinearDML", "CausalForestDML"}:
        model.fit(Y=y_train, T=T_train, X=X_train)
    else:
        model.fit(Y=y_train, T=T_train, X=X_train)
    stop = time.time()
    fit_times[model_name] = stop - start
    print(f"{model_name:<16} fitted in {fit_times[model_name]:.3f} seconds")

timing_df = pd.DataFrame({
    "Model": list(fit_times.keys()),
    "TimeSeconds": list(fit_times.values())
}).sort_values("TimeSeconds")

baseline = timing_df["TimeSeconds"].min()
timing_df["RelativeToFastest"] = (timing_df["TimeSeconds"] / baseline).round(1)

timing_df

plt.figure(figsize=(9, 4.5))
plt.bar(timing_df["Model"], timing_df["RelativeToFastest"])
plt.xticks(rotation=30, ha="right")
plt.ylabel("Relative training time")
plt.title("Relative compute cost across CATE estimators")
plt.show()

def cate_predict(model, X_data):
    return model.effect(X_data)

cate_train = {name: cate_predict(model, X_train) for name, model in models.items()}
cate_test  = {name: cate_predict(model, X_test)  for name, model in models.items()}

pd.DataFrame({k: np.asarray(v).ravel()[:5] for k, v in cate_test.items()})

def uplift_by_decile(y, t, uplift, n_bins=10):
    temp = pd.DataFrame({
        "y": np.asarray(y),
        "t": np.asarray(t).astype(int),
        "uplift": np.asarray(uplift).ravel()
    }).sort_values("uplift", ascending=False).reset_index(drop=True)

    temp["decile"] = pd.qcut(
        np.arange(len(temp)),
        q=n_bins,
        labels=False,
        duplicates="drop"
    )

    rows = []
    for decile, group in temp.groupby("decile"):
        treated = group.loc[group["t"] == 1, "y"]
        control = group.loc[group["t"] == 0, "y"]
        uplift_obs = treated.mean() - control.mean() if len(treated) and len(control) else np.nan
        rows.append({
            "decile": int(decile),
            "n": len(group),
            "treated_n": len(treated),
            "control_n": len(control),
            "observed_uplift": uplift_obs
        })
    return pd.DataFrame(rows)

def plot_uplift_by_decile(decile_df, title):
    plt.figure(figsize=(8, 4.5))
    plt.bar(decile_df["decile"], decile_df["observed_uplift"])
    plt.axhline(0, linewidth=1)
    plt.title(title)
    plt.xlabel("Decile (0 = highest predicted uplift)")
    plt.ylabel("Observed uplift")
    plt.show()

# Example: one model on train and test
example_train = uplift_by_decile(y_train, T_train, cate_train["XLearner"])
example_test  = uplift_by_decile(y_test,  T_test,  cate_test["XLearner"])

plot_uplift_by_decile(example_train, "XLearner — uplift by decile (train)")
plot_uplift_by_decile(example_test, "XLearner — uplift by decile (test)")

fig, axes = plt.subplots(len(models), 2, figsize=(12, 3.4 * len(models)))

for row_idx, (name, _) in enumerate(models.items()):
    train_df = uplift_by_decile(y_train, T_train, cate_train[name])
    test_df  = uplift_by_decile(y_test,  T_test,  cate_test[name])

    axes[row_idx, 0].bar(train_df["decile"], train_df["observed_uplift"])
    axes[row_idx, 0].axhline(0, linewidth=1)
    axes[row_idx, 0].set_title(f"{name} — Train")
    axes[row_idx, 0].set_xlabel("Decile")
    axes[row_idx, 0].set_ylabel("Obs uplift")

    axes[row_idx, 1].bar(test_df["decile"], test_df["observed_uplift"])
    axes[row_idx, 1].axhline(0, linewidth=1)
    axes[row_idx, 1].set_title(f"{name} — Test")
    axes[row_idx, 1].set_xlabel("Decile")
    axes[row_idx, 1].set_ylabel("Obs uplift")

plt.tight_layout()
plt.show()

def expected_response_binary(y, t, uplift_scores, p_treat=None):
    y = np.asarray(y)
    t = np.asarray(t).astype(int)
    uplift_scores = np.asarray(uplift_scores).ravel()

    if p_treat is None:
        p_treat = t.mean()

    p_control = 1 - p_treat
    recommended_treatment = (uplift_scores > 0).astype(int)

    weight = np.where(recommended_treatment == 1, 1 / p_treat, 1 / p_control)
    observed_if_followed = (recommended_treatment == t).astype(int)

    return np.mean(y * observed_if_followed * weight)

metric_rows = []
for name in models:
    train_er = expected_response_binary(y_train, T_train, cate_train[name], p_treat=T_train.mean())
    test_er  = expected_response_binary(y_test,  T_test,  cate_test[name],  p_treat=T_test.mean())
    metric_rows.append({
        "Model": name,
        "ExpectedResponse_Train": train_er,
        "ExpectedResponse_Test": test_er
    })

expected_response_df = pd.DataFrame(metric_rows).sort_values("ExpectedResponse_Test", ascending=False)
expected_response_df

plt.figure(figsize=(9, 4.5))
plt.bar(expected_response_df["Model"], expected_response_df["ExpectedResponse_Test"])
plt.xticks(rotation=30, ha="right")
plt.ylabel("Expected response (test)")
plt.title("Policy value comparison on the test set")
plt.show()

lb, ub = models["LinearDML"].effect_interval(X_test, T0=0, T1=1, alpha=0.05)

print("Lower bounds (first 5):", lb[:5])
print("Upper bounds (first 5):", ub[:5])

intervals = np.column_stack([lb, ub])
contains_zero = np.mean(np.sign(intervals[:, 0]) != np.sign(intervals[:, 1]))

print("Fraction of test observations whose 95% CI contains 0:", round(float(contains_zero), 4))

plt.figure(figsize=(8, 4.5))
sample_idx = np.arange(min(50, len(lb)))
plt.errorbar(
    sample_idx,
    cate_test["LinearDML"][:len(sample_idx)],
    yerr=[
        cate_test["LinearDML"][:len(sample_idx)] - lb[:len(sample_idx)],
        ub[:len(sample_idx)] - cate_test["LinearDML"][:len(sample_idx)]
    ],
    fmt='o'
)
plt.axhline(0, linewidth=1)
plt.title("LinearDML: point estimates with 95% confidence intervals")
plt.xlabel("Sample index")
plt.ylabel("Estimated treatment effect")
plt.show()

conservative_treat = (
    (np.asarray(cate_test["LinearDML"]).ravel() > 0) &
    (lb > 0)
).astype(int)

pd.Series(conservative_treat).value_counts(normalize=True).rename("share")

comparison = expected_response_df.merge(
    timing_df[["Model", "RelativeToFastest"]],
    on="Model",
    how="left"
).rename(columns={"RelativeToFastest": "RelativeComputeCost"})

comparison

# Create a recommendation label from the LinearDML uplift estimates
train_recommend = (np.asarray(cate_train["LinearDML"]).ravel() > 0).astype(int)
test_recommend  = (np.asarray(cate_test["LinearDML"]).ravel() > 0).astype(int)

recommend_train_df = X_train.copy()
recommend_train_df["recommend_treatment"] = train_recommend

recommend_test_df = X_test.copy()
recommend_test_df["recommend_treatment"] = test_recommend

recommend_train_df.head()

# Train a small interpretable recommendation model for DiCE
recommendation_model = LogisticRegression(max_iter=5000)
recommendation_model.fit(X_train, train_recommend)

print("Share recommended for treatment in train:", train_recommend.mean().round(3))
print("Share recommended for treatment in test :", test_recommend.mean().round(3))

# DiCE setup
import dice_ml
from dice_ml import Dice

dice_data = dice_ml.Data(
    dataframe=recommend_train_df,
    continuous_features=["age", "educ", "re74", "re75"],
    outcome_name="recommend_treatment",
)

dice_model = dice_ml.Model(model=recommendation_model, backend="sklearn", model_type="classifier")
dice = Dice(dice_data, dice_model, method="random")

# Pick one person the policy currently does NOT recommend for treatment
candidate_pool = recommend_test_df.copy()
candidate_pool["pred"] = recommendation_model.predict(X_test)

query_idx = candidate_pool.index[candidate_pool["pred"] == 0][0]
query_instance = X_test.loc[[query_idx]]

query_instance

cf = dice.generate_counterfactuals(
    query_instance,
    total_CFs=3,
    desired_class="opposite",
    features_to_vary=["age", "educ", "re74", "re75", "married", "nodegree"]
)

cf.visualize_as_dataframe(show_only_changes=True)