# src/data/02_preprocessing/MLproject
name: preprocessing_and_imputation

conda_env: conda.yaml

entry_points:
  main:
    parameters:
      gcs_input_path: {type: string}
      gcs_output_path: {type: string}
    command: "python main.py --gcs_input_path={gcs_input_path} --gcs_output_path={gcs_output_path}"

python src/data/01_download_data/main.py
python src/data/02_preprocessing/main.py
python src/data/03_feature_engineering/main.py
# ...

"""
MLOps Pipeline Orchestrator
Executes steps sequentially using MLflow + Hydra
"""
import os
import sys
import mlflow
import hydra
from omegaconf import DictConfig
from pathlib import Path
import time

def validate_environment_variables() -> None:
    """Fail fast if critical secrets are missing."""
    required_vars = {
        "GCP_PROJECT_ID": "Google Cloud Project ID",
        "GCS_BUCKET_NAME": "GCS Bucket name",
        "WANDB_API_KEY": "Weights & Biases API Key",
    }

    missing = []
    for var, description in required_vars.items():
        value = os.getenv(var)
        if not value or value in ["your-project-id", "your-key"]:
            missing.append(f"  ERROR: {var}: {description}")

    if missing:
        print("\n" + "="*70)
        print("ERROR: MISSING REQUIRED ENVIRONMENT VARIABLES")
        print("="*70)
        print("\n".join(missing))
        print("\nCreate .env file with:")
        print("  GCP_PROJECT_ID=your-project-id")
        print("  GCS_BUCKET_NAME=your-bucket")
        print("  WANDB_API_KEY=your-key")
        sys.exit(1)

def get_steps_to_execute(config: DictConfig) -> list[str]:
    """Converts execute_steps from config to list."""
    steps = config['main']['execute_steps']
    if isinstance(steps, str):
        return [s.strip() for s in steps.split(',')]
    return list(steps)

def run_step(step_name: str, step_path: Path, entry_point: str, parameters: dict):
    """Executes a step as MLflow project."""
    print(f"\n{'='*70}")
    print(f"EXECUTING: {step_name}")
    print(f"{'='*70}")

    mlflow.run(
        uri=str(step_path),
        entry_point=entry_point,
        env_manager="local",
        parameters=parameters
    )

@hydra.main(config_path='.', config_name="config", version_base="1.3")
def go(config: DictConfig) -> None:
    """Main pipeline entry point."""

    validate_environment_variables()

    mlflow.set_experiment(config['main']['experiment_name'])
    steps_to_execute = get_steps_to_execute(config)
    root_path = Path(__file__).parent

    start_time = time.time()

    try:
        # Step 01: Download Data
        if "01_download_data" in steps_to_execute:
            run_step(
                "01 - Download Data",
                root_path / "src" / "data" / "01_download_data",
                "main",
                {
                    "file_url": config["download_data"]["file_url"],
                    "gcs_output_path": config["download_data"]["gcs_output_path"],
                }
            )

        # ... Steps 02-07 similar pattern ...

        elapsed = time.time() - start_time
        print(f"\nSUCCESS: PIPELINE COMPLETED ({elapsed:.1f}s)")

    except Exception as e:
        print(f"\nERROR: PIPELINE FAILED: {e}")
        raise

if __name__ == "__main__":
    go()

main:
  project_name: "housing-mlops-gcp"
  experiment_name: "end_to_end_pipeline"
  execute_steps:
    - "01_download_data"
    - "02_preprocessing_and_imputation"
    - "03_feature_engineering"
    - "04_segregation"
    - "05_model_selection"
    - "06_sweep"
    - "07_registration"

download_data:
  file_url: "https://raw.githubusercontent.com/ageron/handson-ml2/master/datasets/housing/housing.tgz"
  gcs_output_path: "data/01-raw/housing.parquet"

preprocessing:
  gcs_input_path: "data/01-raw/housing.parquet"
  gcs_output_path: "data/02-processed/housing_processed.parquet"
  imputation_strategy: "auto"  # Will compare 4 strategies

feature_engineering:
  gcs_input_path: "data/02-processed/housing_processed.parquet"
  gcs_output_path: "data/03-features/housing_features.parquet"
  n_clusters: 10
  optimize_hyperparams: true  # Find best K

