import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import joblib
import re

# Шаг 1: Загрузка данных
df = pd.read_csv('../laptops.csv')

# Шаг 2: Проверка и очистка имен столбцов
print("Имена столбцов до очистки:")
print(df.columns.tolist())

# Приведение имен столбцов к нижнему регистру и удаление пробелов
df.columns = df.columns.str.strip().str.lower()
print("\nИмена столбцов после очистки:")
print(df.columns.tolist())

# Шаг 3: Переименование столбцов (если необходимо)
df = df.rename(columns={
    'processor': 'processor',
    'ram': 'ram',
    'os': 'os',
    'ssd': 'ssd',
    'display': 'display',
    'price': 'price'
    # Другие столбцы можно оставить без изменений или переименовать по необходимости
})

# Шаг 4: Проверка наличия необходимых столбцов
required_columns = ['processor', 'ram', 'os', 'ssd', 'display', 'price']
missing_columns = [col for col in required_columns if col not in df.columns]

if missing_columns:
    print(f"\nОтсутствуют следующие столбцы: {missing_columns}")
    # Здесь можно добавить дополнительную обработку или завершить выполнение
    raise Exception(f"Отсутствуют столбцы: {missing_columns}")
else:
    print("\nВсе необходимые столбцы присутствуют.")

# Шаг 5: Удаление строк с пропущенными значениями
df = df.dropna(subset=required_columns)
print(f"\nКоличество строк после удаления пропусков: {df.shape[0]}")

# Шаг 6: Очистка и преобразование колонок

# Функция для очистки числовых колонок, содержащих символы
def clean_numeric_column(column, remove_chars=['₹', ',', ' ']):
    for char in remove_chars:
        column = column.str.replace(char, '', regex=False)
    return pd.to_numeric(column, errors='coerce')

# Очистка колонки 'price'
df['price'] = clean_numeric_column(df['price'])

# Очистка колонки 'ram' (например, '16 GB DDR4 RAM' -> 16)
def extract_numeric_ram(ram_str):
    match = re.search(r'(\d+)', ram_str)
    if match:
        return int(match.group(1))
    else:
        return None

df['ram'] = df['ram'].apply(extract_numeric_ram)

# Очистка колонки 'ssd' (например, '512 GB SSD' -> 512)
def extract_numeric_ssd(ssd_str):
    match = re.search(r'(\d+)', ssd_str)
    if match:
        return int(match.group(1))
    else:
        return None

df['ssd'] = df['ssd'].apply(extract_numeric_ssd)

# Очистка колонки 'display' (убираем лишние символы, если есть)
def clean_display(display_str):
    match = re.search(r'([\d.]+)', display_str)
    if match:
        return float(match.group(1))
    else:
        return None

df['display'] = df['display'].apply(clean_display)

# Проверка на пропущенные значения после очистки
df = df.dropna(subset=['price', 'ram', 'ssd', 'display'])
print(f"\nКоличество строк после очистки числовых колонок: {df.shape[0]}")

# Шаг 7: Выбор необходимых столбцов
df = df[required_columns]
print("\nПример данных после предобработки:")
print(df.head())

# Шаг 8: Преобразование категориальных переменных с помощью One-Hot Encoding
df = pd.get_dummies(df, columns=['processor', 'os'], drop_first=True)
print("\nИмена колонок после One-Hot Encoding:")
print(df.columns.tolist())

# Шаг 9: Разделение данных на обучающую и тестовую выборки
X = df.drop('price', axis=1)
y = df['price']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

print(f"\nРазмер обучающей выборки: {X_train.shape}")
print(f"Размер тестовой выборки: {X_test.shape}")

# Шаг 10: Обучение модели
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Обучение моделей
gbr = GradientBoostingRegressor(n_estimators=100, random_state=42)
gbr.fit(X_train, y_train)

lr = LinearRegression()
lr.fit(X_train, y_train)
print("\nМодели успешно обучена.")

# Оценка моделей

models = {'Random Forest': model, 'Gradient Boosting': gbr, 'Linear Regression': lr}

for name, mdl in models.items():
    y_pred = mdl.predict(X_test)
    mae = mean_absolute_error(y_test, y_pred)
    rmse = mean_squared_error(y_test, y_pred, squared=False)
    r2 = r2_score(y_test, y_pred)
    print(f"{name} - MAE: {mae}, RMSE: {rmse}, R²: {r2}")

# Шаг 12: Сохранение модели
joblib.dump(model, '../laptop_price_model.pkl')
print("\nМодель сохранена как 'laptop_price_model.pkl'.")

# Дополнительно: Сохранение колонок, полученных после One-Hot Encoding, для использования в бэкенде
feature_columns = X.columns.tolist()
joblib.dump(feature_columns, '../feature_columns.pkl')
print("Сохранены названия признаков в 'feature_columns.pkl'.")