# Source: content/notes/ml/types-of-ml.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.tree import DecisionTreeRegressor

rng = np.random.default_rng(101)
n = 20000
X = rng.integers(0, 2, size=(n, 1))
treatment = rng.integers(0, 2, size=n)
untreated_probability = np.where(X[:, 0] == 0, 0.80, 0.30)
treated_probability = np.where(X[:, 0] == 0, 0.78, 0.10)
observed_probability = np.where(treatment == 1, treated_probability, untreated_probability)
outcome = rng.binomial(1, observed_probability)
control_model = DecisionTreeRegressor(max_depth=1, random_state=101).fit(
    X[treatment == 0], outcome[treatment == 0])
treated_model = DecisionTreeRegressor(max_depth=1, random_state=101).fit(
    X[treatment == 1], outcome[treatment == 1])
groups = np.array([[0], [1]])
risk = control_model.predict(groups)
benefit = risk - treated_model.predict(groups)
incremental_value = 100 * benefit - 5
np.testing.assert_allclose(risk, [0.80, 0.30], atol=0.025)
np.testing.assert_allclose(benefit, [0.02, 0.20], atol=0.035)
assert risk.argmax() == 0 and incremental_value.argmax() == 1
assert incremental_value[0] < 0 < incremental_value[1]
print("Untreated risk:", risk)
print("Estimated average churn reduction:", benefit)
print("Offer value:", incremental_value)