from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
import numpy as np
import matplotlib.pyplot as plt

X, y, w = make_regression(n_samples=200, n_features=200, coef=True,
                          random_state=1, bias=0, noise=8, tail_strength=0.9, effective_rank=10)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.4, random_state=2)

alphas = 10.**np.arange(-4, 10, 0.1)

scores_test = []
scores_train = []
for a in alphas:
    print(a)
    clf = Ridge(alpha=a, solver='auto')
    clf.fit(X_train, y_train)
    scores_train.append(mean_squared_error(y_train, clf.predict(X_train)))
    scores_test.append(mean_squared_error(y_test, clf.predict(X_test)))

fig, ax = plt.subplots(figsize=(8, 6))
ax.set_xscale('log')
ax.plot(alphas, scores_test, 'r', label='test')
ax.plot(alphas, scores_train, 'b', label='train')
ax.set_ylabel('MSE', fontsize=16)
ax.set_xlabel('alpha (regularisation strength)', fontsize=16)
ax.legend(fontsize=18)
plt.savefig('bias_variance_tradeoff.png')