lambdas = np.logspace(-4, 5, 10)

train_errors = {
    "ols_own": np.zeros(lambdas.size),
    "ols_sk": np.zeros(lambdas.size),
    "ridge_own": np.zeros(lambdas.size),
    "ridge_sk": np.zeros(lambdas.size),
    "lasso_sk": np.zeros(lambdas.size)
}

test_errors = {
    "ols_own": np.zeros(lambdas.size),
    "ols_sk": np.zeros(lambdas.size),
    "ridge_own": np.zeros(lambdas.size),
    "ridge_sk": np.zeros(lambdas.size),
    "lasso_sk": np.zeros(lambdas.size)
}

plot_counter = 1

fig = plt.figure(figsize=(32, 54))

for i, _lambda in enumerate(tqdm.tqdm(lambdas)):
    omega = get_ols_weights(X_train_own, y_train)
    y_hat_train = X_train_own @ omega
    y_hat_test = X_test_own @ omega

    train_errors["ols_own"][i] = r_squared(y_train, y_hat_train)
    test_errors["ols_own"][i] = r_squared(y_test, y_hat_test)

    plt.subplot(10, 5, plot_counter)
    plt.imshow(omega[1:].reshape(L, L), **cmap_args)
    plt.title("Home made OLS")
    plot_counter += 1

    omega = get_ridge_weights(X_train_own, y_train, _lambda)
    y_hat_train = X_train_own @ omega
    y_hat_test = X_test_own @ omega

    train_errors["ridge_own"][i] = r_squared(y_train, y_hat_train)
    test_errors["ridge_own"][i] = r_squared(y_test, y_hat_test)

    plt.subplot(10, 5, plot_counter)
    plt.imshow(omega[1:].reshape(L, L), **cmap_args)
    plt.title(r"Home made ridge, $\lambda = %.4f$" % _lambda)
    plot_counter += 1

    for key, method in zip(
        ["ols_sk", "ridge_sk", "lasso_sk"],
        [skl.LinearRegression(), skl.Ridge(alpha=_lambda), skl.Lasso(alpha=_lambda)]
    ):
        method = method.fit(X_train, y_train)

        train_errors[key][i] = method.score(X_train, y_train)
        test_errors[key][i] = method.score(X_test, y_test)

        omega = method.coef_.reshape(L, L)

        plt.subplot(10, 5, plot_counter)
        plt.imshow(omega, **cmap_args)
        plt.title(r"%s, $\lambda = %.4f$" % (key, _lambda))
        plot_counter += 1

plt.show()