"""Pipeline runner - Main entry point for pipeline execution.

This module provides the PipelineRunner class, which serves as the main interface
for executing ML pipelines on spectroscopic datasets. It delegates execution to
PipelineOrchestrator and provides prediction/explanation capabilities via
Predictor and Explainer classes.
"""
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, Optional, Union

import numpy as np

from nirs4all.core.logging import configure_logging
from nirs4all.data.config import DatasetConfigs
from nirs4all.data.dataset import SpectroDataset
from nirs4all.data.predictions import Predictions
from nirs4all.pipeline.config.pipeline_config import PipelineConfigs
from nirs4all.pipeline.execution.orchestrator import PipelineOrchestrator
from nirs4all.pipeline.explainer import Explainer
from nirs4all.pipeline.predictor import Predictor
from nirs4all.pipeline.retrainer import ExtractedPipeline, Retrainer, StepMode

if TYPE_CHECKING:
    from nirs4all.pipeline.storage.library import PipelineLibrary

def _get_default_workspace_path() -> Path:
    """Get the default workspace path.

    Uses the workspace module's get_active_workspace() which checks:
    1. Explicitly set workspace via set_active_workspace()
    2. NIRS4ALL_WORKSPACE environment variable
    3. ./workspace in current working directory

    Returns:
        Default workspace path.
    """
    from nirs4all.workspace import get_active_workspace
    return get_active_workspace()

def init_global_random_state(seed: int | None = None):
    """Initialize global random state for reproducibility.

    Sets random seeds for numpy, Python's random module, TensorFlow, PyTorch, and sklearn
    to ensure reproducible results across runs.

    Args:
        seed: Random seed value. If None, uses default seed of 42 for TensorFlow and PyTorch.
    """
    import os
    import random

    if seed is not None:
        np.random.seed(seed)
        random.seed(seed)
        os.environ['PYTHONHASHSEED'] = str(seed)

    try:
        import tensorflow as tf
        tf.random.set_seed(seed if seed is not None else 42)
    except ImportError:
        pass

    try:
        import torch
        torch.manual_seed(seed if seed is not None else 42)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed if seed is not None else 42)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
    except ImportError:
        pass

    try:
        from sklearn.utils import check_random_state
        _ = check_random_state(seed)
    except ImportError:
        pass

