import numpy as np

from ergodic import DomainKnowledge, EdgeProbabilities, decode

# 1. Independent edge probabilities; three edges above 0.5 form a cycle.
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 default decoder: greedy pair-MAP.
result = decode(probs)
print(result)
print(result.graph)
print(result.diagnostics["rejected_cycles"], result.diagnostics["note"])

# 3. A lower threshold nominates more edges.
print(decode(probs, threshold=0.3).graph)

# 4. Exact MAP and the expected-count decoder.
print(decode(probs, "exact_map").graph)
print(decode(probs, "expected_count").graph)

# 5. Candidate DAGs from a seeded sampler, and their consensus.
sampled = decode(probs, "pair_sampler", samples=200, seed=0)
print(sampled)
print(sampled.graph)

# 6. Knowledge: nothing causes the sprinkler.
knowledge = DomainKnowledge().root("sprinkler")
print(decode(probs, threshold=0.3, knowledge=knowledge).graph)