malware_classification.py

"""
Malware Classification using Hyperopt (Bayesian Optimization via TPE).
Enhanced version with comprehensive reporting and visualization.

Supports two classification modes:
- maliciousness: Binary classification (benign=0 vs malicious=1)
- family: Multi-class classification by malware family

Features:
- Multiple ML models with hyperparameter tuning (including ResNet-18)
- Stratified cross-validation
- Optional grouped CV by family (for maliciousness mode)
- Sample weighting for imbalanced classes
- Comprehensive reporting with plots and saved results
"""
import numpy as np
import warnings
import os
import json
from datetime import datetime
# Am Anfang des Skripts, VOR jedem TensorFlow-Import:

import multiprocessing
multiprocessing.set_start_method('spawn', force=True)

from matplotlib import pyplot as plt
from sklearn.decomposition import PCA
from sklearn.exceptions import ConvergenceWarning

warnings.filterwarnings("ignore", category=ConvergenceWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings('ignore')
os.environ["PYTHONWARNINGS"] = "ignore"

from typing import Dict, List, Tuple, Any, Optional, Literal, Set
import weaviate

from sklearn.model_selection import (
    cross_val_score, StratifiedKFold, StratifiedGroupKFold, train_test_split, GroupShuffleSplit
)
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.dummy import DummyClassifier
from sklearn.metrics import (
    classification_report, confusion_matrix, ConfusionMatrixDisplay,
    f1_score, precision_score, recall_score, accuracy_score
)
from sklearn.base import BaseEstimator, ClassifierMixin
import joblib

import xgboost as xgb
import lightgbm as lgb

# Hyperopt
from hyperopt import fmin, tpe, hp, STATUS_OK, Trials, space_eval

import pandas as pd

WEAVIATE_HOST = "weaviate.malwareuniverse.org"
WEAVIATE_HTTP_PORT = 443
WEAVIATE_GRPC_PORT = 50051
COLLECTION_NAME = "pre_traind_distilberd3_3"
SPLIT_KEY_SHA_FIELDS = ("sha256_hash", "sha256", "sha256sum", "shasum")

# Valid classification modes
ClassificationMode = Literal["maliciousness", "malware_family", "threat_actor"]


def _extract_split_key(properties: Dict[str, Any], fallback_id: str) -> str:
    for field_name in SPLIT_KEY_SHA_FIELDS:
        value = properties.get(field_name)
        if value not in (None, ""):
            return str(value).strip().lower()

    sample_key = properties.get("sample_id") or properties.get("id")
    if sample_key not in (None, ""):
        return str(sample_key).strip()
    return str(fallback_id).strip()


def _normalize_split_keys(split_keys: List[str]) -> np.ndarray:
    return np.array([str(key).strip().lower() for key in split_keys], dtype=object)


def _build_holdout_indices(
        y: np.ndarray,
        split_keys: List[str],
        holdout_size: float,
        cv_groups: Optional[np.ndarray],
        fixed_holdout_keys: Optional[Set[str]] = None,
        random_state: int = 42,
) -> Tuple[np.ndarray, np.ndarray, str, Set[str]]:
    normalized_split_keys = _normalize_split_keys(split_keys)

    unique_split_key_count = len(set(normalized_split_keys.tolist()))
    duplicate_split_key_count = len(normalized_split_keys) - unique_split_key_count
    if duplicate_split_key_count > 0:
        print(
            f"  Warning: found {duplicate_split_key_count} duplicate split keys. "
            "Cross-collection holdout identity can differ for tied keys."
        )

    if fixed_holdout_keys is not None:
        normalized_target_keys = {str(key).strip().lower() for key in fixed_holdout_keys}
        available_keys = set(normalized_split_keys.tolist())
        missing_target_keys = normalized_target_keys - available_keys
        if missing_target_keys:
            raise ValueError(
                f"Collection is missing {len(missing_target_keys)} required holdout keys. "
                "Cannot enforce identical holdout split."
            )

        holdout_mask = np.array([key in normalized_target_keys for key in normalized_split_keys], dtype=bool)
        holdout_idx = np.where(holdout_mask)[0]
        train_idx = np.where(~holdout_mask)[0]

        if len(train_idx) == 0 or len(holdout_idx) == 0:
            raise ValueError("Invalid fixed holdout split: empty train or holdout partition.")

        split_strategy = f"Fixed holdout by split key ({len(normalized_target_keys)} keys)"
        holdout_key_set = set(normalized_split_keys[holdout_idx].tolist())
        return train_idx, holdout_idx, split_strategy, holdout_key_set

    stable_order = np.argsort(normalized_split_keys, kind="mergesort")
    ordered_indices = np.arange(len(y))[stable_order]
    y_ordered = y[stable_order]
    cv_groups_ordered = cv_groups[stable_order] if cv_groups is not None else None

    if cv_groups_ordered is not None:
        splitter = GroupShuffleSplit(n_splits=1, test_size=holdout_size, random_state=random_state)
        train_pos, holdout_pos = next(
            splitter.split(np.arange(len(y_ordered)), y_ordered, groups=cv_groups_ordered)
        )
        split_strategy = "GroupShuffleSplit on stable split_key order"
    else:
        try:
            train_pos, holdout_pos = train_test_split(
                np.arange(len(y_ordered)),
                test_size=holdout_size,
                random_state=random_state,
                stratify=y_ordered
            )
            split_strategy = "Stratified train_test_split on stable split_key order"
        except ValueError:
            train_pos, holdout_pos = train_test_split(
                np.arange(len(y_ordered)),
                test_size=holdout_size,
                random_state=random_state
            )
            split_strategy = "Unstratified train_test_split on stable split_key order"

    train_idx = ordered_indices[train_pos]
    holdout_idx = ordered_indices[holdout_pos]
    holdout_key_set = set(normalized_split_keys[holdout_idx].tolist())
    return train_idx, holdout_idx, split_strategy, holdout_key_set


def _filter_dataset_by_keys(
        X: np.ndarray,
        y: np.ndarray,
        split_keys: List[str],
        groups: Optional[np.ndarray],
        allowed_keys: Set[str],
) -> Tuple[np.ndarray, np.ndarray, List[str], Optional[np.ndarray], int]:
    normalized_split_keys = _normalize_split_keys(split_keys)
    keep_mask = np.array([key in allowed_keys for key in normalized_split_keys], dtype=bool)
    dropped_count = int((~keep_mask).sum())

    if dropped_count == 0:
        return X, y, split_keys, groups, dropped_count

    X_filtered = X[keep_mask]
    y_filtered = y[keep_mask]
    groups_filtered = groups[keep_mask] if groups is not None else None
    split_keys_filtered = [key for key, keep in zip(split_keys, keep_mask) if keep]
    return X_filtered, y_filtered, split_keys_filtered, groups_filtered, dropped_count


def _scalar_holdout_metrics(metrics: Dict[str, Any]) -> Dict[str, float]:
    report = metrics.get("classification_report", {})
    macro_avg = report.get("macro avg", {}) if isinstance(report, dict) else {}
    weighted_avg = report.get("weighted avg", {}) if isinstance(report, dict) else {}
    return {
        "accuracy": float(metrics["accuracy"]),
        "f1_weighted": float(metrics["f1_weighted"]),
        "f1_macro": float(metrics["f1_macro"]),
        "precision_macro": float(macro_avg.get("precision", 0.0)),
        "recall_macro": float(macro_avg.get("recall", 0.0)),
        "precision_weighted": float(metrics["precision_weighted"]),
        "recall_weighted": float(metrics["recall_weighted"]),
    }


def _build_holdout_details_payload(
        tuner: Any,
        metrics: Dict[str, Any],
        split_strategy: str,
        holdout_size: float,
        holdout_key_set: Set[str],
) -> Dict[str, Any]:
    report = metrics.get("classification_report", {})
    y_true = np.asarray(metrics.get("y_true", []))
    y_pred = np.asarray(metrics.get("y_pred", []))

    true_labels, true_counts = np.unique(y_true, return_counts=True) if y_true.size else ([], [])
    pred_labels, pred_counts = np.unique(y_pred, return_counts=True) if y_pred.size else ([], [])

    def _label_name(raw_label: Any) -> str:
        try:
            return str(tuner.classes_[int(raw_label)])
        except Exception:
            return str(raw_label)

    true_distribution = [
        {
            "label_encoded": int(label),
            "label_name": _label_name(label),
            "count": int(count),
        }
        for label, count in zip(true_labels, true_counts)
    ]
    predicted_distribution = [
        {
            "label_encoded": int(label),
            "label_name": _label_name(label),
            "count": int(count),
        }
        for label, count in zip(pred_labels, pred_counts)
    ]

    return {
        "best_model_name": tuner.best_model_name,
        "best_model_cv_score": float(tuner.results[tuner.best_model_name]["cv_score_mean"]),
        "best_model_cv_std": float(tuner.results[tuner.best_model_name]["cv_score_std"]),
        "best_model_params": tuner.best_params,
        "split_strategy": split_strategy,
        "requested_holdout_size": float(holdout_size),
        "holdout_split_key_count": int(len(holdout_key_set)),
        "holdout_sample_count": int(y_true.size),
        "overall_metrics": _scalar_holdout_metrics(metrics),
        "classification_report": report,
        "macro_avg": report.get("macro avg", {}) if isinstance(report, dict) else {},
        "weighted_avg": report.get("weighted avg", {}) if isinstance(report, dict) else {},
        "confusion_matrix": metrics.get("confusion_matrix"),
        "true_label_distribution": true_distribution,
        "predicted_label_distribution": predicted_distribution,
    }

class ResultsManager:
    """Manages saving and organizing all experiment results."""

    def __init__(self, output_dir: str = "results"):
        self.output_dir = output_dir
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_dir = os.path.join(output_dir, f"run_{self.timestamp}")
        self.plots_dir = os.path.join(self.run_dir, "plots")
        self.models_dir = os.path.join(self.run_dir, "models")
        self.reports_dir = os.path.join(self.run_dir, "reports")

        # Create directories
        for d in [self.run_dir, self.plots_dir, self.models_dir, self.reports_dir]:
            os.makedirs(d, exist_ok=True)

        print(f"✓ Results will be saved to: {self.run_dir}")

    def save_figure(self, fig, name: str, dpi: int = 150):
        """Save a matplotlib figure."""
        filepath = os.path.join(self.plots_dir, f"{name}.png")
        fig.savefig(filepath, dpi=dpi, bbox_inches='tight', facecolor='white', edgecolor='none')
        plt.close(fig)
        print(f"  ✓ Saved: {filepath}")
        return filepath

    def save_json(self, data: Dict, name: str):
        """Save data as JSON."""
        filepath = os.path.join(self.reports_dir, f"{name}.json")

        # Convert numpy types to Python types
        def convert(obj):
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, (np.int64, np.int32)):
                return int(obj)
            elif isinstance(obj, (np.float64, np.float32)):
                return float(obj)
            elif isinstance(obj, dict):
                return {k: convert(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [convert(v) for v in obj]
            return obj

        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(convert(data), f, indent=2)
        print(f"  ✓ Saved: {filepath}")
        return filepath

    def save_csv(self, df: pd.DataFrame, name: str):
        """Save DataFrame as CSV."""
        filepath = os.path.join(self.reports_dir, f"{name}.csv")
        df.to_csv(filepath, index=True)
        print(f"  ✓ Saved: {filepath}")
        return filepath

    def save_text(self, text: str, name: str):
        """Save text content."""
        filepath = os.path.join(self.reports_dir, f"{name}.txt")
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(text)
        print(f"  ✓ Saved: {filepath}")
        return filepath

    def save_model(self, model, name: str):
        """Save a model using joblib."""
        filepath = os.path.join(self.models_dir, f"{name}.joblib")
        joblib.dump(model, filepath)
        print(f"  ✓ Saved: {filepath}")
        return filepath



class ResNet18Classifier(BaseEstimator, ClassifierMixin):
    """1D ResNet-18 Classifier for embedding vectors."""

    def __init__(
            self,
            base_filters: int = 64,
            kernel_size: int = 3,
            dropout_rate: float = 0.3,
            learning_rate: float = 0.001,
            batch_size: int = 32,
            epochs: int = 100,
            use_dropout: bool = True,
            random_state: int = 42,
    ):
        self.base_filters = base_filters
        self.kernel_size = kernel_size
        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.epochs = epochs
        self.use_dropout = bool(use_dropout)
        self.random_state = random_state

        self.model_ = None
        self.classes_ = None
        self.n_classes_ = None
        self.input_dim_ = None
        self.history_ = None

    def _residual_block(self, x, filters, kernel_size, downsample=False):
        from keras import layers, Model
        stride = 2 if downsample else 1
        shortcut = x

        x = layers.Conv1D(filters=filters, kernel_size=kernel_size, strides=stride, padding='same', use_bias=False)(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('relu')(x)

        x = layers.Conv1D(filters=filters, kernel_size=kernel_size, strides=1, padding='same', use_bias=False)(x)
        x = layers.BatchNormalization()(x)

        if downsample or shortcut.shape[-1] != filters:
            shortcut = layers.Conv1D(filters=filters, kernel_size=1, strides=stride, padding='same', use_bias=False)(
                shortcut)
            shortcut = layers.BatchNormalization()(shortcut)

        x = layers.Add()([x, shortcut])
        x = layers.Activation('relu')(x)

        return x

    def _build_model(self):
        import tensorflow as tf
        from tensorflow import keras
        from keras import layers, Model

        tf.random.set_seed(self.random_state)

        inputs = layers.Input(shape=(self.input_dim_, 1))

        x = layers.Conv1D(filters=self.base_filters, kernel_size=7, strides=2, padding='same', use_bias=False)(inputs)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('relu')(x)
        x = layers.MaxPooling1D(pool_size=3, strides=2, padding='same')(x)

        # Stage 1
        x = self._residual_block(x, self.base_filters, self.kernel_size, downsample=False)
        x = self._residual_block(x, self.base_filters, self.kernel_size, downsample=False)
        if self.use_dropout:
            x = layers.Dropout(self.dropout_rate)(x)

        # Stage 2
        x = self._residual_block(x, self.base_filters * 2, self.kernel_size, downsample=True)
        x = self._residual_block(x, self.base_filters * 2, self.kernel_size, downsample=False)
        if self.use_dropout:
            x = layers.Dropout(self.dropout_rate)(x)

        # Stage 3
        x = self._residual_block(x, self.base_filters * 4, self.kernel_size, downsample=True)
        x = self._residual_block(x, self.base_filters * 4, self.kernel_size, downsample=False)
        if self.use_dropout:
            x = layers.Dropout(self.dropout_rate)(x)

        # Stage 4
        x = self._residual_block(x, self.base_filters * 8, self.kernel_size, downsample=True)
        x = self._residual_block(x, self.base_filters * 8, self.kernel_size, downsample=False)
        if self.use_dropout:
            x = layers.Dropout(self.dropout_rate)(x)

        x = layers.GlobalAveragePooling1D()(x)

        if self.n_classes_ == 2:
            outputs = layers.Dense(1, activation='sigmoid')(x)
        else:
            outputs = layers.Dense(self.n_classes_, activation='softmax')(x)

        model = Model(inputs, outputs)

        optimizer = keras.optimizers.Adam(learning_rate=self.learning_rate)

        if self.n_classes_ == 2:
            model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])
        else:
            model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        return model

    def fit(self, X, y, sample_weight=None):
        from tensorflow import keras
        from keras.callbacks import EarlyStopping

        self.classes_ = np.unique(y)
        self.n_classes_ = len(self.classes_)
        self.input_dim_ = X.shape[1]

        X_reshaped = X.reshape(X.shape[0], X.shape[1], 1)

        self.model_ = self._build_model()

        callbacks = [
            EarlyStopping(monitor='val_loss', patience=15, restore_best_weights=True, verbose=0),
            keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5, min_lr=1e-6, verbose=0)
        ]

        if self.n_classes_ == 2:
            y_train = y.astype(np.float32)
        else:
            y_train = y.astype(np.int32)

        self.history_ = self.model_.fit(
            X_reshaped, y_train,
            batch_size=self.batch_size,
            epochs=self.epochs,
            validation_split=0.15,
            callbacks=callbacks,
            sample_weight=sample_weight,
            verbose=0
        )

        return self

    def predict(self, X):
        if self.model_ is None:
            raise ValueError("Model not fitted yet.")

        X_reshaped = X.reshape(X.shape[0], X.shape[1], 1)

        if self.n_classes_ == 2:
            proba = self.model_.predict(X_reshaped, verbose=0)
            return (proba.flatten() > 0.5).astype(int)
        else:
            proba = self.model_.predict(X_reshaped, verbose=0)
            return np.argmax(proba, axis=1)

    def predict_proba(self, X):
        if self.model_ is None:
            raise ValueError("Model not fitted yet.")

        X_reshaped = X.reshape(X.shape[0], X.shape[1], 1)
        proba = self.model_.predict(X_reshaped, verbose=0)

        if self.n_classes_ == 2:
            proba = proba.flatten()
            return np.column_stack([1 - proba, proba])
        else:
            return proba

    def get_params(self, deep=True):
        return {
            'base_filters': self.base_filters,
            'kernel_size': self.kernel_size,
            'dropout_rate': self.dropout_rate,
            'learning_rate': self.learning_rate,
            'batch_size': self.batch_size,
            'epochs': self.epochs,
            'use_dropout': self.use_dropout,
            'random_state': self.random_state,
        }

    def set_params(self, **params):
        for key, value in params.items():
            setattr(self, key, value)
        return self


