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

import numpy as np
from scipy.special import expit
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, brier_score_loss, log_loss
from sklearn.model_selection import GridSearchCV
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from threadpoolctl import threadpool_limits

rng = np.random.default_rng(121)
groups = np.repeat(np.arange(100), 20)
time = groups + np.tile(np.arange(20), 100) / 100
X = rng.normal(size=(len(groups), 5))
group_effect = rng.normal(scale=0.3, size=100)
probability = expit(-1.7 + 1.8 * X[:, 0] - X[:, 1] + group_effect[groups])
y = rng.binomial(1, probability)
development = np.flatnonzero(groups < 55)
calibration = np.flatnonzero((groups >= 55) & (groups < 70))
policy = np.flatnonzero((groups >= 70) & (groups < 85))
test = np.flatnonzero(groups >= 85)
partitions = [development, calibration, policy, test]
for before, after in zip(partitions[:-1], partitions[1:]):
    assert set(groups[before]).isdisjoint(groups[after])
    assert time[before].max() + 0.3 < time[after].min()

inner_folds = []
for train_end, val_start, val_end in ((25, 26, 35), (35, 36, 45), (45, 46, 55)):
    train_ids = np.flatnonzero(groups[development] < train_end)
    val_ids = np.flatnonzero((groups[development] >= val_start) &
                             (groups[development] < val_end))
    assert set(groups[train_ids]).isdisjoint(groups[val_ids])
    assert time[train_ids].max() + 0.3 < time[val_ids].min()
    assert np.unique(y[val_ids]).size == 2
    inner_folds.append((train_ids, val_ids))
pipeline = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000, random_state=121))
search = GridSearchCV(pipeline, {"logisticregression__C": [0.1, 1.0, 10.0]},
                      cv=inner_folds, scoring="neg_log_loss", n_jobs=1)
with threadpool_limits(limits=1):
    search.fit(X[development], y[development])
    model = search.best_estimator_
    calibrator = LogisticRegression(C=100.0, max_iter=1000).fit(
        model.decision_function(X[calibration])[:, None], y[calibration])
def calibrated(indices):
    score = model.decision_function(X[indices])[:, None]
    return calibrator.predict_proba(score)[:, 1]
policy_probability = calibrated(policy)
fp_cost, fn_cost = 1.0, 6.0
thresholds = np.r_[0., np.unique(policy_probability), np.nextafter(1., 2.)]
costs = []
for threshold in thresholds:
    decisions = policy_probability >= threshold
    costs.append(fp_cost * np.sum(decisions & (y[policy] == 0)) +
                 fn_cost * np.sum(~decisions & (y[policy] == 1)))
threshold = thresholds[np.argmin(costs)]
test_probability = calibrated(test)
test_decisions = test_probability >= threshold
assert np.all((test_probability > 0) & (test_probability < 1))
assert average_precision_score(y[test], test_probability) > y[test].mean()
print("Chosen C:", search.best_params_, "policy threshold:", threshold)
print("Test prevalence/AP:", y[test].mean(), average_precision_score(y[test], test_probability))
print("Test log loss/Brier:", log_loss(y[test], test_probability),
      brier_score_loss(y[test], test_probability))
print("Test error cost:", fp_cost * np.sum(test_decisions & (y[test] == 0)) +
      fn_cost * np.sum(~test_decisions & (y[test] == 1)))