from ergodic import discover
from ergodic.datasets import load_asia
from ergodic.graph import cpdag_of_dag, dag

# 1. Sample 5,000 patients from the Asia network.
asia = load_asia(n=5_000, seed=0)
print(asia.data.variables)

# 2. Learn the graph with greedy equivalence search (GES).
result = discover(asia.data, method="ges")
print(result.method, result.score, result.n)
print(result.graph)

# 3. Build the class of graphs the data could reveal: the truth's CPDAG.
truth = cpdag_of_dag(dag(asia.target["edges"]))
found = {edge.glyph() for edge in result.graph.edges}
expected = {edge.glyph() for edge in truth.edges}

# 4. Compare edge by edge, marks included.
print("right:", len(found & expected), "of", len(expected))
print("missed:", sorted(expected - found))
print("extra:", sorted(found - expected))