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


numeric_features = ["age", "monthly_spend"]
categorical_features = ["plan"]

train = pd.DataFrame(
    {
        "age": [25, 32, np.nan, 45, 51, 28, np.nan, 39],
        "monthly_spend": [120, 210, 180, np.nan, 320, 150, 170, 260],
        "plan": [
            "basic",
            "premium",
            "basic",
            np.nan,
            "premium",
            "basic",
            "standard",
            np.nan,
        ],
        "converted": [0, 1, 0, 1, 1, 0, 0, 1],
    }
)

holdout = pd.DataFrame(
    {
        "age": [np.nan, 42],
        "monthly_spend": [190, np.nan],
        "plan": [np.nan, "premium"],
    }
)

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

preprocessor = ColumnTransformer(
    transformers=[
        ("numeric", numeric_imputer, numeric_features),
        ("categorical", categorical_transformer, categorical_features),
    ],
    verbose_feature_names_out=False,
)

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

X_train = train[numeric_features + categorical_features]
y_train = train["converted"]

model.fit(X_train, y_train)

fitted_preprocessor = model.named_steps["preprocess"]
fitted_numeric = fitted_preprocessor.named_transformers_["numeric"]
fitted_categorical = fitted_preprocessor.named_transformers_["categorical"]
fitted_category_imputer = fitted_categorical.named_steps["imputer"]

transformed_train = fitted_preprocessor.transform(X_train)
transformed_holdout = fitted_preprocessor.transform(holdout)
predictions = model.predict(holdout)

print(f"scikit-learn {sklearn.__version__}")
print("Input missing values:")
print(X_train.isna().sum().to_string())

print()
print("Numeric fill values:")
for name, value in zip(numeric_features, fitted_numeric.statistics_):
    print(f"- {name}: {value:.1f}")

print()
print("Categorical fill values:")
for name, value in zip(categorical_features, fitted_category_imputer.statistics_):
    print(f"- {name}: {value}")

print()
print("Transformed feature names:")
print(", ".join(fitted_preprocessor.get_feature_names_out()))

print()
print(
    "Transformed training shape: "
    f"{transformed_train.shape[0]} rows x {transformed_train.shape[1]} columns"
)
print(f"Missing values after transform: {int(np.isnan(transformed_train).sum())}")

print()
print("Transformed holdout rows:")
holdout_table = pd.DataFrame(
    transformed_holdout,
    columns=fitted_preprocessor.get_feature_names_out(),
).round(2)
print(holdout_table.to_string(index=False))
print(f"Holdout predictions: {predictions.tolist()}")