class MalwareDataLoader:
    """Loads malware embeddings and labels from Weaviate."""

    def __init__(
            self,
            host: str = WEAVIATE_HOST,
            http_port: int = WEAVIATE_HTTP_PORT,
            grpc_port: int = WEAVIATE_GRPC_PORT,
            collection_name: str = COLLECTION_NAME
    ):
        self.host = host
        self.http_port = http_port
        self.grpc_port = grpc_port
        self.collection_name = collection_name
        self.weaviate_client = None

    def _connect_weaviate(self):
        self.weaviate_client = weaviate.connect_to_custom(
            http_host=self.host,
            http_port=self.http_port,
            http_secure=True,
            grpc_host=self.host,
            grpc_port=self.grpc_port,
            grpc_secure=False,
        )

    def _close_weaviate(self):
        if self.weaviate_client:
            self.weaviate_client.close()

    def load_dataset(
            self,
            mode: ClassificationMode = "maliciousness",
            min_family_samples: int = 5,
    ) -> Tuple[np.ndarray, np.ndarray, List[str], Optional[np.ndarray], Dict[str, Any]]:

        print(f"\n{'=' * 50}")
        print(f"Loading malware dataset")
        print(f"  Mode: {mode}")
        print(f"  Collection: {self.collection_name}")
        print(f"{'=' * 50}")

        self._connect_weaviate()

        try:
            collection = self.weaviate_client.collections.get(self.collection_name)
            embeddings = []
            labels = []
            split_keys = []
            families = []
            is_malicious_list = []

            cursor = None
            batch_size = 100

            while True:
                try:
                    if cursor is None:
                        response = collection.query.fetch_objects(limit=batch_size, include_vector=True)
                    else:
                        response = collection.query.fetch_objects(limit=batch_size, include_vector=True, after=cursor)
                except Exception as e:
                    print(f"  Warning: Query failed, trying alternative method...")
                    raise e

                if not response.objects:
                    break

                for obj in response.objects:
                    properties = obj.properties or {}
                    sample_id = properties.get("sample_id") or properties.get("id") or str(obj.uuid)
                    split_key = _extract_split_key(properties, fallback_id=sample_id)
                    family = properties.get("malware_family", "unknown")

                    vec = None

                    if obj.vector:
                        if isinstance(obj.vector, dict):
                            vec = obj.vector.get("default") or list(obj.vector.values())[0]
                        else:
                            vec = obj.vector

                    if vec is None:
                        continue

                    embeddings.append(np.array(vec))
                    families.append(str(family) if family else "unknown")
                    split_keys.append(split_key)

                cursor = response.objects[-1].uuid
                print(f"  Loaded {len(embeddings)} samples...", end='\r')

                if len(response.objects) < batch_size:
                    break

            print(f"\n✓ Loaded {len(embeddings)} samples from Weaviate")

            X = np.array(embeddings)
            families_arr = np.array(families)

            family_encoder = LabelEncoder()
            family_groups = family_encoder.fit_transform(families_arr)

            if mode == "maliciousness":
                y = is_malicious_arr
                groups = family_groups

                label_names = {0: "benign", 1: "malicious"}
                unique, counts = np.unique(y, return_counts=True)

                info = {
                    "mode": mode,
                    "embedding_dim": X.shape[1],
                    "n_samples": len(y),
                    "n_classes": 2,
                    "label_names": label_names,
                    "class_distribution": dict(zip(unique.tolist(), counts.tolist())),
                    "n_families": len(family_encoder.classes_),
                    "family_encoder": family_encoder,
                }

                print(f"\n  Binary classification (maliciousness)")
                print(f"  Embedding dimension: {X.shape[1]}")
                for cls, cnt in zip(unique, counts):
                    print(f"    {label_names[cls]}: {cnt} samples ({cnt / len(y) * 100:.1f}%)")
                print(f"  Families for grouping: {len(family_encoder.classes_)} unique")

            else:  # family mode
                unique_families, family_counts = np.unique(families_arr, return_counts=True)
                valid_families = unique_families[family_counts >= min_family_samples]

                if len(valid_families) < 2:
                    raise ValueError(
                        f"Not enough families with >= {min_family_samples} samples. "
                        f"Found: {len(valid_families)}. Try lowering min_family_samples."
                    )

                valid_mask = np.isin(families_arr, valid_families)
                X = X[valid_mask]
                families_filtered = families_arr[valid_mask]
                split_keys = [key for key, m in zip(split_keys, valid_mask) if m]

                label_encoder = LabelEncoder()
                y = label_encoder.fit_transform(families_filtered)
                groups = None
                label_names = {i: name for i, name in enumerate(label_encoder.classes_)}

                unique, counts = np.unique(y, return_counts=True)

                info = {
                    "mode": mode,
                    "embedding_dim": X.shape[1],
                    "n_samples": len(y),
                    "n_classes": len(label_encoder.classes_),
                    "label_names": label_names,
                    "label_encoder": label_encoder,
                    "class_distribution": dict(zip(unique.tolist(), counts.tolist())),
                    "dropped_families": len(unique_families) - len(valid_families),
                    "min_family_samples": min_family_samples,
                }

                print(f"\n  Multi-class classification (family)")
                print(f"  Embedding dimension: {X.shape[1]}")
                print(f"  Classes: {len(label_encoder.classes_)} families")
                print(
                    f"  Dropped {len(unique_families) - len(valid_families)} families with < {min_family_samples} samples")
                print(f"  Class distribution (top 10):")
                sorted_indices = np.argsort(counts)[::-1][:10]
                for idx in sorted_indices:
                    cls = unique[idx]
                    cnt = counts[idx]
                    print(f"    {label_names[cls]}: {cnt} samples ({cnt / len(y) * 100:.1f}%)")
                if len(unique) > 10:
                    print(f"    ... and {len(unique) - 10} more families")

            return X, y, split_keys, groups, info

        finally:
            self._close_weaviate()

    @staticmethod
    def get_sample_weights(y: np.ndarray, strategy: str = "balanced") -> np.ndarray:
        unique, counts = np.unique(y, return_counts=True)
        n_samples = len(y)

        if strategy == "balanced":
            class_weights = {cls: n_samples / (len(unique) * cnt) for cls, cnt in zip(unique, counts)}
        elif strategy == "sqrt_balanced":
            class_weights = {cls: np.sqrt(n_samples / (len(unique) * cnt)) for cls, cnt in zip(unique, counts)}
        else:
            return np.ones(len(y))

        weights = np.array([class_weights[label] for label in y])
        return weights


