"""
Growth strategies — different approaches to expanding model architecture.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from model_osmosis.core.config import OsmosisConfig
from model_osmosis.growers.capacity import GrowthPlan


class WiderStrategy:
    """Growth strategy that only adds neurons to existing layers.

    Safer and simpler — preserves the existing layer structure.
    Uses Net2WiderNet: clones existing neurons and adds noise
    to break symmetry.
    """

    name = "wider"

    @staticmethod
    def filter_plan(plan: GrowthPlan) -> GrowthPlan:
        """Remove deeper insertions from the plan, keeping only wider targets."""
        return GrowthPlan(
            should_grow=plan.should_grow,
            wider_targets=plan.wider_targets,
            deeper_insertions=[],  # no deeper growth
            estimated_new_params=plan.estimated_new_params,
            growth_ratio=plan.growth_ratio,
        )


class DeeperStrategy:
    """Growth strategy that only inserts new layers.

    More expressive — adds depth for complex function approximation.
    Uses Net2DeeperNet: inserts identity-initialized layers.
    """

    name = "deeper"

    @staticmethod
    def filter_plan(plan: GrowthPlan) -> GrowthPlan:
        """Remove wider targets from the plan, keeping only deeper insertions."""
        return GrowthPlan(
            should_grow=plan.should_grow,
            wider_targets={},  # no wider growth
            deeper_insertions=plan.deeper_insertions,
            estimated_new_params=plan.estimated_new_params,
            growth_ratio=plan.growth_ratio,
        )


class HybridStrategy:
    """Growth strategy combining wider and deeper growth.

    First widens existing layers, then adds depth if more capacity is needed.
    This is the default and most flexible strategy.
    """

    name = "hybrid"

    @staticmethod
    def filter_plan(plan: GrowthPlan) -> GrowthPlan:
        """Keep the full plan — both wider and deeper."""
        return plan
