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 RidgeCV
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)

"""
sklearn.linear_model.RidgeCV(alphas=(0.1, 1.0, 10.0), fit_intercept=True, normalize=False, scoring=None, cv=None, gcv_mode=None, store_cv_values=False)"""

clf = RidgeCV(alphas=(0.001, 0.01, 0.2, 0.5, 1, 2, 5), store_cv_values=True)
clf.fit(x_train, y_train)


clf.fit(x_train, y_train)
print('Coefficients: \n', clf.coef_)
print('Intercept: ', clf.intercept_)
print('Alpha: ', clf.alpha_)
#print(clf.cv_values_)
#print(len(y_test))

train_pred=clf.predict(x_train)
test_pred=clf.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()