"""
.. _example_skd001_skd002_overfitting_underfitting:

SKD001 & SKD002 - Overfitting and underfitting
==============================================

:ref:`SKD002 <skd002-underfitting>` and :ref:`SKD001 <skd001-overfitting>`
describe opposite ends of the same problem: model *expressiveness*.

- Too little expressiveness: train and test scores are close to those of a
  model that guesses randomly (often called a dummy model) (SKD002).
- Too much expressiveness: train scores pull far ahead of test scores
  (SKD001).

A performant model learns from the training data without memorizing its
specificities, so that learned knowledge can be generalized to new data. This
notebook walks the underfitting to overfitting path while showing different
mitigations techniques:

- tweaking model capacity (model family / complexity)
- feature engineering
- tuning regularization
- using early stopping
- using more data

We use California housing (median house value). Half of the rows feed the main
walkthrough; the other half is reserved to show the effect of adding training
data.
"""

# %%
# Load the California housing dataset
# ===================================
#
# Each row is a census block group in a geographical region whose coordinates
# we have (``Latitude``, ``Longitude``). The features mix house characteristics
# (age, rooms, occupancy) with regional aggregates (median income, population).
# The target ``MedHouseVal`` is the median house value for the block group.
# The feature set is small but rich: later we combine columns such as
# ``AveRooms / AveOccup`` into a rooms-per-person marker.

import pandas as pd
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split

housing = fetch_california_housing(as_frame=True)

X, X_heldout, y, y_heldout = train_test_split(
    housing.data,
    housing.target,
    train_size=0.5,
    random_state=42,
)

# %%
# Let us glance at the table. Pay attention to ``AveRooms`` and ``AveOccup``:
# we will combine them into a rooms-per-person feature below.

from skrub import TableReport

TableReport(X)

# %%
# ``MedHouseVal`` is capped at 5.0 (500k USD). Click the column in the report to
# see the histogram and the pile-up at the ceiling.

TableReport(y)

# %%
# We use the same split for every model comparison, with a fixed seed, so that
# we are sure that we are comparing models on the same test data.

from skore import TrainTestSplit

splitter = TrainTestSplit(test_size=0.2, random_state=42)

# %%
# Underfitting: SKD002 fires
# ==========================
#
# We start with a linear model that is overly regularized.
# :class:`~sklearn.linear_model.Ridge` with a very large ``alpha`` shrinks
# coefficients toward zero, so predictions stay close to a dummy baseline.

import skore
from sklearn.linear_model import Ridge

report_underfit = skore.evaluate(
    Ridge(alpha=1e6),
    X=X,
    y=y,
    splitter=splitter,
)
report_underfit.metrics.summarize(data_source="both").frame()

# %%
# For a model that learns so little, SKD002 fires.

report_underfit.checks.summarize(fast_mode=True)

# %%
# Increase model expressiveness
# =============================
#
# Moving away from underfitting means giving the model enough expressiveness to
# use the inputs. Let us drop to a :class:`~sklearn.linear_model.Ridge` with a
# much smaller ``alpha``: it already learns useful weights on each feature and
# is enough to clear SKD002 on this table.

report_ridge = skore.evaluate(Ridge(alpha=1.0), X=X, y=y, splitter=splitter)
report_ridge.metrics.summarize(data_source="both").frame()

# %%
report_ridge.checks.summarize(fast_mode=True)

# %%
# Feature engineering
# ===================
#
# Better features help when the model family is fine but the representation
# is weak (still an underfitting problem). Here, ``AveRooms / AveOccup`` is a
# rooms-per-person signal that a linear model can use more easily than the raw
# counts.
#
# The same step has an overfitting side effect: every new feature also increases
# expressiveness. Mild engineering can close an underfit gap; aggressive
# expansions (very high-degree interactions, huge one-hot spaces) can reopen an
# overfit gap. Think of features as capacity you add to the inputs, not only
# to the estimator.

X_fe = X.assign(RoomsPerPerson=X["AveRooms"] / X["AveOccup"].clip(lower=0.1))

report_ridge_fe = skore.evaluate(Ridge(alpha=1.0), X=X_fe, y=y, splitter=splitter)
report_ridge_fe.metrics.summarize(data_source="both").frame()

# %%
# The test score rises while the train score stays close: we reduced underfitting
# without creating overfitting by adding a signal-rich feature.

report_ridge_fe.checks.summarize(fast_mode=True)

# %%
# Overfitting: SKD001 fires
# =========================
#
# Hand-crafted features are one way to capture nonlinearity for a linear model.
# Another is to switch to a model family that learns nonlinear structure on its
# own. A default :class:`~sklearn.ensemble.RandomForestRegressor` does that,
# but unrestricted leaves can also memorize the training data: train metrics
# look excellent, test metrics improve but not as much as train ones, and
# SKD001 flags the gap.

from sklearn.ensemble import RandomForestRegressor

report_rf = skore.evaluate(
    RandomForestRegressor(random_state=42),
    X=X,
    y=y,
    splitter=splitter,
)
report_rf.metrics.summarize(data_source="both").frame()

# %%
report_rf.checks.summarize(fast_mode=True)

