"""
OsmosisPipeline — the main user-facing API for ModelOsmosis.

Usage::

    from model_osmosis import OsmosisPipeline

    pipeline = OsmosisPipeline()

    # Full transfer in one call
    target = pipeline.absorb(
        source=trained_model_A,
        target=model_B,
        sample_data=X_samples,
        auto_grow=True,
    )

    # Or step by step
    knowledge = pipeline.extract(source=trained_model_A, sample_data=X_samples)
    grown = pipeline.grow(target=model_B, knowledge=knowledge)
    final = pipeline.inject(target=grown, knowledge=knowledge)
"""

from __future__ import annotations

import copy
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from model_osmosis.core.config import OsmosisConfig
from model_osmosis.core.exceptions import (
    AbsorptionError,
    ExtractionError,
    GrowthError,
    ObfuscationError,
    OsmosisError,
)
from model_osmosis.growers.neural_grower import NeuralGrower
from model_osmosis.knowledge.representation import UniversalKnowledge
from model_osmosis.utils.model_utils import (
    auto_select_absorber,
    auto_select_extractor,
    detect_framework,
)
from model_osmosis.utils.determinism import set_global_seed


class OsmosisPipeline:
    """Main pipeline for model-to-model knowledge transfer.

    Orchestrates extraction, capacity estimation, growth, and absorption
    into a single cohesive workflow.

    Args:
        config: Pipeline configuration. If None, uses default settings.
    """

    def __init__(self, config: Optional[OsmosisConfig] = None) -> None:
        self.config = config or OsmosisConfig()
        self._grower = NeuralGrower()

    def _auto_obfuscate(self, result: Any, sample_data: Optional[np.ndarray] = None) -> Any:
        """Automatically obfuscate if auto_obfuscate is enabled in config."""
        if not self.config.auto_obfuscate:
            return result

        if self.config.verbose:
            print()
            print("[Pipeline] Auto-obfuscating output...")

        return self.obfuscate(result, sample_data=sample_data)

    def _count_params(self, model: Any) -> int:
        """Count model parameters for expansion planning/reporting."""
        try:
            import torch.nn as nn

            if isinstance(model, nn.Module):
                return int(sum(p.numel() for p in model.parameters()))
        except ImportError:
            pass

        if isinstance(model, dict):
            total = 0
            for value in model.values():
                arr = np.asarray(value)
                total += int(arr.size)
            return int(total)
        return 0

    def _resolve_target_params(self, target_params: Optional[Union[int, str]], baseline_params: int) -> int:
        """Resolve target parameter envelope from int or ratio spec (`x1.5`)."""
        if target_params is None:
            return int(max(baseline_params, int(baseline_params * 1.2)))
        if isinstance(target_params, (int, np.integer)):
            return int(target_params)

        raw = str(target_params).strip().lower()
        if raw.startswith("x"):
            multiplier = float(raw[1:])
            if multiplier <= 1.0:
                raise ValueError("target_params multiplier must be > 1.0, e.g. x1.5")
            return int(max(baseline_params + 1, baseline_params * multiplier))

        return int(raw)

    def _evaluate_quality_metric(self, model: Any, sample_data: Any, task: str = "auto", seed: int = 42) -> float:
        """Task-aware deterministic quality proxy for expansion gate decisions."""
        from model_osmosis.services.expand_backends import evaluate_metric

        return float(evaluate_metric(model, sample_data, str(task or "auto"), seed=int(seed)))

    def expand(
        self,
        model: Any,
        *,
        sample_data: Any = None,
        task: str = "auto",
        target_params: Optional[Union[int, str]] = None,
        budget_minutes: int = 5,
        budget_ram_gb: float = 1.5,
        seed: int = 42,
        optimizer_mode: str = "auto",
        fallback_mode: str = "off",
        strict_mode: bool = True,
        route: str = "",
        runtime_source: str = "",
    ) -> Tuple[Any, Dict[str, Any]]:
        """Grow and adapt a model with deterministic, quality-gated expansion."""
        from model_osmosis.optimizer.expand_engine import (
            finalize_expand_profile,
            select_expand_profile,
        )
        from model_osmosis.services.expand_backends import (
            quality_gate,
            select_expand_backend,
        )
        from model_osmosis.utils.determinism import set_global_seed

        set_global_seed(int(seed))
        baseline_params = self._count_params(model)
        target_param_count = self._resolve_target_params(target_params, baseline_params)
        normalized_task = str(task or "auto")
        normalized_fallback_mode = str(fallback_mode or "off").strip().lower()
        if normalized_fallback_mode not in {"off", "distill"}:
            raise ValueError("Unsupported fallback_mode. Expected one of: off, distill")

        context = {
            "task": normalized_task,
            "baseline_params": int(baseline_params),
            "target_params": int(target_param_count),
            "budget_minutes": int(max(1, int(budget_minutes))),
            "budget_ram_gb": float(max(0.25, float(budget_ram_gb))),
            "model_type": str(type(model).__name__),
            "route": str(route or "unknown"),
            "runtime_source": str(runtime_source or "unknown"),
            "strict_mode": bool(strict_mode),
            "fallback_mode": normalized_fallback_mode,
        }
        selection = select_expand_profile(
            context=context,
            optimizer_mode=optimizer_mode,
            seed=int(seed),
            record_selection=False,
        )

        baseline_metric = self._evaluate_quality_metric(
            model,
            sample_data,
            task=normalized_task,
            seed=int(seed),
        )
        expanded_model = model
        backend = select_expand_backend(normalized_task)
        backend_id = str(getattr(backend, "backend_id", "unknown"))
        can_expand, can_expand_reason = backend.can_expand(model, sample_data, normalized_task)
        structural_applied = False
        adaptation_applied = False
        degraded = False
        expand_error = ""
        fallback_reason = ""

        if can_expand:
            try:
                execution = backend.execute(
                    model=model,
                    sample_data=sample_data,
                    target_params=int(target_param_count),
                    budget_minutes=int(max(1, int(budget_minutes))),
                    seed=int(seed),
                )
                expanded_model = execution.expanded_model
                structural_applied = bool(execution.structural_applied)
                adaptation_applied = bool(execution.adaptation_applied)
                if str(execution.structural_reason or "") and execution.structural_reason != "ok":
                    fallback_reason = str(execution.structural_reason)
                if str(execution.adaptation_reason or "") and execution.adaptation_reason not in {"", "ok"}:
                    if fallback_reason:
                        fallback_reason = f"{fallback_reason};{execution.adaptation_reason}"
                    else:
                        fallback_reason = str(execution.adaptation_reason)
            except Exception as exc:
                expand_error = str(exc)
                fallback_reason = f"expand_backend_failed:{type(exc).__name__}"
                expanded_model = model
        else:
            fallback_reason = str(can_expand_reason or "strict_expand_not_supported")
            expand_error = str(can_expand_reason or "strict_expand_not_supported")

        if (
            not structural_applied
            and normalized_fallback_mode == "distill"
            and bool(strict_mode)
        ):
            try:
                degraded = True
                distill_model = copy.deepcopy(model)
                if sample_data is not None:
                    knowledge = self.extract(model, sample_data=sample_data)
                    local_config = self.config
                    original_epochs = int(local_config.distillation_epochs)
                    bounded_epochs = max(1, min(original_epochs, int(max(1, int(budget_minutes)))))
                    local_config.distillation_epochs = bounded_epochs
                    try:
                        distill_model = self.inject(distill_model, knowledge)
                    finally:
                        local_config.distillation_epochs = original_epochs
                    adaptation_applied = True
                expanded_model = distill_model
                expand_error = ""
                if not fallback_reason:
                    fallback_reason = "distill_fallback_applied"
            except Exception as exc:
                expand_error = str(exc)
                fallback_reason = f"distill_fallback_failed:{type(exc).__name__}"
                expanded_model = model

        post_metric = self._evaluate_quality_metric(
            expanded_model,
            sample_data,
            task=normalized_task,
            seed=int(seed),
        )
        gate = quality_gate(normalized_task, baseline_metric, post_metric)
        quality_delta = float(post_metric - baseline_metric)
        if degraded:
            accepted = bool(gate.get("passed", False) and not expand_error)
        else:
            accepted = bool(structural_applied and gate.get("passed", False) and not expand_error)
        hard_failure = bool(strict_mode and not accepted)

        finalized = finalize_expand_profile(
            selection,
            quality_ok=accepted,
            fallback_reason=fallback_reason,
            quality_score=post_metric,
        )

        report = {
            "mode": "expand",
            "task": normalized_task,
            "target_params": int(target_param_count),
            "budget_minutes": int(max(1, int(budget_minutes))),
            "budget_ram_gb": float(max(0.25, float(budget_ram_gb))),
            "baseline_metric": float(baseline_metric),
            "post_expand_metric": float(post_metric),
            "quality_delta": float(quality_delta),
            "quality_gate": dict(gate),
            "accepted": bool(accepted),
            "baseline_params": int(baseline_params),
            "post_expand_params": int(self._count_params(expanded_model)),
            "structural_applied": bool(structural_applied),
            "adaptation_applied": bool(adaptation_applied),
            "strict_mode": bool(strict_mode),
            "fallback_mode": normalized_fallback_mode,
            "degraded": bool(degraded),
            "status": "degraded" if degraded and accepted else ("accepted" if accepted else "failed"),
            "hard_failure": bool(hard_failure),
            "route": str(route or "unknown"),
            "runtime_source": str(runtime_source or "unknown"),
            "backend_id": backend_id,
            "can_expand": bool(can_expand),
            "can_expand_reason": str(can_expand_reason or ""),
            "error": expand_error,
            "self_check_status": "pending",
            "self_check_level": "smoke",
            "self_check_checks": [],
            "self_check_skips": [],
        }
        report.update(finalized.to_dict())
        return expanded_model, report

    def absorb(
        self,
        source: Any,
        target: Any,
        sample_data: Optional[np.ndarray] = None,
        auto_grow: bool = True,
    ) -> Any:
        """Full pipeline: extract from source → grow target → inject.

        This is the primary one-call API for knowledge transfer.

        Args:
            source: The trained source model to extract knowledge from.
            target: The target model to inject knowledge into.
            sample_data: Optional sample input data for functional extraction.
                If None, synthetic data is generated.
            auto_grow: Whether to automatically grow the target model to
                accommodate the source's knowledge.

        Returns:
            The target model with absorbed knowledge (grown if auto_grow=True).
        """
        if self.config.verbose:
            src_fw = detect_framework(source)
            tgt_fw = detect_framework(target)
            print(f"[Pipeline] Source: {type(source).__name__} ({src_fw})")
            print(f"[Pipeline] Target: {type(target).__name__} ({tgt_fw})")
            print(f"[Pipeline] Auto-grow: {auto_grow}")
            print()

        # Step 1: Extract
        knowledge = self.extract(source, sample_data)

        if self.config.verbose:
            print()
            print(knowledge.summary())
            print()

        # Step 2: Grow (if requested)
        if auto_grow:
            target = self.grow(target, knowledge)
            if self.config.verbose:
                print()

        # Step 3: Inject
        result = self.inject(target, knowledge)

        if self.config.verbose:
            print()
            print("[Pipeline] ✓ Knowledge transfer complete!")

        return self._auto_obfuscate(result, sample_data)

    def extract(
        self,
        source: Any,
        sample_data: Optional[np.ndarray] = None,
    ) -> UniversalKnowledge:
        """Extract knowledge from a source model.

        Args:
            source: The trained source model.
            sample_data: Optional sample input data.

        Returns:
            UniversalKnowledge containing all extracted knowledge.

        Raises:
            ExtractionError: If extraction fails.
        """
        try:
            extractor = auto_select_extractor(source)
            if self.config.verbose:
                print(
                    f"[Pipeline] Extracting with {type(extractor).__name__}..."
                )
            knowledge = extractor.extract(
                source,
                sample_data=sample_data,
                num_samples=self.config.functional_samples,
            )
            return knowledge
        except Exception as e:
            raise ExtractionError(
                f"Failed to extract knowledge from "
                f"{type(source).__name__}: {e}"
            ) from e

    def grow(
        self,
        target: Any,
        knowledge: UniversalKnowledge,
    ) -> Any:
        """Grow the target model to accommodate incoming knowledge.

        Args:
            target: The target model to expand.
            knowledge: The knowledge it needs to absorb.

        Returns:
            The grown model.

        Raises:
            GrowthError: If growth fails.
        """
        try:
            if self._grower.can_handle(target):
                if self.config.verbose:
                    print("[Pipeline] Growing target architecture...")
                return self._grower.grow(target, knowledge, self.config)
            else:
                if self.config.verbose:
                    print(
                        "[Pipeline] Target model type does not support "
                        "dynamic growth — skipping"
                    )
                return target
        except Exception as e:
            raise GrowthError(
                f"Failed to grow {type(target).__name__}: {e}"
            ) from e

    def inject(
        self,
        target: Any,
        knowledge: UniversalKnowledge,
    ) -> Any:
        """Inject extracted knowledge into the target model.

        Args:
            target: The target model (optionally pre-grown).
            knowledge: The knowledge to inject.

        Returns:
            The target model with absorbed knowledge.

        Raises:
            AbsorptionError: If injection fails.
        """
        try:
            absorber = auto_select_absorber(target)
            if self.config.verbose:
                print(
                    f"[Pipeline] Absorbing with {type(absorber).__name__}..."
                )
            return absorber.absorb(target, knowledge, self.config)
        except Exception as e:
            raise AbsorptionError(
                f"Failed to absorb into {type(target).__name__}: {e}"
            ) from e

    def merge(
        self,
        models: List[Any],
        strategy: str = "auto",
        eval_data: Optional[np.ndarray] = None,
        model_weights: Optional[List[float]] = None,
    ) -> Any:
        """Merge multiple models into one — RAM-efficient with auto-optimization.

        Convenience method that wraps MultiModelMerger.

        Args:
            models: List of models to merge.
            strategy: Merge strategy ('weighted_average', 'slerp', 'ties',
                'dare', 'breadcrumbs', 'consensus', 'auto').
            eval_data: Optional evaluation data for auto-optimization.
            model_weights: Optional per-model importance weights.

        Returns:
            The merged model.
        """
        from model_osmosis.merging.merger import MultiModelMerger

        set_global_seed(self.config.seed)
        merger = MultiModelMerger(
            strategy=strategy,
            model_weights=model_weights,
            config=self.config,
            verbose=self.config.verbose,
        )
        result = merger.merge(models=models, eval_data=eval_data)
        return self._auto_obfuscate(result, eval_data)

    def compare(
        self,
        model_a: Any,
        model_b: Any,
        sample_data: np.ndarray,
    ) -> dict:
        """Compare two models' outputs and return similarity metrics.

        Runs both models on the same sample data and computes:
        - cosine_similarity: Mean cosine similarity across samples
        - mse: Mean squared error between outputs
        - correlation: Mean Pearson correlation per output dimension
        - max_error: Maximum absolute difference

        Args:
            model_a: First model.
            model_b: Second model.
            sample_data: numpy array of input samples.

        Returns:
            Dict with similarity metrics.
        """
        import torch

        x = torch.from_numpy(sample_data).float()

        with torch.no_grad():
            if hasattr(model_a, '__call__'):
                out_a = model_a(x)
                if isinstance(out_a, torch.Tensor):
                    out_a = out_a.numpy()
            else:
                out_a = model_a.predict(sample_data)

            if hasattr(model_b, '__call__'):
                out_b = model_b(x)
                if isinstance(out_b, torch.Tensor):
                    out_b = out_b.numpy()
            else:
                out_b = model_b.predict(sample_data)

        out_a = np.asarray(out_a, dtype=np.float64)
        out_b = np.asarray(out_b, dtype=np.float64)

        # Cosine similarity
        flat_a = out_a.flatten()
        flat_b = out_b.flatten()
        norm_a = np.linalg.norm(flat_a)
        norm_b = np.linalg.norm(flat_b)
        cosine = (
            float(np.dot(flat_a, flat_b) / (norm_a * norm_b))
            if norm_a > 0 and norm_b > 0 else 0.0
        )

        # MSE
        mse = float(np.mean((out_a - out_b) ** 2))

        # Max error
        max_error = float(np.max(np.abs(out_a - out_b)))

        # Per-dimension correlation
        if out_a.ndim >= 2:
            correlations = []
            for i in range(out_a.shape[-1]):
                col_a = out_a[..., i].flatten()
                col_b = out_b[..., i].flatten()
                if np.std(col_a) > 0 and np.std(col_b) > 0:
                    correlations.append(float(np.corrcoef(col_a, col_b)[0, 1]))
            correlation = float(np.mean(correlations)) if correlations else 0.0
        else:
            if np.std(out_a) > 0 and np.std(out_b) > 0:
                correlation = float(np.corrcoef(out_a.flatten(), out_b.flatten())[0, 1])
            else:
                correlation = 0.0

        result = {
            "cosine_similarity": cosine,
            "mse": mse,
            "correlation": correlation,
            "max_error": max_error,
            "num_samples": len(sample_data),
        }

        if self.config.verbose:
            print(f"[Pipeline] Model comparison:")
            print(f"  Cosine similarity: {cosine:.6f}")
            print(f"  MSE: {mse:.2e}")
            print(f"  Correlation: {correlation:.6f}")
            print(f"  Max error: {max_error:.2e}")

        return result

    def benchmark(
        self,
        models: List[Any],
        eval_data: np.ndarray,
        strategies: Optional[List[str]] = None,
    ) -> List[dict]:
        """Benchmark all merge strategies and return ranked results.

        Runs each strategy with auto-optimized weights and evaluates
        the merged model's output consistency.

        Args:
            models: List of models to merge.
            eval_data: Evaluation data to score merged models.
            strategies: List of strategy names to try. Defaults to all
                available strategies.

        Returns:
            List of dicts sorted by score (best first), each containing:
            - strategy: Strategy name
            - score: Consistency score (higher is better)
            - weights: Optimal weights found
        """
        from model_osmosis.merging.strategies import STRATEGY_REGISTRY
        from model_osmosis.merging.evaluator import MergeEvaluator

        if strategies is None:
            strategies = list(STRATEGY_REGISTRY.keys())

        evaluator = MergeEvaluator()
        results = []

        if self.config.verbose:
            print(f"[Pipeline] Benchmarking {len(strategies)} strategies "
                  f"on {len(models)} models...")

        for name in strategies:
            try:
                from model_osmosis.merging.strategies import get_strategy
                strat = get_strategy(name)
                weights, score = evaluator.find_best_weights(
                    models=models,
                    eval_data=eval_data,
                    strategy=strat,
                    n_trials=8,
                    verbose=False,
                )
                results.append({
                    "strategy": name,
                    "score": score,
                    "weights": weights,
                })
                if self.config.verbose:
                    print(f"  {name}: score={score:.6f} "
                          f"weights={[f'{w:.3f}' for w in weights]}")
            except Exception as e:
                if self.config.verbose:
                    print(f"  {name}: FAILED ({e})")

        # Sort by score descending
        results.sort(key=lambda x: x["score"], reverse=True)

        if self.config.verbose and results:
            print(f"  → Best: {results[0]['strategy']} "
                  f"(score: {results[0]['score']:.6f})")

        return results

    def stack(
        self,
        models: List[Any],
        sample_data: Optional[np.ndarray] = None,
        expert_names: Optional[List[str]] = None,
        top_k: int = 2,
        train_gate: bool = True,
    ) -> Any:
        """Stack multiple models into one mega-model (MoE).

        Keeps ALL parameters from every model. A learned gating network
        routes inputs to the best expert(s). Nothing is lost.

        Args:
            models: List of PyTorch models to stack.
            sample_data: Optional data for auto-detection and gate training.
            expert_names: Optional names for each expert.
            top_k: Number of active experts per input.
            train_gate: Whether to train the gating network.

        Returns:
            A StackedModel containing all experts.
        """
        from model_osmosis.stacking.stacker import ModelStacker

        set_global_seed(self.config.seed)
        stacker = ModelStacker(verbose=self.config.verbose)
        result = stacker.stack(
            models=models,
            sample_data=sample_data,
            expert_names=expert_names,
            top_k=top_k,
            train_gate=train_gate,
            seed=self.config.seed,
        )
        return self._auto_obfuscate(result, sample_data)

    def cross_merge(
        self,
        models: List[Any],
        sample_data: Union[np.ndarray, List[str], None] = None,
        output_dim: Optional[int] = None,
        weights: Optional[List[float]] = None,
        student_hidden: int = 512,
        student_model_id: Optional[str] = None,
        epochs: int = 100,
        seed: Optional[int] = None,
        **kwargs,
    ) -> Any:
        """Merge models of ANY architecture into one.

        Unlike merge() which needs same architecture, this works with
        any combination of models by merging their functional knowledge
        in a shared representation space.

        For HuggingFace LLMs, pass model IDs as strings and sample_data
        as a list of text strings. The merger will use knowledge distillation
        to combine the models into a single student LLM.

        Args:
            models: List of models. Can be:
                - nn.Module objects (uses classic CrossMerger)
                - HuggingFace model ID strings (uses LLMCrossMerger)
            sample_data: Sample data for knowledge extraction.
                - np.ndarray for classic CrossMerger
                - List[str] for LLM cross-merger
            output_dim: Output dimension (classic merger only).
            weights: Per-model importance weights (default: equal).
            student_hidden: Hidden size (classic merger only).
            student_model_id: HF model ID for the student base
                (LLM merger only). Defaults to first teacher.
            epochs: Training epochs for the merged model.
            **kwargs: Additional config passed to LLMCrossMergeConfig.

        Returns:
            For classic: a single nn.Module.
            For LLM: a (model, tokenizer) tuple.
        """
        effective_seed = self.config.seed if seed is None else int(seed)
        set_global_seed(effective_seed)
        # Detect if this is an LLM cross-merge
        is_llm = (
            (isinstance(models[0], str))
            or _is_hf_causal_lm(models[0])
        )

        if is_llm:
            from model_osmosis.merging.llm_cross_merger import (
                LLMCrossMerger,
                LLMCrossMergeConfig,
            )

            # Build config
            config_kwargs = {
                "student_model_id": student_model_id,
                "distillation_epochs": min(epochs, 10),
                "verbose": self.config.verbose,
            }
            config_kwargs.update(kwargs)
            config = LLMCrossMergeConfig(**config_kwargs)

            merger = LLMCrossMerger(config)

            # Resolve teacher IDs
            teacher_ids = []
            for m in models:
                if isinstance(m, str):
                    teacher_ids.append(m)
                elif hasattr(m, "config") and hasattr(m.config, "_name_or_path"):
                    teacher_ids.append(m.config._name_or_path)
                else:
                    raise ValueError(
                        f"Cannot determine model ID for {type(m)}. "
                        "Pass HuggingFace model ID strings instead."
                    )

            # Resolve sample texts
            if sample_data is None:
                sample_texts = _default_code_samples()
            elif isinstance(sample_data, list):
                sample_texts = sample_data
            else:
                raise ValueError(
                    "For LLM cross-merge, sample_data must be a list of strings."
                )

            model, tokenizer = merger.merge(
                teacher_ids=teacher_ids,
                sample_texts=sample_texts,
                weights=weights,
            )
            return model, tokenizer
        else:
            from model_osmosis.merging.cross_merger import CrossMerger

            classic_common_dim = kwargs.get("common_dim")
            merger = CrossMerger(
                common_dim=int(classic_common_dim) if classic_common_dim is not None else 256,
                verbose=self.config.verbose,
            )
            result = merger.merge(
                models=models,
                sample_data=sample_data,
                output_dim=output_dim,
                weights=weights,
                student_hidden=student_hidden,
                student_layers=int(kwargs.get("student_layers", 3) or 3),
                epochs=epochs,
                lr=float(kwargs.get("lr", 1e-3) or 1e-3),
                batch_size=int(kwargs.get("batch_size", 32) or 32),
                projection_epochs=kwargs.get("projection_epochs"),
                common_dim=kwargs.get("common_dim"),
                seed=effective_seed,
            )
            return self._auto_obfuscate(result, sample_data)

    def bridge_merge(
        self,
        models: List[Any],
        paired_data: List[np.ndarray],
        modality_names: Optional[List[str]] = None,
        output_dim: Optional[int] = None,
        epochs: int = 100,
    ) -> Any:
        """Bridge-merge models of different modalities using paired data.

        Each model operates on a different data type. Paired data provides
        examples where sample[i] describes the same concept in all modalities.

        Args:
            models: List of models (one per modality).
            paired_data: Paired numpy arrays (same index = same concept).
            modality_names: Names for each modality.
            output_dim: Output dimension.
            epochs: Training epochs.

        Returns:
            A BridgedModel with aligned modality encoders.
        """
        from model_osmosis.merging.modality_bridge import ModalityBridge

        bridge = ModalityBridge(verbose=self.config.verbose)
        result = bridge.merge(
            models=models,
            paired_data=paired_data,
            modality_names=modality_names,
            output_dim=output_dim,
            epochs=epochs,
        )
        return self._auto_obfuscate(result)

    def unsupervised_merge(
        self,
        models: List[Any],
        unpaired_data: List[np.ndarray],
        modality_names: Optional[List[str]] = None,
        output_dim: Optional[int] = None,
        epochs: int = 100,
    ) -> Any:
        """Bridge-merge models of different modalities WITHOUT paired data.

        Uses adversarial alignment — no dictionary needed. Each model
        gets its own independent data.

        Args:
            models: List of models (one per modality).
            unpaired_data: Independent numpy arrays (NOT paired).
            modality_names: Names for each modality.
            output_dim: Output dimension.
            epochs: Training epochs.

        Returns:
            A BridgedModel with adversarially aligned encoders.
        """
        from model_osmosis.merging.modality_bridge import ModalityBridge

        bridge = ModalityBridge(verbose=self.config.verbose)
        result = bridge.unsupervised_merge(
            models=models,
            unpaired_data=unpaired_data,
            modality_names=modality_names,
            output_dim=output_dim,
            epochs=epochs,
        )
        return self._auto_obfuscate(result)

    def obfuscate(
        self,
        model: Any,
        sample_data: Optional[np.ndarray] = None,
        obfuscation_config: Optional[Any] = None,
        target_framework: Optional[str] = None,
    ) -> Any:
        """Obfuscate model weights to prevent source tracing.

        Applies function-preserving transforms (permutation, rotation, scaling)
        that completely change raw weight values while keeping the model's
        input-output behavior identical.

        Works with ANY input format and can output to ANY target framework:
        - OsmoModel, PyTorch, TensorFlow, JAX, ONNX, HuggingFace, raw dicts

        Args:
            model: The model to obfuscate (any supported format).
            sample_data: Optional input data for verification.
            obfuscation_config: ObfuscationConfig instance. If None,
                uses self.config.obfuscation or default settings.
            target_framework: Output format. One of 'pytorch', 'tensorflow',
                'jax', 'onnx', 'huggingface', 'osmo', or None.
                If None, returns in the same format as the input.

        Returns:
            The obfuscated model in the requested framework format.

        Raises:
            ObfuscationError: If obfuscation fails.
        """
        from model_osmosis.obfuscation import (
            WeightObfuscator,
            ObfuscationConfig,
        )

        config = obfuscation_config or self.config.obfuscation or ObfuscationConfig()

        try:
            obfuscator = WeightObfuscator(config)

            if self.config.verbose:
                print("[Pipeline] Obfuscating model weights...")

            from model_osmosis.format.osmo_format import OsmoModel as OsmoModelClass
            from model_osmosis.format.converters import to_osmo, from_osmo

            # --- Step 1: Convert input to OsmoModel and obfuscate ---

            source_fw = None

            if isinstance(model, OsmoModelClass):
                obfuscated_osmo = obfuscator.obfuscate(model, sample_data)
                source_fw = "osmo"
            elif isinstance(model, dict):
                # Raw dict — obfuscate directly
                obfuscated_dict = obfuscator.obfuscate_state_dict(model)
                if target_framework and target_framework != "dict":
                    # Wrap in OsmoModel for conversion
                    from model_osmosis.format.osmo_format import (
                        OsmoManifest, ArchitectureSpec, LayerSpec, LayerType,
                    )
                    obfuscated_osmo = OsmoModelClass(
                        manifest=OsmoManifest(),
                        architecture=ArchitectureSpec(layers=[], hidden_size=0, num_layers=0),
                        parameters=obfuscated_dict,
                    )
                else:
                    return obfuscated_dict
                source_fw = "dict"
            else:
                # Any other model type — detect framework, convert to osmo
                source_fw = detect_framework(model)

                if source_fw == "pytorch" and not target_framework:
                    # Fast path: obfuscate PyTorch model in-place
                    return obfuscator.obfuscate_torch_model(model, sample_data)

                if self.config.verbose:
                    print(
                        f"[Pipeline] Converting {type(model).__name__} "
                        f"to OsmoFormat for obfuscation..."
                    )

                osmo = to_osmo(model)
                obfuscated_osmo = obfuscator.obfuscate(osmo, sample_data)

            # --- Step 2: Convert to target framework ---

            output_fw = target_framework or source_fw

            if output_fw in (None, "osmo"):
                return obfuscated_osmo

            if self.config.verbose:
                print(f"[Pipeline] Converting obfuscated model to {output_fw}...")

            return from_osmo(obfuscated_osmo, target_framework=output_fw)

        except ObfuscationError:
            raise
        except Exception as e:
            raise ObfuscationError(
                f"Failed to obfuscate {type(model).__name__}: {e}"
            ) from e


