models/anomaly-detection/src/components/model_trainer.py

"""
models/anomaly-detection/src/components/model_trainer.py
Model training with Optuna hyperparameter tuning for clustering/anomaly detection
"""
import os
import logging
import joblib
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List
import numpy as np

from ..entity import ModelTrainerConfig, ModelTrainerArtifact
from ..utils import calculate_clustering_metrics, calculate_optuna_objective, format_metrics_report

logger = logging.getLogger("model_trainer")

# MLflow
try:
    import mlflow
    import mlflow.sklearn
    MLFLOW_AVAILABLE = True
except ImportError:
    MLFLOW_AVAILABLE = False
    logger.warning("MLflow not available. Install with: pip install mlflow")

# Optuna
try:
    import optuna
    from optuna.samplers import TPESampler
    OPTUNA_AVAILABLE = True
except ImportError:
    OPTUNA_AVAILABLE = False
    logger.warning("Optuna not available. Install with: pip install optuna")

# Clustering algorithms
try:
    from sklearn.cluster import DBSCAN, KMeans
    from sklearn.ensemble import IsolationForest
    from sklearn.neighbors import LocalOutlierFactor
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False

try:
    import hdbscan
    HDBSCAN_AVAILABLE = True
except ImportError:
    HDBSCAN_AVAILABLE = False
    logger.warning("HDBSCAN not available. Install with: pip install hdbscan")


