import numpy as np

from ergodic import EdgeProbabilities, dag, decode, evaluate_graph, graph_posterior

# 1. The same four variables as the decoder guide, read as independent marginals.
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],
    ]
)
probs = EdgeProbabilities(names, values, kind="independent")

# 2. The posterior over DAGs under this model: exact, every DAG enumerated.
posterior = graph_posterior(probs)
print(posterior)
print(posterior.edge_probabilities().to_frame().round(3))

# 3. Evaluate the greedy decode.
decoded = decode(probs).graph
result = evaluate_graph(decoded, posterior)
print(result)
print(round(result.credible_level, 3), round(result.log_ratio_to_best, 3))
print(round(result.expected_shd, 2), round(result.expected_shd_between_draws, 2))
print(round(result.p_acyclic, 3))
print(result.claims_frame()[["claim", "probability"]].round(3))

# 4. A graph you might have drawn yourself, with slipperiness causing rain.
mine = dag(["rain -> wet", "sprinkler -> wet", "slippery -> rain"], nodes=names)
other = evaluate_graph(mine, posterior)
print(round(other.credible_level, 3), round(other.log_ratio_to_best, 2))