# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_utils.ipynb.

# %% auto 0
__all__ = ['store_variables', 'load_variables', 'store_patch_variables', 'load_patch_variables', 'print_colab_gpu_info',
           'ModelTrackingCallback', 'create_mlflow_callback', 'MLflowUIManager']

# %% ../nbs/07_utils.ipynb 1
import pickle
import torch
from pathlib import Path
import mlflow
import mlflow.pytorch
import os
import tempfile
import json
from datetime import datetime
from fastai.callback.core import Callback
from fastcore.foundation import L
from typing import Any

# %% ../nbs/07_utils.ipynb 3
def store_variables(pkl_fn: str | Path, size: list, apply_reorder: bool, target_spacing: int | list):
    """Save variable values in a pickle file."""
    
    var_vals = [size, apply_reorder, target_spacing]
    
    with open(pkl_fn, 'wb') as f:
        pickle.dump(var_vals, f)

# %% ../nbs/07_utils.ipynb 4
def load_variables(pkl_fn: (str, Path)):
    """Loads stored variable values from a pickle file.

    Args:
        pkl_fn: File path of the pickle file to be loaded.

    Returns:
        The deserialized value of the pickled data.
    """
    with open(pkl_fn, 'rb') as f:
        return pickle.load(f)

# %% ../nbs/07_utils.ipynb 6
def store_patch_variables(
    pkl_fn: str | Path,
    patch_size: list,
    patch_overlap: int | float | list,
    aggregation_mode: str,
    apply_reorder: bool = False,
    target_spacing: list = None,
    sampler_type: str = 'uniform',
    label_probabilities: dict = None,
    samples_per_volume: int = 8,
    queue_length: int = 300,
    queue_num_workers: int = 4,
    keep_largest_component: bool = False
):
    """Save patch-based training and inference configuration to a pickle file.

    Args:
        pkl_fn: Path to save the pickle file.
        patch_size: Size of patches [x, y, z].
        patch_overlap: Overlap for inference (int, float 0-1, or list).
        aggregation_mode: GridAggregator mode ('crop', 'average', 'hann').
        apply_reorder: Whether to reorder to canonical (RAS+) orientation.
        target_spacing: Target voxel spacing [x, y, z].
        sampler_type: Type of sampler used during training.
        label_probabilities: Label probabilities for LabelSampler.
        samples_per_volume: Number of patches extracted per volume during training.
        queue_length: Maximum number of patches in queue buffer.
        queue_num_workers: Number of workers for parallel patch extraction.
        keep_largest_component: If True, keep only the largest connected component
            in binary segmentation predictions during inference.

    Example:
        >>> store_patch_variables(
        ...     'patch_settings.pkl',
        ...     patch_size=[96, 96, 96],
        ...     patch_overlap=0.5,
        ...     aggregation_mode='hann',
        ...     apply_reorder=True,
        ...     target_spacing=[1.0, 1.0, 1.0],
        ...     samples_per_volume=16,
        ...     keep_largest_component=True
        ... )
    """
    config = {
        'patch_size': patch_size,
        'patch_overlap': patch_overlap,
        'aggregation_mode': aggregation_mode,
        'apply_reorder': apply_reorder,
        'target_spacing': target_spacing,
        'sampler_type': sampler_type,
        'label_probabilities': label_probabilities,
        'samples_per_volume': samples_per_volume,
        'queue_length': queue_length,
        'queue_num_workers': queue_num_workers,
        'keep_largest_component': keep_largest_component
    }

    with open(pkl_fn, 'wb') as f:
        pickle.dump(config, f)

