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

print('Model sa wo')
lr=LinearRegression()
lr.fit(x_train, y_train)

print('Koeficijenti: ', lr.coef_)
print('w0: ', lr.intercept_)

print('f(x)=', lr.intercept_[0], '+',
      lr.coef_[0][0], '*populacija' )

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)

r2_train = r2_score(y_train, train_pred)
r2_test = r2_score(y_test, test_pred)


print('Model bez wo')

lr2=LinearRegression(fit_intercept=False)
lr2.fit(x_train, y_train)

print('Koeficijenti: ', lr2.coef_)

print('f(x)=', lr2.coef_[0][0], '*populacija' )

train_pred2 = lr2.predict(x_train)
test_pred2 = lr2.predict(x_test)

mse_train2 = mean_squared_error(y_train, train_pred2)
mse_test2 = mean_squared_error(y_test, test_pred2)

r2_train2 = r2_score(y_train, train_pred2)
r2_test2 = r2_score(y_test, test_pred2)

print('MSE trening skup:', mse_train2)
print('MSE test skup:', mse_test2)
print('R2 trening skup:', r2_train2)
print('R2 test skup:', r2_test2)

plt.subplot(211)
plt.scatter(x_test, y_test,  c='k', s=10, label='Test skup')
plt.plot(x_test, test_pred, c='r',
label='Model sa $w_0$: MSE=%.2f, $R^2$=%.2f'%(mse_test, r2_test))
plt.plot(x_test, test_pred2, c='b',
label='Model bez $w_0$: MSE=%.2f, $R^2$=%.2f'%(mse_test2, r2_test2))
plt.xlabel('populacija*$10^4$')
plt.ylabel('dobit*$10^4\$$')
plt.legend(loc='upper left', fontsize=6)
plt.title('Linearna regresija')

plt.subplot(212)

plt.scatter(test_pred, y_test-test_pred, c='b',s=10, label='Test skup MSE=%.2f, $R^2$=%.2f'%(mse_test, r2_test))
plt.scatter(train_pred, y_train-train_pred, c='r', marker='x', s=10, label='Trening skup MSE=%.2f, $R^2$=%.2f'%(mse_train, r2_train))
plt.hlines(y=0, xmin=0, xmax=25)
plt.xlabel('$\hat{y}_i$')
plt.ylabel('$y_i-\hat{y}_i$')
plt.legend(loc='upper left', fontsize=6)
plt.title('Prikaz reziduala')

plt.tight_layout()

plt.show()