"""
ModelOsmosis — Universal Model Knowledge Transfer Framework

Extract knowledge from any trained ML model and inject it into another,
with dynamic architecture growth so the target can always accommodate more.

## OsmoFormat - Universal Model Interchange

The .osmo format enables converting, merging, and transferring models
across ANY framework (PyTorch, TensorFlow, JAX, ONNX, etc.):

    from model_osmosis.format import to_osmo, merge_osmo, save_osmo

    # Convert any model to universal format
    osmo = to_osmo(my_model)

    # Merge models from different frameworks
    merged = merge_osmo([osmo1, osmo2], strategy="slerp")

    # Save to disk
    save_osmo(merged, "merged.osmo")

Supports:
- Text LLMs (GPT, LLaMA, Mistral, etc.)
- Vision Models (ViT, ResNet, ConvNeXT, timm, torchvision)
- Audio Models (Whisper, Wav2Vec2, HuBERT)
- Diffusion Models (Stable Diffusion UNet, VAE, text encoders)
- Multimodal Models (CLIP, LLaVA, BLIP)
- LoRA/PEFT Adapters
- ONNX and MLX models
"""

__version__ = "0.8.4"

# Core pipeline
from model_osmosis.pipeline import OsmosisPipeline

# Merging
from model_osmosis.merging.merger import MultiModelMerger
from model_osmosis.merging.cross_merger import CrossMerger
from model_osmosis.merging.modality_bridge import ModalityBridge

# Stacking
from model_osmosis.stacking.stacker import ModelStacker
from model_osmosis.stacking.stacked_model import StackedModel
from model_osmosis.stacking.safe_stacker import SafeStacker

# Core HuggingFace adapter (always available)
from model_osmosis.adapters.huggingface import (
    HuggingFaceAdapter,
    HuggingFaceExpert,
    ModelModality,
)

# Configuration
from model_osmosis.core.config import OsmosisConfig, GrowthStrategyType

# Exceptions
from model_osmosis.core.exceptions import (
    OsmosisError,
    ExtractionError,
    AbsorptionError,
    GrowthError,
    IncompatibleModelError,
    ObfuscationError,
    VerificationError,
)

# Knowledge representations
from model_osmosis.knowledge.representation import (
    UniversalKnowledge,
    StructuralKnowledge,
    FunctionalKnowledge,
    MetaKnowledge,
)

# Utilities
from model_osmosis.utils.model_card import (
    ModelCardGenerator,
    MergeCardConfig,
    StackCardConfig,
)

# Obfuscation
from model_osmosis.obfuscation import (
    WeightObfuscator,
    ObfuscationConfig,
    obfuscate_osmo,
)

# OsmoFormat - Universal Model Interchange
from model_osmosis.format import (
    OsmoModel,
    OsmoManifest,
    to_osmo,
    from_osmo,
    save_osmo,
    load_osmo,
    merge_osmo,
    stack_osmo,
)

__all__ = [
    # Version
    "__version__",
    # Pipeline
    "OsmosisPipeline",
    # Merging
    "MultiModelMerger",
    "CrossMerger",
    "ModalityBridge",
    # Stacking
    "ModelStacker",
    "StackedModel",
    "SafeStacker",
    # Adapters (core)
    "HuggingFaceAdapter",
    "HuggingFaceExpert",
    "ModelModality",
    # Config
    "OsmosisConfig",
    "GrowthStrategyType",
    # Exceptions
    "OsmosisError",
    "ExtractionError",
    "AbsorptionError",
    "GrowthError",
    "IncompatibleModelError",
    "ObfuscationError",
    "VerificationError",
    # Knowledge
    "UniversalKnowledge",
    "StructuralKnowledge",
    "FunctionalKnowledge",
    "MetaKnowledge",
    # Utilities
    "ModelCardGenerator",
    "MergeCardConfig",
    "StackCardConfig",
    # Obfuscation
    "WeightObfuscator",
    "ObfuscationConfig",
    "obfuscate_osmo",
    # OsmoFormat
    "OsmoModel",
    "OsmoManifest",
    "to_osmo",
    "from_osmo",
    "save_osmo",
    "load_osmo",
    "merge_osmo",
    "stack_osmo",
]


def get_adapter(adapter_type: str):
    """
    Get an adapter by type name.

    Args:
        adapter_type: One of 'huggingface', 'vision', 'audio', 'diffusion',
                     'multimodal', 'lora', 'onnx', 'mlx'

    Returns:
        The adapter class

    Example:
        >>> adapter = get_adapter('vision')
        >>> expert = adapter.load('google/vit-base-patch16-224')
    """
    from model_osmosis import adapters

    adapter_map = {
        "huggingface": "HuggingFaceAdapter",
        "hf": "HuggingFaceAdapter",
        "vision": "VisionAdapter",
        "audio": "AudioAdapter",
        "diffusion": "DiffusionAdapter",
        "multimodal": "MultimodalAdapter",
        "lora": "LoRAAdapter",
        "peft": "LoRAAdapter",
        "onnx": "ONNXAdapter",
        "mlx": "MLXAdapter",
        "tensorflow": "TensorFlowAdapter",
        "tf": "TensorFlowAdapter",
        "jax": "JAXAdapter",
        "flax": "JAXAdapter",
    }

    if adapter_type.lower() not in adapter_map:
        raise ValueError(
            f"Unknown adapter type: {adapter_type}. "
            f"Available: {list(adapter_map.keys())}"
        )

    return getattr(adapters, adapter_map[adapter_type.lower()])
