import numpy as np
import pandas as pd
from ergodic import SCM, admg, dag, tabular

# 1. Simulate 500 units. A hidden U drives X and Y; X acts on Y only through M.
rng = np.random.default_rng(0)
n = 500
u = rng.normal(size=n)
x = u + rng.normal(size=n)
m = 0.8 * x + rng.normal(size=n)
y = 1.5 * m + 2.0 * u + rng.normal(size=n)
data = tabular(pd.DataFrame({"X": x, "M": m, "Y": y}))
print(f"true effect of X on Y: {0.8 * 1.5:.2f}")
print(f"slope of Y on X:       {np.polyfit(x, y, 1)[0]:.2f}")

# 2. Draw the graph. X <-> Y says that a hidden common cause links X and Y.
graph = admg(["X -> M", "M -> Y", "X <-> Y"])

# 3. Fit one mechanism per node.
scm = SCM(graph, data, draws=500, tune=500, chains=1, random_state=0).fit()
for node, mechanism in scm.mechanisms.items():
    print(f"{node}: {type(mechanism).__name__}, parents {mechanism.parents}")

# 4. Ask for the average effect of a one-unit increase in X.
effect = scm.ate("X", "Y")
print(effect)

# 5. Drop the mediator from the graph. Now no formula identifies the effect.
bow = dag(["U -> X", "U -> Y", "X -> Y"], latent=["U"])
bounds = SCM(bow, data, draws=500, tune=500, chains=1, random_state=0).fit().ate("X", "Y")
low, high = bounds.range()
print(f"{type(bounds).__name__}: effect between {low:.2f} and {high:.2f}")
print(f"effect if the hidden path carries covariance 2.0: {bounds.at(2.0):.2f}")