# %% ../nbs/07_utils.ipynb 7
def load_patch_variables(pkl_fn: str | Path) -> dict:
    """Load patch-based training and inference configuration from a pickle file.

    Args:
        pkl_fn: Path to the pickle file.

    Returns:
        Dictionary with patch configuration including:
        - patch_size, patch_overlap, aggregation_mode
        - apply_reorder, target_spacing, sampler_type, label_probabilities
        - samples_per_volume, queue_length, queue_num_workers

    Example:
        >>> config = load_patch_variables('patch_settings.pkl')
        >>> from fastMONAI.vision_patch import PatchConfig
        >>> patch_config = PatchConfig(**config)
    """
    with open(pkl_fn, 'rb') as f:
        return pickle.load(f)

# %% ../nbs/07_utils.ipynb 8
def print_colab_gpu_info(): 
    """Check if we have a GPU attached to the runtime."""
    
    colab_gpu_msg =(f"{'#'*80}\n"
                    "Remember to attach a GPU to your Colab Runtime:"
                    "\n1. From the **Runtime** menu select **Change Runtime Type**"
                    "\n2. Choose **GPU** from the drop-down menu"
                    "\n3. Click **'SAVE'**\n"
                    f"{'#'*80}")
    
    if torch.cuda.is_available(): print('GPU attached.')
    else: print(colab_gpu_msg)

# %% ../nbs/07_utils.ipynb 9
def _detect_patch_workflow(dls) -> bool:
    """Detect if DataLoaders are patch-based (MedPatchDataLoaders).
    
    Args:
        dls: DataLoaders instance
        
    Returns:
        True if dls is a MedPatchDataLoaders instance
    """
    return hasattr(dls, 'patch_config') or hasattr(dls, '_patch_config')


def _extract_size_from_transforms(tfms) -> list | None:
    """Extract target size from PadOrCrop transform if present.
    
    Args:
        tfms: List of transforms
        
    Returns:
        Target size as list, or None if not found
    """
    if tfms is None:
        return None
    for tfm in tfms:
        if hasattr(tfm, 'pad_or_crop') and hasattr(tfm.pad_or_crop, 'target_shape'):
            return list(tfm.pad_or_crop.target_shape)
    return None


def _extract_standard_config(learn) -> dict:
    """Extract config from standard MedDataBlock workflow.
    
    Args:
        learn: fastai Learner instance
        
    Returns:
        Dictionary with extracted configuration
    """
    from fastMONAI.vision_core import MedBase
    dls = learn.dls

    # Get preprocessing from MedBase class attributes
    apply_reorder = MedBase.apply_reorder
    target_spacing = MedBase.target_spacing

    # Extract item_tfms from DataLoaders pipeline
    item_tfms = []
    if hasattr(dls, 'after_item') and dls.after_item:
        item_tfms = list(dls.after_item.fs)

    # Extract size from PadOrCrop transform
    size = _extract_size_from_transforms(item_tfms)

    return {
        'apply_reorder': apply_reorder,
        'target_spacing': target_spacing,
        'size': size,
        'item_tfms': item_tfms,
        'batch_size': dls.bs,
        'patch_config': None,
    }


def _extract_patch_config(learn) -> dict:
    """Extract config from MedPatchDataLoaders workflow.
    
    Args:
        learn: fastai Learner instance
        
    Returns:
        Dictionary with extracted configuration including patch-specific params
    """
    dls = learn.dls
    patch_config = getattr(dls, '_patch_config', None) or getattr(dls, 'patch_config', None)

    config = {
        'apply_reorder': getattr(dls, '_apply_reorder', patch_config.apply_reorder if patch_config else False),
        'target_spacing': getattr(dls, '_target_spacing', patch_config.target_spacing if patch_config else None),
        'size': patch_config.patch_size if patch_config else None,
        'item_tfms': getattr(dls, '_pre_patch_tfms', []) or [],
        'batch_size': dls.bs,
    }

    # Add patch-specific params for logging
    if patch_config:
        config['patch_config'] = {
            'patch_size': patch_config.patch_size,
            'patch_overlap': patch_config.patch_overlap,
            'samples_per_volume': patch_config.samples_per_volume,
            'sampler_type': patch_config.sampler_type,
            'label_probabilities': str(patch_config.label_probabilities) if patch_config.label_probabilities else None,
            'queue_length': patch_config.queue_length,
            'aggregation_mode': patch_config.aggregation_mode,
            'padding_mode': patch_config.padding_mode,
            'keep_largest_component': patch_config.keep_largest_component,
        }
    else:
        config['patch_config'] = None

    return config


