from ergodic import dag
from ergodic.datasets import load_earthquake
from ergodic.graph import cpdag_of_dag, mag_of_dag, pag_of_mag


def show(title, graph):
    print(f"{title} ({graph.kind.value}):", ", ".join(sorted(e.glyph() for e in graph.edges)))


# 1. Build the burglar-alarm DAG from its published edges.
alarm = load_earthquake(n=100, seed=0)
graph = dag(alarm.target["edges"])
show("network", graph)

# 2. Ask separation questions.
print(graph.d_separated("Burglary", "Earthquake"))
print(graph.d_separated("Burglary", "Earthquake", "Alarm"))
print(graph.d_separated("JohnCalls", "MaryCalls", "Alarm"))

# 3. Intervene on the alarm: the edges into it are cut.
show("do(Alarm)", graph.do("Alarm"))

# 4. Read what data alone can tell apart: the equivalence class.
show("class", cpdag_of_dag(graph))
calls = graph.induced_subgraph({"Alarm", "JohnCalls", "MaryCalls"})
show("calls class", cpdag_of_dag(calls))

# 5. Hide the alarm, then read the class of what stays visible.
hidden = mag_of_dag(dag(alarm.target["edges"], latent=["Alarm"]))
show("alarm hidden", hidden)
show("its class", pag_of_mag(hidden))