import numpy as np
import pandas as pd

from ergodic import dag, tabular
from ergodic.identification import identify_mediation
from ergodic.inference import estimate_mediation


def expit(z):
    return 1 / (1 + np.exp(-z))


# 1. Simulate 2,000 workers. The direct effect of training is 1.0.
rng = np.random.default_rng(0)
n = 2_000
experience = rng.normal(size=n)
training = rng.binomial(1, expit(0.5 * experience))
certified = rng.binomial(1, expit(-1 + 2 * training + 0.5 * experience))
wage = 20 + 1.0 * training + 3.0 * certified + 1.5 * experience + rng.normal(0, 2, n)
frame = pd.DataFrame(
    {"experience": experience, "training": training, "certified": certified, "wage": wage}
)
data = tabular(frame, types={"training": "categorical", "certified": "categorical"})

# 2. Draw the graph and check that the split is identifiable.
graph = dag([
    "experience -> training", "experience -> certified", "experience -> wage",
    "training -> certified", "training -> wage", "certified -> wage",
])
print(identify_mediation(graph, "training", "wage", "certified"))

# 3. Estimate the direct and indirect effects.
split = estimate_mediation(
    graph, data, "training", "wage", "certified", treated=1, random_state=0, n_boot=200
)
for name, effect in [("direct", split.nde), ("indirect", split.nie), ("total", split.total)]:
    low, high = effect.interval()
    print(f"{name:8} {effect.value:.2f}  95% interval [{low:.2f}, {high:.2f}]")
print(f"proportion mediated: {split.proportion_mediated:.2f}")

# 4. The planted indirect effect: 3.0 times the rise in certification.
planted = 3.0 * np.mean(expit(1 + 0.5 * experience) - expit(-1 + 0.5 * experience))
print(f"planted: direct 1.00, indirect {planted:.2f}")

# 5. A confounder of certification and wage that training itself affects.
confounded = dag([*graph.edge_list(), "training -> motivation",
                  "motivation -> certified", "motivation -> wage"])
print(identify_mediation(confounded, "training", "wage", "certified"))