utils.py

"""
Utility functions for MLflow tutorial.

This module contains common functions used across all tutorial components.
"""

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # Use non-GUI backend for matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.datasets import make_classification, make_regression
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, classification_report
import os


def print_header(title, width=80):
    """
    Print a formatted header.
    
    Args:
        title: Title text
        width: Width of the header
    """
    print("\n" + "=" * width)
    print(f"  {title}")
    print("=" * width + "\n")


def print_subheader(title, width=80):
    """
    Print a formatted subheader.
    
    Args:
        title: Title text
        width: Width of the subheader
    """
    print("\n" + "-" * width)
    print(f"  {title}")
    print("-" * width + "\n")


def create_classification_dataset(n_samples=1000, n_features=20, n_classes=2, 
                                 test_size=0.2, random_state=42):
    """
    Create a synthetic classification dataset.
    
    Args:
        n_samples: Number of samples
        n_features: Number of features
        n_classes: Number of classes
        test_size: Proportion of test set
        random_state: Random seed
    
    Returns:
        X_train, X_test, y_train, y_test
    """
    X, y = make_classification(
        n_samples=n_samples,
        n_features=n_features,
        n_informative=int(n_features * 0.7),
        n_redundant=int(n_features * 0.2),
        n_classes=n_classes,
        random_state=random_state
    )
    
    return train_test_split(X, y, test_size=test_size, random_state=random_state)


def create_regression_dataset(n_samples=1000, n_features=20, 
                              test_size=0.2, random_state=42):
    """
    Create a synthetic regression dataset.
    
    Args:
        n_samples: Number of samples
        n_features: Number of features
        test_size: Proportion of test set
        random_state: Random seed
    
    Returns:
        X_train, X_test, y_train, y_test
    """
    X, y = make_regression(
        n_samples=n_samples,
        n_features=n_features,
        n_informative=int(n_features * 0.7),
        random_state=random_state
    )
    
    return train_test_split(X, y, test_size=test_size, random_state=random_state)


def plot_confusion_matrix(y_true, y_pred, labels=None, save_path=None):
    """
    Plot confusion matrix.
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        labels: Class labels
        save_path: Path to save the plot
    
    Returns:
        Path to saved plot (if save_path provided)
    """
    cm = confusion_matrix(y_true, y_pred)
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=labels, yticklabels=labels)
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        plt.close()
        return save_path
    else:
        plt.show()


def plot_feature_importance(feature_importance, feature_names=None, 
                           top_n=10, save_path=None):
    """
    Plot feature importance.
    
    Args:
        feature_importance: Array of feature importances
        feature_names: Names of features
        top_n: Number of top features to show
        save_path: Path to save the plot
    
    Returns:
        Path to saved plot (if save_path provided)
    """
    # Get top N features
    indices = np.argsort(feature_importance)[::-1][:top_n]
    values = feature_importance[indices]
    
    # Create labels
    if feature_names is None:
        labels = [f"Feature {i}" for i in indices]
    else:
        labels = [feature_names[i] for i in indices]
    
    plt.figure(figsize=(10, 6))
    plt.bar(range(len(values)), values)
    plt.xticks(range(len(values)), labels, rotation=45, ha='right')
    plt.title(f'Top {top_n} Feature Importances')
    plt.xlabel('Feature')
    plt.ylabel('Importance')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        plt.close()
        return save_path
    else:
        plt.show()


def plot_training_history(history, metrics=['loss', 'accuracy'], save_path=None):
    """
    Plot training history.
    
    Args:
        history: Dictionary with metric names as keys and lists of values
        metrics: List of metrics to plot
        save_path: Path to save the plot
    
    Returns:
        Path to saved plot (if save_path provided)
    """
    n_metrics = len(metrics)
    fig, axes = plt.subplots(1, n_metrics, figsize=(6*n_metrics, 5))
    
    if n_metrics == 1:
        axes = [axes]
    
    for i, metric in enumerate(metrics):
        if metric in history:
            axes[i].plot(history[metric])
            axes[i].set_title(f'Training {metric.capitalize()}')
            axes[i].set_xlabel('Epoch')
            axes[i].set_ylabel(metric.capitalize())
            axes[i].grid(True)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        plt.close()
        return save_path
    else:
        plt.show()


def create_temp_dir(base_dir="temp"):
    """
    Create a temporary directory for artifacts.
    
    Args:
        base_dir: Base directory name
    
    Returns:
        Path to created directory
    """
    os.makedirs(base_dir, exist_ok=True)
    return base_dir


