from ergodic import dag, estimate_effect
from ergodic.datasets import load_birthweight

# 1. Load 4,642 births: the mother's smoking, the birthweight and 17 confounders.
births = load_birthweight()
smoking, weight = births.treatment, births.outcome
print(smoking, "->", weight, "with", len(births.confounders), "confounders")

# 2. Draw the graph: each confounder affects both smoking and birthweight.
edges = [f"{c} -> {smoking}" for c in births.confounders]
edges += [f"{c} -> {weight}" for c in births.confounders]
graph = dag([*edges, f"{smoking} -> {weight}"])

# 3. Compare the raw difference in mean birthweight.
rows = births.data.to_pandas()
means = rows.groupby(smoking, observed=True)[weight].mean()
print(f"raw difference: {means['smoker'] - means['nonsmoker']:.1f} g")

# 4. Estimate the average effect, adjusting for the confounders.
estimate = estimate_effect(
    graph, births.data, smoking, weight, treated=births.treated, random_state=0
)
print(estimate)
low, high = estimate.interval()
print(f"95% interval: {low:.0f} g to {high:.0f} g")