class MalwareHyperoptTuner:
    """Hyperparameter tuning using Hyperopt for malware classification."""

    def __init__(
            self,
            random_state: int = 42,
            scoring: str = "f1_weighted"
    ):
        self.random_state = random_state
        self.scoring = scoring
        self.best_model = None
        self.best_params = None
        self.best_model_name = None
        self.results = {}
        self.label_encoder = LabelEncoder()
        self.classes_ = None
        self.info = None
        self.results_manager = None

    def get_hyperopt_spaces(self) -> Dict[str, Dict[str, Any]]:
        """Define Hyperopt search spaces for classification models."""
        spaces = {
            "resnet18": {
                "model_class": ResNet18Classifier,
                "static_params": {"random_state": self.random_state, "epochs": 100, "n_jobs": 1},
                "space": {
                    "base_filters": hp.choice("resnet_base_filters", [32, 64, 128]),
                    "kernel_size": hp.choice("resnet_kernel_size", [3, 5, 7]),
                    "dropout_rate": hp.uniform("resnet_dropout_rate", 0.1, 0.5),
                    "learning_rate": hp.loguniform("resnet_learning_rate", np.log(1e-5), np.log(1e-2)),
                    "batch_size": hp.choice("resnet_batch_size", [16, 32, 64, 128]),
                    "use_dropout": hp.randint("resnet_use_dropout", 2),
                },
                "supports_sample_weight": True,
                "skip_pca": True,
            },
            "xgboost": {
                "model_class": xgb.XGBClassifier,
                "static_params": {"random_state": self.random_state, "verbosity": 0, "use_label_encoder": False,
                                  "eval_metric": "mlogloss", "device": "cpu"},
                "space": {
                    "n_estimators": hp.choice("xgb_n_estimators", [50, 100, 200, 300, 500, 750, 1000]),
                    "max_depth": hp.choice("xgb_max_depth", [3, 4, 5, 6, 7, 8, 10, 12, 15]),
                    "learning_rate": hp.loguniform("xgb_learning_rate", np.log(0.005), np.log(0.3)),
                    "subsample": hp.uniform("xgb_subsample", 0.5, 1.0),
                    "colsample_bytree": hp.uniform("xgb_colsample_bytree", 0.5, 1.0),
                    "min_child_weight": hp.choice("xgb_min_child_weight", [1, 3, 5, 7, 10]),
                    "gamma": hp.uniform("xgb_gamma", 0, 0.5),
                    "reg_alpha": hp.loguniform("xgb_reg_alpha", np.log(1e-6), np.log(10.0)),
                    "reg_lambda": hp.loguniform("xgb_reg_lambda", np.log(1e-3), np.log(100.0)),
                },
                "supports_sample_weight": True
            },
            "lightgbm": {
                "model_class": lgb.LGBMClassifier,
                "static_params": {"random_state": self.random_state, "verbosity": -1, "device": "cpu"},
                "space": {
                    "n_estimators": hp.choice("lgb_n_estimators", [50, 100, 200, 300, 500, 750, 1000]),
                    "max_depth": hp.choice("lgb_max_depth", [-1, 3, 5, 7, 10, 15, 20]),
                    "learning_rate": hp.loguniform("lgb_learning_rate", np.log(0.005), np.log(0.3)),
                    "num_leaves": hp.choice("lgb_num_leaves", [15, 31, 63, 127, 255]),
                    "subsample": hp.uniform("lgb_subsample", 0.5, 1.0),
                    "colsample_bytree": hp.uniform("lgb_colsample_bytree", 0.5, 1.0),
                    "min_child_samples": hp.choice("lgb_min_child_samples", [5, 10, 20, 50, 100]),
                    "reg_alpha": hp.loguniform("lgb_reg_alpha", np.log(1e-6), np.log(10.0)),
                    "reg_lambda": hp.loguniform("lgb_reg_lambda", np.log(1e-6), np.log(10.0)),
                    "class_weight": hp.choice("lgb_class_weight", [None, "balanced"]),
                },
                "supports_sample_weight": True
            },
            "random_forest": {
                "model_class": RandomForestClassifier,
                "static_params": {"random_state": self.random_state, "n_jobs": -1},
                "space": {
                    "n_estimators": hp.choice("rf_n_estimators", [50, 100, 200, 300, 500, 750]),
                    "max_depth": hp.choice("rf_max_depth", [None, 10, 20, 30, 50, 100]),
                    "min_samples_split": hp.choice("rf_min_samples_split", [2, 5, 10, 20]),
                    "min_samples_leaf": hp.choice("rf_min_samples_leaf", [1, 2, 4, 8]),
                    "max_features": hp.choice("rf_max_features", ["sqrt", "log2", 0.3, 0.5, 0.7]),
                    "class_weight": hp.choice("rf_class_weight", [None, "balanced", "balanced_subsample"]),
                },
                "supports_sample_weight": True
            },
            "svc": {
                "model_class": SVC,
                "static_params": {"random_state": self.random_state, "probability": True},
                "space": {
                    "C": hp.loguniform("svc_C", np.log(0.01), np.log(1000)),
                    "kernel": hp.choice("svc_kernel", ["rbf", "linear", "poly"]),
                    "gamma": hp.choice("svc_gamma", ["scale", "auto", 0.001, 0.01, 0.1]),
                    "degree": hp.choice("svc_degree", [2, 3, 4]),
                    "class_weight": hp.choice("svc_class_weight", [None, "balanced"]),
                },
                "supports_sample_weight": True
            },
            "logistic_regression": {
                "model_class": LogisticRegression,
                "static_params": {"random_state": self.random_state, "max_iter": 2000, "n_jobs": -1},
                "space": {
                    "C": hp.loguniform("lr_C", np.log(1e-4), np.log(100)),
                    "solver": hp.choice("lr_solver", ["lbfgs", "saga"]),
                    "class_weight": hp.choice("lr_class_weight", [None, "balanced"]),
                    "penalty": hp.choice("lr_penalty", ["l2", None]),
                },
                "supports_sample_weight": True
            },
            "knn": {
                "model_class": KNeighborsClassifier,
                "static_params": {"n_jobs": -1},
                "space": {
                    "n_neighbors": hp.choice("knn_n_neighbors", [1, 3, 5, 7, 11, 15, 21, 31]),
                    "weights": hp.choice("knn_weights", ["uniform", "distance"]),
                    "metric": hp.choice("knn_metric", ["euclidean", "cosine", "manhattan"]),
                    "leaf_size": hp.choice("knn_leaf_size", [20, 30, 50]),
                },
                "supports_sample_weight": False
            },
        }
        return spaces

    def _build_preprocessor(self, pca_components: Any = None) -> Pipeline:
        steps = [("scaler", StandardScaler())]
        if pca_components is not None:
            steps.append(("pca", PCA(n_components=pca_components)))
        return Pipeline(steps)

    def _create_objective(
            self,
            X: np.ndarray,
            y: np.ndarray,
            model_class,
            static_params: Dict,
            cv,
            groups: Optional[np.ndarray] = None,
            sample_weights: Optional[np.ndarray] = None,
            supports_sample_weight: bool = True,
            skip_pca: bool = False,
    ):
        def objective(params):
            if skip_pca:
                pca_components = None
                model_params = params
            else:
                pca_components = params.get("pca_components", None)
                model_params = {k: v for k, v in params.items() if k != "pca_components"}

            all_params = {**static_params, **model_params}
            model = model_class(**all_params)

            if skip_pca:
                preprocessor = Pipeline([("scaler", StandardScaler())])
            else:
                preprocessor = self._build_preprocessor(pca_components)

            pipeline = Pipeline([("preprocessor", preprocessor), ("model", model)])

            params = {}
            if sample_weights is not None and supports_sample_weight:
                params["model__sample_weight"] = sample_weights

            if groups is not None:
                scores = cross_val_score(
                    pipeline, X, y, cv=cv, scoring=self.scoring,
                    n_jobs=1 if model_class == ResNet18Classifier else -1,
                    groups=groups, params=params if params else None
                )
            else:
                scores = cross_val_score(
                    pipeline, X, y, cv=cv, scoring=self.scoring,
                    n_jobs=1 if model_class == ResNet18Classifier else -1,
                    params=params if params else None
                )

            return {
                "loss": -scores.mean(),
                "status": STATUS_OK,
                "score_mean": scores.mean(),
                "score_std": scores.std(),
                "fold_scores": scores.tolist(),
                "pca_components": pca_components
            }

        return objective

    def tune_model(
            self,
            X: np.ndarray,
            y: np.ndarray,
            model_name: str,
            max_evals: int = 100,
            cv=None,
            groups: Optional[np.ndarray] = None,
            sample_weights: Optional[np.ndarray] = None,
            use_pca: bool = True,
    ) -> Dict[str, Any]:
        """Tune a single model using Hyperopt."""

        spaces = self.get_hyperopt_spaces()
        if model_name not in spaces:
            raise ValueError(f"Unknown model: {model_name}. Available: {list(spaces.keys())}")

        config = spaces[model_name]
        supports_sample_weight = config.get("supports_sample_weight", True)
        skip_pca = config.get("skip_pca", False)

        space = config["space"].copy()
        if use_pca and not skip_pca:
            space["pca_components"] = hp.choice(f"{model_name}_pca", [None, 0.90, 0.95, 0.99, 50, 100, 200])

        if cv is None:
            if groups is not None:
                cv = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=self.random_state)
            else:
                cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=self.random_state)

        cv_type = "StratifiedGroupKFold" if groups is not None else "StratifiedKFold"

        print(f"\n{'=' * 60}")
        print(f"Hyperopt Tuning: {model_name}")
        print(f"  Scoring: {self.scoring}")
        print(f"  Max evaluations: {max_evals}")
        print(f"  CV strategy: {cv_type} with {cv.n_splits} folds")
        print(f"  Using sample weights: {sample_weights is not None and supports_sample_weight}")
        print(f"  PCA search: {use_pca and not skip_pca}")
        print(f"{'=' * 60}")

        objective = self._create_objective(
            X, y, config["model_class"], config["static_params"], cv,
            groups=groups, sample_weights=sample_weights,
            supports_sample_weight=supports_sample_weight, skip_pca=skip_pca,
        )

        trials = Trials()
        best_indices = fmin(
            fn=objective, space=space, algo=tpe.suggest, max_evals=max_evals,
            trials=trials, rstate=np.random.default_rng(self.random_state),
            show_progressbar=True
        )

        best_params = space_eval(space, best_indices)
        best_trial = trials.best_trial
        best_score = best_trial["result"]["score_mean"]
        best_score_std = best_trial["result"]["score_std"]

        pca_components = best_params.pop("pca_components", None) if not skip_pca else None
        final_params = {**config["static_params"], **best_params}
        final_model = config["model_class"](**final_params)

        if skip_pca:
            preprocessor = Pipeline([("scaler", StandardScaler())])
        else:
            preprocessor = self._build_preprocessor(pca_components)

        final_pipeline = Pipeline([("preprocessor", preprocessor), ("model", final_model)])

        if sample_weights is not None and supports_sample_weight:
            final_pipeline.fit(X, y, model__sample_weight=sample_weights)
        else:
            final_pipeline.fit(X, y)

        # Extract trial history for plotting
        trial_history = []
        for i, trial in enumerate(trials.trials):
            trial_history.append({
                "trial": i + 1,
                "score": trial["result"]["score_mean"],
                "std": trial["result"]["score_std"],
                "loss": trial["result"]["loss"]
            })

        result = {
            "model_name": model_name,
            "best_params": best_params,
            "best_params_with_static": final_params,
            "pca_components": pca_components,
            "cv_score_mean": best_score,
            "cv_score_std": best_score_std,
            "scoring_metric": self.scoring,
            "best_estimator": final_pipeline,
            "trials": trials,
            "trial_history": trial_history,
            "n_evals": len(trials.trials),
            "used_sample_weights": sample_weights is not None and supports_sample_weight,
            "cv_strategy": cv_type
        }

        print(f"\n  ✓ Best params: {best_params}")
        if not skip_pca:
            print(f"  ✓ PCA components: {pca_components}")
        print(f"  ✓ CV {self.scoring}: {best_score:.4f} (+/- {best_score_std:.4f})")

        return result

    def tune_all_models(
            self,
            X: np.ndarray,
            y: np.ndarray,
            info: Dict[str, Any],
            models_to_tune: Optional[List[str]] = None,
            max_evals: int = 100,
            k_folds: int = 5,
            groups: Optional[np.ndarray] = None,
            sample_weights: Optional[np.ndarray] = None,
            use_pca: bool = True,
    ) -> Dict[str, Dict]:
        """Tune all specified models."""

        self.info = info
        spaces = self.get_hyperopt_spaces()

        if models_to_tune is None:
            models_to_tune = list(spaces.keys())
        else:
            invalid = set(models_to_tune) - set(spaces.keys())
            if invalid:
                raise ValueError(f"Unknown models: {invalid}. Available: {list(spaces.keys())}")

        if groups is not None:
            cv = StratifiedGroupKFold(n_splits=k_folds, shuffle=True, random_state=self.random_state)
            cv_type = "StratifiedGroupKFold"
        else:
            cv = StratifiedKFold(n_splits=k_folds, shuffle=True, random_state=self.random_state)
            cv_type = "StratifiedKFold"

        self.classes_ = np.unique(y)
        y_encoded = self.label_encoder.fit_transform(y)

        print(f"\n{'#' * 60}")
        print(f"# Hyperopt Classification Tuning - {len(models_to_tune)} models")
        print(f"# Mode: {info['mode']}")
        print(f"# Scoring: {self.scoring}")
        print(f"# CV Strategy: {cv_type} with {k_folds} folds, {max_evals} evals/model")
        print(f"# Embedding dimension: {info['embedding_dim']}")
        print(f"# Classes: {info['n_classes']}")
        print(f"# Samples: {info['n_samples']}")
        if groups is not None:
            print(f"# Groups (families): {len(np.unique(groups))}")
        print(f"# Using sample weights: {sample_weights is not None}")
        print(f"{'#' * 60}")

        best_score = -float('inf')

        for model_name in models_to_tune:
            result = self.tune_model(
                X, y_encoded, model_name, max_evals=max_evals, cv=cv,
                groups=groups, sample_weights=sample_weights, use_pca=use_pca,
            )
            self.results[model_name] = result

            if result["cv_score_mean"] > best_score:
                best_score = result["cv_score_mean"]
                self.best_model = result["best_estimator"]
                self.best_params = result["best_params"]
                self.best_model_name = model_name

        return self.results

    def print_summary(self):
        """Print summary of all results."""

        print(f"\n{'=' * 80}")
        print(f"HYPEROPT RESULTS (sorted by {self.scoring})")
        print(f"{'=' * 80}")

        sorted_results = sorted(self.results.items(), key=lambda x: x[1]["cv_score_mean"], reverse=True)

        print(f"\n{'Model':<22} {self.scoring:<16} {'Std':<10} {'PCA':<12} {'CV Type':<20}")
        print("-" * 80)

        for name, result in sorted_results:
            marker = " ★" if name == self.best_model_name else ""
            pca_str = str(result.get("pca_components", "None"))[:10]
            cv_type = result.get("cv_strategy", "Unknown")
            print(
                f"{name:<22} {result['cv_score_mean']:<16.4f} {result['cv_score_std']:<10.4f} {pca_str:<12} {cv_type:<20}{marker}")

        print(f"\n✓ Best model: {self.best_model_name}")
        print(f"  Best params: {self.best_params}")

    def predict(self, X: np.ndarray) -> np.ndarray:
        if self.best_model is None:
            raise ValueError("No model trained yet.")
        y_pred = self.best_model.predict(X)
        return self.label_encoder.inverse_transform(y_pred)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        if self.best_model is None:
            raise ValueError("No model trained yet.")
        return self.best_model.predict_proba(X)

    def evaluate(self, X: np.ndarray, y: np.ndarray, groups: Optional[np.ndarray] = None) -> Dict[str, Any]:
        """Evaluate the selected best model on a holdout set."""
        if self.best_model is None:
            raise ValueError("No model trained yet.")

        y_encoded = self.label_encoder.transform(y)
        y_pred = self.best_model.predict(X)

        if self.info and "label_names" in self.info:
            target_names = [str(self.info["label_names"].get(i, i)) for i in range(len(self.classes_))]
        else:
            target_names = [str(c) for c in self.classes_]

        report = classification_report(y_encoded, y_pred, target_names=target_names, output_dict=True)

        metrics = {
            "accuracy": accuracy_score(y_encoded, y_pred),
            "f1_weighted": f1_score(y_encoded, y_pred, average="weighted"),
            "f1_macro": f1_score(y_encoded, y_pred, average="macro"),
            "precision_weighted": precision_score(y_encoded, y_pred, average="weighted"),
            "recall_weighted": recall_score(y_encoded, y_pred, average="weighted"),
            "classification_report": report,
            "confusion_matrix": confusion_matrix(y_encoded, y_pred),
            "y_true": y_encoded,
            "y_pred": y_pred,
        }

        print(f"\n{'=' * 60}")
        print("EVALUATION METRICS (Holdout Set)")
        print(f"{'=' * 60}")
        print(f"  Holdout samples:   {len(y_encoded)}")
        print(f"  Accuracy:          {metrics['accuracy']:.4f}")
        print(f"  F1 (weighted):     {metrics['f1_weighted']:.4f}")
        print(f"  F1 (macro):        {metrics['f1_macro']:.4f}")
        print(f"  Precision (weighted): {metrics['precision_weighted']:.4f}")
        print(f"  Recall (weighted):    {metrics['recall_weighted']:.4f}")
        print(f"\n  Classification Report:")
        print(classification_report(y_encoded, y_pred, target_names=target_names))

        return metrics

    def generate_all_reports(
            self,
            X: np.ndarray,
            y: np.ndarray,
            groups: Optional[np.ndarray] = None,
            baseline_score: float = None,
            results_manager: ResultsManager = None
    ):
        """Generate comprehensive reports and visualizations."""

        if results_manager is None:
            results_manager = ResultsManager()

        self.results_manager = results_manager

        print(f"\n{'=' * 60}")
        print("GENERATING COMPREHENSIVE REPORTS")
        print(f"{'=' * 60}")

        # 1. Model Comparison Bar Chart
        self._plot_model_comparison()

        # 2. Hyperopt Convergence Plots
        self._plot_convergence_curves()

        # 3. Confusion Matrix for Best Model
        self._plot_confusion_matrix(X, y, groups)

        # 4. Per-class Performance
        self._plot_per_class_performance(X, y, groups)

        # 5. Class Distribution
        self._plot_class_distribution(y)

        # 6. Hyperparameter Importance (for tree-based models)
        self._plot_hyperparameter_analysis()

        # 7. Save all model results as JSON
        self._save_results_json(baseline_score)

        # 8. Save summary table as CSV
        self._save_summary_csv()

        # 9. Generate text report
        self._save_text_report(baseline_score)

        # 10. Save all trained models
        self._save_all_models()

        print(f"\n✓ All reports saved to: {results_manager.run_dir}")

        return results_manager.run_dir

    def _plot_model_comparison(self):
        """Create bar chart comparing all models."""
        fig, ax = plt.subplots(figsize=(12, 6))

        models = []
        scores = []
        stds = []
        colors = []

        sorted_results = sorted(self.results.items(), key=lambda x: x[1]["cv_score_mean"], reverse=True)

        for name, result in sorted_results:
            models.append(name)
            scores.append(result["cv_score_mean"])
            stds.append(result["cv_score_std"])
            colors.append('#2ecc71' if name == self.best_model_name else '#3498db')

        bars = ax.barh(models, scores, xerr=stds, color=colors, edgecolor='black', capsize=3)

        ax.set_xlabel(f'{self.scoring} Score', fontsize=12)
        ax.set_title('Model Comparison (Cross-Validated)', fontsize=14, fontweight='bold')
        ax.set_xlim(0, 1)

        # Add value labels
        for bar, score, std in zip(bars, scores, stds):
            ax.text(score + std + 0.02, bar.get_y() + bar.get_height() / 2,
                    f'{score:.4f}', va='center', fontsize=10)

        ax.axvline(x=max(scores), color='red', linestyle='--', alpha=0.5, label='Best Score')
        ax.legend(loc='lower right')

        plt.tight_layout()
        self.results_manager.save_figure(fig, "model_comparison")

    def _plot_convergence_curves(self):
        """Plot hyperopt convergence for each model."""
        n_models = len(self.results)
        n_cols = min(3, n_models)
        n_rows = (n_models + n_cols - 1) // n_cols

        fig, axes = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))
        if n_models == 1:
            axes = np.array([axes])
        axes = axes.flatten()

        for idx, (name, result) in enumerate(self.results.items()):
            ax = axes[idx]
            history = result.get("trial_history", [])

            if history:
                trials = [h["trial"] for h in history]
                scores = [h["score"] for h in history]

                # Running best
                running_best = []
                current_best = -float('inf')
                for s in scores:
                    current_best = max(current_best, s)
                    running_best.append(current_best)

                ax.scatter(trials, scores, alpha=0.5, s=20, label='Trial Score')
                ax.plot(trials, running_best, 'r-', linewidth=2, label='Best So Far')

            ax.set_xlabel('Trial')
            ax.set_ylabel(self.scoring)
            ax.set_title(f'{name}', fontweight='bold')
            ax.legend(loc='lower right', fontsize=8)
            ax.grid(True, alpha=0.3)

        # Hide empty subplots
        for idx in range(n_models, len(axes)):
            axes[idx].set_visible(False)

        plt.suptitle('Hyperopt Convergence Curves', fontsize=14, fontweight='bold', y=1.02)
        plt.tight_layout()
        self.results_manager.save_figure(fig, "convergence_curves")

    def _plot_confusion_matrix(self, X, y, groups):
        """Plot holdout confusion matrix for best model."""
        y_encoded = self.label_encoder.transform(y)
        y_pred = self.best_model.predict(X)

        if self.info and "label_names" in self.info:
            display_labels = [str(self.info["label_names"].get(i, i)) for i in range(len(self.classes_))]
        else:
            display_labels = [str(c) for c in self.classes_]

        # Normalize if many classes
        normalize = "true" if len(display_labels) > 10 else None
        cm = confusion_matrix(y_encoded, y_pred, normalize=normalize)

        n_classes = len(display_labels)
        size = max(8, min(20, n_classes * 0.6))
        fig, ax = plt.subplots(figsize=(size, size))

        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=display_labels)
        disp.plot(ax=ax, cmap="Blues", values_format=".2f" if normalize else "d")

        ax.set_title(f'Confusion Matrix (Holdout): {self.best_model_name}', fontsize=14, fontweight='bold')
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()

        self.results_manager.save_figure(fig, "confusion_matrix")

    def _plot_per_class_performance(self, X, y, groups):
        """Plot holdout per-class precision, recall, and F1."""
        y_encoded = self.label_encoder.transform(y)
        y_pred = self.best_model.predict(X)

        if self.info and "label_names" in self.info:
            target_names = [str(self.info["label_names"].get(i, i)) for i in range(len(self.classes_))]
        else:
            target_names = [str(c) for c in self.classes_]

        report = classification_report(y_encoded, y_pred, target_names=target_names, output_dict=True)

        # Extract per-class metrics
        classes = [c for c in target_names if c in report]
        precision = [report[c]['precision'] for c in classes]
        recall = [report[c]['recall'] for c in classes]
        f1 = [report[c]['f1-score'] for c in classes]
        support = [report[c]['support'] for c in classes]

        # Sort by F1 score
        sorted_indices = np.argsort(f1)[::-1]

        # Limit to top 30 classes for readability
        max_classes = 30
        if len(classes) > max_classes:
            sorted_indices = sorted_indices[:max_classes]

        classes = [classes[i] for i in sorted_indices]
        precision = [precision[i] for i in sorted_indices]
        recall = [recall[i] for i in sorted_indices]
        f1_sorted = [f1[i] for i in sorted_indices]

        fig, ax = plt.subplots(figsize=(12, max(6, len(classes) * 0.3)))

        y_pos = np.arange(len(classes))
        width = 0.25

        ax.barh(y_pos - width, precision, width, label='Precision', color='#3498db')
        ax.barh(y_pos, recall, width, label='Recall', color='#e74c3c')
        ax.barh(y_pos + width, f1_sorted, width, label='F1-Score', color='#2ecc71')

        ax.set_yticks(y_pos)
        ax.set_yticklabels(classes)
        ax.set_xlabel('Score')
        ax.set_title(f'Per-Class Performance (Holdout): {self.best_model_name}', fontsize=14, fontweight='bold')
        ax.legend(loc='lower right')
        ax.set_xlim(0, 1.1)
        ax.grid(True, alpha=0.3, axis='x')

        plt.tight_layout()
        self.results_manager.save_figure(fig, "per_class_performance")

    def _plot_class_distribution(self, y):
        """Plot class distribution."""
        unique, counts = np.unique(y, return_counts=True)

        if self.info and "label_names" in self.info:
            labels = [str(self.info["label_names"].get(i, i)) for i in unique]
        else:
            labels = [str(c) for c in unique]

        # Sort by count
        sorted_indices = np.argsort(counts)[::-1]
        labels = [labels[i] for i in sorted_indices]
        counts = counts[sorted_indices]

        # Limit for readability
        max_show = 30
        if len(labels) > max_show:
            labels = labels[:max_show] + [f"Other ({len(labels) - max_show} classes)"]
            other_count = sum(counts[max_show:])
            counts = np.append(counts[:max_show], other_count)

        fig, ax = plt.subplots(figsize=(12, max(6, len(labels) * 0.3)))

        colors = plt.cm.viridis(np.linspace(0, 1, len(labels)))
        bars = ax.barh(labels, counts, color=colors)

        # Add count labels
        for bar, count in zip(bars, counts):
            ax.text(count + max(counts) * 0.01, bar.get_y() + bar.get_height() / 2,
                    f'{count}', va='center', fontsize=9)

        ax.set_xlabel('Number of Samples')
        ax.set_title('Class Distribution', fontsize=14, fontweight='bold')
        ax.invert_yaxis()

        plt.tight_layout()
        self.results_manager.save_figure(fig, "class_distribution")

    def _plot_hyperparameter_analysis(self):
        """Analyze hyperparameter importance for each model."""
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        axes = axes.flatten()

        plot_idx = 0
        for name, result in self.results.items():
            if plot_idx >= 4:
                break

            trials = result.get("trials")
            if trials is None:
                continue

            # Extract trial data
            trial_data = []
            for trial in trials.trials:
                if trial['result']['status'] == STATUS_OK:
                    row = {'score': trial['result']['score_mean']}
                    # Get hyperparameters from misc
                    vals = trial['misc']['vals']
                    for k, v in vals.items():
                        if v:  # non-empty
                            row[k.replace(f'{name}_', '')] = v[0]
                    trial_data.append(row)

            if not trial_data:
                continue

            df = pd.DataFrame(trial_data)

            # Find numeric columns (excluding score)
            numeric_cols = [c for c in df.columns if c != 'score' and df[c].dtype in ['int64', 'float64']]

            if numeric_cols:
                ax = axes[plot_idx]

                # Correlation with score
                correlations = {}
                for col in numeric_cols[:6]:  # Limit columns
                    corr = df['score'].corr(df[col])
                    if not np.isnan(corr):
                        correlations[col] = corr

                if correlations:
                    sorted_corr = sorted(correlations.items(), key=lambda x: abs(x[1]), reverse=True)
                    params = [c[0] for c in sorted_corr]
                    corrs = [c[1] for c in sorted_corr]

                    colors = ['#2ecc71' if c > 0 else '#e74c3c' for c in corrs]
                    ax.barh(params, corrs, color=colors)
                    ax.axvline(x=0, color='black', linewidth=0.5)
                    ax.set_xlabel('Correlation with Score')
                    ax.set_title(f'{name}', fontweight='bold')
                    ax.set_xlim(-1, 1)

                plot_idx += 1

        # Hide unused subplots
        for idx in range(plot_idx, 4):
            axes[idx].set_visible(False)

        plt.suptitle('Hyperparameter Correlation with Score', fontsize=14, fontweight='bold', y=1.02)
        plt.tight_layout()
        self.results_manager.save_figure(fig, "hyperparameter_analysis")

    def _save_results_json(self, baseline_score):
        """Save all results as JSON."""
        results_data = {
            "timestamp": self.results_manager.timestamp,
            "mode": self.info.get("mode", "unknown"),
            "scoring_metric": self.scoring,
            "best_model": self.best_model_name,
            "best_score": self.results[self.best_model_name]["cv_score_mean"],
            "baseline_score": baseline_score,
            "dataset_info": {
                "n_samples": self.info.get("n_samples"),
                "n_classes": self.info.get("n_classes"),
                "embedding_dim": self.info.get("embedding_dim"),
                "class_distribution": self.info.get("class_distribution"),
            },
            "models": {}
        }

        for name, result in self.results.items():
            results_data["models"][name] = {
                "cv_score_mean": result["cv_score_mean"],
                "cv_score_std": result["cv_score_std"],
                "best_params": result["best_params"],
                "pca_components": result.get("pca_components"),
                "n_evals": result["n_evals"],
                "cv_strategy": result["cv_strategy"],
                "trial_history": result.get("trial_history", [])
            }

        self.results_manager.save_json(results_data, "experiment_results")

    def _save_summary_csv(self):
        """Save model comparison as CSV."""
        rows = []
        for name, result in self.results.items():
            rows.append({
                "Model": name,
                "CV_Score_Mean": result["cv_score_mean"],
                "CV_Score_Std": result["cv_score_std"],
                "PCA_Components": result.get("pca_components"),
                "N_Evaluations": result["n_evals"],
                "CV_Strategy": result["cv_strategy"],
                "Is_Best": name == self.best_model_name
            })

        df = pd.DataFrame(rows)
        df = df.sort_values("CV_Score_Mean", ascending=False)
        self.results_manager.save_csv(df, "model_comparison")

    def _save_text_report(self, baseline_score):
        """Generate and save a comprehensive text report."""
        report = []
        report.append("=" * 80)
        report.append("MALWARE CLASSIFICATION EXPERIMENT REPORT")
        report.append("=" * 80)
        report.append(f"\nTimestamp: {self.results_manager.timestamp}")
        report.append(f"Mode: {self.info.get('mode', 'unknown')}")
        report.append(f"Scoring Metric: {self.scoring}")

        report.append("\n" + "-" * 40)
        report.append("DATASET INFORMATION")
        report.append("-" * 40)
        report.append(f"  Samples: {self.info.get('n_samples', 'N/A')}")
        report.append(f"  Classes: {self.info.get('n_classes', 'N/A')}")
        report.append(f"  Embedding Dimension: {self.info.get('embedding_dim', 'N/A')}")

        report.append("\n" + "-" * 40)
        report.append("MODEL RESULTS (sorted by score)")
        report.append("-" * 40)

        sorted_results = sorted(self.results.items(), key=lambda x: x[1]["cv_score_mean"], reverse=True)

        for rank, (name, result) in enumerate(sorted_results, 1):
            marker = " ★ BEST" if name == self.best_model_name else ""
            report.append(f"\n  {rank}. {name}{marker}")
            report.append(f"     Score: {result['cv_score_mean']:.4f} (+/- {result['cv_score_std']:.4f})")
            report.append(f"     PCA: {result.get('pca_components', 'None')}")
            report.append(f"     CV Strategy: {result['cv_strategy']}")
            report.append(f"     Hyperparameters:")
            for k, v in result['best_params'].items():
                report.append(f"       - {k}: {v}")

        report.append("\n" + "-" * 40)
        report.append("COMPARISON TO BASELINE")
        report.append("-" * 40)

        if baseline_score is not None:
            best_score = self.results[self.best_model_name]["cv_score_mean"]
            report.append(f"  Baseline Score: {baseline_score:.4f}")
            report.append(f"  Best Model Score: {best_score:.4f}")

            if best_score > baseline_score:
                improvement = abs(best_score - baseline_score) / abs(baseline_score) * 100
                report.append(f"  ✓ Model beats baseline by {improvement:.1f}%")
            else:
                report.append(f"  ✗ Model does NOT beat baseline")

        report.append("\n" + "=" * 80)
        report.append("END OF REPORT")
        report.append("=" * 80)

        self.results_manager.save_text("\n".join(report), "experiment_report")

    def _save_all_models(self):
        """Save all trained models."""
        print("\n  Saving trained models...")

        for name, result in self.results.items():
            model_data = {
                "model": result["best_estimator"],
                "model_name": name,
                "params": result["best_params"],
                "cv_score": result["cv_score_mean"],
                "cv_score_std": result["cv_score_std"],
                "scoring_metric": self.scoring,
                "label_encoder": self.label_encoder,
                "classes": self.classes_,
                "info": self.info,
            }
            self.results_manager.save_model(model_data, f"model_{name}")

        # Also save best model separately with clear naming
        best_result = self.results[self.best_model_name]
        best_model_data = {
            "model": best_result["best_estimator"],
            "model_name": self.best_model_name,
            "params": best_result["best_params"],
            "cv_score": best_result["cv_score_mean"],
            "scoring_metric": self.scoring,
            "label_encoder": self.label_encoder,
            "classes": self.classes_,
            "info": self.info,
        }
        self.results_manager.save_model(best_model_data, "BEST_MODEL")

    def save_model(self, filepath: str = None):
        """Save the best model and metadata."""
        if self.best_model is None:
            raise ValueError("No model to save.")

        if filepath is None:
            mode = self.info.get("mode", "unknown") if self.info else "unknown"
            filepath = f"malware_{mode}_classifier.joblib"

        save_data = {
            "model": self.best_model,
            "model_name": self.best_model_name,
            "params": self.best_params,
            "cv_score": self.results[self.best_model_name]["cv_score_mean"],
            "scoring_metric": self.scoring,
            "label_encoder": self.label_encoder,
            "classes": self.classes_,
            "info": self.info,
        }

        joblib.dump(save_data, filepath)
        print(f"✓ Model saved to {filepath}")

    def load_model(self, filepath: str):
        """Load a saved model."""
        data = joblib.load(filepath)
        self.best_model = data["model"]
        self.best_model_name = data["model_name"]
        self.best_params = data["params"]
        self.scoring = data["scoring_metric"]
        self.label_encoder = data["label_encoder"]
        self.classes_ = data["classes"]
        self.info = data.get("info", {})

        print(f"✓ Loaded {data['model_name']} from {filepath}")
        print(f"  CV {data['scoring_metric']}: {data['cv_score']:.4f}")
        print(f"  Mode: {self.info.get('mode', 'unknown')}")
        print(f"  Classes: {len(self.classes_)}")

        return self.best_model