segregation:
  gcs_input_path: "data/03-features/housing_features.parquet"
  gcs_train_output_path: "data/04-split/train/train.parquet"
  gcs_test_output_path: "data/04-split/test/test.parquet"
  test_size: 0.2
  target_column: "median_house_value"

model_selection:
  gcs_train_path: "data/04-split/train/train.parquet"
  gcs_test_path: "data/04-split/test/test.parquet"

sweep:
  sweep_count: 50  # 50 Bayesian optimization runs
  metric_name: "mape"
  metric_goal: "minimize"

registration:
  registered_model_name: "housing_price_model"
  model_stage: "Staging"  # Or "Production"

execute_steps: ["03_feature_engineering", "05_model_selection"]

"""
Imputation Analyzer - Automatically compares strategies
Author: Carlos Daniel Jiménez
"""
from dataclasses import dataclass
from typing import Dict, Tuple
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from sklearn.impute import SimpleImputer, KNNImputer, IterativeImputer
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

@dataclass
class ImputationResult:
    """Result of an imputation strategy."""
    method_name: str
    rmse: float
    imputed_values: np.ndarray
    imputer: object

class ImputationAnalyzer:
    """
    Analyzes and compares imputation strategies.
    Automatically selects the best based on RMSE.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        target_column: str = "total_bedrooms",
        test_size: float = 0.2,
        random_state: int = 42
    ):
        self.df = df
        self.target_column = target_column
        self.test_size = test_size
        self.random_state = random_state
        self.results: Dict[str, ImputationResult] = {}
        self.best_method: str = None
        self.best_imputer: object = None

    def prepare_validation_set(self) -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
        """
        Creates masked validation set to compare strategies.

        Strategy:
        1. Remove rows with missing target (can't validate against NaN)
        2. Split into train/val
        3. Mask target in val set (simulate missing values)
        4. Save ground truth

        Returns:
            (train_set, val_set_missing, y_val_true)
        """
        housing_numeric = self.df.select_dtypes(include=[np.number])
        housing_known = housing_numeric.dropna(subset=[self.target_column])

        train_set, val_set = train_test_split(
            housing_known,
            test_size=self.test_size,
            random_state=self.random_state
        )

        # Mask target in val
        val_set_missing = val_set.copy()
        val_set_missing[self.target_column] = np.nan

        # Ground truth
        y_val_true = val_set[self.target_column].copy()

        return train_set, val_set_missing, y_val_true

    def evaluate_simple_imputer(
        self,
        train_set: pd.DataFrame,
        val_set_missing: pd.DataFrame,
        y_val_true: pd.Series,
        strategy: str = "median"
    ) -> ImputationResult:
        """Evaluates SimpleImputer with given strategy."""
        imputer = SimpleImputer(strategy=strategy)
        imputer.fit(train_set)

        val_imputed = imputer.transform(val_set_missing)

        target_col_idx = train_set.columns.get_loc(self.target_column)
        y_val_pred = val_imputed[:, target_col_idx]

        rmse = np.sqrt(mean_squared_error(y_val_true, y_val_pred))

        return ImputationResult(
            method_name=f"Simple Imputer ({strategy})",
            rmse=rmse,
            imputed_values=y_val_pred,
            imputer=imputer
        )

    def evaluate_knn_imputer(
        self,
        train_set: pd.DataFrame,
        val_set_missing: pd.DataFrame,
        y_val_true: pd.Series,
        n_neighbors: int = 5
    ) -> ImputationResult:
        """
        Evaluates KNNImputer with scaling.

        CRITICAL: KNN requires scaled features or explodes with overflow.
        """
        import warnings
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore', category=RuntimeWarning)

            # Scale data
            scaler = StandardScaler()
            train_scaled = scaler.fit_transform(train_set)
            val_scaled = scaler.transform(val_set_missing)

            # KNN imputation
            imputer = KNNImputer(n_neighbors=n_neighbors)
            imputer.fit(train_scaled)
            val_imputed_scaled = imputer.transform(val_scaled)

            # Inverse scale
            val_imputed = scaler.inverse_transform(val_imputed_scaled)

        target_col_idx = train_set.columns.get_loc(self.target_column)
        y_val_pred = val_imputed[:, target_col_idx]

        rmse = np.sqrt(mean_squared_error(y_val_true, y_val_pred))

        return ImputationResult(
            method_name=f"KNN Imputer (k={n_neighbors})",
            rmse=rmse,
            imputed_values=y_val_pred,
            imputer=(scaler, imputer)  # Store tuple!
        )

    def evaluate_iterative_imputer(
        self,
        train_set: pd.DataFrame,
        val_set_missing: pd.DataFrame,
        y_val_true: pd.Series
    ) -> ImputationResult:
        """Evaluates IterativeImputer with RandomForest estimator."""
        estimator = RandomForestRegressor(
            n_jobs=-1,
            random_state=self.random_state
        )
        imputer = IterativeImputer(
            estimator=estimator,
            random_state=self.random_state
        )

        imputer.fit(train_set)
        val_imputed = imputer.transform(val_set_missing)

        target_col_idx = train_set.columns.get_loc(self.target_column)
        y_val_pred = val_imputed[:, target_col_idx]

        rmse = np.sqrt(mean_squared_error(y_val_true, y_val_pred))

        return ImputationResult(
            method_name="Iterative Imputer (RF)",
            rmse=rmse,
            imputed_values=y_val_pred,
            imputer=imputer
        )

    def compare_all_methods(self) -> Dict[str, ImputationResult]:
        """Compares all strategies and selects the best."""
        train_set, val_set_missing, y_val_true = self.prepare_validation_set()

        # Evaluate all
        self.results['simple_median'] = self.evaluate_simple_imputer(
            train_set, val_set_missing, y_val_true, strategy="median"
        )

        self.results['simple_mean'] = self.evaluate_simple_imputer(
            train_set, val_set_missing, y_val_true, strategy="mean"
        )

        self.results['knn'] = self.evaluate_knn_imputer(
            train_set, val_set_missing, y_val_true, n_neighbors=5
        )

        self.results['iterative_rf'] = self.evaluate_iterative_imputer(
            train_set, val_set_missing, y_val_true
        )

        # Select best
        best_key = min(self.results, key=lambda k: self.results[k].rmse)
        self.best_method = best_key
        self.best_imputer = self.results[best_key].imputer

        # Print summary
        print("\n" + "="*70)
        print("IMPUTATION METHODS COMPARISON")
        print("="*70)
        for key, result in sorted(self.results.items(), key=lambda x: x[1].rmse):
            status = "[BEST]" if key == best_key else ""
            print(f"  {result.method_name:30s} RMSE: {result.rmse:8.4f} {status}")
        print("="*70)

        return self.results

    def apply_best_imputer(self, df: pd.DataFrame) -> pd.DataFrame:
        """Applies the best imputer to the complete dataset."""
        if self.best_imputer is None:
            raise ValueError("Run compare_all_methods() first")

        df_out = df.copy()
        numeric_df = df_out.select_dtypes(include=[np.number])

        import warnings
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore', category=RuntimeWarning)

            # Check if it's a tuple (KNN with scaler)
            if isinstance(self.best_imputer, tuple):
                scaler, imputer = self.best_imputer
                numeric_scaled = scaler.transform(numeric_df)
                imputed_scaled = imputer.transform(numeric_scaled)
                imputed_array = scaler.inverse_transform(imputed_scaled)
            else:
                imputed_array = self.best_imputer.transform(numeric_df)

        target_col_idx = numeric_df.columns.get_loc(self.target_column)
        df_out[self.target_column] = imputed_array[:, target_col_idx]

        return df_out

    def create_comparison_plot(self) -> plt.Figure:
        """Creates bar plot comparing RMSE of methods."""
        methods = [r.method_name for r in self.results.values()]
        rmses = [r.rmse for r in self.results.values()]

        fig, ax = plt.subplots(figsize=(10, 6))
        colors = ['green' if i == np.argmin(rmses) else 'skyblue'
                  for i in range(len(rmses))]
        bars = ax.bar(methods, rmses, color=colors)

        ax.set_xlabel('Imputation Method', fontweight='bold')
        ax.set_ylabel('RMSE', fontweight='bold')
        ax.set_title('Comparison of Imputation Strategies', fontsize=14)
        ax.grid(axis='y', alpha=0.3)

        # Value labels
        for bar, rmse in zip(bars, rmses):
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width()/2.,
                height,
                f'{rmse:.4f}',
                ha='center',
                va='bottom'
            )

        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()

        return fig

train_set, val_set = train_test_split(housing_known, test_size=0.2)
val_set_missing = val_set.copy()
val_set_missing[self.target_column] = np.nan
y_val_true = val_set[self.target_column].copy()

scaler = StandardScaler()
train_scaled = scaler.fit_transform(train_set)

if isinstance(self.best_imputer, tuple):
    scaler, imputer = self.best_imputer
    # ... apply both

# In Step 02 main.py
import wandb
import mlflow

analyzer = ImputationAnalyzer(df, target_column="total_bedrooms")
results = analyzer.compare_all_methods()

# Log to W&B
comparison_plot = analyzer.create_comparison_plot()
wandb.log({
    "imputation/comparison": wandb.Image(comparison_plot),
    "imputation/best_method": analyzer.best_method,
    "imputation/best_rmse": results[analyzer.best_method].rmse,
})

# Apply to complete dataset
housing_clean = analyzer.apply_best_imputer(housing_df)

# Save imputer
import joblib
joblib.dump(analyzer.best_imputer, "artifacts/imputer.pkl")
mlflow.log_artifact("artifacts/imputer.pkl")

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.cluster import KMeans
import numpy as np

class ClusterSimilarity(BaseEstimator, TransformerMixin):
    """
    Custom transformer for geographic clustering.

    Design: Scikit-learn transformer to integrate into Pipeline.
    """

    def __init__(self, n_clusters=10, gamma=1.0, random_state=None):
        self.n_clusters = n_clusters
        self.gamma = gamma  # Placeholder for RBF kernel (not currently used)
        self.random_state = random_state

    def fit(self, X, y=None, sample_weight=None):
        """Fit KMeans on geographic coordinates."""
        self.kmeans_ = KMeans(
            self.n_clusters,
            n_init=10,
            random_state=self.random_state
        )
        self.kmeans_.fit(X, sample_weight=sample_weight)
        return self

    def transform(self, X):
        """Transform coordinates to cluster labels."""
        cluster_labels = self.kmeans_.predict(X)
        return np.expand_dims(cluster_labels, axis=1)

    def get_feature_names_out(self, names=None):
        """Returns feature names for Pipeline."""
        return ["cluster_label"]

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder, StandardScaler

def create_preprocessing_pipeline(n_clusters=10):
    """
    Creates pipeline that processes:
    - Numeric: impute + scale
    - Categorical: impute + one-hot
    - Geo: clustering
    """

    num_pipeline = Pipeline([
        ("impute", SimpleImputer(strategy="median")),
        ("standardize", StandardScaler()),
    ])

    cat_pipeline = Pipeline([
        ("impute", SimpleImputer(strategy="most_frequent")),
        ("onehot", OneHotEncoder(handle_unknown="ignore")),
    ])

    num_attribs = [
        "longitude", "latitude", "housing_median_age",
        "total_rooms", "total_bedrooms", "population",
        "households", "median_income"
    ]

    cat_attribs = ["ocean_proximity"]

    preprocessing = ColumnTransformer([
        ("num", num_pipeline, num_attribs),
        ("cat", cat_pipeline, cat_attribs),
        ("geo", ClusterSimilarity(n_clusters=n_clusters),
         ["latitude", "longitude"]),
    ])

    return preprocessing

from sklearn.metrics import silhouette_score, davies_bouldin_score

def optimize_n_clusters(
    df: pd.DataFrame,
    min_clusters=2,
    max_clusters=20
) -> Tuple[int, Dict]:
    """
    Finds the best K for KMeans using silhouette score.

    Metrics:
    - Silhouette score (0 to 1): Cluster separation. Maximize.
    - Davies-Bouldin index: Internal dispersion vs separation. Minimize.
    """
    geo_features = df[["latitude", "longitude"]].values

    cluster_range = range(min_clusters, max_clusters + 1)
    silhouette_scores = []
    davies_bouldin_scores = []
    inertias = []

    for n in cluster_range:
        kmeans = KMeans(n_clusters=n, n_init=10, random_state=42)
        labels = kmeans.fit_predict(geo_features)

        silhouette_scores.append(silhouette_score(geo_features, labels))
        davies_bouldin_scores.append(davies_bouldin_score(geo_features, labels))
        inertias.append(kmeans.inertia_)

    # Select K with best silhouette
    optimal_n = cluster_range[np.argmax(silhouette_scores)]

    metrics = {
        "optimal_n_clusters": optimal_n,
        "best_silhouette": max(silhouette_scores),
        "cluster_range": list(cluster_range),
        "silhouette_scores": silhouette_scores,
        "davies_bouldin_scores": davies_bouldin_scores,
        "inertias": inertias,
    }

    return optimal_n, metrics

def create_optimization_plots(metrics: Dict) -> Dict[str, plt.Figure]:
    """Creates K optimization plots."""

    # Plot 1: Elbow Method (Inertia)
    fig1, ax1 = plt.subplots(figsize=(10, 6))
    ax1.plot(metrics["cluster_range"], metrics["inertias"], 'bo-')
    ax1.axvline(
        metrics["optimal_n_clusters"],
        color='r',
        linestyle='--',
        label=f'Optimal K={metrics["optimal_n_clusters"]}'
    )
    ax1.set_xlabel('Number of Clusters (K)')
    ax1.set_ylabel('Inertia')
    ax1.set_title('Elbow Method - KMeans Optimization')
    ax1.legend()
    ax1.grid(True)

    # Plot 2: Silhouette Score
    fig2, ax2 = plt.subplots(figsize=(10, 6))
    ax2.plot(metrics["cluster_range"], metrics["silhouette_scores"], 'go-')
    ax2.axvline(
        metrics["optimal_n_clusters"],
        color='r',
        linestyle='--'
    )
    ax2.set_xlabel('Number of Clusters (K)')
    ax2.set_ylabel('Silhouette Score')
    ax2.set_title('Silhouette Score vs K')
    ax2.grid(True)

    return {"elbow_method": fig1, "silhouette_scores": fig2}

# In Step 03 main.py
import wandb
import mlflow
import joblib

# Download data from GCS
df = download_from_gcs(bucket, "data/02-processed/housing_processed.parquet")

# Optimize K
optimal_k, metrics = optimize_n_clusters(df, min_clusters=5, max_clusters=15)

print(f"Optimal K: {optimal_k}")
print(f"   Silhouette: {metrics['best_silhouette']:.4f}")

# Create plots
plots = create_optimization_plots(metrics)

# Log to W&B
wandb.log({
    "optimization/optimal_k": optimal_k,
    "optimization/silhouette": metrics["best_silhouette"],
    "optimization/elbow_plot": wandb.Image(plots["elbow_method"]),
    "optimization/silhouette_plot": wandb.Image(plots["silhouette_scores"]),
})

# Create pipeline with optimal K
preprocessing_pipeline = create_preprocessing_pipeline(n_clusters=optimal_k)

# Fit pipeline
target_column = "median_house_value"
y = df[target_column]
X = df.drop(columns=[target_column])

preprocessing_pipeline.fit(X, y)

# Transform data
X_transformed = preprocessing_pipeline.transform(X)

# Reconstruct DataFrame with target
df_transformed = pd.DataFrame(
    X_transformed,
    columns=preprocessing_pipeline.get_feature_names_out()
)
df_transformed[target_column] = y.values

# Upload to GCS
upload_to_gcs(df_transformed, bucket, "data/03-features/housing_features.parquet")

# Save pipeline
joblib.dump(preprocessing_pipeline, "artifacts/preprocessing_pipeline.pkl")
mlflow.log_artifact("artifacts/preprocessing_pipeline.pkl")

for k in range(5, 15):
    pipeline = create_preprocessing_pipeline(n_clusters=k)
    X_train_transformed = pipeline.fit_transform(X_train, y_train)
    X_test_transformed = pipeline.transform(X_test)

    model = RandomForestRegressor()
    model.fit(X_train_transformed, y_train)

    mape = calculate_mape(model, X_test_transformed, y_test)
    # Select K with best MAPE

OneHotEncoder(handle_unknown="ignore")

joblib.dump(preprocessing_pipeline, "artifacts/preprocessing_pipeline.pkl")

# =================================================================
# W&B Sweep Configuration for Random Forest
# Author: Carlos Daniel Jiménez
# =================================================================

program: main.py
method: bayes  # Bayesian optimization, not random, not grid

metric:
  name: wmape  # Weighted MAPE (less biased than MAPE)
  goal: minimize

parameters:
  n_estimators:
    min: 50
    max: 500

  max_depth:
    min: 5
    max: 30

  min_samples_split:
    min: 2
    max: 20

  min_samples_leaf:
    min: 1
    max: 10

  max_features:
    values: ['sqrt', 'log2']

# Early stopping: eliminates poor runs early
early_terminate:
  type: hyperband
  min_iter: 10   # Minimum 10 runs before terminating
  eta: 3         # Eliminates 1/3 of poor runs
  s: 2

name: housing-rf-sweep-improved
description: "Optimize Random Forest with wmape + feature tracking"

"""
W&B Sweep for Random Forest Hyperparameter Optimization.
Author: Carlos Daniel Jiménez
"""
import argparse
import yaml
import wandb
import logging
from pathlib import Path

from utils import (
    download_data_from_gcs,
    prepare_data,
    train_random_forest,
    evaluate_model,
    log_feature_importances
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Module-level data cache (loaded once, reused in all runs)
_data_cache = {
    "X_train": None,
    "X_test": None,
    "y_train": None,
    "y_test": None,
    "feature_names": None
}

def train():
    """
    Training function called by W&B Sweep agent.
    Executed for each hyperparameter combination.

    Uses module-level cache to avoid reloading data on each run.
    """
    run = wandb.init()
    config = wandb.config

    logger.info("="*70)
    logger.info(f"SWEEP RUN: {run.name}")
    logger.info("="*70)

    try:
        # Prepare parameters
        params = {
            'n_estimators': int(config.n_estimators),
            'max_depth': int(config.max_depth) if config.max_depth else None,
            'min_samples_split': int(config.min_samples_split),
            'min_samples_leaf': int(config.min_samples_leaf),
            'max_features': config.max_features,
            'random_state': 42
        }

        # Train model using cached data
        model = train_random_forest(
            _data_cache["X_train"],
            _data_cache["y_train"],
            params
        )

        # Evaluate model
        metrics = evaluate_model(
            model,
            _data_cache["X_test"],
            _data_cache["y_test"]
        )

        # Log feature importances
        feature_importances = log_feature_importances(
            model,
            _data_cache["feature_names"]
        )

        # Log everything to W&B
        wandb.log({
            **params,
            **metrics,
            **{f"feature_importance_{k}": v
               for k, v in list(feature_importances.items())[:10]}
        })

        logger.info(f"SUCCESS: MAPE={metrics['mape']:.2f}% | "
                   f"WMAPE={metrics['wmape']:.2f}%")

    except Exception as e:
        logger.error(f"ERROR: Run failed: {str(e)}")
        wandb.log({
            "error": str(e),
            "mape": 999.9,
            "wmape": 999.9
        })
        raise

    finally:
        run.finish()

def main():
    """Main function to initialize and run the sweep."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--gcs_train_path", type=str, required=True)
    parser.add_argument("--gcs_test_path", type=str, required=True)
    parser.add_argument("--bucket_name", type=str, required=True)
    parser.add_argument("--wandb_project", type=str, required=True)
    parser.add_argument("--target_column", type=str, default="median_house_value")
    parser.add_argument("--sweep_count", type=int, default=50)
    parser.add_argument("--sweep_config", type=str, default="sweep_config.yaml")
    args = parser.parse_args()

    logger.info("="*70)
    logger.info("W&B SWEEP - HYPERPARAMETER OPTIMIZATION")
    logger.info("="*70)

    # Load data ONCE into module-level cache
    logger.info("\nLoading data into cache...")
    train_df = download_data_from_gcs(args.bucket_name, args.gcs_train_path)
    X_train, y_train = prepare_data(train_df, args.target_column)

    test_df = download_data_from_gcs(args.bucket_name, args.gcs_test_path)
    X_test, y_test = prepare_data(test_df, args.target_column)

    # Store in cache
    _data_cache["X_train"] = X_train
    _data_cache["X_test"] = X_test
    _data_cache["y_train"] = y_train
    _data_cache["y_test"] = y_test
    _data_cache["feature_names"] = X_train.columns.tolist()

    logger.info(f"Data cached: Train {X_train.shape}, Test {X_test.shape}")

    # Load sweep configuration
    sweep_config_path = Path(__file__).parent / args.sweep_config

    with open(sweep_config_path, 'r') as f:
        sweep_config = yaml.safe_load(f)

    logger.info(f"\nSweep config:")
    logger.info(f"  Method: {sweep_config['method']}")
    logger.info(f"  Metric: {sweep_config['metric']['name']} ({sweep_config['metric']['goal']})")

    # Initialize sweep
    sweep_id = wandb.sweep(
        sweep=sweep_config,
        project=args.wandb_project
    )

    logger.info(f"\nSweep created: {sweep_id}")
    logger.info(f"  View at: https://wandb.ai/{args.wandb_project}/sweeps/{sweep_id}")

    # Run sweep agent
    logger.info(f"\nStarting sweep agent ({args.sweep_count} runs)...")
    wandb.agent(
        sweep_id,
        function=train,
        count=args.sweep_count,
        project=args.wandb_project
    )

    logger.info("\n" + "="*70)
    logger.info("SWEEP COMPLETED")
    logger.info("="*70)

    # Save best params
    api = wandb.Api()
    sweep = api.sweep(f"{args.wandb_project}/{sweep_id}")
    best_run = sweep.best_run()

    best_params = {
        "hyperparameters": {
            "n_estimators": int(best_run.config.get('n_estimators')),
            "max_depth": int(best_run.config.get('max_depth')) if best_run.config.get('max_depth') else None,
            "min_samples_split": int(best_run.config.get('min_samples_split')),
            "min_samples_leaf": int(best_run.config.get('min_samples_leaf')),
            "max_features": best_run.config.get('max_features'),
        },
        "metrics": {
            "mape": float(best_run.summary.get('mape')),
            "wmape": float(best_run.summary.get('wmape')),
            "r2": float(best_run.summary.get('r2')),
        },
        "sweep_id": sweep_id,
        "best_run_id": best_run.id
    }

    # Save to YAML
    best_params_path = Path(__file__).parent / "best_params.yaml"
    with open(best_params_path, 'w') as f:
        yaml.dump(best_params, f)

    logger.info(f"\nBest params saved to: {best_params_path}")
    logger.info(f"   MAPE: {best_params['metrics']['mape']:.2f}%")