class PipelineRunner:
    """Main pipeline execution interface.

    Orchestrates pipeline execution on datasets, providing a simplified interface for
    training, prediction, and explanation workflows. Delegates actual execution to
    PipelineOrchestrator, Predictor, and Explainer.

    Attributes:
        workspace_path (Path): Root workspace directory
        verbose (int): Verbosity level (0=quiet, 1=info, 2=debug, 3=trace)
        mode (str): Execution mode ('train', 'predict', 'explain')
        save_artifacts (bool): Whether to save binary artifacts (models, transformers)
        save_charts (bool): Whether to save charts and visual outputs
        enable_tab_reports (bool): Whether to generate tabular reports
        continue_on_error (bool): Whether to continue on step failures
        show_spinner (bool): Whether to show progress spinners
        keep_datasets (bool): Whether to keep raw/preprocessed data snapshots
        max_preprocessed_snapshots_per_dataset (int): Maximum number of
            preprocessed snapshots retained per dataset when keep_datasets=True.
        plots_visible (bool): Whether to display plots interactively
        orchestrator (PipelineOrchestrator): Underlying orchestrator for execution
        predictor (Predictor): Handler for prediction mode
        explainer (Explainer): Handler for explanation mode
        raw_data (Dict[str, np.ndarray]): Raw dataset snapshots (if keep_datasets=True)
        pp_data (Dict[str, Dict[str, np.ndarray]]): Preprocessed data snapshots

    Example:
        >>> # Training workflow
        >>> runner = PipelineRunner(workspace_path="./workspace", verbose=1)
        >>> pipeline = [{"preprocessing": StandardScaler()}, {"model": SVC()}]
        >>> X, y = load_data()
        >>> predictions, dataset_preds = runner.run(pipeline, (X, y))

        >>> # Prediction workflow
        >>> runner = PipelineRunner(mode="predict")
        >>> y_pred, preds = runner.predict(best_model, X_new)

        >>> # Explanation workflow
        >>> runner = PipelineRunner(mode="explain")
        >>> shap_results, out_dir = runner.explain(best_model, X_test)
    """

    def __init__(
        self,
        workspace_path: str | Path | None = None,
        verbose: int = 0,
        mode: str = "train",
        save_artifacts: bool = True,
        save_charts: bool = True,
        enable_tab_reports: bool = True,
        continue_on_error: bool = False,
        show_spinner: bool = True,
        keep_datasets: bool = False,
        max_preprocessed_snapshots_per_dataset: int = 3,
        plots_visible: bool = False,
        random_state: int | None = None,
        n_jobs: int = 1,
        report_naming: Literal["nirs", "ml", "auto"] = "nirs",
        # Logging configuration
        log_file: bool = True,
        log_format: str = "pretty",
        use_unicode: bool | None = None,
        use_colors: bool | None = None,
        show_progress_bar: bool = True,
        json_output: bool = False,
    ):
        """Initialize pipeline runner.

        Args:
            workspace_path: Workspace root directory. Defaults to './workspace'
            verbose: Verbosity level (0=quiet, 1=info, 2=debug, 3=trace)
            mode: Execution mode ('train', 'predict', 'explain')
            save_artifacts: Whether to save binary artifacts (models, transformers)
            save_charts: Whether to save charts and visual outputs
            enable_tab_reports: Whether to generate tabular reports
            continue_on_error: Whether to continue on step failures
            show_spinner: Whether to show progress spinners
            keep_datasets: Whether to keep data snapshots (raw/preprocessed)
            max_preprocessed_snapshots_per_dataset: Max number of
                preprocessed snapshots retained per dataset
            plots_visible: Whether to display plots interactively
            random_state: Random seed for reproducibility
            n_jobs: Number of parallel workers for pipeline variants (sweeps).
                - 1: Sequential execution (default)
                - -1: Use all available CPU cores
                - N > 1: Use N parallel workers
                Note: This parallelizes different pipeline variants generated by
                _cartesian_, _or_, etc. Branch parallelization (branch keyword)
                is controlled separately via branch-level parallel config.
            report_naming: Naming convention for metrics in reports ("nirs", "ml", or "auto").
                Defaults to "nirs" (chemometrics terminology).
            log_file: Whether to write logs to workspace/logs/ directory
            log_format: Output format: "pretty" (default), "minimal", or "json"
            use_unicode: Use Unicode symbols (auto-detected if None). Set to False
                for HPC/cluster environments without Unicode support.
            use_colors: Use ANSI colors (auto-detected if None). Set to False to
                disable colored output.
            show_progress_bar: Whether to show TTY-aware progress bars
            json_output: Also write JSON Lines log file for machine parsing
        """
        if random_state is not None:
            init_global_random_state(random_state)
        self.random_state = random_state

        if workspace_path is None:
            workspace_path = _get_default_workspace_path()
        self.workspace_path = Path(workspace_path)

        self.verbose = verbose
        self.mode = mode
        self.save_artifacts = save_artifacts
        self.save_charts = save_charts
        self.enable_tab_reports = enable_tab_reports
        self.continue_on_error = continue_on_error
        self.show_spinner = show_spinner
        self.keep_datasets = keep_datasets
        self.max_preprocessed_snapshots_per_dataset = max(0, int(max_preprocessed_snapshots_per_dataset))
        self.plots_visible = plots_visible
        self.n_jobs = n_jobs
        self.report_naming = report_naming

        # Store logging configuration
        self.log_file = log_file
        self.log_format = log_format
        self.use_unicode = use_unicode
        self.use_colors = use_colors
        self.show_progress_bar = show_progress_bar
        self.json_output = json_output

        # Configure logging system
        log_dir = self.workspace_path / "logs" if log_file else None
        configure_logging(
            verbose=verbose,
            log_file=log_file,
            log_dir=log_dir,
            log_format=log_format,
            use_unicode=use_unicode,
            use_colors=use_colors,
            show_progress=True,
            show_progress_bar=show_progress_bar,
            json_output=json_output,
        )

        # Create orchestrator
        self.orchestrator = PipelineOrchestrator(
            workspace_path=self.workspace_path,
            verbose=verbose,
            mode=mode,
            save_artifacts=save_artifacts,
            save_charts=save_charts,
            enable_tab_reports=enable_tab_reports,
            continue_on_error=continue_on_error,
            show_spinner=show_spinner,
            keep_datasets=keep_datasets,
            max_preprocessed_snapshots_per_dataset=self.max_preprocessed_snapshots_per_dataset,
            n_jobs=n_jobs,
            plots_visible=plots_visible,
            random_state=random_state,
            report_naming=report_naming,
        )

        # Create predictor and explainer
        self.predictor = Predictor(self)
        self.explainer = Explainer(self)
        self.retrainer = Retrainer(self)

        # Expose orchestrator state for convenience
        self.raw_data = self.orchestrator.raw_data
        self.pp_data = self.orchestrator.pp_data
        self._figure_refs = self.orchestrator._figure_refs

        # Model capture support for explainer
        self._capture_model: bool = False

        # Last run aggregate settings (for visualization integration)
        self._last_aggregate_column: str | None = None
        self._last_aggregate_method: str | None = None
        self._last_aggregate_exclude_outliers: bool = False

        # Execution state (synchronized from executor during execution)
        self.step_number: int = 0
        self.substep_number: int = -1
        self.operation_count: int = 0

        # Runtime components (set during execution)
        self.artifact_loader: Any = None  # ArtifactLoader for predict/explain modes
        self.pipeline_uid: str | None = None  # Current pipeline UID
        self.target_model: dict | None = None  # Target model for predict/explain modes
        self.last_execution_trace: Any = None  # ExecutionTrace from last run

        # Cache configuration (set via nirs4all.run(cache=...))
        self.cache_config: Any = None  # CacheConfig instance

        # Library for template management
        self._library: PipelineLibrary | None = None

    def run(
        self,
        pipeline: PipelineConfigs | list[Any] | dict | str,
        dataset: DatasetConfigs | SpectroDataset | list[SpectroDataset] | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
        pipeline_name: str = "",
        dataset_name: str = "dataset",
        max_generation_count: int = 10000,
        refit: bool | dict[str, Any] | list[dict[str, Any]] | None = True,
        store_run_id: str | None = None,
        manage_store_run: bool = True,
    ) -> tuple[Predictions, dict[str, Any]]:
        """Execute pipeline on dataset(s).

        Main entry point for training workflows. Executes one or more pipeline
        configurations on one or more datasets, tracking predictions and artifacts.

        Args:
            pipeline: Pipeline definition (PipelineConfigs, list of steps, dict, or path)
            dataset: Dataset definition (see DatasetConfigs for supported formats)
            pipeline_name: Optional pipeline name for identification
            dataset_name: Name for array-based datasets
            max_generation_count: Max pipeline combinations to generate
            refit: Refit configuration.
                - ``True``: Refit top 1 by RMSECV (default).
                - ``False`` or ``None``: Disable refit.
                - ``dict``: Single criterion, e.g. ``{"top_k": 3}``.
                - ``list[dict]``: Multiple criteria for union selection.
            store_run_id: Optional existing store run ID to join.
            manage_store_run: Whether to manage the store run lifecycle.

        Returns:
            Tuple of (run_predictions, datasets_predictions)
        """
        # Propagate cache_config to orchestrator for RuntimeContext injection
        self.orchestrator.cache_config = getattr(self, 'cache_config', None)

        run_predictions, dataset_predictions = self.orchestrator.execute(
            pipeline=pipeline,
            dataset=dataset,
            pipeline_name=pipeline_name,
            dataset_name=dataset_name,
            max_generation_count=max_generation_count,
            artifact_loader=self.artifact_loader,
            target_model=self.target_model,
            explainer=self.explainer,
            refit=refit,
            store_run_id=store_run_id,
            manage_store_run=manage_store_run,
        )

        # Sync state
        if self.keep_datasets:
            self.raw_data = self.orchestrator.raw_data
            self.pp_data = self.orchestrator.pp_data
        self._figure_refs = self.orchestrator._figure_refs

        # Sync pipeline UID from last executed pipeline
        if self.orchestrator.last_pipeline_uid is not None:
            self.pipeline_uid = self.orchestrator.last_pipeline_uid

        # Sync execution state from last executor (via orchestrator)
        if hasattr(self.orchestrator, 'last_executor') and self.orchestrator.last_executor:
            self.step_number = self.orchestrator.last_executor.step_number
            self.substep_number = self.orchestrator.last_executor.substep_number
            self.operation_count = self.orchestrator.last_executor.operation_count

        # Sync aggregate column from last dataset for visualization integration
        if hasattr(self.orchestrator, 'last_aggregate_column'):
            self._last_aggregate_column = self.orchestrator.last_aggregate_column
        if hasattr(self.orchestrator, 'last_aggregate_method'):
            self._last_aggregate_method = self.orchestrator.last_aggregate_method
        if hasattr(self.orchestrator, 'last_aggregate_exclude_outliers'):
            self._last_aggregate_exclude_outliers = self.orchestrator.last_aggregate_exclude_outliers

        # Sync execution trace for post-run visualization
        if hasattr(self.orchestrator, 'last_execution_trace'):
            self.last_execution_trace = self.orchestrator.last_execution_trace

        return run_predictions, dataset_predictions

    def predict(
        self,
        prediction_obj: dict[str, Any] | str,
        dataset: DatasetConfigs | SpectroDataset | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
        dataset_name: str = "prediction_dataset",
        all_predictions: bool = False,
        verbose: int = 0
    ) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
        """Run prediction using a saved model on new dataset.

        Delegates to Predictor class for actual execution.

        Args:
            prediction_obj: Model identifier (dict with config_path or prediction ID)
            dataset: New dataset to predict on
            dataset_name: Name for the dataset
            all_predictions: If True, return all predictions; if False, return single best
            verbose: Verbosity level

        Returns:
            If all_predictions=False: (y_pred, predictions)
            If all_predictions=True: (predictions_dict, predictions)
        """
        return self.predictor.predict(
            prediction_obj=prediction_obj,
            dataset=dataset,
            dataset_name=dataset_name,
            all_predictions=all_predictions,
            verbose=verbose
        )

    def explain(
        self,
        prediction_obj: dict[str, Any] | str,
        dataset: DatasetConfigs | SpectroDataset | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
        dataset_name: str = "explain_dataset",
        shap_params: dict[str, Any] | None = None,
        verbose: int = 0,
        plots_visible: bool = True
    ) -> tuple[dict[str, Any], str]:
        """Generate SHAP explanations for a saved model.

        Delegates to Explainer class for actual execution.

        Args:
            prediction_obj: Model identifier (dict with config_path or prediction ID)
            dataset: Dataset to explain on
            dataset_name: Name for the dataset
            shap_params: SHAP configuration parameters
            verbose: Verbosity level
            plots_visible: Whether to display plots interactively

        Returns:
            Tuple of (shap_results_dict, output_directory_path)
        """
        return self.explainer.explain(
            prediction_obj=prediction_obj,
            dataset=dataset,
            dataset_name=dataset_name,
            shap_params=shap_params,
            verbose=verbose,
            plots_visible=plots_visible
        )

    @property
    def store(self) -> Any:
        """Get the WorkspaceStore from the orchestrator.

        Returns:
            WorkspaceStore instance for DuckDB-backed persistence
        """
        return self.orchestrator.store

    def close(self) -> None:
        """Close the underlying WorkspaceStore to release DuckDB resources."""
        if hasattr(self, "orchestrator") and self.orchestrator is not None:
            store = self.orchestrator.store
            if store is not None:
                store.close()

    @property
    def runs_dir(self) -> Path:
        """Get runs directory (legacy compatibility).

        Returns:
            Path to workspace root (runs are now in DuckDB, not filesystem)
        """
        return self.workspace_path

    @property
    def last_aggregate(self) -> str | None:
        """Get aggregate column from the last executed dataset.

        Returns the aggregation setting from the last dataset processed by run().
        This can be used to create a PredictionAnalyzer with matching defaults.

        Returns:
            Aggregate column name ('y' for y-based aggregation, column name for
            metadata-based aggregation, or None if no aggregation was set).

        Example:
            >>> runner = PipelineRunner()
            >>> predictions, _ = runner.run(pipeline, DatasetConfigs({"train_x": path, "aggregate": "sample_id"}))
            >>> # Create analyzer with same aggregate setting
            >>> analyzer = PredictionAnalyzer(predictions, default_aggregate=runner.last_aggregate)
        """
        return self._last_aggregate_column

    @property
    def last_aggregate_method(self) -> str | None:
        """Get aggregate method from the last executed dataset.

        Returns:
            Aggregate method ('mean', 'median', 'vote') or None for default.
        """
        return self._last_aggregate_method

    @property
    def last_aggregate_exclude_outliers(self) -> bool:
        """Get aggregate exclude_outliers setting from the last executed dataset.

        Returns:
            True if T-squared outlier exclusion was enabled, False otherwise.
        """
        return self._last_aggregate_exclude_outliers

    @property
    def library(self) -> "PipelineLibrary":
        """Get pipeline library for template management.

        Returns:
            PipelineLibrary instance for managing pipeline templates
        """
        if self._library is None:
            from nirs4all.pipeline.storage.library import PipelineLibrary
            self._library = PipelineLibrary(self.workspace_path)
        return self._library

    def next_op(self) -> int:
        """Get the next operation ID (for controller compatibility).

        Returns:
            Next operation counter value
        """
        self.operation_count += 1
        return self.operation_count

    def export(
        self,
        source: dict[str, Any] | str | Path,
        output_path: str | Path,
        format: str = "n4a",
        include_metadata: bool = True,
        compress: bool = True
    ) -> Path:
        """Export a trained pipeline to a standalone bundle.

        Creates a self-contained prediction bundle that can be used for
        deployment, sharing, or archival without requiring the original
        workspace or full nirs4all installation.

        Supported formats:
            - 'n4a': Full bundle (ZIP archive with artifacts and metadata)
            - 'n4a.py': Portable Python script with embedded artifacts

        Args:
            source: Prediction source to export. Can be:
                - prediction dict: From a previous run's Predictions object
                - folder path: Path to a pipeline directory
                - Run object: Best prediction from a Run
            output_path: Path for the output bundle file
            format: Bundle format ('n4a' or 'n4a.py')
            include_metadata: Whether to include full metadata in bundle
            compress: Whether to compress artifacts (for .n4a format)

        Returns:
            Path to the created bundle file

        Raises:
            ValueError: If format is not supported
            FileNotFoundError: If source cannot be resolved

        Example:
            >>> runner = PipelineRunner()
            >>> predictions, _ = runner.run(pipeline, dataset)
            >>> best_pred = predictions.top(n=1)[0]
            >>>
            >>> # Export to .n4a bundle
            >>> runner.export(best_pred, "exports/wheat_model.n4a")
            >>>
            >>> # Export to portable Python script
            >>> runner.export(best_pred, "exports/wheat_model.n4a.py", format='n4a.py')
            >>>
            >>> # Later, predict from bundle
            >>> y_pred, _ = runner.predict("exports/wheat_model.n4a", X_new)
        """
        from nirs4all.pipeline.bundle import BundleGenerator

        generator = BundleGenerator(
            workspace_path=self.workspace_path,
            verbose=self.verbose,
            store=self.store,
        )

        return generator.export(
            source=source,
            output_path=output_path,
            format=format,
            include_metadata=include_metadata,
            compress=compress
        )

    def export_model(
        self,
        source: dict[str, Any] | str | Path,
        output_path: str | Path,
        format: str | None = None,
        fold: int | None = None
    ) -> Path:
        """Export only the model artifact from a trained pipeline.

        Unlike `export()` which creates a full bundle with all preprocessing
        artifacts and metadata, this method exports just the model binary.
        This is useful when you want a lightweight model file that can be
        loaded directly into other pipelines or used with external tools.

        Args:
            source: Prediction source to export from. Can be:
                - prediction dict: From a previous run's Predictions object
                - folder path: Path to a pipeline directory
                - bundle path: Path to a .n4a bundle
            output_path: Path for the output model file. Extension determines
                format: .joblib, .pkl, .h5, .keras, .pt
            format: Optional explicit format ('joblib', 'pickle', 'keras_h5').
                If None, determined from output_path extension.
            fold: Optional fold index to export. If None, exports fold 0 or
                the primary model artifact.

        Returns:
            Path to the created model file

        Raises:
            ValueError: If no model artifact found
            FileNotFoundError: If source cannot be resolved

        Example:
            >>> runner = PipelineRunner()
            >>> predictions, _ = runner.run(pipeline, dataset)
            >>> best_pred = predictions.top(n=1)[0]
            >>>
            >>> # Export just the model
            >>> runner.export_model(best_pred, "exports/pls_model.joblib")
            >>>
            >>> # Later, use in new pipeline
            >>> new_pipeline = [
            ...     MinMaxScaler(),
            ...     {"model": "exports/pls_model.joblib", "name": "pretrained"}
            ... ]
        """
        from nirs4all.pipeline.resolver import PredictionResolver
        from nirs4all.pipeline.storage.artifacts.artifact_persistence import to_bytes

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Resolve the prediction source
        resolver = PredictionResolver(
            workspace_path=self.workspace_path,
            runs_dir=self.runs_dir,
            store=self.store,
        )
        resolved = resolver.resolve(source, verbose=self.verbose)

        if resolved.model_step_index is None:
            raise ValueError("No model step found in the resolved prediction")

        # Get the model artifact
        if resolved.artifact_provider is None:
            raise ValueError("No artifact provider available for this source")

        artifacts = resolved.artifact_provider.get_artifacts_for_step(
            resolved.model_step_index
        )

        if not artifacts:
            raise ValueError(
                f"No model artifacts found at step {resolved.model_step_index}"
            )

        # Select the fold
        if fold is not None:
            # Find artifact for specific fold
            model = None
            for artifact_id, artifact in artifacts:
                if f":{fold}" in str(artifact_id) or artifact_id.endswith(f"_{fold}"):
                    model = artifact
                    break
            if model is None:
                raise ValueError(f"No artifact found for fold {fold}")
        else:
            # Use first artifact (fold 0 or primary)
            _, model = artifacts[0]

        # Determine format from extension if not specified
        if format is None:
            ext = output_path.suffix.lower()
            format_map = {
                '.joblib': 'joblib',
                '.pkl': 'cloudpickle',
                '.pickle': 'cloudpickle',
                '.h5': 'keras_h5',
                '.hdf5': 'keras_h5',
                '.keras': 'tensorflow_keras',
                '.pt': 'pytorch_state_dict',
                '.pth': 'pytorch_state_dict',
            }
            format = format_map.get(ext, 'joblib')

        # Serialize and write
        data, actual_format = to_bytes(model, format)
        with open(output_path, 'wb') as f:
            f.write(data)

        if self.verbose > 0:
            print(f"Exported model to {output_path} (format: {actual_format})")

        return output_path

    def retrain(
        self,
        source: dict[str, Any] | str | Path,
        dataset: DatasetConfigs | SpectroDataset | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
        mode: str = "full",
        dataset_name: str = "retrain_dataset",
        new_model: Any | None = None,
        epochs: int | None = None,
        step_modes: list[StepMode] | None = None,
        verbose: int = 0,
        **kwargs
    ) -> tuple[Predictions, dict[str, Any]]:
        """Retrain a pipeline on new data.

        Enables retraining trained pipelines with various modes:
        - full: Train from scratch with same pipeline structure
        - transfer: Use existing preprocessing artifacts, train new model
        - finetune: Continue training existing model with new data

        Args:
            source: Prediction source to retrain from.
            dataset: New dataset to train on. Supports same formats as run()
            mode: Retrain mode ('full', 'transfer', 'finetune')
            dataset_name: Name for the dataset if array-based
            new_model: Optional new model for transfer mode (replaces original)
            epochs: Optional epochs for fine-tuning
            step_modes: Optional per-step mode overrides for fine-grained control
            verbose: Verbosity level
            **kwargs: Additional parameters

        Returns:
            Tuple of (run_predictions, datasets_predictions)

        Example:
            >>> runner = PipelineRunner()
            >>> predictions, _ = runner.run(pipeline, dataset)
            >>> best_pred = predictions.top(n=1)[0]
            >>>
            >>> # Full retrain on new data
            >>> new_preds, _ = runner.retrain(best_pred, new_data, mode='full')
        """
        return self.retrainer.retrain(
            source=source,
            dataset=dataset,
            mode=mode,
            dataset_name=dataset_name,
            new_model=new_model,
            epochs=epochs,
            step_modes=step_modes,
            verbose=verbose,
            **kwargs
        )

    def extract(
        self,
        source: dict[str, Any] | str | Path
    ) -> ExtractedPipeline:
        """Extract a trained pipeline for inspection or modification.

        Loads a trained pipeline from a prediction source and returns an
        ExtractedPipeline object that can be inspected, modified, and then
        executed with runner.run().

        Args:
            source: Prediction source to extract.

        Returns:
            ExtractedPipeline object with steps, trace, artifact_provider,
            model_step_index, and preprocessing_chain.

        Example:
            >>> runner = PipelineRunner()
            >>> predictions, _ = runner.run(pipeline, dataset)
            >>> best_pred = predictions.top(n=1)[0]
            >>>
            >>> # Extract for inspection
            >>> extracted = runner.extract(best_pred)
            >>> print(f"Steps: {len(extracted.steps)}")
        """
        return self.retrainer.extract(source, verbose=self.verbose)
