# Source: content/notes/ml/linear-models.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.linear_model import (HuberRegressor, LinearRegression,
                                  QuantileRegressor, RANSACRegressor, Ridge,
                                  TheilSenRegressor)
from sklearn.metrics import mean_absolute_error, mean_pinball_loss
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, SplineTransformer

rng = np.random.default_rng(8)
X = rng.uniform(-2, 2, size=(160, 1))
y = 1 + 2 * X[:, 0] + rng.normal(0, 0.3, size=160)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=8)
y_corrupt = y_train.copy()
y_corrupt[::12] += 12
models = {
    "OLS": LinearRegression(), "Huber": HuberRegressor(),
    "RANSAC": RANSACRegressor(random_state=8),
    "Theil-Sen": TheilSenRegressor(random_state=8, max_subpopulation=300),
    "polynomial": make_pipeline(PolynomialFeatures(2, include_bias=False), Ridge()),
    "spline": make_pipeline(SplineTransformer(n_knots=4), Ridge()),
}
for name, model in models.items():
    model.fit(X_train, y_corrupt)
    pred = model.predict(X_test)
    assert np.isfinite(pred).all()
    print(name, "clean-test MAE", mean_absolute_error(y_test, pred))
quantile = QuantileRegressor(quantile=0.9, alpha=0.01, solver="highs")
quantile.fit(X_train, y_train)
upper = quantile.predict(X_test)
assert np.isfinite(upper).all()
print("Pinball loss:", mean_pinball_loss(y_test, upper, alpha=0.9))
print("Observed upper-quantile coverage:", np.mean(y_test <= upper))