if __name__ == "__main__":
    main()

method: bayes  # Not random, not grid

metric:
  name: wmape  # Weighted MAPE

_data_cache = {
    "X_train": None,
    "X_test": None,
    # ...
}

early_terminate:
  type: hyperband
  min_iter: 10
  eta: 3

feature_importances = log_feature_importances(model, feature_names)
wandb.log({
    **{f"feature_importance_{k}": v
       for k, v in list(feature_importances.items())[:10]}
})

sweep_id: f73ao31m
best_run_id: 5q1840qa
best_run_name: dry-sweep-5

hyperparameters:
  n_estimators: 128
  max_depth: 23
  min_samples_split: 9
  min_samples_leaf: 9
  max_features: log2
  random_state: 42

metrics:
  # Primary metrics
  mae: 38901.45
  rmse: 53277.02
  r2: 0.78339

  # Percentage error metrics
  mape: 20.4002
  smape: 18.85
  wmape: 19.12
  median_ape: 16.73

  # Accuracy within thresholds
  within_5pct: 12.4
  within_10pct: 36.2
  within_15pct: 52.8

sweep_url: https://wandb.ai/danieljimenez88m-carlosdanieljimenez-com/housing-mlops-gcp/sweeps/f73ao31m

import joblib
joblib.dump(model, "best_model.pkl")

