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

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import roc_auc_score

rng = np.random.default_rng(12)
n = 400
age = rng.normal(40, 11, n)
income = rng.lognormal(3.8, 0.4, n)
region = rng.choice(["north", "south", "west"], n).astype(object)
score = 0.10 * (age - 40) + 0.025 * (income - 45) + (region == "north")
y = (score + rng.normal(0, 0.5, n) > 0).astype(int)
age[rng.random(n) < 0.08] = np.nan
region[rng.random(n) < 0.05] = np.nan
X = pd.DataFrame({"age": age, "income": income, "region": region})
train, test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=4)
test = test.copy()
test.loc[test.index[:3], "region"] = "unseen"
numeric = Pipeline([("impute", SimpleImputer(strategy="median", add_indicator=True)),
                    ("scale", StandardScaler())])
categorical = Pipeline([("impute", SimpleImputer(strategy="most_frequent")),
    ("encode", OneHotEncoder(handle_unknown="ignore", sparse_output=True))])
preprocess = ColumnTransformer([
    ("numeric", numeric, ["age", "income"]),
    ("categorical", categorical, ["region"])])
model = Pipeline([("features", preprocess),
                  ("classifier", LogisticRegression(max_iter=500, random_state=4))])
cv_auc = cross_val_score(model, train, y_train, cv=3, scoring="roc_auc", n_jobs=1)
model.fit(train, y_train)
probability = model.predict_proba(test)[:, 1]
features = model.named_steps["features"]
names = features.get_feature_names_out()
assert len(names) == features.transform(test).shape[1]
assert np.isfinite(probability).all()
assert np.allclose(probability, model.predict_proba(test.copy())[:, 1])
assert roc_auc_score(y_test, probability) > 0.80
encoder = features.named_transformers_["categorical"].named_steps["encode"]
assert "unseen" not in encoder.categories_[0]
print("CV AUC", cv_auc, "holdout AUC", roc_auc_score(y_test, probability))
print("feature columns", names.tolist())