import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt

from sklearn.pipeline import make_pipeline

df = pd.read_csv('kamion_prodavnica.csv')

X=df[['populacija']]
Y=df[['dobit']]

x_train, x_test, y_train, y_test = train_test_split(X, Y, train_size=0.8)


colors = ['red', 'green', 'gold']
plt.scatter(x_test, y_test, color='navy', s=30, marker='o', label="test skup")

for i, degree in enumerate([2,3,4]):
    model = make_pipeline(PolynomialFeatures(degree), LinearRegression())
    model.fit(x_train, y_train)
    y_plot = model.predict(x_test)
    plt.plot(x_test, y_plot, color=colors[i], linewidth=0.5, label="stepen %d, MSE=%.2f, $R^2$=%.2f" % (degree,
                                                                                                        mean_squared_error(y_test, y_plot),
                                                                                                        r2_score(y_test, y_plot)))

plt.legend(loc='lower right')

plt.show()