#Step 1: Import packages and classes
import numpy as np
from sklearn.linear_model import LinearRegression
import scipy.stats
import matplotlib.pyplot as plt

#Step 2: Provide data

#x = np.array([5, 15, 25, 35, 45, 55]).reshape((-1, 1))
#y = np.array([5, 20, 14, 32, 22, 38])

x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).reshape((-1, 1))
y = np.array([1, 3, 2, 5, 7, 8, 8, 9, 10, 12])

#Step 3: Create a model and fit it
model = LinearRegression()
model.fit(x, y)

#Step 4: Get results (optional)
r_sq = model.score(x, y)
print(f"coefficient of determination: {r_sq}")
print(f"intercept: {model.intercept_}")
print(f"slope: {model.coef_}")

x_new = np.array([[16]])  # New value of X
y_new = model.predict(x_new) #Prediction of the new Y, using the new X value

#y_new = model.predict(np.array([[16]]))

print(y_new)




# Adding the plot with scipy.stats.linregress
slope, intercept, r, *__ = scipy.stats.linregress(x.ravel(), y) #slope,y-intersept,corr. coeff. "r"
line = f'Regression line: y={intercept:.2f}+{slope:.2f}x, r={r:.2f}'
fig, ax = plt.subplots()
ax.plot(x, y, linewidth=0, marker='s', label='Data points')
ax.plot(x, intercept + slope * x, label=line)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(facecolor='yellow')
plt.show()