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 Ridge
import matplotlib.pyplot as plt

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

print(df.columns.tolist())

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


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

lr=Ridge()
lr.fit(x_train, y_train)
print('Coefficients: \n', lr.coef_)
print('Intercept: ', lr.intercept_)

train_pred=lr.predict(x_train)
test_pred=lr.predict(x_test)

mse_train = mean_squared_error(y_train, train_pred)
mse_test = mean_squared_error(y_test, test_pred)
print("MSE na trening skupu: %.2f"% mse_train)
print("MSE na test skupu: %.2f"% mse_test)
print()
r2_train=r2_score(y_train, train_pred)
r2_test= r2_score(y_test, test_pred)
print("R^2 na trening skupu: %.2f"% r2_train)
print("R^2 na test skupu: %.2f"% r2_test)


plt.scatter(x_train, y_train, color='red',  marker='x',  label='Trening skup')
plt.plot(x_train, train_pred, color='red', label='Trening skup, $R^2$=%.2f'%r2_train)


plt.scatter(x_test, y_test, color='blue', label='Test skup')
plt.plot(x_test, test_pred, color='blue', label='Test skup, $R^2$=%.2f'%r2_test)
plt.xlabel('populacija* $10^4$')
plt.ylabel('dobit*$10^4\$$')
plt.legend(loc='lower right')
plt.title('Grebena linearna regresija')

plt.show()