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

X, y = make_classification(
    n_features=2,
    n_redundant=0,
    n_informative=2,
    random_state=0,
    n_clusters_per_class=1
)

rng = np.random.RandomState(2)
X += 2 * rng.uniform(size=X.shape)
linearly_dataset = (X, y)
moon_dataset = make_moons(noise=0.3, random_state=0)
circles_dataset = make_circles(noise=0.2, factor=0.5, random_state=1)
datasets = [moon_dataset, circles_dataset, linearly_dataset]

"""
Данные:
· moon_dataset
· circles_dataset
· linearly_dataset
"""
for ds_cnt, ds in enumerate(datasets):
    X, y = ds
    X = StandardScaler().fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=.4, random_state=42
    )
    """
    Модели:
    · Линейную регрессию
    · Полиномиальную регрессию (со степенью 3)
    · Гребневую полиномиальную регрессию (со степенью 3, alpha = 1.0)
    """

    # Линейная
    linear_regression = LinearRegression()
    linear_regression.fit(X_train, y_train)
    linear_predictions = linear_regression.predict(X_test)
    linear_mse = mean_squared_error(y_test, linear_predictions)

    # Полиномиальная (degree=3)
    poly_regression = make_pipeline(PolynomialFeatures(degree=3), LinearRegression())
    poly_regression.fit(X_train, y_train)
    poly_predictions = poly_regression.predict(X_test)
    poly_mse = mean_squared_error(y_test, poly_predictions)

    # Гребневая (degree=3, alpha=1.0)
    poly_regression_alpha = make_pipeline(PolynomialFeatures(degree=3), Ridge(alpha=1.0))
    poly_regression_alpha.fit(X_train, y_train)
    poly_alpha_predictions = poly_regression_alpha.predict(X_test)
    poly_alpha_mse = mean_squared_error(y_test, poly_alpha_predictions)

    # График данных
    plt.figure(figsize=(10, 6))
    plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap='coolwarm')
    plt.title('Датасет №' + str(ds_cnt))
    plt.xlabel('X')
    plt.ylabel('Y')

    # График линейной модели
    plt.figure(figsize=(10, 6))
    plt.scatter(X_test[:, 0], X_test[:, 1], c=linear_predictions, cmap='coolwarm')
    plt.title('Линейная ds'+ str(ds_cnt))
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.show()

    # График полиномиальной модели (degree=3)
    plt.figure(figsize=(10, 6))
    plt.scatter(X_test[:, 0], X_test[:, 1], c=poly_predictions, cmap='coolwarm')
    plt.title('Полиномиальная (degree=3) ds' + str(ds_cnt))
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.show()

    # График гребневой модели (degree=3, alpha=1.0)
    plt.figure(figsize=(10, 6))
    plt.scatter(X_test[:, 0], X_test[:, 1], c=poly_alpha_predictions, cmap='coolwarm')
    plt.title('Гребневая (degree=3, alpha=1.0) ds' + str(ds_cnt))
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.show()

    # Сравнение качества
    print('Линейная MSE:', linear_mse)
    print('Полиномиальная (degree=3) MSE:', poly_mse)
    print('Гребневая (degree=3, alpha=1.0) MSE:', poly_alpha_mse)