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

# 1. Load the births and draw the graph: every confounder affects both columns.
births = load_birthweight()
smoking, weight = births.treatment, births.outcome
edges = [f"{c} -> {smoking}" for c in births.confounders]
edges += [f"{c} -> {weight}" for c in births.confounders]
graph = dag([*edges, f"{smoking} -> {weight}"])

# 2. Estimate the effect once.
estimate = estimate_effect(
    graph, births.data, smoking, weight, treated=births.treated, random_state=0
)
print(estimate)

# 3. Run the three refutations. Pass the same options you used for the estimate.
for method in ("placebo_treatment", "random_common_cause", "data_subset"):
    check = refute_effect(
        graph,
        births.data,
        smoking,
        weight,
        method=method,
        treated=births.treated,
        estimand_value=estimate.value,
        random_state=0,
    )
    print(check)
    print(f"  spread of the perturbed estimates: {check.refuted_std:.1f} g")

# 4. Refute an estimate that leaves out every confounder.
naive = dag([f"{smoking} -> {weight}"])
check = refute_effect(
    naive, births.data, smoking, weight, treated=births.treated, random_state=0
)
print(check)