import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

# загрузка и обработка данных
tutors_data = pd.read_csv('tutors.csv')
tutors_data = tutors_data.dropna()

status_coding = {'Student': 0, 'School teacher': 1, 'Postgraduate student': 2, 'Private tutor': 3, 'University professor': 4, 'Native speaker': 5}
tutors_data['Status'] = tutors_data['Status'].replace(status_coding)

photo_coding = {'No': 0, 'Yes': 1}
tutors_data['Photo'] = tutors_data['Photo'].replace(photo_coding)
tutors_data['Video_presentation'] = tutors_data['Video_presentation'].replace(photo_coding)

# выбор признаков и целевой переменной
X = tutors_data[['Score', 'Reviews_number', 'Experience', 'Status', 'Video_presentation', 'Photo']]
y = tutors_data['Price']

# разделение данных на обучащий и тестовый наборы
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.05, random_state=42)

# создание модели линейной регрессии и обучение на тренировочных данных
model = LinearRegression()
model.fit(X_train, y_train)
# предсказание стоимости занятия
predictions = model.predict(X_test)
# оценка качества модели
score = model.score(X_test, y_test)
print(f"Качество модели линейной регрессии = {score:.2f}")

# вывод результатов в консоль
test_data = y_test.values
print("------------ СТОИМОСТИ ЗАНЯТИЙ ------------")
for i in range(len(predictions)):
    print(f"Фактическая: {test_data[i]}\t\tПредсказанная: {int(predictions[i])}")

# отображение результатов на графике
plt.scatter(y_test, predictions, color='b')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], color='r')
plt.xlabel("Фактическая стоимость занятия")
plt.ylabel("Предсказанная стоимость занятия")
plt.title("Линейная регрессия")
plt.show()