from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn.datasets import make_circles
from sklearn.linear_model import LinearRegression, 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

# Нелинейный генератор - позволяет сгенерировать такие классы-признаки,
# что признаки одного класса геометрически окружают признаки другого класса
X, y = make_circles(noise=0.2, factor=0.5, random_state=1)
X = StandardScaler().fit_transform(X)

# Разделение на обучающую и тестовую выборки (40% данных будет использовано для тестов)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=.4, random_state=42)

# Создание необходимых для оценки моделей
def models():
    # Линейная регрессия
    linear_regression = LinearRegression()
    linear_regression.fit(X_train, y_train)

    # Полиномиальная регрессия (degree=3)
    poly_regression = make_pipeline(PolynomialFeatures(degree=3), LinearRegression())
    poly_regression.fit(X_train, y_train)

    # Гребневая рекрессия (degree=3, alpha=1.0)
    ridge_regression = make_pipeline(PolynomialFeatures(degree=3), Ridge(alpha=1.0))
    ridge_regression.fit(X_train, y_train)

    models = [linear_regression, poly_regression, ridge_regression]

    # Предсказанные y
    linear_predict = linear_regression.predict(X_test)
    poly_predict = poly_regression.predict(X_test)
    ridge_predict = ridge_regression.predict(X_test)

    pred = [linear_predict, poly_predict, ridge_predict]

    # Среднеквадратичные ошибки
    lin_mse = mean_squared_error(y_test, linear_predict)
    poly_mse = mean_squared_error(y_test, poly_predict)
    rr_mse = mean_squared_error(y_test, ridge_predict)

    mse = [lin_mse, poly_mse, rr_mse]

    grafics(pred, mse, models)

# Графики
def grafics(pred, mse, models):
    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(8, 8))
    cm_color1 = ListedColormap(['r', 'g'])

    plt.suptitle('Лабораторная работа 1. Вариант 2.', fontweight='bold')

    # График данных
    plt.subplot(2, 2, 1)
    plt.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_color1, marker='o', label='тренировочные данные')
    plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm_color1, marker='x', label='тестовые данные')
    plt.title('Датасет circles', fontsize=10, loc='left')
    plt.legend(loc='upper left')

    # График линейной модели
    plt.subplot(2, 2, 2)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=pred[0], cmap=cm_color1)
    plt.title('Линейная регрессия', fontsize=10, loc='left')

    # График полиномиальной модели (degree=3)
    plt.subplot(2, 2, 3)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=pred[1], cmap=cm_color1)
    plt.title('Полиномиальная регрессия (degree=3)', fontsize=10, loc='left')
    plt.xlabel('X')
    plt.ylabel('Y')

    # График гребневой модели (degree=3, alpha=1.0)
    plt.subplot(2, 2, 4)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=pred[2], cmap=cm_color1)
    plt.title('Гребневая регрессия (degree=3, alpha=1.0)', fontsize=10, loc='left')
    plt.show()

    # Сравнение качества
    print('Линейная MSE:', mse[0])
    print('Полиномиальная (degree=3) MSE:', mse[1])
    print('Гребневая (degree=3, alpha=1.0) MSE:', mse[2])

models()