"""第 10 章：在相同划分上比较直线、树和集成；合成数据。"""

import numpy as np
from sklearn.dummy import DummyRegressor
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import KFold, cross_val_score, train_test_split
from sklearn.tree import DecisionTreeRegressor


def make_data():
    rng = np.random.default_rng(2026)
    x = rng.uniform(0, 1, (240, 1))
    y = 2*x[:, 0] + 3*(x[:, 0] > 0.55) + rng.normal(0, 0.12, 240)
    return train_test_split(x, y, test_size=0.25, random_state=2026)


def models():
    return {
        "均值基线": DummyRegressor(),
        "线性回归": LinearRegression(),
        "深度3的树": DecisionTreeRegressor(max_depth=3, min_samples_leaf=5,
                                           random_state=2026),
        "随机森林": RandomForestRegressor(n_estimators=100, max_depth=6,
                                          min_samples_leaf=3, random_state=2026,
                                          n_jobs=1),
        "梯度提升": GradientBoostingRegressor(n_estimators=100, learning_rate=0.05,
                                             max_depth=2, min_samples_leaf=5,
                                             random_state=2026),
    }


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


def main():
    x_train, x_test, y_train, y_test = make_data()
    results = compare(x_train, y_train)
    for name, train_mae, scores, _ in results:
        print(f"{name}：训练 MAE={train_mae:.3f}，五折 MAE={scores.mean():.3f}")
    name, _, _, selected = min(results, key=lambda row: row[2].mean())
    errors = selected.predict(x_test)-y_test
    print(f"固定模型：{name}；测试 MAE={np.abs(errors).mean():.3f}")
    for label, mask in [("阈值附近", np.abs(x_test[:, 0]-0.55) < 0.08),
                        ("其余位置", np.abs(x_test[:, 0]-0.55) >= 0.08)]:
        print(f"{label}：{mask.sum()} 条，MAE={np.abs(errors[mask]).mean():.3f}")


if __name__ == "__main__":
    main()