from ergodic import decode
from ergodic.datasets import load_birthweight
from ergodic.welt import discover

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

# 2. Run CDFM once and decode its edge scores with greedy pair-MAP.
result = discover(data, model="cdfm", decoder="greedy", seed=0)
for edge in sorted(edge.glyph() for edge in result.graph.edges):
    print(edge)
print(result.diagnostics["decoded"], result.diagnostics["threshold_source"])
print(result.diagnostics["parameters"]["threshold"], result.diagnostics["rejected_cycles"])

# 3. The edges CDFM drew itself, cycles and two-way pairs included.
print(sorted(result.diagnostics["welt_native_edges"]))

# 4. Decode the same scores another way, locally: no new run.
probabilities = result.diagnostics["edge_probabilities"]
print(probabilities)
print(decode(probabilities, "expected_count").graph)