def main(
        mode: ClassificationMode = "maliciousness",
        use_sample_weights: bool = True,
        weight_strategy: str = "balanced",
        use_grouped_cv: bool = True,
        min_family_samples: int = 5,
        max_evals: int = 100,
        models_to_tune: Optional[List[str]] = None,
        scoring: str = "f1_weighted",
        use_pca: bool = True,
        k_folds: int = 5,
        holdout_size: float = 0.2,
        output_dir: str = "results",
        collection_name: str = COLLECTION_NAME,
        fixed_holdout_keys: Optional[Set[str]] = None,
        preloaded_data: Optional[Tuple[np.ndarray, np.ndarray, List[str], Optional[np.ndarray], Dict[str, Any]]] = None,
):
    """
    Main training function for malware classification.
    """
    print("=" * 60)
    print(f"Malware Classification - Hyperopt Bayesian Optimization")
    print("=" * 60)
    print(f"  Mode: {mode}")
    print(f"  Scoring: {scoring}")
    print(f"  Use sample weights: {use_sample_weights} ({weight_strategy})")
    print(f"  Use grouped CV: {use_grouped_cv}")
    print(f"  Use PCA search: {use_pca}")
    print(f"  K-folds: {k_folds}")
    print(f"  Holdout size: {holdout_size:.0%}")
    print(f"  Max evals per model: {max_evals}")
    print(f"  Collection: {collection_name}")
    print(f"  Output directory: {output_dir}")

    if not (0.0 < holdout_size < 1.0):
        raise ValueError(f"holdout_size must be in (0, 1), got {holdout_size}")

    # Initialize results manager
    results_manager = ResultsManager(output_dir)

    # Load data
    if preloaded_data is None:
        loader = MalwareDataLoader(collection_name=collection_name)
        X, y, split_keys, groups, info = loader.load_dataset(
            mode=mode,
            min_family_samples=min_family_samples,
        )
    else:
        X, y, split_keys, groups, info = preloaded_data

    # Groups for CV
    cv_groups = None
    if use_grouped_cv and groups is not None:
        cv_groups = groups
        print(f"  Grouped CV by: family ({len(np.unique(groups))} groups)")

    train_idx, holdout_idx, split_strategy, holdout_key_set = _build_holdout_indices(
        y=y,
        split_keys=split_keys,
        holdout_size=holdout_size,
        cv_groups=cv_groups,
        fixed_holdout_keys=fixed_holdout_keys,
        random_state=42,
    )

    X_train, X_holdout = X[train_idx], X[holdout_idx]
    y_train, y_holdout = y[train_idx], y[holdout_idx]
    cv_groups_train = cv_groups[train_idx] if cv_groups is not None else None
    holdout_groups = cv_groups[holdout_idx] if cv_groups is not None else None

    # Sample weights are computed only on train split (avoid holdout leakage)
    sample_weights_train = None
    if use_sample_weights:
        sample_weights_train = MalwareDataLoader.get_sample_weights(y_train, strategy=weight_strategy)
        print(f"\n  Train sample weights range: [{sample_weights_train.min():.3f}, {sample_weights_train.max():.3f}]")

    print(f"\n  Holdout split strategy: {split_strategy}")
    print(f"  Train samples: {len(y_train)}")
    print(f"  Holdout samples: {len(y_holdout)}")

    # Default models
    if models_to_tune is None:
        models_to_tune = ["resnet18", "xgboost", "lightgbm", "random_forest", "svc", "logistic_regression", "knn"]

    # Baseline
    print(f"\n{'=' * 50}")
    print("BASELINE")
    print(f"{'=' * 50}")

    if cv_groups is not None:
        cv = StratifiedGroupKFold(n_splits=k_folds, shuffle=True, random_state=42)
        print(f"  (Using StratifiedGroupKFold with {k_folds} folds)")
    else:
        cv = StratifiedKFold(n_splits=k_folds, shuffle=True, random_state=42)
        print(f"  (Using StratifiedKFold with {k_folds} folds)")

    baseline_results = {}
    for strategy in ["most_frequent", "stratified"]:
        dummy = DummyClassifier(strategy=strategy, random_state=42)
        if cv_groups_train is not None:
            baseline_scores = cross_val_score(dummy, X_train, y_train, cv=cv, scoring=scoring, groups=cv_groups_train)
        else:
            baseline_scores = cross_val_score(dummy, X_train, y_train, cv=cv, scoring=scoring)
        baseline_results[strategy] = baseline_scores.mean()
        print(f"  {strategy:15} - {scoring}: {baseline_scores.mean():.4f} (+/- {baseline_scores.std():.4f})")

    baseline_score = baseline_results["most_frequent"]

    # Tune models
    tuner = MalwareHyperoptTuner(random_state=42, scoring=scoring)

    info = dict(info)
    info["collection_name"] = collection_name
    info["holdout_size"] = holdout_size
    info["train_samples"] = int(len(y_train))
    info["holdout_samples"] = int(len(y_holdout))
    info["holdout_split_strategy"] = split_strategy
    info["holdout_split_key_count"] = int(len(holdout_key_set))
    info["holdout_split_keys"] = sorted(holdout_key_set)

    tuner.tune_all_models(
        X_train, y_train, info=info, models_to_tune=models_to_tune, max_evals=max_evals,
        k_folds=k_folds, groups=cv_groups_train, sample_weights=sample_weights_train, use_pca=use_pca,
    )

    tuner.print_summary()

    # Evaluate best model
    metrics = tuner.evaluate(X_holdout, y_holdout, groups=holdout_groups)
    info["holdout_metrics"] = _scalar_holdout_metrics(metrics)
    holdout_details = _build_holdout_details_payload(
        tuner=tuner,
        metrics=metrics,
        split_strategy=split_strategy,
        holdout_size=holdout_size,
        holdout_key_set=holdout_key_set,
    )
    holdout_report_path = results_manager.save_json(
        holdout_details,
        "holdout_evaluation_detailed",
    )
    info["holdout_report_path"] = holdout_report_path
    info["holdout_macro_avg"] = holdout_details.get("macro_avg", {})
    info["holdout_weighted_avg"] = holdout_details.get("weighted_avg", {})

    # Compare to baseline
    best_score = tuner.results[tuner.best_model_name]["cv_score_mean"]

    print(f"\n{'=' * 50}")
    print("VS BASELINE")
    print(f"{'=' * 50}")
    print(f"  Baseline ({scoring}): {baseline_score:.4f}")
    print(f"  Best model ({scoring}): {best_score:.4f}")

    if best_score > baseline_score:
        improvement = abs(best_score - baseline_score) / abs(baseline_score) * 100
        print(f"  ✓ Model beats baseline by {improvement:.1f}%")
    else:
        print(f"  ✗ Model does NOT beat baseline")

    # Generate comprehensive reports
    tuner.generate_all_reports(
        X_holdout, y_holdout, groups=holdout_groups,
        baseline_score=baseline_score,
        results_manager=results_manager
    )

    return tuner, info, results_manager


