"""sagellm-compression: Model Compression & Acceleration Module for sageLLM."""

from __future__ import annotations

from sagellm_compression._version import __version__
from sagellm_compression.cot import (
    ChainShorteningConfig,
    CompressionResult,
    CoTAccelerationResult,
    CoTAccelerator,
    CoTAcceleratorConfig,
    CoTDetectionResult,
    CoTDetector,
    CoTMetrics,
    CoTStrategyType,
    CoTTemplateManager,
    HeuristicPromptCompressor,
    ParallelChainsConfig,
    PromptCompressionConfig,
    PromptCompressor,
)
from sagellm_compression.fusion import (
    AttentionFusion,
    AttentionFusionConfig,
    FusionConfig,
    FusionResult,
    FusionType,
    GatedMLPFusionConfig,
    KernelFusion,
    MLPFusion,
    MLPFusionConfig,
    StubKernelFusion,
)
from sagellm_compression.quantization import (
    Calibrator,
    Exporter,
    FP8Config,
    QuantizationConfig,
    Quantizer,
    W4A16Config,
    W8A8Config,
)
from sagellm_compression.sparsity import (
    BlockSparsityConfig,
    NMSparsityConfig,
    SparseModel,
    SparseWeight,
    Sparsifier,
    SparsityConfig,
    SparsityType,
    StubSparsifier,
    UnstructuredSparsityConfig,
)
from sagellm_compression.speculative import (
    DraftModel,
    SpeculativeController,
    SpeculativeMetrics,
    SpeculativeResult,
    VerificationResult,
    Verifier,
)

# Public API
__all__ = [
    "__version__",
    # Kernel Fusion (Task 3.4)
    "KernelFusion",
    "AttentionFusion",
    "MLPFusion",
    "StubKernelFusion",
    "FusionConfig",
    "FusionType",
    "AttentionFusionConfig",
    "MLPFusionConfig",
    "GatedMLPFusionConfig",
    "FusionResult",
    # Quantization (Task 3.1)
    "Calibrator",
    "Quantizer",
    "Exporter",
    "QuantizationConfig",
    "W8A8Config",
    "W4A16Config",
    "FP8Config",
    # Speculative decoding (Task 3.3)
    "DraftModel",
    "Verifier",
    "VerificationResult",
    "SpeculativeController",
    "SpeculativeResult",
    "SpeculativeMetrics",
    # Sparsity (Task 3.2)
    "SparsityConfig",
    "SparsityType",
    "UnstructuredSparsityConfig",
    "NMSparsityConfig",
    "BlockSparsityConfig",
    "Sparsifier",
    "StubSparsifier",
    "SparseWeight",
    "SparseModel",
    # CoT acceleration (Task 3.5)
    "CoTAcceleratorConfig",
    "CoTStrategyType",
    "PromptCompressionConfig",
    "ChainShorteningConfig",
    "ParallelChainsConfig",
    "CoTDetector",
    "CoTDetectionResult",
    "PromptCompressor",
    "HeuristicPromptCompressor",
    "CompressionResult",
    "CoTAccelerator",
    "CoTAccelerationResult",
    "CoTMetrics",
    "CoTTemplateManager",
]