def _extract_loss_name(learn) -> str:
    """Extract loss function name from Learner.
    
    Args:
        learn: fastai Learner instance
        
    Returns:
        Name of the loss function
    """
    loss_func = learn.loss_func
    # Handle CustomLoss wrapper
    if hasattr(loss_func, 'loss_func'):
        inner = loss_func.loss_func
        return inner._get_name() if hasattr(inner, '_get_name') else inner.__class__.__name__
    return loss_func._get_name() if hasattr(loss_func, '_get_name') else loss_func.__class__.__name__


def _extract_model_name(learn) -> str:
    """Extract model architecture name from Learner.
    
    Args:
        learn: fastai Learner instance
        
    Returns:
        Name of the model architecture
    """
    model = learn.model
    return model._get_name() if hasattr(model, '_get_name') else model.__class__.__name__

# %% ../nbs/07_utils.ipynb 10
class ModelTrackingCallback(Callback):
    """
    A FastAI callback for comprehensive MLflow experiment tracking.
    
    This callback automatically logs hyperparameters, metrics, model artifacts,
    and configuration to MLflow during training. If a SaveModelCallback is present,
    the best model checkpoint will also be logged as an artifact.
    
    Supports auto-managed runs when created via `create_mlflow_callback()`.
    """
    
    def __init__(
        self, 
        model_name: str, 
        loss_function: str, 
        item_tfms: list[Any],
        size: list[int], 
        target_spacing: list[float], 
        apply_reorder: bool,
        experiment_name: str = None,
        run_name: str = None,
        auto_start: bool = False,
        patch_config: dict = None,
        extra_params: dict = None,
        extra_tags: dict = None
    ):
        """
        Initialize the MLflow tracking callback.
        
        Args:
            model_name: Name of the model architecture for registration
            loss_function: Name of the loss function being used
            item_tfms: List of item transforms
            size: Model input dimensions
            target_spacing: Resampling dimensions
            apply_reorder: Whether reordering augmentation is applied
            experiment_name: MLflow experiment name (used with auto_start)
            run_name: MLflow run name (auto-generated if None)
            auto_start: If True, auto-starts/stops MLflow run
            patch_config: Patch configuration dict for logging (from MedPatchDataLoaders)
            extra_params: Additional parameters to log
            extra_tags: MLflow tags to set on the run
        """
        self.model_name = model_name
        self.loss_function = loss_function
        self.item_tfms = item_tfms
        self.size = size
        self.target_spacing = target_spacing
        self.apply_reorder = apply_reorder
        
        # New auto-management fields
        self.experiment_name = experiment_name
        self.run_name = run_name
        self.auto_start = auto_start
        self.patch_config = patch_config
        self.extra_params = extra_params or {}
        self.extra_tags = extra_tags or {}
        self._auto_started = False
        
        self.config = self._build_config()
        
    def extract_all_params(self, tfm):
        """
        Extract all parameters from a transform object for detailed logging.
        
        Args:
            tfm: Transform object to extract parameters from
            
        Returns:
            dict: Dictionary with 'name' and 'params' keys containing transform details
        """
        class_name = tfm.__class__.__name__
        params = {}
        
        for key, value in tfm.__dict__.items():
            if not key.startswith('_') and key != '__signature__':
                if hasattr(value, '__dict__') and hasattr(value, 'target_shape'):
                    params['target_shape'] = value.target_shape
                elif hasattr(value, '__dict__') and not key.startswith('_'):
                    nested_params = {k: v for k, v in value.__dict__.items() 
                                   if not k.startswith('_') and isinstance(v, (int, float, str, bool, tuple, list))}
                    params.update(nested_params)
                elif isinstance(value, (int, float, str, bool, tuple, list)):
                    params[key] = value
        
        return {
            'name': class_name,
            'params': params
        }
        
    def _build_config(self) -> dict[str, Any]:
        """Build configuration dictionary from initialization parameters."""
        # Extract detailed transform information
        transform_details = [self.extract_all_params(tfm) for tfm in self.item_tfms]
        
        return {
            "model_name": self.model_name,
            "loss_function": self.loss_function,
            "transform_details": transform_details,
            "size": self.size,
            "target_spacing": self.target_spacing,
            "apply_reorder": self.apply_reorder,
        }
    
    def _extract_training_params(self) -> dict[str, Any]:
        """Extract training hyperparameters from the learner."""
        params = {}
        
        params["epochs"] = self.learn.n_epoch
        params["learning_rate"] = float(self.learn.lr)
        params["optimizer"] = self.learn.opt_func.__name__
        params["batch_size"] = self.learn.dls.bs
        
        params["loss_function"] = self.config["loss_function"]
        params["size"] = self.config["size"]
        params["target_spacing"] = self.config["target_spacing"]
        params["apply_reorder"] = self.config["apply_reorder"]
        
        params["transformations"] = json.dumps(
            self.config["transform_details"], 
            indent=2, 
            separators=(',', ': ')
        )
        
        # Add patch-specific params if present
        if self.patch_config:
            for key, value in self.patch_config.items():
                if value is not None:
                    params[f"patch_{key}"] = value if isinstance(value, (int, float, str, bool)) else str(value)
        
        return params
    
    def _extract_epoch_metrics(self) -> dict[str, float]:
        """Extract metrics from the current epoch."""
        recorder = self.learn.recorder
        
        # Get custom metric names and values (skip 'epoch' and 'time')
        metric_names = recorder.metric_names[2:]
        raw_metric_values = recorder.log[2:]
        
        metrics = {}
        
        # Process each metric, handling both scalars and tensors
        for name, val in zip(metric_names, raw_metric_values):
            if val is None:
                continue  # Skip None values during inference
            if isinstance(val, torch.Tensor):
                if val.numel() == 1:
                    # Single value tensor (like binary dice score)
                    metrics[name] = float(val)
                else:
                    # Multi-element tensor (like multiclass dice scores)
                    val_list = val.tolist() if hasattr(val, 'tolist') else list(val)
                    # Log individual class scores
                    for i, class_score in enumerate(val_list):
                        metrics[f"{name}_class_{i+1}"] = float(class_score)
                    # Log mean across classes
                    metrics[f"{name}_mean"] = float(torch.mean(val))
            else:
                metrics[name] = float(val)
        
        # Handle loss values
        if len(recorder.log) >= 2:
            if recorder.log[1] is not None:
                metrics['train_loss'] = float(recorder.log[1])
            if len(recorder.log) >= 3 and recorder.log[2] is not None:
                metrics['valid_loss'] = float(recorder.log[2])
                
        return metrics
    
    def _save_model_artifacts(self, temp_dir: Path) -> None:
        """Save model weights, learner, and configuration as artifacts."""
        weights_path = temp_dir / "weights"
        self.learn.save(str(weights_path))
        
        weights_file = f"{weights_path}.pth"
        if os.path.exists(weights_file):
            mlflow.log_artifact(weights_file, "model")
        
        # Auto-detect SaveModelCallback and log best model
        from fastai.callback.tracker import SaveModelCallback
        best_model_cb = None
        for cb in self.learn.cbs:
            if isinstance(cb, SaveModelCallback):
                best_model_path = self.learn.path/self.learn.model_dir/f'{cb.fname}.pth'
                if best_model_path.exists():
                    mlflow.log_artifact(str(best_model_path), "model")
                    print(f"Logged best model artifact: {cb.fname}.pth")
                    best_model_cb = cb
                break

        # Load best model weights before export if SaveModelCallback was used
        if best_model_cb is not None:
            self.learn.load(best_model_cb.fname)
            print(f"Loaded best model weights ({best_model_cb.fname}) for learner export")

        # Remove MLflow callbacks before exporting learner for inference
        # This prevents the callback from being triggered during inference
        original_cbs = self.learn.cbs.copy()  # Save original callbacks
        
        # Remove ModelTrackingCallback instances from learner using proper collection type
        filtered_cbs = L([cb for cb in self.learn.cbs if not isinstance(cb, ModelTrackingCallback)])
        self.learn.cbs = filtered_cbs
        
        # Export clean learner without MLflow callbacks
        learner_path = temp_dir / "learner.pkl"
        self.learn.export(str(learner_path))
        mlflow.log_artifact(str(learner_path), "model")
        
        # Restore original callbacks for current session
        self.learn.cbs = original_cbs
        
        config_path = temp_dir / "inference_settings.pkl"
        store_variables(config_path, self.size, self.apply_reorder, self.target_spacing)
        mlflow.log_artifact(str(config_path), "config")
    
    def _register_pytorch_model(self) -> None:
        """Register the PyTorch model with MLflow."""
        mlflow.pytorch.log_model(
            pytorch_model=self.learn.model,
            registered_model_name=self.model_name
        )
    
    def before_fit(self) -> None:
        """Log hyperparameters before training starts. Auto-start run if configured."""
        # Auto-start run if requested
        if self.auto_start:
            if self.experiment_name:
                mlflow.set_experiment(self.experiment_name)
            self.run_name = self.run_name or f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            mlflow.start_run(run_name=self.run_name)
            self._auto_started = True
        
        # Set tags
        if self.extra_tags:
            mlflow.set_tags(self.extra_tags)
        
        # Log params
        params = self._extract_training_params()
        params.update(self.extra_params)
        mlflow.log_params(params)
    
    def after_epoch(self) -> None:
        """Log metrics after each epoch."""
        metrics = self._extract_epoch_metrics()
        if metrics:
            mlflow.log_metrics(metrics, step=self.learn.epoch)
    
    def after_fit(self) -> None:
        """Log model artifacts after training completion."""
        print("\nTraining finished. Logging model artifacts to MLflow...")
        
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            
            self._save_model_artifacts(temp_path)
            
            self._register_pytorch_model()
        
        run_id = mlflow.active_run().info.run_id
        print(f"MLflow run completed. Run ID: {run_id}")
        
        # End run if auto-started
        if self._auto_started:
            mlflow.end_run()
            self._auto_started = False