"""
Model registration in MLflow Model Registry.
Author: Carlos Daniel Jiménez
"""
import mlflow
import mlflow.sklearn
from mlflow.tracking import MlflowClient

def register_model_to_mlflow(
    model,
    model_name: str,
    model_stage: str,
    params: dict,
    metrics: dict,
    feature_columns: list,
    target_column: str,
    gcs_train_path: str,
    gcs_test_path: str
) -> tuple:
    """
    Registers model in MLflow Model Registry with rich metadata.

    Args:
        model: Trained sklearn model
        model_name: Name for registered model
        model_stage: Stage (Staging/Production)
        params: Hyperparameters
        metrics: Evaluation metrics
        feature_columns: List of features
        target_column: Target column name
        gcs_train_path: Path to training data
        gcs_test_path: Path to test data

    Returns:
        (model_uri, model_version, run_id)
    """
    logger.info("="*70)
    logger.info("REGISTERING MODEL TO MLFLOW")
    logger.info("="*70)

    client = MlflowClient()
    run_id = mlflow.active_run().info.run_id
    model_uri = f"runs:/{run_id}/model"

    # Log model to MLflow
    mlflow.sklearn.log_model(model, "model")
    logger.info(f"Model logged: {model_uri}")

    # Create or get registered model
    try:
        client.create_registered_model(
            name=model_name,
            description="Housing price prediction - Random Forest"
        )
        logger.info(f"Created new registered model: {model_name}")
    except Exception as e:
        if "already exists" in str(e):
            logger.info(f"Model already exists: {model_name}")
        else:
            raise

    # Create model version
    model_version = client.create_model_version(
        name=model_name,
        source=model_uri,
        run_id=run_id
    )
    logger.info(f"Version created: {model_version.version}")

    # Transition to stage
    client.transition_model_version_stage(
        name=model_name,
        version=model_version.version,
        stage=model_stage  # "Staging" or "Production"
    )
    logger.info(f"Transitioned to: {model_stage}")

    # Create comprehensive description (MARKDOWN)
    description = f"""
# Housing Price Prediction Model

**Algorithm:** Random Forest Regressor

## Hyperparameters
- n_estimators: {params['n_estimators']}
- max_depth: {params.get('max_depth', 'None')}
- min_samples_split: {params['min_samples_split']}
- min_samples_leaf: {params['min_samples_leaf']}
- max_features: {params.get('max_features', 'sqrt')}

## Performance Metrics
- MAPE: {metrics['mape']:.2f}%
- Median APE: {metrics['median_ape']:.2f}%
- Within 10%: {metrics['within_10pct']:.1f}%
- RMSE: {metrics['rmse']:.2f}
- R²: {metrics['r2']:.4f}

## Features
- Number of features: {len(feature_columns)}
- Target: {target_column}

## Data Sources
- Training: {gcs_train_path}
- Testing: {gcs_test_path}
"""

    client.update_model_version(
        name=model_name,
        version=model_version.version,
        description=description
    )

    # Add searchable tags
    tags = {
        "algorithm": "RandomForest",
        "framework": "sklearn",
        "mape": f"{metrics['mape']:.2f}",
        "within_10pct": f"{metrics['within_10pct']:.1f}",
        "rmse": f"{metrics['rmse']:.2f}",
        "r2": f"{metrics['r2']:.4f}",
        "n_features": str(len(feature_columns)),
        "target": target_column,
    }

    for key, value in tags.items():
        client.set_model_version_tag(
            model_name,
            model_version.version,
            key,
            value
        )

    logger.info("Tags added to model version")
    logger.info("="*70)

    return model_uri, model_version.version, run_id

mlflow.sklearn.log_model(model, "model")  # Only artifact
# Usage: mlflow.sklearn.load_model(f"runs://{run_id}/model")

client.create_model_version(name="housing_price_model", source=model_uri)
# Usage: mlflow.pyfunc.load_model("models:/housing_price_model/Production")

description = f"""
# Housing Price Prediction Model

**Algorithm:** Random Forest

## Hyperparameters
- n_estimators: {params['n_estimators']}
...

## Performance Metrics
- MAPE: {metrics['mape']:.2f}%
...
"""

tags = {
    "algorithm": "RandomForest",
    "mape": f"{metrics['mape']:.2f}",
    "r2": f"{metrics['r2']:.4f}",
}

for key, value in tags.items():
    client.set_model_version_tag(model_name, model_version.version, key, value)

model_config = {
    'model': {
        'name': model_name,
        'version': str(model_version),
        'stage': model_stage,
        'parameters': params,
        'metrics': metrics,
        'feature_columns': feature_columns,
        'mlflow_run_id': run_id,
        'sweep_id': sweep_id
    }
}

config_path = Path("configs/model_config.yaml")
with open(config_path, 'w') as f:
    yaml.dump(model_config, f)

mlflow.log_artifact(str(config_path), artifact_path="config")