from ergodic import estimate_effect, identify_effect
from ergodic.datasets import load_birthweight
from ergodic.welt import WeltLearner, discover

# 1. Take 400 births and five continuous columns.
births = load_birthweight()
columns = ["mage", "medu", "nprenatal", "bweight", "fage"]
data = births.data.select(columns).take(range(400))
print(data.shape)

# 2. Learn a DAG with Welt's Arrow model: one paid discovery job.
result = discover(data, model="arrow", seed=0)
for edge in sorted(edge.glyph() for edge in result.graph.edges):
    print(edge)
print(result.diagnostics["welt_model_version"])

# 3. Ask whether the graph identifies the effect of prenatal visits on birthweight.
identification = identify_effect(result.graph, "nprenatal", "bweight")
print(identification)

# 4. Estimate it with DML and a Welt tabular model in each nuisance slot.
learner = WeltLearner("tabicl-v2", max_fits=4)
estimate = estimate_effect(
    result.graph, data, "nprenatal", "bweight",
    strategy="dml", learner=learner, n_folds=2, random_state=0,
)
print(estimate)
print("paid fits:", learner.fits_used)