"""模型对照：pip install scikit-learn；运行 python 04_models.py。"""
from random import Random
from time import perf_counter

from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import log_loss, roc_auc_score
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.tree import DecisionTreeClassifier

rng = Random(32)
X = [[rng.gauss(0, 1), rng.gauss(0, 1), rng.gauss(0, 1)] for _ in range(1200)]
y = [int((a > 0.2) == (b > -0.2)) ^ int(rng.random() < 0.08) for a, b, _ in X]
development, holdout = train_test_split(list(range(len(y))), random_state=4, stratify=y)
splits = list(StratifiedKFold(4, shuffle=True, random_state=5)
              .split(development, [y[i] for i in development]))
specs = [("linear", lambda: LogisticRegression(C=1)),
         ("shallow tree", lambda: DecisionTreeClassifier(max_depth=3, random_state=5)),
         ("histogram GBDT", lambda: HistGradientBoostingClassifier(max_iter=70, max_leaf_nodes=8,
                                                                  early_stopping=False, random_state=5))]

def fit(factory, ids):
    return factory().fit([X[i] for i in ids], [y[i] for i in ids])

for name, factory in specs:
    elapsed = perf_counter()
    scores = []
    for fit_idx, valid_idx in splits:
        train_ids = [development[i] for i in fit_idx]
        valid_ids = [development[i] for i in valid_idx]
        p = fit(factory, train_ids).predict_proba([X[i] for i in valid_ids])[:, 1]
        scores.append(roc_auc_score([y[i] for i in valid_ids], p))
    p = fit(factory, development).predict_proba([X[i] for i in holdout])[:, 1]
    true = [y[i] for i in holdout]
    near = [j for j, i in enumerate(holdout) if abs(X[i][0] - 0.2) < 0.4]
    print(name, "CV", [round(s, 3) for s in scores],
          "holdout AUC", round(roc_auc_score(true, p), 3), "logloss", round(log_loss(true, p), 3),
          "boundary AUC", round(roc_auc_score([true[j] for j in near], p[near]), 3),
          "seconds", round(perf_counter() - elapsed, 3))