import random
from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn.datasets import make_moons
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import Pipeline

rs = random.randrange(50)

X, y = make_moons(n_samples=250, noise=0.3, random_state=rs)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=42)

figure = plt.figure(1, figsize=(16, 9))
axis = figure.subplots(4, 3)
cm = ListedColormap(['#FF0000', "#0000FF"])
arr_res = list(range(len(y_test)))
X_scale = list(range(len(y_test)))


def test(col, model):
    global axis
    global arr_res
    global X_test
    global X_train
    global y_train
    global y_test

    model.fit(X_train, y_train)
    res_y = model.predict(X_test)
    print(model.score(X_test, y_test))

    axis[0, col].scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm)
    axis[1, col].scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm)
    axis[2, col].scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm)
    axis[2, col].scatter(X_test[:, 0], X_test[:, 1], c=res_y, cmap=cm)
    axis[3, col].plot([i for i in range(len(res_y))], y_test, c="g")
    axis[3, col].plot([i for i in range(len(res_y))], res_y, c="r")


def start():
    lin = LinearRegression()
    poly = Pipeline([('poly', PolynomialFeatures(degree=3)),
                     ('linear', LinearRegression())])
    ridge = Pipeline([('poly', PolynomialFeatures(degree=3)),
                      ('ridge', Ridge(alpha=1.0))])

    test(0, lin)
    test(1, poly)
    test(2, ridge)

    plt.show()


start()