# %% ../nbs/07_utils.ipynb 11
def create_mlflow_callback(
    learn,
    experiment_name: str = None,
    run_name: str = None,
    auto_start: bool = True,
    model_name: str = None,
    extra_params: dict = None,
    extra_tags: dict = None,
) -> ModelTrackingCallback:
    """Create MLflow tracking callback with auto-extracted configuration.
    
    This factory function automatically extracts configuration from the Learner,
    eliminating the need to manually specify parameters like size, transforms,
    loss function, etc.
    
    Auto-extracts from Learner:
    - Preprocessing: apply_reorder, target_spacing, size/patch_size
    - Transforms: item_tfms or pre_patch_tfms
    - Training: loss_func, model architecture
    
    Args:
        learn: fastai Learner instance
        experiment_name: MLflow experiment name. If None, uses model name.
        run_name: MLflow run name. If None, auto-generates with timestamp.
        auto_start: If True, auto-starts/stops MLflow run in before_fit/after_fit.
        model_name: Override auto-extracted model name for registration.
        extra_params: Additional parameters to log (e.g., {'dropout': 0.5}).
        extra_tags: MLflow tags to set on the run.
    
    Returns:
        ModelTrackingCallback ready to use with learn.fit()
    
    Example:
        >>> # Instead of this (6 manual params):
        >>> # mlflow_callback = ModelTrackingCallback(
        >>> #     model_name=f"{task}_{model._get_name()}",
        >>> #     loss_function=loss_func.loss_func._get_name(),
        >>> #     item_tfms=item_tfms,
        >>> #     size=size,
        >>> #     target_spacing=target_spacing,
        >>> #     apply_reorder=True,
        >>> # )
        >>> # with mlflow.start_run(run_name="training"):
        >>> #     learn.fit_one_cycle(30, lr, cbs=[mlflow_callback])
        >>> 
        >>> # Do this (zero manual params):
        >>> callback = create_mlflow_callback(learn, experiment_name="Task02_Heart")
        >>> learn.fit_one_cycle(30, lr, cbs=[callback, save_best])
    """
    # Detect workflow and extract config
    if _detect_patch_workflow(learn.dls):
        config = _extract_patch_config(learn)
    else:
        config = _extract_standard_config(learn)
    
    # Extract model/loss info
    _model_name = model_name or _extract_model_name(learn)
    _loss_name = _extract_loss_name(learn)
    
    # Set experiment name
    if experiment_name is None:
        experiment_name = _model_name
    
    # Validate size was extracted
    if config['size'] is None:
        raise ValueError(
            "Could not auto-extract 'size'. Either:\n"
            "1. Add PadOrCrop to item_tfms, or\n"
            "2. Use MedPatchDataLoaders with PatchConfig, or\n"
            "3. Use ModelTrackingCallback directly with manual params"
        )
    
    return ModelTrackingCallback(
        model_name=_model_name,
        loss_function=_loss_name,
        item_tfms=config['item_tfms'],
        size=config['size'],
        target_spacing=config['target_spacing'],
        apply_reorder=config['apply_reorder'],
        experiment_name=experiment_name,
        run_name=run_name,
        auto_start=auto_start,
        patch_config=config['patch_config'],
        extra_params=extra_params,
        extra_tags=extra_tags,
    )

