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

rng = np.random.default_rng(7)
n = 160
X = pd.DataFrame({
    "hours": rng.normal(100, 20, n),
    "temperature": rng.normal(35, 5, n),
    "model": rng.choice(["A", "B", "C"], n),
})
# A synthetic outcome, generated before measurements go missing.
signal = (X["hours"] + 3 * X["temperature"]
          + 10 * (X["model"] == "B") + rng.normal(0, 20, n))
y = (signal > 210).astype(int)
X.loc[::11, "hours"] = np.nan
X.loc[::17, "model"] = np.nan
print("Missing values:", X.isna().sum().to_dict())

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42, stratify=y
)
numeric = Pipeline([
    ("impute", SimpleImputer(strategy="median")),
    ("scale", StandardScaler()),
])
categorical = Pipeline([
    ("impute", SimpleImputer(strategy="most_frequent")),
    ("encode", OneHotEncoder(handle_unknown="ignore",
                             sparse_output=False)),
])
prepare = ColumnTransformer([
    ("numeric", numeric, ["hours", "temperature"]),
    ("categorical", categorical, ["model"]),
])
model = Pipeline([
    ("prepare", prepare),
    ("classifier", LogisticRegression(max_iter=1000)),
])
model.fit(X_train, y_train)
fitted = model.named_steps["prepare"]
train_ready = fitted.transform(X_train)
test_ready = fitted.transform(X_test)
assert train_ready.shape[1] == test_ready.shape[1]
assert np.isfinite(train_ready).all()
assert np.isfinite(test_ready).all()

print("Shapes:", train_ready.shape, test_ready.shape)
print("Features:", fitted.get_feature_names_out())
print("Test balanced accuracy:", round(balanced_accuracy_score(
    y_test, model.predict(X_test)), 3))

new_row = pd.DataFrame({"hours": [np.nan],
                        "temperature": [36.0], "model": ["D"]})
print("New row prediction:", model.predict(new_row))