# %%
# Regularization
# ==============
#
# Once SKD001 appears, we pull expressiveness back with capacity limits on the
# estimator (leaf size, feature fraction, depth, learning rate, …).
#
# Here, we set the hyperparameters of the model by hand. In practice, it is
# rather hard to know in advance which combination of hyperparameters leads to
# generalization. One should lean towards tuning hyperparameters. We would
# advocate for randomized search, with successive halving when the amount of
# samples is large. See `tuning the hyper-parameters of an estimator
# <https://scikit-learn.org/stable/modules/grid_search.html#tuning-the-hyper-parameters-of-an-estimator>`_.
# In this particular example, we are not implementing the search in order to
# keep the execution time of the example short.

model_rf_reg = RandomForestRegressor(
    min_samples_leaf=20,
    max_features=0.5,
    random_state=42,
)
report_rf_reg = skore.evaluate(
    model_rf_reg,
    X=X,
    y=y,
    splitter=splitter,
)
report_rf_reg.metrics.summarize(data_source="both").frame()

# %%
report_rf_reg.checks.summarize(fast_mode=True)

# %%
# Early stopping
# ==============
#
# For iterative learners, early stopping is another way to limit expressiveness
# that does not require a hyperparameter search. We hold out a validation set,
# monitor a metric, and stop when that metric stops improving. Further
# iterations are assumed to overfit. We can either pass a validation set to
# ``.fit`` or use the ``validation_fraction`` parameter. See scikit-learn's
# example on `gradient boosting with early stopping
# <https://scikit-learn.org/stable/auto_examples/ensemble/plot_gradient_boosting_early_stopping.html>`_.
#
# Let us compare a :class:`~sklearn.ensemble.HistGradientBoostingRegressor`
# with and without early stopping on the same split to see the effect. We bump
# ``max_iter`` so that early stopping has room to act. With
# ``validation_fraction``, the early-stopped model trains on only 90% of each
# train fold; a small test drop relative to the unstopped run can come from
# that, rather than from stopping itself. Passing an external validation set
# (for example ``X_heldout`` / ``y_heldout``) would avoid that, but wiring fit
# parameters through skore currently needs a skrub DataOp, so we keep
# ``validation_fraction`` here.

from sklearn.ensemble import HistGradientBoostingRegressor

report_hgbr = skore.evaluate(
    HistGradientBoostingRegressor(max_iter=1000, random_state=42),
    X=X,
    y=y,
    splitter=splitter,
)

model_hgbr_es = HistGradientBoostingRegressor(
    early_stopping=True,
    validation_fraction=0.1,
    n_iter_no_change=10,
    max_iter=1000,
    random_state=42,
)
report_hgbr_es = skore.evaluate(
    model_hgbr_es,
    X=X,
    y=y,
    splitter=splitter,
)

skore.compare(
    {
        "hgbr": report_hgbr,
        "hgbr_early_stopping": report_hgbr_es,
    }
).metrics.summarize(data_source="both").frame()

# %%
report_hgbr.checks.summarize(fast_mode=True)

# %%
report_hgbr_es.checks.summarize(fast_mode=True)

# %%
# More training data
# ==================
#
# Extra labeled rows help both regimes, but they are especially useful against
# overfitting: the model has fewer opportunities to memorize a small sample.
# Let us keep one fixed test fold, fit on the original train fold, then refit
# after concatenating the held-out half of the dataset. We use
# :func:`~skore.evaluate` with ``splitter="prefit"`` once the estimator is
# fitted. A side-effect is that fit time is unavailable (skore did not time
# ``.fit``), so that metric appears as NaN in the reported tables.

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

model_less = HistGradientBoostingRegressor(
    early_stopping=True,
    validation_fraction=0.1,
    n_iter_no_change=10,
    random_state=42,
).fit(X_train, y_train)
report_less = skore.evaluate(model_less, X_test, y_test, splitter="prefit")

model_more = HistGradientBoostingRegressor(
    early_stopping=True,
    validation_fraction=0.1,
    n_iter_no_change=10,
    random_state=42,
).fit(pd.concat([X_train, X_heldout]), pd.concat([y_train, y_heldout]))
report_more = skore.evaluate(model_more, X_test, y_test, splitter="prefit")

skore.compare(
    {"less_training_data": report_less, "more_training_data": report_more}
).metrics.summarize(data_source="test").frame()

# %%
# In production, a `learning curve
# <https://scikit-learn.org/stable/modules/learning_curve.html#learning-curve>`_
# (many refits on growing subsets) shows whether more labels are still worth
# the cost. skore does not wrap that API yet; we can treat it as a
# scikit-learn-side diagnostic next to these checks.

# %%
# Summary comparison
# ==================
#
# Here is a side-by-side comparison of the impact of the different techniques we
# applied, measured on the same test set.

skore.compare(
    {
        "ridge_large_alpha": report_underfit,
        "ridge": report_ridge,
        "ridge_feature_engineering": report_ridge_fe,
        "default_rf": report_rf,
        "regularized_rf": report_rf_reg,
        "hgbr": report_hgbr,
        "hgbr_early_stopping": report_hgbr_es,
        "more_data": report_more,
    }
).metrics.summarize(data_source="test").frame()

# %%
# Conclusion
# ==========
#
# SKD002 and SKD001 are two different consequences of a unsuited expressiveness.
# We relax regularization / add capacity and use informative features until the
# model beats a weak baseline; then we regularize, stop early, and add data if
# train scores have a noticeable gap compared to the test scores.
#
# Feature engineering helps the underfit side of the path; capacity control
# and more data help the overfit side. Hyperparameter choices play a dual
# role. In practice we combine several of these levers and re-run the checks
# until the results are satisfying enough.