def _is_hf_causal_lm(model: Any) -> bool:
    """Check if a model is a HuggingFace causal language model."""
    try:
        from transformers import PreTrainedModel
        return isinstance(model, PreTrainedModel) and hasattr(model, "generate")
    except ImportError:
        return False


def _default_code_samples() -> List[str]:
    """Default code sample texts for LLM cross-merging."""
    return [
        "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)",
        "class BinaryTree:\n    def __init__(self, value):\n        self.value = value\n        self.left = None\n        self.right = None",
        "import asyncio\n\nasync def fetch_data(url):\n    async with aiohttp.ClientSession() as session:\n        async with session.get(url) as response:\n            return await response.json()",
        "def quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quicksort(left) + middle + quicksort(right)",
        "def binary_search(arr, target):\n    low, high = 0, len(arr) - 1\n    while low <= high:\n        mid = (low + high) // 2\n        if arr[mid] == target:\n            return mid\n        elif arr[mid] < target:\n            low = mid + 1\n        else:\n            high = mid - 1\n    return -1",
        "class LinkedList:\n    def __init__(self):\n        self.head = None\n    def append(self, data):\n        new_node = Node(data)\n        if not self.head:\n            self.head = new_node\n            return\n        current = self.head\n        while current.next:\n            current = current.next\n        current.next = new_node",
        "from flask import Flask, jsonify, request\n\napp = Flask(__name__)\n\n@app.route('/api/users', methods=['GET'])\ndef get_users():\n    return jsonify({'users': users})",
        "import torch\nimport torch.nn as nn\n\nclass Transformer(nn.Module):\n    def __init__(self, d_model, nhead, num_layers):\n        super().__init__()\n        self.encoder = nn.TransformerEncoder(\n            nn.TransformerEncoderLayer(d_model, nhead), num_layers\n        )",
        "def merge_sort(arr):\n    if len(arr) <= 1:\n        return arr\n    mid = len(arr) // 2\n    left = merge_sort(arr[:mid])\n    right = merge_sort(arr[mid:])\n    return merge(left, right)",
        "class Graph:\n    def __init__(self):\n        self.adjacency_list = {}\n    def add_edge(self, u, v):\n        if u not in self.adjacency_list:\n            self.adjacency_list[u] = []\n        self.adjacency_list[u].append(v)",
        "async function fetchData(url) {\n    const response = await fetch(url);\n    const data = await response.json();\n    return data;\n}",
        "SELECT u.name, COUNT(o.id) as order_count\nFROM users u\nLEFT JOIN orders o ON u.id = o.user_id\nGROUP BY u.name\nHAVING COUNT(o.id) > 5\nORDER BY order_count DESC;",
        "fn main() {\n    let mut vec = Vec::new();\n    vec.push(1);\n    vec.push(2);\n    for item in &vec {\n        println!(\"{}\", item);\n    }\n}",
        "package main\n\nimport \"fmt\"\n\nfunc main() {\n    ch := make(chan int)\n    go func() {\n        ch <- 42\n    }()\n    fmt.Println(<-ch)\n}",
        "The capital of France is Paris.",
        "Machine learning is a subset of artificial intelligence.",
        "To implement a REST API, you need to define endpoints for CRUD operations.",
        "The time complexity of binary search is O(log n).",
        "Docker containers provide isolated environments for running applications.",
        "Git is a distributed version control system used for tracking changes in code.",
    ]

