import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
import seaborn as sns

# === Загружаем данные ===
columns = ['Position', 'Height(inches)', 'Weight(pounds)', 'Age']
df = pd.read_csv('SOCR_Data_MLB_HeightsWeights.txt', sep='\t', nrows=115, usecols=columns)

# Кодируем категориальный признак "Position"
df['Position_num'] = pd.factorize(df['Position'])[0]

# === Признаки и целевая переменная ===
X = df[['Position_num', 'Height(inches)', 'Age']]
y = df['Weight(pounds)']

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

# === Обучение модели ===
model = LinearRegression()
model.fit(X_train, y_train)

# === Предсказание ===
y_pred = model.predict(X_test)

# === Оценка точности ===
mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_test, y_pred)

print("=== Оценка качества модели ===")
print(f"СКО (RMSE): {rmse:.4f}")
print(f"Коэффициент детерминации (R^2): {r2:.4f}")

# === График ошибок ===
errors = y_test - y_pred

plt.figure(figsize=(10, 5))
sns.histplot(errors, bins=15, kde=True, color='lightcoral', edgecolor='black')
plt.title('Гистограмма ошибок предсказания')
plt.xlabel('Ошибка (y_true - y_pred)')
plt.ylabel('Частота')
plt.show()

# === Визуализация предсказанного vs реального значения ===
plt.figure(figsize=(8, 6))
plt.scatter(y_test, y_pred, color='blue', edgecolor='k')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
plt.xlabel('Истинные значения')
plt.ylabel('Предсказанные значения')
plt.title('Предсказанные vs Истинные значения')
plt.show()