import numpy as np

from ergodic import (
    EdgeProbabilities,
    adjustment_probability,
    average_effect,
    graph_posterior,
    tabular,
)

# 1. Data where c confounds a binary treatment t and the outcome y; the effect is 2.
rng = np.random.default_rng(0)
c = rng.normal(size=2000)
t = (0.9 * c + rng.normal(size=2000) > 0).astype(float)
y = 2.0 * t + 1.5 * c + rng.normal(size=2000)
data = tabular({"c": c, "t": t, "y": y})

# 2. A model's pair-state probabilities: sure about t -> y and c -> y, unsure about c, t.
names = ("c", "t", "y")
values = np.array(
    [
        [0.00, 0.60, 0.95],
        [0.15, 0.00, 0.90],
        [0.02, 0.05, 0.00],
    ]
)
posterior = graph_posterior(EdgeProbabilities(names, values, kind="pair_state"))

# 3. How likely is each adjustment set to be valid?
print(round(adjustment_probability(posterior, "t", "y", ["c"]).probability, 3))
print(round(adjustment_probability(posterior, "t", "y", []).probability, 3))

# 4. The effect, estimated once per adjustment set and averaged by posterior mass.
result = average_effect(posterior, data, "t", "y")
print(result.to_frame()[["adjustment_set", "mass", "n_dags", "estimate", "std_error"]].round(3))
print(round(result.mean, 3), round(result.std, 3))
print(round(result.zero_effect_mass, 3), result.unidentified_mass, result.truncated_mass)