"""
Capacity estimator — determine how much a model needs to grow.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Tuple

import numpy as np

from model_osmosis.core.config import GrowthStrategyType, OsmosisConfig
from model_osmosis.knowledge.representation import UniversalKnowledge


@dataclass
class GrowthPlan:
    """Specification for how a model should be grown.

    Attributes:
        should_grow: Whether growth is needed at all.
        wider_targets: Dict mapping layer names to new widths.
            e.g., {'layer1': 128} means layer1 should be expanded to 128 units.
        deeper_insertions: List of (after_layer_name, new_layer_width) tuples
            specifying where to insert new layers and their width.
        estimated_new_params: Estimated total params after growth.
        growth_ratio: Ratio of new capacity to current capacity.
    """

    should_grow: bool = False
    wider_targets: Dict[str, int] = field(default_factory=dict)
    deeper_insertions: List[Tuple[str, int]] = field(default_factory=list)
    estimated_new_params: int = 0
    growth_ratio: float = 1.0


class CapacityEstimator:
    """Estimate how much a target model needs to grow to absorb source knowledge.

    The estimator compares the source model's parameter count and complexity
    against the target model's current capacity, then produces a GrowthPlan.
    """

    def estimate(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> GrowthPlan:
        """Estimate required growth and produce a plan.

        Args:
            target_model: The current target model.
            knowledge: Extracted knowledge from the source.
            config: Pipeline configuration.

        Returns:
            A GrowthPlan specifying what needs to be expanded.
        """
        source_params = knowledge.structural.total_params
        target_params = self._count_params(target_model)

        # If target already has enough capacity, no growth needed
        if target_params >= source_params:
            return GrowthPlan(should_grow=False, growth_ratio=1.0)

        # Calculate required growth
        growth_ratio = max(config.growth_factor, source_params / max(target_params, 1))

        plan = GrowthPlan(
            should_grow=True,
            growth_ratio=growth_ratio,
        )

        # Build layer-specific growth targets
        try:
            import torch.nn as nn

            if isinstance(target_model, nn.Module):
                plan = self._plan_pytorch_growth(
                    target_model, knowledge, config, growth_ratio
                )
        except ImportError:
            pass

        return plan

    def _count_params(self, model: Any) -> int:
        """Count total parameters in a model."""
        try:
            import torch.nn as nn

            if isinstance(model, nn.Module):
                return sum(p.numel() for p in model.parameters())
        except ImportError:
            pass

        # sklearn models
        total = 0
        if hasattr(model, "coef_"):
            total += np.asarray(model.coef_).size
        if hasattr(model, "intercept_"):
            total += np.asarray(model.intercept_).size
        return total

    def _plan_pytorch_growth(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
        growth_ratio: float,
    ) -> GrowthPlan:
        """Create a detailed growth plan for a PyTorch model."""
        import torch.nn as nn

        wider_targets: Dict[str, int] = {}
        deeper_insertions: List[Tuple[str, int]] = []
        strategy = config.growth_strategy

        # Collect linear/conv layers that can be widened
        linear_layers = []
        for name, module in target_model.named_modules():
            if isinstance(module, nn.Linear):
                linear_layers.append((name, module))

        if strategy in (GrowthStrategyType.WIDER, GrowthStrategyType.HYBRID):
            # Expand hidden layers (not the output layer)
            for name, module in linear_layers[:-1]:  # skip last layer
                new_width = max(
                    module.out_features + 1,
                    int(module.out_features * growth_ratio),
                )
                wider_targets[name] = new_width

        if strategy in (GrowthStrategyType.DEEPER, GrowthStrategyType.HYBRID):
            # Insert new layers after existing hidden layers
            if len(linear_layers) >= 2:
                # Pick the widest hidden layer to insert after
                widest_name = max(
                    linear_layers[:-1],
                    key=lambda x: x[1].out_features,
                )
                width = widest_name[1].out_features
                if strategy == GrowthStrategyType.HYBRID:
                    # Use the grown width if we're also widening
                    width = wider_targets.get(widest_name[0], width)
                deeper_insertions.append((widest_name[0], width))

        # Estimate new param count
        estimated = self._estimate_new_params(
            target_model, wider_targets, deeper_insertions
        )

        return GrowthPlan(
            should_grow=True,
            wider_targets=wider_targets,
            deeper_insertions=deeper_insertions,
            estimated_new_params=estimated,
            growth_ratio=growth_ratio,
        )

    def _estimate_new_params(
        self,
        model: Any,
        wider_targets: Dict[str, int],
        deeper_insertions: List[Tuple[str, int]],
    ) -> int:
        """Rough estimate of parameter count after growth."""
        import torch.nn as nn

        total = sum(p.numel() for p in model.parameters())

        for name, new_width in wider_targets.items():
            for n, m in model.named_modules():
                if n == name and isinstance(m, nn.Linear):
                    # Additional params from widening
                    added = (new_width - m.out_features) * (m.in_features + 1)
                    total += added

        for after_name, width in deeper_insertions:
            # New linear layer: width x width + bias
            total += width * width + width

        return total
