from ergodic.graph import cpdag, dag
from ergodic.identification import (
    adjustment_set,
    frontdoor_set,
    ida,
    instruments,
    is_instrument,
    is_valid_adjustment_set,
    minimal_adjustment_set,
)

# 1. Confounders Z1, Z2, Z3, a mediator W and an instrument I.
graph = dag([
    "Z1 -> X", "Z1 -> Z3", "Z2 -> Z3", "Z2 -> Y", "Z3 -> X", "Z3 -> Y",
    "X -> W", "W -> Y", "I -> X",
])

# 2. The canonical adjustment set, and a minimal one.
print("canonical:", sorted(adjustment_set(graph, "X", "Y")))
print("minimal:", sorted(minimal_adjustment_set(graph, "X", "Y")))

# 3. Check sets of your own.
for candidate in [[], ["Z3"], ["Z1", "Z3"], ["Z3", "W"]]:
    print(candidate, is_valid_adjustment_set(graph, "X", "Y", candidate))

# 4. Instruments, then a check with a larger conditioning set.
print("instruments:", instruments(graph, "X", "Y"))
print("Z1 given Z2, Z3:", is_instrument(graph, "Z1", "X", "Y", given=["Z2", "Z3"]))

# 5. With the confounder hidden, adjustment fails but the front door is open.
hidden = dag(["U -> X", "U -> Y", "X -> M", "M -> Y"], latent=["U"])
print("adjustment:", adjustment_set(hidden, "X", "Y"))
print("front-door:", sorted(frontdoor_set(hidden, "X", "Y")))

# 6. On a CPDAG, identify the effect in each member graph.
result = ida(cpdag(["Z --- X", "X --- Y", "Z --- Y"]), "X", "Y")
print(result)
print(sorted(member.strategy for member in result.members))