def cleanup_temp_dir(temp_dir="temp"):
    """
    Remove temporary directory.
    
    Args:
        temp_dir: Directory to remove
    """
    import shutil
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)


def format_metrics(metrics_dict, precision=4):
    """
    Format metrics dictionary for display.
    
    Args:
        metrics_dict: Dictionary of metrics
        precision: Number of decimal places
    
    Returns:
        Formatted string
    """
    lines = []
    for key, value in metrics_dict.items():
        if isinstance(value, float):
            lines.append(f"  {key}: {value:.{precision}f}")
        else:
            lines.append(f"  {key}: {value}")
    return "\n".join(lines)


def generate_classification_report_text(y_true, y_pred, save_path=None):
    """
    Generate and optionally save classification report.
    
    Args:
        y_true: True labels
        y_pred: Predicted labels
        save_path: Path to save the report
    
    Returns:
        Classification report as string
    """
    report = classification_report(y_true, y_pred)
    
    if save_path:
        with open(save_path, 'w') as f:
            f.write(report)
    
    return report


def save_dict_as_json(data_dict, file_path):
    """
    Save dictionary as JSON file.
    
    Args:
        data_dict: Dictionary to save
        file_path: Path to save JSON file
    """
    import json
    with open(file_path, 'w') as f:
        json.dump(data_dict, f, indent=2)


def load_json_as_dict(file_path):
    """
    Load JSON file as dictionary.
    
    Args:
        file_path: Path to JSON file
    
    Returns:
        Dictionary
    """
    import json
    with open(file_path, 'r') as f:
        return json.load(f)


def print_run_info(run):
    """
    Print MLflow run information.
    
    Args:
        run: MLflow run object
    """
    print(f"Run ID: {run.info.run_id}")
    print(f"Experiment ID: {run.info.experiment_id}")
    print(f"Status: {run.info.status}")
    print(f"Start Time: {run.info.start_time}")
    print(f"End Time: {run.info.end_time}")
    print(f"Artifact URI: {run.info.artifact_uri}")


def print_experiment_info(experiment):
    """
    Print MLflow experiment information.
    
    Args:
        experiment: MLflow experiment object
    """
    print(f"Experiment Name: {experiment.name}")
    print(f"Experiment ID: {experiment.experiment_id}")
    print(f"Artifact Location: {experiment.artifact_location}")
    print(f"Lifecycle Stage: {experiment.lifecycle_stage}")


def compare_runs(runs, metrics=['accuracy', 'f1_score']):
    """
    Compare multiple MLflow runs.
    
    Args:
        runs: List of MLflow run objects
        metrics: List of metrics to compare
    
    Returns:
        DataFrame with comparison
    """
    data = []
    for run in runs:
        row = {
            'run_id': run.info.run_id,
            'run_name': run.data.tags.get('mlflow.runName', 'N/A'),
            'status': run.info.status
        }
        
        # Add parameters
        for param_key, param_value in run.data.params.items():
            row[f'param_{param_key}'] = param_value
        
        # Add metrics
        for metric in metrics:
            row[metric] = run.data.metrics.get(metric, None)
        
        data.append(row)
    
    return pd.DataFrame(data)


def create_sample_features_dataframe(n_samples=100, n_features=5):
    """
    Create a sample features DataFrame.
    
    Args:
        n_samples: Number of samples
        n_features: Number of features
    
    Returns:
        DataFrame with sample features
    """
    data = np.random.randn(n_samples, n_features)
    columns = [f'feature_{i}' for i in range(n_features)]
    return pd.DataFrame(data, columns=columns)


# Color schemes for visualizations
COLORS = {
    'primary': '#1f77b4',
    'secondary': '#ff7f0e',
    'success': '#2ca02c',
    'danger': '#d62728',
    'warning': '#ff9896',
    'info': '#17becf'
}


def set_plotting_style():
    """Set a consistent plotting style for all visualizations."""
    plt.style.use('seaborn-v0_8-darkgrid')
    sns.set_palette("husl")


if __name__ == "__main__":
    # Test utilities
    print_header("Testing Utility Functions")
    
    # Test data generation
    print("Creating classification dataset...")
    X_train, X_test, y_train, y_test = create_classification_dataset()
    print(f"✓ Train: {X_train.shape}, Test: {X_test.shape}")
    
    # Test plotting
    print("\nTesting plotting functions...")
    set_plotting_style()
    print("✓ Plotting style set")
    
    print("\n✓ All utility functions working correctly!")