from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Ridge
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PolynomialFeatures


def lin(x_train, x_test, y_train, y_test):
    plt.scatter(x_test, y_test)
    model = LinearRegression().fit(x_train, y_train)
    y_predict = model.intercept_ + model.coef_ * x_test
    plt.plot(x_test, y_predict, color='red')
    plt.title('Линейная регрессия')
    plt.savefig('images/linear.png')
    plt.show()

    print('Линейная регрессия')
    print('Оценка качества:', model.score(x_train, y_train))


def polynom(x_train, y_train):
    plt.scatter(x_train, y_train)
    x_poly = PolynomialFeatures(degree=4).fit_transform(x_train)
    pol_reg = LinearRegression()
    model = pol_reg.fit(x_poly, y_train)
    y_predict = pol_reg.predict(x_poly)
    plt.plot(x_train, y_predict, color='green')
    plt.title('Полиномиальная регрессия')
    plt.savefig('images/polynomial.png')
    plt.show()

    print('Полиномиальная регрессия')
    print('Оценка качества:', model.score(x_poly, y_train))


def greb_polynom(x_train, x_test, y_train, y_test):
    plt.scatter(x_test, y_test)
    pipeline = Pipeline([("polynomial_features", PolynomialFeatures(degree=4)), ("ridge", Ridge(alpha=1.0))])
    model = pipeline.fit(x_train, y_train)
    y_predict = pipeline.predict(x_test)
    plt.plot(x_test, y_predict, color='blue')
    plt.title('Гребневая полиномиальная регрессия')
    plt.savefig('images/greb_polynom.png')
    plt.show()

    print('Гребневая полиномиальная регрессия')
    print('Оценка качества:', model.score(x_train, y_train))