"""第 9 章：把数值、缺失值和类别一起放进训练流程。"""

import numpy as np
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler


def data():
    # 列依次为页数、队列长度、打印模式；全部为教学构造数据。
    train_x = np.array([
        [5, 0, "mono"], [10, 2, "color"], [15, np.nan, "mono"],
        [20, 1, "color"], [np.nan, 3, "mono"], [30, 2, "color"],
    ], dtype=object)
    train_y = np.array([35, 64, 69, 92, 101, 124])
    validation_x = np.array([[12, 1, "mono"], [18, np.nan, "color"],
                             [24, 2, "draft"]], dtype=object)
    validation_y = np.array([58, 88, 96])
    return train_x, validation_x, train_y, validation_y


def make_model():
    numeric = make_pipeline(
        SimpleImputer(strategy="median", add_indicator=True),
        StandardScaler(),
    )
    categorical = make_pipeline(
        SimpleImputer(strategy="most_frequent"),
        OneHotEncoder(handle_unknown="ignore", sparse_output=False),
    )
    preprocessing = ColumnTransformer([
        ("numeric", numeric, [0, 1]),
        ("category", categorical, [2]),
    ])
    return make_pipeline(preprocessing, Ridge(alpha=1.0))


def main():
    train_x, validation_x, train_y, validation_y = data()
    model = make_model().fit(train_x, train_y)
    preprocessing = model.steps[0][1]
    numeric = preprocessing.named_transformers_["numeric"]
    medians = numeric.steps[0][1].statistics_
    transformed = preprocessing.transform(validation_x)
    predictions = model.predict(validation_x)
    print("训练集填补中位数：" + ", ".join(f"{x:.1f}" for x in medians))
    print(f"原始列数={train_x.shape[1]}，转换后列数={transformed.shape[1]}")
    print("新类别 draft 的独热部分：" + str(transformed[-1, -2:].tolist()))
    print("验证预测：" + ", ".join(f"{x:.2f}" for x in predictions))
    print(f"验证 MAE={mean_absolute_error(validation_y, predictions):.3f} 秒")


if __name__ == "__main__":
    main()