import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import ListedColormap
from sklearn.datasets import make_circles
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

rs = 42

X, y = make_circles(noise=0.2, factor=0.5, random_state=rs)
X = StandardScaler().fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=.4, random_state=rs)

# Инициализируем модели
linear_model = LinearRegression()
poly_model = make_pipeline(PolynomialFeatures(4), LinearRegression())
ridge_model = make_pipeline(PolynomialFeatures(4), Ridge(alpha=1.0))

# Обучаем модели
linear_model.fit(X_train, y_train)
poly_model.fit(X_train, y_train)
ridge_model.fit(X_train, y_train)

# Предсказываем значения для тестового набора
y_pred_linear = linear_model.predict(X_test)
y_pred_poly = poly_model.predict(X_test)
y_pred_ridge = ridge_model.predict(X_test)

# Качество моделей
mse_linear = mean_squared_error(y_test, y_pred_linear)
mse_poly = mean_squared_error(y_test, y_pred_poly)
mse_ridge = mean_squared_error(y_test, y_pred_ridge)
models_data = [(linear_model, "Линейная", mse_linear), (poly_model, "Полиномиальная", mse_poly),
               (ridge_model, "Греб. полиномиальная", mse_ridge)]
# Печатаем MSE
print(f"mse линейная: {mse_linear}")
print(f"mse полиномиальная: {mse_poly}")
print(f"mse гребневая полиномиальная: {mse_ridge}")

cm = plt.cm.RdBu
cm_bright = ListedColormap(['#FF0000', '#0000FF'])

fig, axs = plt.subplots(1, 3, figsize=(15, 7))
fig.suptitle('Сравнение регрессионных моделей')


# Функция отрисовки моделей
def draw_plot(model_data, i):
    h = .02  # шаг регулярной сетки
    x0_min, x0_max = X[:, 0].min() - .5, X[:, 0].max() + .5
    x1_min, x1_max = X[:, 1].min() - .5, X[:, 1].max() + .5
    xx0, xx1 = np.meshgrid(np.arange(x0_min, x0_max, h), np.arange(x1_min, x1_max, h))
    Z = model_data[0].predict(np.c_[xx0.ravel(), xx1.ravel()])  # 3

    Z = Z.reshape(xx0.shape)
    axs[i].contourf(xx0, xx1, Z, cmap=cm, alpha=.8)
    axs[i].set_xlim(xx0.min(), xx0.max())
    axs[i].set_ylim(xx0.min(), xx1.max())
    axs[i].set_xticks(())
    axs[i].set_yticks(())
    axs[i].set_title(model_data[1])
    axs[i].text(xx0.max() - .3, xx1.min() + .3, ('%.2f' % model_data[2]).lstrip('0'),
                size=15, horizontalalignment='right')
    axs[i].scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm_bright, alpha=0.6)


for i, model_data in enumerate(models_data):
    draw_plot(model_data, i)
plt.savefig("plots.png")
plt.show()