"""第 8 章：交叉验证选择复杂度，最后才使用测试集。"""

import numpy as np
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import KFold, cross_val_score, train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler


def make_data():
    rng = np.random.default_rng(2026)
    x = rng.uniform(-3, 3, size=(60, 1))
    y = 2 + 0.8*x[:, 0] + 1.2*x[:, 0]**2 + rng.normal(0, 1, 60)
    return train_test_split(x, y, test_size=0.3, random_state=2026)


def make_model(degree):
    return make_pipeline(
        PolynomialFeatures(degree, include_bias=False),
        StandardScaler(), Ridge(alpha=1e-6),
    )


def compare(x_train, y_train):
    folds = KFold(n_splits=5, shuffle=True, random_state=2026)
    results = []
    for degree in [1, 2, 10]:
        model = make_model(degree)
        scores = -cross_val_score(model, x_train, y_train, cv=folds,
                                  scoring="neg_mean_absolute_error")
        model.fit(x_train, y_train)
        train_mae = mean_absolute_error(y_train, model.predict(x_train))
        results.append((degree, train_mae, scores, model))
    return results


def main():
    x_train, x_test, y_train, y_test = make_data()
    results = compare(x_train, y_train)
    print(f"开发集 {len(x_train)} 条；保留测试集 {len(x_test)} 条")
    for degree, train_mae, scores, _ in results:
        print(f"次数={degree}：训练 MAE={train_mae:.3f}，"
              f"五折 MAE 均值={scores.mean():.3f}，折间标准差={scores.std():.3f}")
    degree, _, _, selected = min(results, key=lambda row: row[2].mean())
    print(f"选择次数 {degree}，在完整开发集重新拟合")
    print(f"测试 MAE={mean_absolute_error(y_test, selected.predict(x_test)):.3f}")
    # 模拟部署输入范围改变；没有根据这组诊断数据重新选择模型。
    rng = np.random.default_rng(2029)
    shifted_x = rng.uniform(4, 6, (100, 1))
    shifted_y = (2 + 0.8*shifted_x[:, 0] + 1.2*shifted_x[:, 0]**2
                 + rng.normal(0, 1, 100))
    print(f"范围变为 [4,6] 后的诊断 MAE="
          f"{mean_absolute_error(shifted_y, selected.predict(shifted_x)):.3f}")


if __name__ == "__main__":
    main()