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

# 1. Simulate 4,000 people. The effect is +0.8 when need > 0 and age > 0, else negative.
rng = np.random.default_rng(0)
n = 4000
age = rng.normal(0, 1, n)
need = rng.normal(0, 1, n)
income = rng.normal(0, 1, n)
coached = rng.binomial(1, 1 / (1 + np.exp(-age)))
true_effect = 1.0 * (need > 0) + 1.0 * (age > 0) - 1.2
outcome = age + need + income + true_effect * coached + rng.normal(0, 1, n)
frame = pd.DataFrame(
    {"age": age, "need": need, "income": income, "coached": coached, "outcome": outcome}
)
data = tabular(frame, types={"coached": "categorical"})
graph = dag(
    [
        "age -> coached",
        "income -> coached",
        "age -> outcome",
        "need -> outcome",
        "income -> outcome",
        "coached -> outcome",
    ]
)

# 2. Split the rows in half: one half learns the rule, the other evaluates it.
learn = np.arange(n) < n // 2
train, test = data.take(np.flatnonzero(learn)), data.take(np.flatnonzero(~learn))

# 3. Learn the best depth-2 rule from the per-unit effects on the first half.
effects = cate(graph, train, "coached", "outcome", random_state=0, uncertainty=False)
tree = policy_tree(effects, depth=2)
print(tree.describe())
print(f"gain over treating no one, in sample: {tree.value:.3f}")
print(f"best possible gain, learning half: {np.mean(np.maximum(true_effect[learn], 0)):.3f}")

# 4. Evaluate the rule on the held-out half, with fresh doubly robust scores.
held_out = cate(graph, test, "coached", "outcome", random_state=0, uncertainty=False)
value = tree.evaluate(held_out)
low, high = value.interval(0.95)
print(f"gain over treating no one, held out: {value.value:.3f} [{low:.3f}, {high:.3f}]")
x_test = test.to_pandas()[list(tree.covariates)].to_numpy()
treat = tree.predict(x_test)
print(f"held-out share treated: {treat.mean():.2f}")

# 5. Check the rule against the truth on the held-out half.
truth = true_effect[~learn]
print(f"true gain: {np.mean(treat * truth):.3f}")
print(f"best possible gain, held-out half: {np.mean(np.maximum(truth, 0)):.3f}")