mlflow

RUN_NAME = "model_bayesian_search"
STUDY_DB_NAME = "sqlite:///local.study.db"
STUDY_NAME = "churn_model"


experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
if not experiment:
    experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
else:
    experiment_id = experiment.experiment_id

if mlflow.active_run() is not None:
    mlflow.end_run()


def objective(trial: optuna.Trial) -> float:
    print("Optimizing...")
    # Вложенные запуски
    # with mlflow.start_run(run_id=run_id, nested=True) as nested_run:
    print("Start nested_run")

    param = {
        "learning_rate": trial.suggest_float("learning_rate", 0.001, 0.1, log=True),
        "depth": trial.suggest_int("depth", 1, 12),
        "l2_leaf_reg": trial.suggest_float("l2_leaf_reg", 0.1, 5),
        "random_strength": trial.suggest_float("random_strength", 0.1, 5),
        "loss_function": "Logloss",
        "task_type": "CPU",
        "random_seed": 0,
        "iterations": 300,
        "verbose": False,
        "auto_class_weights": 'Balanced',
        "cat_features": cat_cols,
    }

    model = CatBoostClassifier(**param)

    skf = StratifiedKFold(n_splits=2,
                        shuffle=True,
                        random_state=555)

    metrics = defaultdict(list)
    #  train_index и val_index — это массивы индексов, которые указывают, какие строки из исходного DataFrame X_train (или y_train) должны быть включены в обучающую и валидационную выборки для каждого фолда i, соответственно.
    for i, (train_index, val_index) in enumerate(skf.split(X_train, y_train)):
        # Разделяем данные на обучающую и валидационную выборки
        train_x, val_x = X_train.iloc[train_index], X_train.iloc[val_index]
        train_y, val_y = y_train.iloc[train_index], y_train.iloc[val_index]

        # Обучаем модель на обучающей выборке
        model.fit(train_x, train_y, verbose=False) #verbose=False для уменьшения вывода

        # Получаем предсказания и вероятности на валидационной выборке
        prediction = model.predict(val_x)
        probas = model.predict_proba(val_x)[:, 1]

        # Вычисление метрик
        # Вычисление confusion matrix с нормализацией
        cm = confusion_matrix(val_y, prediction, normalize='all')
        # Извлечение значений из нормализованной матрицы
        TN = cm[0, 0]  # True Negative
        FP = cm[0, 1]  # False Positive (ошибка первого рода)
        FN = cm[1, 0]  # False Negative (ошибка второго рода)
        TP = cm[1, 1]  # True Positive
        # err_1 — ошибка первого рода
        # err_2 — ошибка второго рода
        err_1 = FP  # False Positive Rate (доля от общего числа)
        err_2 = FN  # False Negative Rate (доля от общего числа)

        auc = roc_auc_score(val_y, probas)
        precision = precision_score(val_y, prediction)
        recall = recall_score(val_y, prediction)
        f1 = f1_score(val_y, prediction)
        logloss = log_loss(val_y, probas)

        metrics["err1"].append(err_1)
        metrics["err2"].append(err_2)
        metrics["auc"].append(auc)
        metrics["precision"].append(precision)
        metrics["recall"].append(recall)
        metrics["f1"].append(f1)
        metrics["logloss"].append(logloss)

    err_1 = np.median(metrics['err1'])
    err_2 = np.median(metrics['err2'])
    auc = np.median(metrics['auc'])
    precision = np.median(metrics["precision"])
    recall = np.median(metrics["recall"])
    f1 = np.median(metrics["f1"])
    logloss = np.median(metrics["logloss"])

    mlflow.log_metric("err1", err_1)
    mlflow.log_metric("err2", err_2)
    mlflow.log_metric("auc", auc)
    mlflow.log_metric("precision", precision)
    mlflow.log_metric("recall", recall)
    mlflow.log_metric("f1", f1)
    mlflow.log_metric("logloss", logloss)

    # mlflow.end_run()

    return auc


with mlflow.start_run(run_name=RUN_NAME, experiment_id=experiment_id) as run:
    run_id = run.info.run_id

    # MLflow Callback
    mlflc = MLflowCallback(
        tracking_uri=f'http://{TRACKING_SERVER_HOST}:{TRACKING_SERVER_PORT}',
        metric_name="AUC",
        create_experiment=False, # не нужно создавать эксперимент каждый раз
        mlflow_kwargs = {'experiment_id': experiment_id, 'tags': {"mlflow.parentRunId": run_id}}
    )

    study = optuna.create_study(direction='maximize',
                                study_name=STUDY_NAME,
                                storage=STUDY_DB_NAME,
                                load_if_exists=True,
                                sampler=optuna.samplers.TPESampler())

    study.optimize(objective,
                    n_trials=10,
                    callbacks=[mlflc])
    best_params = study.best_trial.params



    # Создаем модель с лучшими параметрами
    best_model = CatBoostClassifier(**best_params,
                                    auto_class_weights='Balanced',
                                    cat_features=cat_cols,
                                    loss_function="Logloss", # нужно указать явно, т.к. его нет в best_params
                                    task_type="CPU",
                                    random_seed=555,
                                    verbose=False)

    # Обучаем модель на ВСЕМ тренировочном наборе
    best_model.fit(X_train, y_train)

    # Логируем модель в MLflow
    # Код логирования модели находится вне цикла оптимизации и вызывается только один раз после того, как Optuna завершит поиск лучших гиперпараметров.
    mlflow.catboost.log_model(
        cb_model=best_model,
        artifact_path="models", # директория для хранения модели
        registered_model_name=REGISTRY_MODEL_NAME,
        pip_requirements=pip_requirements,
        signature=signature,
        input_example=input_example,
        await_registration_for=await_registration_for
        )
    print("Best model logged to MLflow.")
    # Получаем данные о запуске эксперимента по его уникальному идентификатору
    run = mlflow.get_run(run_id)


    print(f"Number of finished trials: {len(study.trials)}")
    print(f"Best params: {best_params}")
    print(f"Данные о запуске\n{run}")