"""
Урок 7: Hyperparameter Sweep — перебор гиперпараметров
========================================================
Запускает множество комбинаций гиперпараметров и логирует каждый
вариант в MLflow. В UI можно сравнить все варианты и найти лучший.

Запуск:
    python src/hyperparam_sweep.py
    python src/hyperparam_sweep.py --max-combos 50  # больше комбинаций
"""
import argparse
import itertools
import random
import time

from sklearn.datasets import load_digits
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
from sklearn.metrics import recall_score, precision_score, f1_score, accuracy_score
from sklearn.model_selection import train_test_split

import mlflow


def main():
    parser = argparse.ArgumentParser(description="MLflow hyperparameter sweep")
    parser.add_argument("--max-combos", type=int, default=20,
                        help="максимум комбинаций для перебора")
    parser.add_argument("--cv-folds", type=int, default=5,
                        help="фолдов кросс-валидации")
    parser.add_argument("--experiment-name", type=str, default="sweep_digits")
    args = parser.parse_args()

    # ─── Сетка гиперпараметров ───────────────────────────────
    param_grid = {
        "n_estimators":   [50, 100, 150, 200, 300],
        "max_depth":       [4, 6, 8, 12, 16, None],
        "min_samples_split": [2, 5, 10],
        "min_samples_leaf":  [1, 2, 4],
        "max_features":    ["sqrt", "log2", None],
    }

    # Генерируем все комбинации, берём случайные max_combos
    all_combos = list(itertools.product(
        param_grid["n_estimators"],
        param_grid["max_depth"],
        param_grid["min_samples_split"],
        param_grid["min_samples_leaf"],
        param_grid["max_features"],
    ))
    random.seed(42)
    random.shuffle(all_combos)
    combos = all_combos[:args.max_combos]

    print(f"🔬 Hyperparameter Sweep")
    print(f"   Всего возможных комбинаций: {len(all_combos)}")
    print(f"   Будет запущено: {len(combos)}")
    print(f"   CV фолдов: {args.cv_folds}")
    print(f"   Эксперимент: {args.experiment_name}")
    print()

    # ─── Данные ──────────────────────────────────────────────
    digits = load_digits()
    X, y = digits.data, digits.target
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    mlflow.set_experiment(args.experiment_name)

    best_acc = 0.0
    best_run_id = None
    best_params = None

    for i, (n_est, depth, min_split, min_leaf, max_feat) in enumerate(combos):
        params = {
            "n_estimators": n_est,
            "max_depth": depth,
            "min_samples_split": min_split,
            "min_samples_leaf": min_leaf,
            "max_features": max_feat,
        }

        # Описание для имени run
        depth_str = str(depth) if depth else "None"
        feat_str = str(max_feat) if max_feat else "None"
        run_name = f"rf_ne{n_est}_d{depth_str}_ms{min_split}_ml{min_leaf}_mf{feat_str}"

        with mlflow.start_run(run_name=run_name) as run:
            # Логируем ВСЕ параметры
            mlflow.log_params(params)

            # Обучаем
            model = RandomForestClassifier(
                random_state=42, n_jobs=-1, **params
            )
            model.fit(X_train, y_train)

            # Кросс-валидация (более надёжная оценка)
            cv_scores = cross_val_score(model, X_train, y_train,
                                         cv=args.cv_folds, scoring="accuracy")

            # Метрики на тесте
            y_pred = model.predict(X_test)
            test_acc = accuracy_score(y_test, y_pred)
            test_recall = recall_score(y_test, y_pred, average="macro")
            test_precision = precision_score(y_test, y_pred, average="macro")
            test_f1 = f1_score(y_test, y_pred, average="macro")

            # Логируем метрики
            mlflow.log_metric("cv_mean_accuracy", cv_scores.mean())
            mlflow.log_metric("cv_std_accuracy", cv_scores.std())
            mlflow.log_metric("test_accuracy", test_acc)
            mlflow.log_metric("test_recall_macro", test_recall)
            mlflow.log_metric("test_precision_macro", test_precision)
            mlflow.log_metric("test_f1_macro", test_f1)

            # Тег — номер комбинации
            mlflow.set_tag("combo_index", str(i))

            # Отслеживаем лучший
            if test_acc > best_acc:
                best_acc = test_acc
                best_run_id = run.info.run_id
                best_params = params

            # Прогресс
            bar_len = 30
            filled = int(bar_len * (i + 1) / len(combos))
            bar = "█" * filled + "░" * (bar_len - filled)
            print(f"\r  [{bar}] {i+1}/{len(combos)} | "
                  f"acc={test_acc:.4f} | {run_name[:40]:<40}", end="")

    print(f"\n\n{'='*60}")
    print(f"🏆 Лучший результат:")
    print(f"   Run ID: {best_run_id}")
    print(f"   Accuracy: {best_acc:.4f}")
    print(f"   Параметры: {best_params}")
    print(f"\n📊 Откройте MLflow UI:")
    print(f"   Эксперимент: {args.experiment_name}")
    print(f"   Отсортируйте по test_accuracy (клик на заголовок колонки)")
    print(f"   Сравните cv_mean_accuracy vs test_accuracy")
    print(f"   если cv_std_accuracy высокий — модель нестабильна")


if __name__ == "__main__":
    main()