import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sqlalchemy import create_engine


def run_prediction_module():
    engine = create_engine('mysql+pymysql://wind:wind@193.124.203.110:3306/wind_towers')

    query = """
    SELECT BarTrend, CRC, DateStamp, DewPoint, HeatIndex, ETDay, HumIn, HumOut,
           Pressure, RainDay, RainMonth, RainRate, RainStorm, RainYear,
           TempIn, TempOut, WindDir, WindSpeed, WindSpeed10Min
    FROM weather_data
    WHERE DateStamp >= '2024-10-14 21:00:00' - INTERVAL 36 HOUR;
    """
    df = pd.read_sql(query, engine)  # Загружаем данные из SQL-запроса в DataFrame

    df['DateStamp'] = pd.to_datetime(df['DateStamp'])  # Преобразуем столбец 'DateStamp' в формат datetime
    df.set_index('DateStamp', inplace=True)  # Устанавливаем 'DateStamp' как индекс
    df.sort_index(inplace=True)  # Сортируем DataFrame по индексу (по дате)

    lags = 3  # Задаем количество временных сдвигов (лагов)
    shifted_dfs = [df]  # Создаем список для хранения исходного DataFrame и его лаговых версий

    for lag in range(1, lags + 1):
        shifted_df = df.shift(lag).add_suffix(f'_t-{lag}')  # Создаем сдвинутый на lag строк DataFrame
        shifted_dfs.append(shifted_df)  # Добавляем его в список

    df_with_lags = pd.concat(shifted_dfs, axis=1)  # Объединяем исходный DataFrame и все лаги по столбцам

    df_with_lags.dropna(inplace=True)  # Удаляем строки с пропущенными значениями
    df_with_lags = df_with_lags.copy()  # Создаем копию DataFrame (для предотвращения предупреждений)

    # Преобразуем столбец 'BarTrend' в числовой формат, используя кодировщик категорий
    le = LabelEncoder()
    df_with_lags['BarTrend_encoded'] = le.fit_transform(df_with_lags['BarTrend'])

    # Оставляем в DataFrame только числовые столбцы
    df_with_lags = df_with_lags.select_dtypes(include=['float64', 'int64'])

    # Создаем словари для хранения моделей и значений MSE
    models = {}
    mse_scores = {}

    # Обучаем модели для каждого целевого столбца
    for target_column in df.columns:
        if target_column not in df_with_lags.columns:  # Пропускаем, если столбец отсутствует в df_with_lags
            continue

        X = df_with_lags.drop(columns=[target_column]).values  # Признаки - все столбцы, кроме целевого
        y = df_with_lags[target_column].values  # Целевой столбец

        # Разделяем данные на обучающую и тестовую выборки без перемешивания (временной ряд)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

        model = RandomForestRegressor()  # Инициализируем модель случайного леса
        model.fit(X_train, y_train)  # Обучаем модель

        y_pred = model.predict(X_test)  # Делаем предсказания на тестовой выборке
        mse = mean_squared_error(y_test, y_pred)  # Вычисляем среднеквадратичную ошибку
        mse_scores[target_column] = mse  # Сохраняем MSE для целевого столбца
        models[target_column] = model  # Сохраняем модель для целевого столбца

        quality = "хорошая" if mse < 1.0 else "плохая"  # Определяем качество модели
        print(f"MSE для {target_column}: {mse} ({quality})")  # Выводим MSE и качество

    # Обучаем модель для столбца 'BarTrend_encoded' отдельно
    X_bartrend = df_with_lags.drop(columns=['BarTrend_encoded']).values  # Признаки
    y_bartrend = df_with_lags['BarTrend_encoded'].values  # Целевой столбец 'BarTrend_encoded'

    # Разделяем данные на обучающую и тестовую выборки без перемешивания
    X_train_bartrend, X_test_bartrend, y_train_bartrend, y_test_bartrend = train_test_split(X_bartrend, y_bartrend,
                                                                                            test_size=0.2,
                                                                                            shuffle=False)

    model_bartrend = RandomForestRegressor()  # Инициализируем модель случайного леса
    model_bartrend.fit(X_train_bartrend, y_train_bartrend)  # Обучаем модель

    y_pred_bartrend = model_bartrend.predict(X_test_bartrend)  # Предсказания на тестовой выборке для 'BarTrend_encoded'
    mse_bartrend = mean_squared_error(y_test_bartrend, y_pred_bartrend)  # Вычисляем MSE
    models['BarTrend_encoded'] = model_bartrend  # Сохраняем модель для 'BarTrend_encoded'
    mse_scores['BarTrend_encoded'] = mse_bartrend  # Сохраняем MSE для 'BarTrend_encoded'

    quality_bartrend = "хорошая" if mse_bartrend < 1.0 else "плохая"  # Определяем качество модели для 'BarTrend_encoded'
    print(f"MSE для BarTrend: {mse_bartrend} ({quality_bartrend})")  # Выводим MSE и качество

    last_data = X[-1].reshape(1, -1)  # Берем последнюю строку данных и преобразуем в формат для предсказания

    predictions = {}  # Создаем словарь для хранения предсказаний
    for target_column, model in models.items():
        prediction = model.predict(last_data)[0]  # Делаем предсказание для последней строки данных
        if target_column == 'BarTrend_encoded':
            prediction = le.inverse_transform([int(prediction)])[0]  # Декодируем категориальное значение
            predictions['BarTrend'] = prediction  # Сохраняем предсказание для 'BarTrend'
            #print(f"Предсказание для BarTrend: {prediction}")  # Выводим предсказание
            continue  # Продолжаем цикл после предсказания для 'BarTrend_encoded'
        predictions[target_column] = prediction  # Сохраняем предсказание для остальных столбцов
        #print(f"Предсказание для {target_column}: {prediction}")  # Выводим предсказание для столбца

    return predictions  # Возвращаем словарь с предсказанными значениями и названиями столбцов