"""
Configuration classes for ModelOsmosis.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional


class GrowthStrategyType(Enum):
    """Strategy for how the target model's architecture should grow."""

    WIDER = "wider"
    """Add more neurons/units to existing layers."""

    DEEPER = "deeper"
    """Insert new layers between existing layers."""

    HYBRID = "hybrid"
    """Combination of wider and deeper — first widen, then deepen if needed."""


@dataclass
class OsmosisConfig:
    """Global configuration for the osmosis pipeline.

    Attributes:
        growth_strategy: How to expand the target model architecture.
        growth_factor: Multiplier for how much to grow (1.5 = 50% more capacity).
        max_growth_iterations: Safety limit on growth iterations.
        distillation_epochs: Number of training epochs for knowledge distillation.
        distillation_lr: Learning rate for knowledge distillation training.
        distillation_batch_size: Batch size during distillation.
        functional_samples: Number of I/O sample pairs to capture for functional knowledge.
        noise_std: Standard deviation of noise added during Net2Net wider growth
            to break symmetry between cloned neurons.
        device: Device for PyTorch operations ('cpu', 'cuda', 'auto').
        verbose: Whether to print progress information.
    """

    growth_strategy: GrowthStrategyType = GrowthStrategyType.HYBRID
    growth_factor: float = 1.5
    max_growth_iterations: int = 10
    distillation_epochs: int = 50
    distillation_lr: float = 1e-3
    distillation_batch_size: int = 64
    functional_samples: int = 1000
    noise_std: float = 1e-2
    device: str = "cpu"
    verbose: bool = True
    auto_obfuscate: bool = False
    """When True, automatically obfuscate every model produced by the pipeline."""
    seed: int = 42
    """Global seed for deterministic behavior in merge/stack/cross-merge flows."""

    obfuscation: Optional[Any] = None
    """Optional ObfuscationConfig for weight obfuscation. Import from
    model_osmosis.obfuscation to avoid circular imports."""

    def __post_init__(self) -> None:
        if self.growth_factor < 1.0:
            raise ValueError(
                f"growth_factor must be >= 1.0, got {self.growth_factor}"
            )
        if self.max_growth_iterations < 1:
            raise ValueError(
                f"max_growth_iterations must be >= 1, got {self.max_growth_iterations}"
            )
        if self.distillation_epochs < 1:
            raise ValueError(
                f"distillation_epochs must be >= 1, got {self.distillation_epochs}"
            )
        if self.device == "auto":
            try:
                import torch

                self.device = "cuda" if torch.cuda.is_available() else "cpu"
            except ImportError:
                self.device = "cpu"
