import numpy as np
import pandas as pd
from ergodic import cate, dag, tabular

# 1. Simulate 2,000 patients. The true effect of the drug is 1 + 2 * need.
rng = np.random.default_rng(0)
n = 2000
age = rng.normal(0, 1, n)
need = rng.normal(0, 1, n)
drug = rng.binomial(1, 1 / (1 + np.exp(-age)))
true_effect = 1 + 2 * need
score = age + need + true_effect * drug + rng.normal(0, 1, n)
frame = pd.DataFrame({"age": age, "need": need, "drug": drug, "score": score})
data = tabular(frame, types={"drug": "categorical"})

# 2. Draw the graph: age confounds; need only changes the outcome.
graph = dag(["age -> drug", "age -> score", "need -> score", "drug -> score"])

# 3. Estimate one effect per patient, with an interval for each.
effects = cate(graph, data, "drug", "score", random_state=0)
print(effects.source, effects.kind, effects.covariates)
print(effects.to_frame().head(3).round(2))
error = np.sqrt(np.mean((effects.values - true_effect) ** 2))
print(f"mean effect {effects.values.mean():.2f}, error against the truth {error:.2f}")

# 4. Test for heterogeneity, then describe it.
summary = effects.explain()
print(f"heterogeneity p-value: {summary.calibration_test().heterogeneity_pvalue:.1e}")
for name, row in summary.blp().frame().items():
    print(f"{name:12} {row['coef']:5.2f}  [{row['ci_low']:.2f}, {row['ci_high']:.2f}]")
print("effect by predicted fifth:", summary.gates().group_effect.round(2))