"""
OsmoFormat - Universal Model Interchange Format

The .osmo format is a universal intermediate representation that allows
any ML model to be converted, merged, stacked, and transferred regardless
of its original framework.

Quick Start:
    # Convert any model to OsmoFormat
    from model_osmosis.format import to_osmo, from_osmo, save_osmo, load_osmo

    # PyTorch -> OsmoFormat
    osmo = to_osmo(my_torch_model)

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

    # Load and convert to TensorFlow
    osmo = load_osmo("my_model.osmo")
    tf_weights = from_osmo(osmo, target_framework="tensorflow")

    # Merge models from different frameworks
    from model_osmosis.format import merge_osmo

    osmo1 = to_osmo(pytorch_model)
    osmo2 = to_osmo(jax_model)
    merged = merge_osmo([osmo1, osmo2], weights=[0.7, 0.3])

Supported Frameworks:
    - PyTorch (torch.nn.Module, state_dict)
    - TensorFlow/Keras (tf.keras.Model)
    - JAX/Flax (params dict)
    - ONNX (onnx.ModelProto)
    - HuggingFace Transformers
"""

# Core format types
from model_osmosis.format.osmo_format import (
    OsmoModel,
    OsmoManifest,
    ArchitectureSpec,
    LayerSpec,
    LayerType,
    ModelModality,
    FrameworkOrigin,
    ModelCapabilities,
    TrainingInfo,
    OSMO_FORMAT_VERSION,
    OSMO_EXTENSION,
)

# Serialization
from model_osmosis.format.serialization import (
    OsmoSerializer,
    OsmoLoader,
    save_osmo,
    load_osmo,
)

# Converters
from model_osmosis.format.converters import (
    UniversalConverter,
    PyTorchConverter,
    TensorFlowConverter,
    JAXConverter,
    ONNXConverter,
    HuggingFaceConverter,
    to_osmo,
    from_osmo,
)

# Operations
from model_osmosis.format.operations import (
    OsmoMerger,
    OsmoStacker,
    OsmoComposer,
    merge_osmo,
    stack_osmo,
)

# Image container
from model_osmosis.format.image import (
    OsmoImage,
    OsmoRuntime,
)

# Adaptive compression
from model_osmosis.format.adaptive_compression import (
    ADAPTIVE_PAYLOAD_CODEC,
    ADAPTIVE_COMPRESSION_ENGINE_VERSION,
)
from model_osmosis.format.parity import (
    ParityResult,
    compare_tensor_maps,
    DEFAULT_TOLERANCE_POLICY,
)
from model_osmosis.format.trust import (
    DEFAULT_SIGNER_ID,
    TRUST_SIGNATURE_ALGORITHM,
    DEFAULT_SIGNER_BACKEND,
    KMS_SIGNER_BACKEND,
    SignerBackend,
    LocalSignerBackend,
    AWSSignerBackend,
    generate_signer,
    list_keys,
    import_trusted_signer,
    revoke_trusted_signer,
    get_signer_backend,
    sign_statement_with_backend,
    verify_trust_block_with_backend,
    publish_transparency_event,
    verify_transparency_event,
    append_audit_event,
)

# File format registration
from model_osmosis.format.file_format import (
    OSMOIMG_EXTENSION,
    OSMOIMG_MAGIC,
    OSMOIMG_MIME_TYPE,
    is_osmoimg,
    register_file_format,
    unregister_file_format,
)

__all__ = [
    # Core types
    "OsmoModel",
    "OsmoManifest",
    "ArchitectureSpec",
    "LayerSpec",
    "LayerType",
    "ModelModality",
    "FrameworkOrigin",
    "ModelCapabilities",
    "TrainingInfo",
    "OSMO_FORMAT_VERSION",
    "OSMO_EXTENSION",
    # Serialization
    "OsmoSerializer",
    "OsmoLoader",
    "save_osmo",
    "load_osmo",
    # Converters
    "UniversalConverter",
    "PyTorchConverter",
    "TensorFlowConverter",
    "JAXConverter",
    "ONNXConverter",
    "HuggingFaceConverter",
    "to_osmo",
    "from_osmo",
    # Operations
    "OsmoMerger",
    "OsmoStacker",
    "OsmoComposer",
    "merge_osmo",
    "stack_osmo",
    # Image
    "OsmoImage",
    "OsmoRuntime",
    # Compression
    "ADAPTIVE_PAYLOAD_CODEC",
    "ADAPTIVE_COMPRESSION_ENGINE_VERSION",
    # Parity validation
    "ParityResult",
    "compare_tensor_maps",
    "DEFAULT_TOLERANCE_POLICY",
    # Trust
    "DEFAULT_SIGNER_ID",
    "TRUST_SIGNATURE_ALGORITHM",
    "DEFAULT_SIGNER_BACKEND",
    "KMS_SIGNER_BACKEND",
    "SignerBackend",
    "LocalSignerBackend",
    "AWSSignerBackend",
    "generate_signer",
    "list_keys",
    "import_trusted_signer",
    "revoke_trusted_signer",
    "get_signer_backend",
    "sign_statement_with_backend",
    "verify_trust_block_with_backend",
    "publish_transparency_event",
    "verify_transparency_event",
    "append_audit_event",
]
