from ergodic import DomainKnowledge, discover, identify_effect
from ergodic.datasets import load_asia
from ergodic.graph import dag

asia = load_asia(n=5_000, seed=0)
truth = dag(asia.target["edges"])

# 1. Search without knowledge. Two edges stay undirected.
plain = discover(asia.data, method="ges").graph
print(plain.kind.value, sum(edge.is_undirected for edge in plain.edges), "undirected")

# 2. Write down what you know: exposures, then diseases, then symptoms.
knowledge = (
    DomainKnowledge()
    .with_tiers([["asia", "smoke"], ["tub", "lung", "bronc"], ["either"], ["xray", "dysp"]])
    .require("asia", "tub")
)
print(knowledge)

# 3. Search again with the knowledge.
result = discover(asia.data, method="ges", knowledge=knowledge)
print(result.graph)
print("same edges as the truth:", result.graph.edges == truth.edges)

# 4. Re-face the graph as a DAG before you identify an effect.
graph = result.graph.as_kind("dag")
print(identify_effect(graph, "smoke", "dysp"))