import numpy as np

from ergodic import DomainKnowledge, EdgeProbabilities, decode, identify_effect

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")

# 1. A soft statement: you're fairly sure nothing causes the sprinkler.
print(decode(probs, threshold=0.3).graph)
soft = DomainKnowledge().root("sprinkler", confidence=0.7)
print(decode(probs, threshold=0.3, knowledge=soft).graph)

# 2. A statement the model disagrees with: rain doesn't wet the ground.
doubtful = soft.forbid("rain", "wet", confidence=0.8)
result = decode(probs, threshold=0.3, knowledge=doubtful)
print(result.graph)

# 3. The audit: what the model makes of each statement.
for entry in result.diagnostics["knowledge_audit"]:
    print(
        f"{entry['statement']:<20} confidence {entry['confidence']}, "
        f"model {entry['probability']:.2f}, holds {entry['holds']}, "
        f"without knowledge {entry['holds_without_knowledge']}"
    )
print([entry["statement"] for entry in result.diagnostics["knowledge_conflicts"]])
print(round(result.diagnostics["knowledge_cost"], 3))

# 4. A hidden common cause of the sprinkler and slipperiness: an ADMG.
confounded = decode(probs, knowledge=DomainKnowledge().confound("sprinkler", "slippery"))
print(confounded)
print(confounded.graph)
print(identify_effect(confounded.graph, "sprinkler", "slippery"))