import numpy as np

from ergodic import EdgeProbabilities, admg, dag, evaluate_graph, graph_posterior
from ergodic.graph import cpdag_of_dag, mag_of_dag, pag_of_mag

names = ("rain", "sprinkler", "wet", "slippery")
values = np.array(
    [
        [0.00, 0.45, 0.91, 0.06],
        [0.40, 0.00, 0.83, 0.05],
        [0.04, 0.07, 0.00, 0.88],
        [0.58, 0.02, 0.10, 0.00],
    ]
)
posterior = graph_posterior(EdgeProbabilities(names, values, kind="independent"))
graph = dag(["rain -> wet", "sprinkler -> wet", "wet -> slippery"], nodes=names)

# 1. A CPDAG: adjacencies, compelled directions and the equivalence class.
cpdag = evaluate_graph(cpdag_of_dag(graph), posterior)
print(cpdag.claims_frame()[["claim", "kind", "probability"]].round(3))

# 2. A PAG: marks read as ancestry; circles claim nothing.
pag = pag_of_mag(mag_of_dag(graph))
print(pag)
for claim in evaluate_graph(pag, posterior).claims:
    if claim.kind in ("ancestor", "not_ancestor"):
        print(f"{claim.claim:<35} {claim.probability:.3f}")

# 3. An ADMG: directed edges are claims; a hidden common cause is not assessable.
confounded = admg(["rain -> wet", "sprinkler -> wet", "wet -> slippery", "rain <-> slippery"])
result = evaluate_graph(confounded, posterior)
for claim in result.claims:
    shown = "not assessable" if claim.probability is None else f"{claim.probability:.3f}"
    print(f"{claim.claim:<20} {shown}")