def run_multi_collection_experiments(
        collections: List[str],
        mode: ClassificationMode = "maliciousness",
        use_sample_weights: bool = True,
        weight_strategy: str = "balanced",
        use_grouped_cv: bool = True,
        min_family_samples: int = 5,
        max_evals: int = 100,
        models_to_tune: Optional[List[str]] = None,
        scoring: str = "f1_weighted",
        use_pca: bool = True,
        k_folds: int = 5,
        holdout_size: float = 0.2,
        output_dir: str = "results",
) -> Tuple[Dict[str, Any], str]:
    collection_names = list(dict.fromkeys(collections))
    if len(collection_names) < 2:
        raise ValueError("Please provide at least 2 distinct collections for multi-collection mode.")

    print("\n" + "=" * 72)
    print("MULTI-COLLECTION PRECHECK")
    print("=" * 72)
    print(f"Collections: {collection_names}")

    loaded_datasets: Dict[str, Tuple[np.ndarray, np.ndarray, List[str], Optional[np.ndarray], Dict[str, Any]]] = {}
    key_sets: Dict[str, Set[str]] = {}

    for collection_name in collection_names:
        loader = MalwareDataLoader(collection_name=collection_name)
        dataset = loader.load_dataset(mode=mode, min_family_samples=min_family_samples)
        loaded_datasets[collection_name] = dataset
        key_set = set(_normalize_split_keys(dataset[2]).tolist())
        key_sets[collection_name] = key_set
        print(f"  {collection_name}: samples={len(dataset[1])}, unique_split_keys={len(key_set)}")

    common_keys = set.intersection(*(key_sets[name] for name in collection_names))
    if not common_keys:
        raise ValueError("No common split keys across collections. Cannot build an identical holdout set.")

    print(f"Common split keys across all collections: {len(common_keys)}")

    aligned_datasets: Dict[str, Tuple[np.ndarray, np.ndarray, List[str], Optional[np.ndarray], Dict[str, Any]]] = {}
    for collection_name in collection_names:
        X, y, split_keys, groups, info = loaded_datasets[collection_name]
        X, y, split_keys, groups, dropped_count = _filter_dataset_by_keys(
            X=X,
            y=y,
            split_keys=split_keys,
            groups=groups,
            allowed_keys=common_keys,
        )
        if dropped_count > 0:
            print(
                f"  {collection_name}: dropped {dropped_count} samples not present in all collections."
            )

        info = dict(info)
        info["n_samples"] = int(len(y))
        info["aligned_to_common_keys"] = True
        info["dropped_for_alignment"] = dropped_count
        aligned_datasets[collection_name] = (X, y, split_keys, groups, info)

    default_holdout_keys_by_collection: Dict[str, Set[str]] = {}
    for collection_name in collection_names:
        _, y, split_keys, groups, _ = aligned_datasets[collection_name]
        cv_groups = groups if use_grouped_cv and groups is not None else None
        _, _, _, holdout_keys = _build_holdout_indices(
            y=y,
            split_keys=split_keys,
            holdout_size=holdout_size,
            cv_groups=cv_groups,
            fixed_holdout_keys=None,
            random_state=42,
        )
        default_holdout_keys_by_collection[collection_name] = holdout_keys

    canonical_holdout_keys = default_holdout_keys_by_collection[collection_names[0]]
    default_holdouts_identical = all(
        default_holdout_keys_by_collection[name] == canonical_holdout_keys
        for name in collection_names
    )

    if default_holdouts_identical:
        print("Default holdout split is already identical across all collections.")
    else:
        print("Default holdout split differs across collections. Enforcing canonical holdout key set.")
    print(f"Canonical holdout key count: {len(canonical_holdout_keys)}")

    collection_results: Dict[str, Any] = {}
    for collection_name in collection_names:
        collection_output_dir = os.path.join(output_dir, collection_name)
        tuner, info, results_manager = main(
            mode=mode,
            use_sample_weights=use_sample_weights,
            weight_strategy=weight_strategy,
            use_grouped_cv=use_grouped_cv,
            min_family_samples=min_family_samples,
            max_evals=max_evals,
            models_to_tune=models_to_tune,
            scoring=scoring,
            use_pca=use_pca,
            k_folds=k_folds,
            holdout_size=holdout_size,
            output_dir=collection_output_dir,
            collection_name=collection_name,
            fixed_holdout_keys=canonical_holdout_keys,
            preloaded_data=aligned_datasets[collection_name],
        )

        collection_results[collection_name] = {
            "best_model_name": tuner.best_model_name,
            "best_cv_score": float(tuner.results[tuner.best_model_name]["cv_score_mean"]),
            "best_cv_std": float(tuner.results[tuner.best_model_name]["cv_score_std"]),
            "holdout_metrics": info.get("holdout_metrics", {}),
            "holdout_macro_avg": info.get("holdout_macro_avg", {}),
            "holdout_weighted_avg": info.get("holdout_weighted_avg", {}),
            "holdout_split_strategy": info.get("holdout_split_strategy"),
            "train_samples": int(info.get("train_samples", 0)),
            "holdout_samples": int(info.get("holdout_samples", 0)),
            "holdout_report_path": info.get("holdout_report_path"),
            "run_dir": results_manager.run_dir,
        }

    os.makedirs(output_dir, exist_ok=True)
    summary_payload = {
        "mode": mode,
        "collections": collection_names,
        "default_holdouts_identical": default_holdouts_identical,
        "canonical_holdout_key_count": len(canonical_holdout_keys),
        "results": collection_results,
    }
    summary_path = os.path.join(
        output_dir,
        f"multi_collection_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
    )
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary_payload, f, indent=2)

    print("\n" + "=" * 72)
    print("MULTI-COLLECTION SUMMARY")
    print("=" * 72)
    for collection_name in collection_names:
        item = collection_results[collection_name]
        print(
            f"  {collection_name}: best={item['best_model_name']} "
            f"cv={item['best_cv_score']:.4f} holdout_f1w={item['holdout_metrics'].get('f1_weighted', float('nan')):.4f}"
        )
        print(f"    run_dir={item['run_dir']}")
    print(f"Summary JSON: {summary_path}")

    return collection_results, summary_path


