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

# Загрузка данных
data = pd.read_csv("F1DriversDataset.csv")

# Выбор признаков и целевой переменной
features = ['Race_Entries', 'Race_Starts', 'Pole_Positions', 'Race_Wins', 'Podiums', 'Fastest_Laps']
target = 'Championships'
X = data[features]
y = data[target]

# Разделение данных на обучающую и тестовую выборки
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)

# Оценка качества модели на тестовой выборке
score = model.score(X_test, y_test)
print("Коэффициент детерминации R^2 на тестовой выборке:", score)

# Предсказание количества чемпионских титулов для гонщика
new_features = []
for feature in features:
    value = float(input(f"Введите значение для {feature}: "))
    new_features.append(value)
new_features = [new_features]
predicted_championships = model.predict(new_features)
print("Предсказанное количество чемпионских титулов:", predicted_championships[0])

# Визуализация модели
y_pred = model.predict(X_test)

plt.scatter(y_test, y_pred)
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2)
plt.xlabel('Фактическое количество чемпионских титулов')
plt.ylabel('Предсказанное количество чемпионских титулов')
plt.title('Предсказания модели линейной регрессии')
plt.show()