"""
Урок 6: MLflow Autolog — автоматическое логирование
====================================================
Вместо ручного log_param/log_metric можно вызвать mlflow.autolog()
ОДИН раз — и MLflow сам залогирует всё: параметры, метрики, модель,
признаки, даже feature importance.

Сравните с train_simple.py — там всё логировалось вручную.

Запуск:
    python src/autolog_demo.py
"""
from sklearn.datasets import load_wine
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, recall_score

import mlflow
import mlflow.sklearn

# Подавляем INFO MLflow о переменных окружения (напр. OPENAI_API_KEY) при логировании модели
import os
os.environ.setdefault("MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING", "false")
mlflow.set_tracking_uri(os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5555"))


def main():
    # ─── КЛЮЧЕВАЯ СТРОКА ──────────────────────────────────────
    # autolog() включает автоматическое логирование для sklearn.
    # Доступно для: sklearn, pytorch, tensorflow, xgboost, lightgbm,
    #               keras, fastai, spark, autogluon, statsmodels
    mlflow.sklearn.autolog(
        log_input_examples=False,   # не логировать примеры входа
        log_model_signatures=False,  # не строить сигнатуру модели
        log_models=True,             # сохранить модель автоматически
        log_datasets=True,           # залогировать датасет
        max_tuning_runs=10,           # лимит для GridSearch
        log_post_training_metrics=True,
    )
    # ─────────────────────────────────────────────────────────

    mlflow.set_experiment("autolog_wine")

    # Датасет wine — классификация вин по 3 сортам
    wine = load_wine()
    X, y = wine.data, wine.target
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42
    )
    print(f"📊 Датасет wine: {X.shape[0]} образцов, {X.shape[1]} признаков, 3 класса")

    # autolog начинает работать автоматически внутри start_run
    with mlflow.start_run(run_name="gb_autolog") as run:
        print(f"MLflow run ID: {run.info.run_id}")

        # Просто обучаем модель — НИКАКИХ log_param/log_metric!
        model = GradientBoostingClassifier(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=3,
            random_state=42,
        )
        model.fit(X_train, y_train)

        # Можно добавить и ручные метрики — они дополнят автолог
        y_pred = model.predict(X_test)
        acc = accuracy_score(y_test, y_pred)
        recall_macro = recall_score(y_test, y_pred, average="macro")
        recall_weighted = recall_score(y_test, y_pred, average="weighted")
        mlflow.log_metric("manual_accuracy", acc)
        mlflow.log_metric("manual_recall_macro", recall_macro)
        mlflow.log_metric("manual_recall_weighted", recall_weighted)

        print(f"\n✅ Готово! Откройте MLflow UI:")
        print(f"   Эксперимент: autolog_wine")
        print(f"   Run: {run.info.run_id}")
        print(f"   Accuracy:         {acc:.4f}")
        print(f"   Recall (macro):   {recall_macro:.4f}")
        print(f"   Recall (weighted):{recall_weighted:.4f}")
        print(f"\n🔍 Что autolog залогировал автоматически:")
        print(f"   • Параметры: n_estimators, learning_rate, max_depth, ...")
        print(f"   • Метрики: training_accuracy, training_log_loss, ...")
        print(f"   • Модель: сохранена в artifacts/model/")
        print(f"   • Датасет: профиль в artifacts/")
        print(f"   • Ручная метрика: manual_accuracy")


if __name__ == "__main__":
    main()