# 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 #ac941446-cf7e-4f5c-ace0-a36fb078022f
import pickle
import torch
from pathlib import Path
import mlflow
import mlflow.pytorch
try:
    _mlruns_db = Path(__file__).resolve().parent.parent / 'mlruns.db'
except NameError:
    _mlruns_db = Path('mlruns.db').resolve()
mlflow.set_tracking_uri(f"sqlite:///{_mlruns_db}")
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 #5f954a64-c9df-4ee8-b70a-fce7f857eaf2
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 #c2db5512-171c-4dfd-a26e-561b773a6069
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 #w143kixa46j
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 #03xsquvh55db
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 #14e1513a-25d6-497a-bec3-6adc008452d9
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 #czquspt567w
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,
            'preprocessed': patch_config.preprocessed,
        }
    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 #030876d2
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 checkpoint callback
    (SaveModelCallback, EMACheckpoint, or any TrackerCallback with fname) 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,
        dataset_version: str = None,
        log_split: bool = True
    ):
        """
        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
            dataset_version: Dataset version hash string for tracking
            log_split: If True, auto-logs train/val split CSV when dls has split_df
        """
        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.dataset_version = dataset_version
        self.log_split = log_split
        self._auto_started = False
        self.run_id = None
        
        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."""
        import shutil

        # Save final epoch weights (without optimizer state to reduce file size)
        weights_path = temp_dir / "final_weights"
        self.learn.save(str(weights_path), with_opt=False)
        weights_file = f"{weights_path}.pth"
        if os.path.exists(weights_file):
            mlflow.log_artifact(weights_file, "model")

        # Auto-detect checkpoint callback (SaveModelCallback, EMACheckpoint, etc.)
        from fastai.callback.tracker import TrackerCallback
        best_model_cb = None
        for cb in self.learn.cbs:
            if isinstance(cb, TrackerCallback) and hasattr(cb, 'fname'):
                best_model_path = self.learn.path / self.learn.model_dir / f'{cb.fname}.pth'
                if best_model_path.exists():
                    best_weights_dest = temp_dir / "best_weights.pth"
                    shutil.copy2(str(best_model_path), str(best_weights_dest))
                    mlflow.log_artifact(str(best_weights_dest), "model")
                    print(f"Logged best model weights: best_weights.pth")
                    best_model_cb = cb
                break

        # Remove ModelTrackingCallback before exporting learners
        original_cbs = self.learn.cbs.copy()
        filtered_cbs = L([cb for cb in self.learn.cbs
                          if not isinstance(cb, ModelTrackingCallback)])
        self.learn.cbs = filtered_cbs

        # Export final epoch learner (before loading best weights)
        final_learner_path = temp_dir / "final_learner.pkl"
        self.learn.export(str(final_learner_path))
        mlflow.log_artifact(str(final_learner_path), "model")
        print("Logged final epoch learner: final_learner.pkl")

        # If best model exists, load best weights and export best learner
        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 best learner export")

            best_learner_path = temp_dir / "best_learner.pkl"
            self.learn.export(str(best_learner_path))
            mlflow.log_artifact(str(best_learner_path), "model")
            print("Logged best epoch learner: best_learner.pkl")

        # Restore original callbacks
        self.learn.cbs = original_cbs

        # Save inference config
        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 _log_split_df(self) -> None:
        """Log train/validation split as CSV artifact if available."""
        dls = self.learn.dls
        if not hasattr(dls, 'split_df'):
            return
        try:
            split_df = dls.split_df
            cols = []
            img_col = getattr(dls, '_img_col', None)
            mask_col = getattr(dls, '_mask_col', None)
            if img_col and img_col in split_df.columns:
                cols.append(img_col)
            if mask_col and mask_col in split_df.columns:
                cols.append(mask_col)
            cols.append('is_valid')
            split_df = split_df[cols]

            mlflow.log_text(split_df.to_csv(index=False), "data/train_val_split.csv")
            print(f"Logged train/val split ({len(split_df)} rows) to MLflow artifacts")
        except Exception as e:
            print(f"Warning: Could not log train/val split: {e}")

    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
        
        # Log dataset version tag if provided
        if self.dataset_version is not None:
            mlflow.set_tag("dataset_version", self.dataset_version)
        
        # 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)

        # Log train/val split
        if self.log_split:
            self._log_split_df()
    
    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()
        
        self.run_id = mlflow.active_run().info.run_id
        print(f"MLflow run completed. Run ID: {self.run_id}")
        
        # End run if auto-started
        if self._auto_started:
            mlflow.end_run()
            self._auto_started = False

    def log_metrics(self, metrics: dict) -> None:
        """Log additional metrics to the completed MLflow run.

        Reopens the run to log metrics, then closes it. Useful for
        post-training evaluation (e.g., validation set patch inference results).
        NaN values are silently skipped (MLflow rejects them).

        Args:
            metrics: Dictionary of metric name -> value pairs to log.
        """
        if self.run_id is None:
            raise RuntimeError("No MLflow run found. Train the model first.")

        import math
        valid = {k: v for k, v in metrics.items()
                 if not (isinstance(v, float) and math.isnan(v))}

        if valid:
            with mlflow.start_run(run_id=self.run_id):
                mlflow.log_metrics(valid)
            print(f"Logged {len(valid)} metric(s) to MLflow run {self.run_id}")

    def log_metrics_table(self, df, metrics=None, key='validation_metrics', display=False):
        """Log a summary table of metrics as an image to the MLflow Images tab.

        Creates a styled matplotlib table with Metric, Mean, and Std columns
        from the given DataFrame and logs it via mlflow.log_image() so it
        appears in the MLflow UI Images tab.

        Args:
            df: DataFrame containing per-sample metric values.
            metrics: List of column names in df to summarize. If None,
                all numeric columns are used automatically.
            key: Image key for MLflow Images tab.
            display: If True, display the table inline in the notebook.
        """
        import math
        import io
        import matplotlib.pyplot as plt
        from PIL import Image

        if self.run_id is None:
            raise RuntimeError("No MLflow run found. Train the model first.")

        if metrics is None:
            metrics = df.select_dtypes(include='number').columns.tolist()

        summary_data = []
        for metric in metrics:
            mean = df[metric].mean()
            std = df[metric].std()
            std_str = f'{std:.4f}' if not math.isnan(std) else 'N/A'
            summary_data.append([metric, f'{mean:.4f}', std_str])

        fig, ax = plt.subplots(figsize=(5, len(metrics) * 0.8 + 0.5))
        ax.axis('off')

        table = ax.table(
            cellText=summary_data,
            colLabels=['Metric', 'Mean', 'Std'],
            cellLoc='center',
            loc='center'
        )
        table.auto_set_font_size(False)
        table.set_fontsize(14)
        table.scale(1.5, 2.2)

        for col in range(3):
            table[0, col].set_facecolor('#4472C4')
            table[0, col].set_text_props(color='white', fontweight='bold')

        fig.tight_layout()

        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight', dpi=250)
        buf.seek(0)
        pil_image = Image.open(buf)

        with mlflow.start_run(run_id=self.run_id):
            mlflow.log_image(pil_image, key=key)

        buf.close()
        if display: plt.show()
        plt.close(fig)
        print(f"Logged metrics table to MLflow run {self.run_id}")

    def log_dataframe(self, df, artifact_name='validation_results.csv', artifact_path='results'):
        """Log a DataFrame as a CSV artifact to the MLflow run.

        Args:
            df: DataFrame to save.
            artifact_name: Filename for the CSV.
            artifact_path: Artifact subdirectory in MLflow.
        """
        if self.run_id is None:
            raise RuntimeError("No MLflow run found. Train the model first.")

        with tempfile.NamedTemporaryFile(suffix='.csv', delete=False, mode='w') as f:
            df.to_csv(f, index=False)
            temp_path = f.name

        with mlflow.start_run(run_id=self.run_id):
            mlflow.log_artifact(temp_path, artifact_path)

        os.remove(temp_path)
        print(f"Logged DataFrame ({len(df)} rows) to MLflow run {self.run_id}")

# %% ../nbs/07_utils.ipynb #38avhd96apq
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,
    dataset_version: str = None,
    log_split: bool = True,
) -> 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.
        dataset_version: Dataset version hash string for tracking.
        log_split: If True, auto-logs train/val split CSV when dls has split_df.
    
    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,
        dataset_version=dataset_version,
        log_split=log_split,
    )

# %% ../nbs/07_utils.ipynb #a3d2d585
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 = f'sqlite:///{_mlruns_db}'
        self._owns_process = False
        self._reusing_external = False
        
    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_running_mlflow(self):
        """Find an existing MLflow UI in our port range. Returns port or None."""
        for port in range(self.port, self.port + 10):
            if not self.is_port_available(port):
                try:
                    response = requests.get(f'http://localhost:{port}', timeout=2)
                    if response.status_code == 200 and 'mlflow' in response.text.lower():
                        return port
                except (requests.RequestException, ConnectionError, OSError):
                    continue
        return None
    
    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
        
        # Check for existing MLflow UI before launching a new one
        existing_port = self._find_running_mlflow()
        if existing_port is not None:
            self.port = existing_port
            self._owns_process = False
            self._reusing_external = True
            if not quiet:
                display(HTML(f'''
                    <div style="background-color: #e3f2fd; border: 2px solid #1976d2; padding: 15px; border-radius: 8px; margin: 10px 0;">
                        <div style="color: #0d47a1; font-weight: bold; font-size: 16px; margin-bottom: 10px;">
                            Reusing existing MLflow UI on port {self.port}
                        </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>
                '''))
            else:
                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}, reused)
                    </a>
                '''))
            return True
        
        # 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():
                self._owns_process = True
                self._reusing_external = False
                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 and self._owns_process:
            self.process.terminate()
            self.process = None
            self._owns_process = False
            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>
            '''))
        elif self._reusing_external:
            self._reusing_external = False
            display(HTML('''
                <div style="background-color: #e3f2fd; border: 2px solid #1976d2; padding: 10px; border-radius: 6px;">
                    <span style="color: #0d47a1; font-weight: bold; font-size: 14px;">MLflow UI was started externally — not stopping it</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>
            '''))
