import numpy as np

from ergodic import DomainKnowledge, EdgeProbabilities, evaluate_graph, graph_posterior

# 1. Edge probabilities over six variables, as a model might return them.
rng = np.random.default_rng(1)
names = ("a", "b", "c", "d", "e", "f")
values = np.zeros((6, 6))
for i in range(6):
    for j in range(i + 1, 6):
        _, forward, backward = rng.dirichlet([2.0, 1.0, 0.5])
        values[i, j], values[j, i] = forward, backward
probs = EdgeProbabilities(names, values, kind="pair_state")

# 2. Without knowledge the target factorizes by pair: exact draws.
exact = graph_posterior(probs, seed=0)
print(exact.method, exact.diagnostics["engine"], exact.n_draws)

# 3. "a before f" with e and b free to sit anywhere is an ancestral constraint: MCMC.
knowledge = DomainKnowledge().order("a", "f").forbid("b", "c", confidence=0.9)
chained = graph_posterior(probs, knowledge=knowledge, seed=0)
diagnostics = chained.diagnostics
print(chained.method, diagnostics["engine"], diagnostics["converged"])
print(round(diagnostics["rhat_log_target"], 2), diagnostics["ess_log_target"] > 2000)
print(diagnostics["target"]["ancestral_constraints"])

# 4. A feature that is not one edge: is a an ancestor of f?
feature = chained.probability(lambda graph: "a" in graph.ancestors("f"))
print(round(feature.probability, 2), feature.mc_error < 0.02)

# 5. The best graph found, and its claims.
result = evaluate_graph(chained.best_graph, chained)
print(result.best_source, round(result.log_ratio_to_best, 3), round(result.credible_level, 3))