# %% ../nbs/07_utils.ipynb 13
import subprocess
import threading
import time
import socket
import os
from IPython.display import display, HTML, clear_output
from IPython.core.magic import register_line_magic
from IPython import get_ipython
import requests
import shutil

class MLflowUIManager:
    def __init__(self):
        self.process = None
        self.thread = None
        self.port = 5001
        self.host = '0.0.0.0'
        self.backend_store_uri = './mlruns'
        
    def is_port_available(self, port):
        """Check if a port is available."""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(('localhost', port))
                return True
            except OSError:
                return False
    
    def is_mlflow_running(self):
        """Check if MLflow UI is actually responding."""
        try:
            response = requests.get(f'http://localhost:{self.port}', timeout=2)
            return response.status_code == 200
        except (requests.RequestException, ConnectionError, OSError):
            return False
    
    def find_available_port(self, start_port=5001):
        """Find an available port starting from start_port."""
        for port in range(start_port, start_port + 10):
            if self.is_port_available(port):
                return port
        return None
    
    def check_mlflow_installed(self):
        """Check if MLflow is installed."""
        return shutil.which('mlflow') is not None
    
    def start_ui(self, auto_open=True, quiet=False):
        """Start MLflow UI with better error handling and user feedback."""
        
        # Check if MLflow is installed
        if not self.check_mlflow_installed():
            if not quiet:
                display(HTML('<div style="color: #d32f2f; font-weight: bold; font-size: 14px;">❌ MLflow not installed. Run: pip install mlflow</div>'))
            return False
        
        # Find available port
        available_port = self.find_available_port(self.port)
        if available_port is None:
            if not quiet:
                display(HTML('<div style="color: #d32f2f; font-weight: bold; font-size: 14px;">❌ No available ports found (5001-5010)</div>'))
            return False
        
        self.port = available_port
        
        # Start MLflow UI in a separate thread
        def run_mlflow():
            try:
                self.process = subprocess.Popen([
                    'mlflow', 'ui', 
                    '--host', self.host,
                    '--port', str(self.port),
                    '--backend-store-uri', self.backend_store_uri
                ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                self.process.wait()
            except Exception as e:
                if not quiet:
                    display(HTML(f'<div style="color: #d32f2f; font-weight: bold; font-size: 14px;">❌ Error: {str(e)}</div>'))
        
        self.thread = threading.Thread(target=run_mlflow, daemon=True)
        self.thread.start()
        
        # Wait and check if server started successfully
        max_wait = 10
        for i in range(max_wait):
            time.sleep(1)
            if self.is_mlflow_running():
                if quiet:
                    # Bright, visible link for quiet mode
                    display(HTML(f'''
                        <a href="http://localhost:{self.port}" target="_blank" 
                           style="color: #1976d2; font-weight: bold; font-size: 16px; text-decoration: underline;">
                           🔗 MLflow UI (Port {self.port})
                        </a>
                    '''))
                else:
                    # Success message with high contrast colors
                    display(HTML(f'''
                        <div style="background-color: #c8e6c9; border: 2px solid #388e3c; padding: 15px; border-radius: 8px; margin: 10px 0;">
                            <div style="color: #1b5e20; font-weight: bold; font-size: 16px; margin-bottom: 10px;">
                                ✅ MLflow UI is running successfully!
                            </div>
                            <a href="http://localhost:{self.port}" target="_blank" 
                               style="background-color: #1976d2; color: white; padding: 12px 24px; text-decoration: none; border-radius: 6px; font-weight: bold; font-size: 14px; display: inline-block; margin: 5px 0;">
                                🔗 Open MLflow UI
                            </a>
                            <div style="margin-top: 10px;">
                                <div style="color: #424242; font-size: 13px;">URL: http://localhost:{self.port}</div>
                            </div>
                        </div>
                    '''))
                return True
        
        # If we get here, server didn't start properly
        if not quiet:
            display(HTML('<div style="color: #d32f2f; font-weight: bold; font-size: 14px;">❌ Failed to start MLflow UI</div>'))
        return False
    
    def stop(self):
        """Stop the MLflow UI server."""
        if self.process:
            self.process.terminate()
            self.process = None
            display(HTML('''
                <div style="background-color: #ffecb3; border: 2px solid #f57c00; padding: 10px; border-radius: 6px;">
                    <span style="color: #e65100; font-weight: bold; font-size: 14px;">🛑 MLflow UI stopped</span>
                </div>
            '''))
        else:
            display(HTML('''
                <div style="background-color: #f0f0f0; border: 2px solid #757575; padding: 10px; border-radius: 6px;">
                    <span style="color: #424242; font-weight: bold; font-size: 14px;">ℹ️ MLflow UI is not currently running</span>
                </div>
            '''))
    
    def status(self):
        """Check MLflow UI status."""
        if self.is_mlflow_running():
            display(HTML(f'''
                <div style="background-color: #c8e6c9; border: 2px solid #388e3c; padding: 10px; border-radius: 6px;">
                    <div style="color: #1b5e20; font-weight: bold; font-size: 14px;">✅ MLflow UI is running</div>
                    <a href="http://localhost:{self.port}" target="_blank" 
                       style="color: #1976d2; font-weight: bold; text-decoration: underline;">
                       http://localhost:{self.port}
                    </a>
                </div>
            '''))
        else:
            display(HTML('''
                <div style="background-color: #ffcdd2; border: 2px solid #d32f2f; padding: 10px; border-radius: 6px;">
                    <div style="color: #b71c1c; font-weight: bold; font-size: 14px;">❌ MLflow UI is not running</div>
                    <div style="color: #424242; font-size: 13px; margin-top: 5px;">
                        Run <code style="background-color: #f5f5f5; padding: 2px 4px; border-radius: 3px;">mlflow_ui.start_ui()</code> to start it.
                    </div>
                </div>
            '''))
