from ergodic import discover, estimate_effect, identify_effect
from ergodic.datasets import load_earthquake

# 1. Load 10,000 rows sampled from the burglar-alarm network.
alarm = load_earthquake(n=10_000, seed=0)
rows = alarm.data.to_pandas()
print(rows.shape)
print((rows == "True").mean().round(3))

# 2. Learn the graph from the data alone.
graph = discover(alarm.data, method="pc").graph
for edge in sorted(edge.glyph() for edge in graph.edges):
    print(edge)

# 3. Ask whether the data can answer the question.
identification = identify_effect(graph, "Burglary", "JohnCalls")
print(identification)

# 4. Estimate the effect of a burglary on John calling.
estimate = estimate_effect(graph, alarm.data, "Burglary", "JohnCalls", treated="True")
print(estimate)