import numpy as np

from ergodic import dag, evaluate_graph, graph_posterior, tabular
from ergodic.graph import cpdag_of_dag

# 1. Linear-Gaussian data from a known DAG: x -> y -> w <- z.
rng = np.random.default_rng(0)
x = rng.normal(size=300)
y = 0.8 * x + rng.normal(size=300)
z = rng.normal(size=300)
w = 0.7 * y - 0.6 * z + rng.normal(size=300)
data = tabular({"x": x, "y": y, "z": z, "w": w})
truth = dag(["x -> y", "y -> w", "z -> w"])

# 2. The BGe posterior: exact, every DAG enumerated. Equivalent DAGs score the same,
#    so add up the mass of each equivalence class and show the heaviest classes.
posterior = graph_posterior(data, score="bge")
print(posterior)
classes = {}
for index, weight in enumerate(posterior.weights):
    cpdag = cpdag_of_dag(posterior.graph(index))
    mass, members = classes.get(cpdag, (0.0, 0))
    classes[cpdag] = (mass + weight, members + 1)
for cpdag, (mass, members) in sorted(classes.items(), key=lambda item: -item[1][0])[:3]:
    print(round(mass, 3), members, cpdag)

# 3. The true DAG ties with its Markov-equivalent twin; the class holds most mass.
result = evaluate_graph(truth, posterior)
print(round(result.credible_level, 3), round(result.mass_above, 3))
print(round(evaluate_graph(cpdag_of_dag(truth), posterior).class_probability.probability, 3))

# 4. BIC only approximates the marginal likelihood, and the result says so.
approximate = graph_posterior(data, score="bic")
print(approximate.diagnostics["approximate"], approximate.diagnostics["note"])