class ModelTrainer:
    """
    Model training component with:
    1. Optuna hyperparameter optimization
    2. Multiple clustering algorithms (DBSCAN, KMeans, HDBSCAN)
    3. Anomaly detection (Isolation Forest, LOF)
    4. MLflow experiment tracking
    """

    def __init__(self, config: Optional[ModelTrainerConfig] = None):
        """
        Initialize model trainer.
        
        Args:
            config: Optional configuration
        """
        self.config = config or ModelTrainerConfig()

        # Ensure output directory exists
        Path(self.config.output_directory).mkdir(parents=True, exist_ok=True)

        # Setup MLflow
        self._setup_mlflow()

        logger.info("[ModelTrainer] Initialized")
        logger.info(f"  Models to train: {self.config.models_to_train}")
        logger.info(f"  Optuna trials: {self.config.n_optuna_trials}")

    def _setup_mlflow(self):
        """Configure MLflow tracking"""
        if not MLFLOW_AVAILABLE:
            logger.warning("[ModelTrainer] MLflow not available")
            return

        try:
            # Set tracking URI
            mlflow.set_tracking_uri(self.config.mlflow_tracking_uri)

            # Set credentials for DagsHub
            if self.config.mlflow_username and self.config.mlflow_password:
                os.environ["MLFLOW_TRACKING_USERNAME"] = self.config.mlflow_username
                os.environ["MLFLOW_TRACKING_PASSWORD"] = self.config.mlflow_password

            # Create or get experiment
            try:
                mlflow.create_experiment(self.config.experiment_name)
            except Exception:
                pass
            mlflow.set_experiment(self.config.experiment_name)

            logger.info(f"[ModelTrainer] MLflow configured: {self.config.mlflow_tracking_uri}")

        except Exception as e:
            logger.warning(f"[ModelTrainer] MLflow setup error: {e}")

    def _train_dbscan(self, X: np.ndarray, trial: Optional['optuna.Trial'] = None) -> Dict[str, Any]:
        """
        Train DBSCAN with optional Optuna tuning.
        """
        if not SKLEARN_AVAILABLE:
            return {"error": "sklearn not available"}

        # Hyperparameters
        if trial:
            eps = trial.suggest_float("eps", 0.1, 2.0)
            min_samples = trial.suggest_int("min_samples", 2, 20)
        else:
            eps = 0.5
            min_samples = 5

        model = DBSCAN(eps=eps, min_samples=min_samples, n_jobs=-1)
        labels = model.fit_predict(X)

        metrics = calculate_clustering_metrics(X, labels)
        metrics["eps"] = eps
        metrics["min_samples"] = min_samples

        return {
            "model": model,
            "labels": labels,
            "metrics": metrics,
            "params": {"eps": eps, "min_samples": min_samples}
        }

    def _train_kmeans(self, X: np.ndarray, trial: Optional['optuna.Trial'] = None) -> Dict[str, Any]:
        """
        Train KMeans with optional Optuna tuning.
        """
        if not SKLEARN_AVAILABLE:
            return {"error": "sklearn not available"}

        # Hyperparameters
        if trial:
            n_clusters = trial.suggest_int("n_clusters", 2, 20)
            n_init = trial.suggest_int("n_init", 5, 20)
        else:
            n_clusters = 5
            n_init = 10

        model = KMeans(n_clusters=n_clusters, n_init=n_init, random_state=42)
        labels = model.fit_predict(X)

        metrics = calculate_clustering_metrics(X, labels)
        metrics["n_clusters"] = n_clusters

        return {
            "model": model,
            "labels": labels,
            "metrics": metrics,
            "params": {"n_clusters": n_clusters, "n_init": n_init}
        }

    def _train_hdbscan(self, X: np.ndarray, trial: Optional['optuna.Trial'] = None) -> Dict[str, Any]:
        """
        Train HDBSCAN with optional Optuna tuning.
        """
        if not HDBSCAN_AVAILABLE:
            return {"error": "hdbscan not available"}

        # Hyperparameters
        if trial:
            min_cluster_size = trial.suggest_int("min_cluster_size", 5, 50)
            min_samples = trial.suggest_int("min_samples", 1, 20)
        else:
            min_cluster_size = 15
            min_samples = 5

        model = hdbscan.HDBSCAN(
            min_cluster_size=min_cluster_size,
            min_samples=min_samples,
            core_dist_n_jobs=-1
        )
        labels = model.fit_predict(X)

        metrics = calculate_clustering_metrics(X, labels)

        return {
            "model": model,
            "labels": labels,
            "metrics": metrics,
            "params": {"min_cluster_size": min_cluster_size, "min_samples": min_samples}
        }

    def _train_isolation_forest(self, X: np.ndarray, trial: Optional['optuna.Trial'] = None) -> Dict[str, Any]:
        """
        Train Isolation Forest for anomaly detection.
        """
        if not SKLEARN_AVAILABLE:
            return {"error": "sklearn not available"}

        # Hyperparameters
        if trial:
            contamination = trial.suggest_float("contamination", 0.01, 0.3)
            n_estimators = trial.suggest_int("n_estimators", 50, 200)
        else:
            contamination = 0.1
            n_estimators = 100

        model = IsolationForest(
            contamination=contamination,
            n_estimators=n_estimators,
            random_state=42,
            n_jobs=-1
        )
        predictions = model.fit_predict(X)
        labels = (predictions == -1).astype(int)  # -1 = anomaly

        n_anomalies = int(np.sum(labels))

        return {
            "model": model,
            "labels": labels,
            "metrics": {
                "n_anomalies": n_anomalies,
                "anomaly_rate": n_anomalies / len(X),
                "contamination": contamination,
                "n_estimators": n_estimators
            },
            "params": {"contamination": contamination, "n_estimators": n_estimators},
            "anomaly_indices": np.where(labels == 1)[0].tolist()
        }

    def _train_lof(self, X: np.ndarray, trial: Optional['optuna.Trial'] = None) -> Dict[str, Any]:
        """
        Train Local Outlier Factor for anomaly detection.
        """
        if not SKLEARN_AVAILABLE:
            return {"error": "sklearn not available"}

        # Hyperparameters
        if trial:
            n_neighbors = trial.suggest_int("n_neighbors", 5, 50)
            contamination = trial.suggest_float("contamination", 0.01, 0.3)
        else:
            n_neighbors = 20
            contamination = 0.1

        model = LocalOutlierFactor(
            n_neighbors=n_neighbors,
            contamination=contamination,
            n_jobs=-1,
            novelty=True  # For prediction on new data
        )
        model.fit(X)
        predictions = model.predict(X)
        labels = (predictions == -1).astype(int)  # -1 = anomaly

        n_anomalies = int(np.sum(labels))

        return {
            "model": model,
            "labels": labels,
            "metrics": {
                "n_anomalies": n_anomalies,
                "anomaly_rate": n_anomalies / len(X),
                "n_neighbors": n_neighbors,
                "contamination": contamination
            },
            "params": {"n_neighbors": n_neighbors, "contamination": contamination},
            "anomaly_indices": np.where(labels == 1)[0].tolist()
        }

    def _optimize_model(self, model_name: str, X: np.ndarray) -> Dict[str, Any]:
        """
        Use Optuna to find best hyperparameters for a model.
        """
        if not OPTUNA_AVAILABLE:
            logger.warning("[ModelTrainer] Optuna not available, using defaults")
            return self._train_model(model_name, X, None)

        train_func = {
            "dbscan": self._train_dbscan,
            "kmeans": self._train_kmeans,
            "hdbscan": self._train_hdbscan,
            "isolation_forest": self._train_isolation_forest,
            "lof": self._train_lof
        }.get(model_name)

        if not train_func:
            return {"error": f"Unknown model: {model_name}"}

        def objective(trial):
            try:
                result = train_func(X, trial)
                if "error" in result:
                    return -1.0

                metrics = result.get("metrics", {})

                # For clustering: use silhouette
                if model_name in ["dbscan", "kmeans", "hdbscan"]:
                    score = metrics.get("silhouette_score", -1)
                    return score if score is not None else -1

                # For anomaly detection: balance anomaly rate
                else:
                    # Target anomaly rate around 5-15%
                    rate = metrics.get("anomaly_rate", 0)
                    target = 0.1
                    return -abs(rate - target)  # Closer to target is better

            except Exception as e:
                logger.debug(f"Trial failed: {e}")
                return -1.0

        # Create and run study
        study = optuna.create_study(
            direction="maximize",
            sampler=TPESampler(seed=42)
        )

        study.optimize(
            objective,
            n_trials=self.config.n_optuna_trials,
            timeout=self.config.optuna_timeout_seconds,
            show_progress_bar=True
        )

        logger.info(f"[ModelTrainer] {model_name} best params: {study.best_params}")
        logger.info(f"[ModelTrainer] {model_name} best score: {study.best_value:.4f}")

        # Train with best params
        best_result = train_func(X, None)  # Use defaults as base
        # Override with best params
        if study.best_params:
            # Re-train with best params would require custom logic
            # For now, we just log the best params
            best_result["best_params"] = study.best_params
            best_result["best_score"] = study.best_value
            best_result["study_name"] = study.study_name

        return best_result

    def _train_model(self, model_name: str, X: np.ndarray, trial=None) -> Dict[str, Any]:
        """Train a single model"""
        train_funcs = {
            "dbscan": self._train_dbscan,
            "kmeans": self._train_kmeans,
            "hdbscan": self._train_hdbscan,
            "isolation_forest": self._train_isolation_forest,
            "lof": self._train_lof
        }

        func = train_funcs.get(model_name)
        if func:
            return func(X, trial)
        return {"error": f"Unknown model: {model_name}"}

    def initiate_model_trainer(self, feature_path: str) -> ModelTrainerArtifact:
        """
        Execute model training pipeline.
        
        Args:
            feature_path: Path to feature matrix (.npy)
            
        Returns:
            ModelTrainerArtifact with results
        """
        logger.info(f"[ModelTrainer] Starting training: {feature_path}")
        start_time = datetime.now()

        # Load features
        X = np.load(feature_path)
        logger.info(f"[ModelTrainer] Loaded features: {X.shape}")

        # Start MLflow run
        mlflow_run_id = ""
        mlflow_experiment_id = ""

        if MLFLOW_AVAILABLE:
            try:
                run = mlflow.start_run()
                mlflow_run_id = run.info.run_id
                mlflow_experiment_id = run.info.experiment_id

                mlflow.log_param("n_samples", X.shape[0])
                mlflow.log_param("n_features", X.shape[1])
                mlflow.log_param("models", self.config.models_to_train)
            except Exception as e:
                logger.warning(f"[ModelTrainer] MLflow run start error: {e}")

        # Train all models
        trained_models = []
        best_model = None
        best_score = -float('inf')

        for model_name in self.config.models_to_train:
            logger.info(f"[ModelTrainer] Training {model_name}...")

            try:
                result = self._optimize_model(model_name, X)

                if "error" in result:
                    logger.warning(f"[ModelTrainer] {model_name} error: {result['error']}")
                    continue

                # Save model
                model_path = Path(self.config.output_directory) / f"{model_name}_model.joblib"
                joblib.dump(result["model"], model_path)

                # Log to MLflow
                if MLFLOW_AVAILABLE:
                    try:
                        mlflow.log_params({f"{model_name}_{k}": v for k, v in result.get("params", {}).items()})
                        mlflow.log_metrics({f"{model_name}_{k}": v for k, v in result.get("metrics", {}).items() if isinstance(v, (int, float))})
                        mlflow.sklearn.log_model(result["model"], model_name)
                    except Exception as e:
                        logger.debug(f"MLflow log error: {e}")

                # Track results
                model_info = {
                    "name": model_name,
                    "path": str(model_path),
                    "params": result.get("params", {}),
                    "metrics": result.get("metrics", {})
                }
                trained_models.append(model_info)

                # Check if best (for clustering models)
                score = result.get("metrics", {}).get("silhouette_score", -1)
                if score and score > best_score:
                    best_score = score
                    best_model = model_info

                logger.info(f"[ModelTrainer] ✓ {model_name} complete")

            except Exception as e:
                logger.error(f"[ModelTrainer] {model_name} failed: {e}")

        # End MLflow run
        if MLFLOW_AVAILABLE:
            try:
                mlflow.end_run()
            except Exception:
                pass

        # Calculate duration
        duration = (datetime.now() - start_time).total_seconds()

        # Get anomaly info from best anomaly detector
        n_anomalies = None
        anomaly_indices = None
        for model_info in trained_models:
            if model_info["name"] in ["isolation_forest", "lof"]:
                n_anomalies = model_info["metrics"].get("n_anomalies")
                break

        # Build artifact
        artifact = ModelTrainerArtifact(
            best_model_name=best_model["name"] if best_model else "",
            best_model_path=best_model["path"] if best_model else "",
            best_model_metrics=best_model["metrics"] if best_model else {},
            trained_models=trained_models,
            mlflow_run_id=mlflow_run_id,
            mlflow_experiment_id=mlflow_experiment_id,
            n_clusters=best_model["metrics"].get("n_clusters") if best_model else None,
            n_anomalies=n_anomalies,
            anomaly_indices=anomaly_indices,
            training_duration_seconds=duration,
            optuna_study_name=None
        )

        logger.info(f"[ModelTrainer] Training complete in {duration:.1f}s")
        logger.info(f"[ModelTrainer] Best model: {best_model['name'] if best_model else 'N/A'}")

        # ============================================
        # TRAIN PER-LANGUAGE MODELS FOR LIVE INFERENCE
        # ============================================
        # Different BERT models produce embeddings in different vector spaces.
        # We train separate Isolation Forest models per language to avoid
        # mixing incompatible embeddings.
        try:
            # Check if features include extra metadata (> 768 dims)
            if X.shape[1] > 768:
                X_embeddings = X[:, :768]  # Extract BERT embeddings only
            else:
                X_embeddings = X

            logger.info(f"[ModelTrainer] Training per-language models on {X_embeddings.shape[0]} samples...")

            # Load language labels from the same directory as features
            feature_dir = Path(feature_path).parent
            lang_files = list(feature_dir.glob("languages_*.npy"))
            
            if lang_files:
                # Get most recent language file
                latest_lang_file = max(lang_files, key=lambda p: p.stem)
                languages = np.load(latest_lang_file, allow_pickle=True)
                logger.info(f"[ModelTrainer] Loaded language labels from {latest_lang_file.name}")
            else:
                # Fallback: try to load from transformed data parquet
                parquet_files = list(feature_dir.glob("transformed_*.parquet"))
                if parquet_files:
                    import pandas as pd
                    latest_parquet = max(parquet_files, key=lambda p: p.stem)
                    df_temp = pd.read_parquet(latest_parquet)
                    if "language" in df_temp.columns:
                        languages = df_temp["language"].values
                        logger.info(f"[ModelTrainer] Loaded {len(languages)} language labels from parquet")
                    else:
                        languages = np.array(["english"] * len(X_embeddings))
                        logger.warning("[ModelTrainer] No language column in parquet, defaulting to english")
                else:
                    languages = np.array(["english"] * len(X_embeddings))
                    logger.warning("[ModelTrainer] No language data found, defaulting to english")

            # Train per-language models
            MIN_SAMPLES_PER_LANGUAGE = 10
            per_lang_models = {}

            for lang in ["english", "sinhala", "tamil"]:
                lang_mask = languages == lang
                X_lang = X_embeddings[lang_mask]

                if len(X_lang) >= MIN_SAMPLES_PER_LANGUAGE:
                    logger.info(f"[ModelTrainer] Training {lang} model on {len(X_lang)} samples...")
                    
                    lang_model = IsolationForest(
                        contamination=0.1,
                        n_estimators=100,
                        random_state=42,
                        n_jobs=-1
                    )
                    lang_model.fit(X_lang)

                    # Save per-language model
                    model_path = Path(self.config.output_directory) / f"isolation_forest_{lang}.joblib"
                    joblib.dump(lang_model, model_path)
                    per_lang_models[lang] = str(model_path)

                    logger.info(f"[ModelTrainer] ✓ Saved: isolation_forest_{lang}.joblib ({len(X_lang)} samples)")
                else:
                    logger.warning(f"[ModelTrainer] Skipping {lang}: only {len(X_lang)} samples (min: {MIN_SAMPLES_PER_LANGUAGE})")

            # Also save a legacy "embeddings_only" model for backward compatibility (trained on English)
            if "english" in per_lang_models:
                import shutil
                english_model_path = Path(per_lang_models["english"])
                legacy_path = Path(self.config.output_directory) / "isolation_forest_embeddings_only.joblib"
                shutil.copy(english_model_path, legacy_path)
                logger.info(f"[ModelTrainer] ✓ Legacy model copied: isolation_forest_embeddings_only.joblib")

            logger.info(f"[ModelTrainer] Per-language training complete: {list(per_lang_models.keys())}")

        except Exception as e:
            logger.warning(f"[ModelTrainer] Per-language model training failed: {e}")
            import traceback
            traceback.print_exc()

        return artifact