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


train = pd.DataFrame(
    {
        "age": [25, 31, 45, 22, 36, 52, 29, 41, 33, 48, 27, 39],
        "income": [
            52000,
            72000,
            68000,
            39000,
            85000,
            58000,
            np.nan,
            91000,
            62000,
            88000,
            45000,
            76000,
        ],
        "visits": [4, 7, 6, 2, 8, 3, 5, 9, 5, np.nan, 3, 6],
        "city": [
            "London",
            "Paris",
            "New York",
            "London",
            "Paris",
            "New York",
            "London",
            "Paris",
            "New York",
            "London",
            "Paris",
            "New York",
        ],
        "plan": [
            "basic",
            "premium",
            "standard",
            "basic",
            "premium",
            "standard",
            "standard",
            "premium",
            "standard",
            "premium",
            "basic",
            "premium",
        ],
        "converted": [0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1],
    }
)

holdout = pd.DataFrame(
    {
        "age": [30, 50, 34],
        "income": [64000, 56000, 79000],
        "visits": [5, 4, 7],
        "city": ["Rome", "London", "Paris"],
        "plan": ["premium", "basic", "standard"],
    }
)

numeric_features = ["age", "income", "visits"]
categorical_features = ["city", "plan"]

numeric_transformer = Pipeline(
    steps=[
        ("imputer", SimpleImputer(strategy="median")),
        ("scaler", StandardScaler()),
    ]
)

categorical_transformer = Pipeline(
    steps=[
        ("imputer", SimpleImputer(strategy="most_frequent")),
        ("encoder", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
    ]
)

preprocessor = ColumnTransformer(
    transformers=[
        ("numeric", numeric_transformer, numeric_features),
        ("categorical", categorical_transformer, categorical_features),
    ],
    remainder="drop",
)

model = Pipeline(
    steps=[
        ("preprocess", preprocessor),
        ("classifier", LogisticRegression(max_iter=1000)),
    ]
)

X_train = train.drop(columns="converted")
y_train = train["converted"]
X_holdout = holdout

model.fit(X_train, y_train)

transformed_holdout = model.named_steps["preprocess"].transform(X_holdout)
feature_names = model.named_steps["preprocess"].get_feature_names_out()
predictions = model.predict(X_holdout)

print(f"scikit-learn {sklearn.__version__}")
print(
    "Transformed holdout shape: "
    f"{transformed_holdout.shape[0]} rows x {transformed_holdout.shape[1]} columns"
)
print()
print("Feature names:")
for name in feature_names:
    print(f"- {name}")
print()
print("Holdout categories:")
print(X_holdout[["city", "plan"]].to_string(index=False))
print()
print(f"Predictions: {predictions.tolist()}")
print(f"Pipeline accepted {len(X_holdout)} holdout rows")