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

# 1. Load the births and name the question: smoking's effect on birthweight.
births = load_birthweight()
smoking, weight = births.treatment, births.outcome


def graph(confounders, latent=()):
    edges = [f"{c} -> {smoking}" for c in confounders]
    edges += [f"{c} -> {weight}" for c in confounders]
    return dag([*edges, f"{smoking} -> {weight}"], latent=latent)


# 2. Assume every confounder is measured, and estimate.
full = estimate_effect(
    graph(births.confounders), births.data, smoking, weight,
    treated=births.treated, random_state=0,
)
print("17 confounders:", full)

# 3. Assume there are none: the same data, another answer.
none = estimate_effect(
    graph([]), births.data, smoking, weight, treated=births.treated, random_state=0
)
print("no confounders:", none)

# 4. Assume a hidden common cause: no adjustment can remove it.
hidden = graph([*births.confounders, "U"], latent=["U"])
print(identify_effect(hidden, smoking, weight))