"""
Модуль для обучения модели и выполнения тренировочного пайплайна.
"""

import os
import json
import logging
import logging.config
import pickle

import hydra
import pandas as pd

from entities import entities
from data import data
from features import features
from models import models

logger = logging.getLogger(name)


def prepare_val_features_for_predict(
    train_features: pd.DataFrame, val_features: pd.DataFrame
) -> pd.DataFrame:
    """
    Подготавливает валидационные признаки для прогнозирования.

    :param train_features: DataFrame с обучающими признаками.
    :param val_features: DataFrame с валидационными признаками.
    :return: DataFrame с подготовленными валидационными признаками.
    ""
    train_features, val_features = train_features.align(val_features, join="left", axis=1)
    val_features = val_features.fillna(0)
    return val_features


def train_pipeline(params: entities.TrainingPipelineConfig) -> None:
    """E2E функция для обучения модели"""
    logger.info("Начало выполнения пайплайна")
    data = data.read_data(params.dataset.input_data_path)
    logger.info("Размер данных: %s", data.shape)

    train_df, test_df = data.split_train_test_data(data, params.split.test_size, params.split.random_state)
    logger.info("Размер обучающего набора данных: %s", train_df.shape)
    logger.info("Размер тестового набора данных: %s", test_df.shape)

    logger.info("Начало преобразований")
    transformer = features.build_transformer(
        categorical_features=params.feature.categorical_features,
        numerical_features=params.feature.numerical_features,
    )
    transformer.fit(train_df)
    train_features = features.make_features(transformer, train_df)
    train_target = features.extract_target(train_df, params.feature.target_col)
    logger.info("Размер обучающих признаков: %s", train_features.shape)

    model = models.train_model(
        params.model.model_params,
        train_features, train_target
    )

    test_features = features.make_features(transformer, test_df)
    test_target = features.extract_target(test_df, params.feature.target_col)

    val_features_prepared = prepare_val_features_for_predict(
        train_features, test_features
    )
    logger.info("Размер валидационных признаков: %s", val_features_prepared.shape)
    predicts = models.predict_model(
        model,
        val_features_prepared
    )

    metrics = models.evaluate_model(
        predicts,
        test_target
    )

    metrics_output_dir = os.path.join(os.getcwd(), ‘metrics’)
    model_output_dir = os.path.join(os.getcwd(), ‘models’)

    os.makedirs(metrics_output_dir, exist_ok=True)
    metrics_filepath = os.path.join(
        metrics_output_dir, f"{params.model.model_name}.json"
    )
    with open(metrics_filepath, "w", encoding="utf-8") as metric_file:
        json.dump(metrics, metric_file)
        logger.info("Метрики: %s", metrics)
    logger.info("Метрики сохранены в файл: %s", metrics_filepath)

    os.makedirs(model_output_dir, exist_ok=True)
    models_filepath = os.path.join(
        model_output_dir, f"{params.model.model_name}.pkl"
    )

    # Сохранение модели в файл
    with open(models_filepath, "wb") as model_file:
        pickle.dump(model, model_file)
    logger.info("Модель сохранена в файл: %s", models_filepath)

    logger.info("Обучение завершено")


@hydra.main(config_path="…/configs", config_name="config", version_base="1.1")
def main(cfg: entities.TrainingPipelineConfig = None) -> None:
    """Основная функция для настройки логгера и запуска train_pipeline"""
    train_pipeline(cfg)


if name == "main":
    main()