if __name__ == "__main__":
    # ============================================================
    # CONFIGURATION
    # ============================================================

    COLLECTIONS = ["pre_traind_distilberd3_3", "distilbert3_3"]
    CLASSIFICATION_MODE: ClassificationMode = "malware_family"
    SCORING = "f1_weighted"
    USE_SAMPLE_WEIGHTS = False
    WEIGHT_STRATEGY = "balanced"
    USE_GROUPED_CV = True
    K_FOLDS = 10
    MIN_FAMILY_SAMPLES = 5
    MAX_EVALS = 30
    USE_PCA = None
    MODELS_TO_TUNE = ["svc", "lightgbm", "random_forest"]
    HOLDOUT_SIZE = 0.2
    OUTPUT_DIR = "results"

    if len(COLLECTIONS) == 1:
        tuner, info, results_manager = main(
            mode=CLASSIFICATION_MODE,
            use_sample_weights=USE_SAMPLE_WEIGHTS,
            weight_strategy=WEIGHT_STRATEGY,
            use_grouped_cv=USE_GROUPED_CV,
            min_family_samples=MIN_FAMILY_SAMPLES,
            max_evals=MAX_EVALS,
            models_to_tune=MODELS_TO_TUNE,
            scoring=SCORING,
            use_pca=USE_PCA,
            k_folds=K_FOLDS,
            holdout_size=HOLDOUT_SIZE,
            output_dir=OUTPUT_DIR,
            collection_name=COLLECTIONS[0],
        )

        print(f"\n{'=' * 60}")
        print(f"EXPERIMENT COMPLETE")
        print(f"{'=' * 60}")
        print(f"All results saved to: {results_manager.run_dir}")
        print(f"\nContents:")
        print(f"  📁 plots/     - Visualizations (PNG)")
        print(f"  📁 models/    - Trained models (joblib)")
        print(f"  📁 reports/   - JSON, CSV, and text reports")
    else:
        run_multi_collection_experiments(
            collections=COLLECTIONS,
            mode=CLASSIFICATION_MODE,
            use_sample_weights=USE_SAMPLE_WEIGHTS,
            weight_strategy=WEIGHT_STRATEGY,
            use_grouped_cv=USE_GROUPED_CV,
            min_family_samples=MIN_FAMILY_SAMPLES,
            max_evals=MAX_EVALS,
            models_to_tune=MODELS_TO_TUNE,
            scoring=SCORING,
            use_pca=USE_PCA,
            k_folds=K_FOLDS,
            holdout_size=HOLDOUT_SIZE,
            output_dir=OUTPUT_DIR,
        )