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

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

X = data[['Archon Picks', 'Archon Win Rate', 'Legend Picks', 'Ancient Picks', 'Ancient Win Rate']]
y = data['Legend Win Rate']
names = data['Name']

# Разбиваем данные на обучающую и тестовую выборки
X_train, X_test, y_train, y_test, names_train, names_test = train_test_split(X, y, names, test_size=0.1, random_state=42)

# Применяем полиномиальные признаки к обучающим данным
poly_features = PolynomialFeatures(degree=2)
X_train_poly = poly_features.fit_transform(X_train)

# Создаем и обучаем модель полиномиальной регрессии
poly_model = LinearRegression()
poly_model.fit(X_train_poly, y_train)

# Применяем полиномиальные признаки к тестовым данным и делаем предсказания
X_test_poly = poly_features.transform(X_test)
y_pred = poly_model.predict(X_test_poly)

# Оценка качества модели на тестовых данных
r2 = poly_model.score(X_test_poly, y_test)
print(f"R-квадрат: {r2}")

# Построение графика с именами персонажей
plt.figure(figsize=(10, 6))
plt.title('Корреляция между выбором персонажей и победами в ранге "Legend"')
plt.grid(True)
plt.scatter(X_train['Legend Picks'], y_train, color='blue', alpha=0.5, label='Обучающая выборка')
plt.scatter(X_test['Legend Picks'], y_test, color='red', alpha=0.5, label='Тестовая выборка')

# Добавляем имена персонажей на график
for i, name in enumerate(names_test):
    plt.annotate(name, (X_test['Legend Picks'].iloc[i], y_pred[i]), fontsize=8, alpha=0.7, color='black')

plt.xlabel('Legend Picks')
plt.ylabel('Legend Win Rate')
plt.legend()
plt.show()