"""
Neural network grower — Net2Net-inspired dynamic architecture expansion.

Implements function-preserving transformations:
- Net2WiderNet: Add neurons by cloning existing ones + noise
- Net2DeeperNet: Insert identity-initialized layers

Reference: Chen, Goodfellow, Shlens — "Net2Net: Accelerating Learning via Knowledge Transfer" (2015)
"""

from __future__ import annotations

import copy
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from model_osmosis.core.config import GrowthStrategyType, OsmosisConfig
from model_osmosis.core.exceptions import GrowthError
from model_osmosis.growers.base import BaseGrower
from model_osmosis.growers.capacity import CapacityEstimator, GrowthPlan
from model_osmosis.growers.strategies import (
    DeeperStrategy,
    HybridStrategy,
    WiderStrategy,
)
from model_osmosis.knowledge.representation import UniversalKnowledge


class NeuralGrower(BaseGrower):
    """Dynamically grow PyTorch neural network architectures.

    Uses Net2Net-inspired function-preserving transforms so the grown
    model produces the **same outputs** as before growth.
    """

    def __init__(self) -> None:
        self.estimator = CapacityEstimator()

    def can_handle(self, model: Any) -> bool:
        try:
            import torch.nn as nn

            return isinstance(model, nn.Module)
        except ImportError:
            return False

    def grow(
        self,
        model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> Any:
        """Grow the PyTorch model according to the growth plan.

        Args:
            model: The PyTorch nn.Module to expand.
            knowledge: Source knowledge (used to determine required capacity).
            config: Pipeline configuration.

        Returns:
            A new, larger model that produces the same outputs as the original.
        """
        import torch.nn as nn

        plan = self.estimator.estimate(model, knowledge, config)

        if not plan.should_grow:
            if config.verbose:
                print("[Grower] Target already has sufficient capacity — no growth needed")
            return model

        # Apply strategy filter
        if config.growth_strategy == GrowthStrategyType.WIDER:
            plan = WiderStrategy.filter_plan(plan)
        elif config.growth_strategy == GrowthStrategyType.DEEPER:
            plan = DeeperStrategy.filter_plan(plan)
        else:
            plan = HybridStrategy.filter_plan(plan)

        if config.verbose:
            print(f"[Grower] Growth plan: ratio={plan.growth_ratio:.2f}")
            if plan.wider_targets:
                print(f"[Grower]   Wider: {plan.wider_targets}")
            if plan.deeper_insertions:
                print(f"[Grower]   Deeper insertions: {plan.deeper_insertions}")

        grown_model = copy.deepcopy(model)

        # Apply wider growth first
        if plan.wider_targets:
            grown_model = self._apply_wider(grown_model, plan, config)

        # Then apply deeper growth
        if plan.deeper_insertions:
            grown_model = self._apply_deeper(grown_model, plan, config)

        if config.verbose:
            old_params = sum(p.numel() for p in model.parameters())
            new_params = sum(p.numel() for p in grown_model.parameters())
            print(
                f"[Grower] Growth complete: {old_params:,} -> {new_params:,} "
                f"parameters ({new_params / max(old_params, 1):.1f}x)"
            )

        return grown_model

    def _apply_wider(
        self, model: Any, plan: GrowthPlan, config: OsmosisConfig
    ) -> Any:
        """Apply Net2WiderNet — add neurons to existing layers.

        For each widened layer:
        1. Select neurons to clone (randomly)
        2. Copy their weights
        3. Add small noise to break symmetry
        4. Update the next layer's input weights accordingly
        """
        import torch
        import torch.nn as nn

        # Build ordered list of linear layers
        linear_layers = []
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                linear_layers.append((name, module))

        for target_name, new_width in plan.wider_targets.items():
            # Find this layer and the next layer
            target_idx = None
            for i, (name, _) in enumerate(linear_layers):
                if name == target_name:
                    target_idx = i
                    break

            if target_idx is None:
                continue

            layer = linear_layers[target_idx][1]
            old_width = layer.out_features

            if new_width <= old_width:
                continue

            extra_neurons = new_width - old_width

            # Select neurons to clone (cycle through existing)
            clone_indices = [i % old_width for i in range(extra_neurons)]

            # Create new weight matrix for this layer
            old_weight = layer.weight.data.clone()
            old_bias = layer.bias.data.clone() if layer.bias is not None else None

            new_weight = torch.zeros(new_width, layer.in_features)
            new_weight[:old_width] = old_weight
            for i, src_idx in enumerate(clone_indices):
                noise = torch.randn_like(old_weight[src_idx]) * config.noise_std
                new_weight[old_width + i] = old_weight[src_idx] + noise

            new_bias = None
            if old_bias is not None:
                new_bias = torch.zeros(new_width)
                new_bias[:old_width] = old_bias
                for i, src_idx in enumerate(clone_indices):
                    new_bias[old_width + i] = old_bias[src_idx]

            # Replace the layer with a wider one
            new_layer = nn.Linear(layer.in_features, new_width, bias=layer.bias is not None)
            new_layer.weight.data = new_weight
            if new_bias is not None:
                new_layer.bias.data = new_bias

            _replace_module(model, target_name, new_layer)

            # Update the next layer's input features
            if target_idx + 1 < len(linear_layers):
                next_name, next_layer = linear_layers[target_idx + 1]
                old_next_weight = next_layer.weight.data.clone()

                new_next_weight = torch.zeros(
                    next_layer.out_features, new_width
                )
                new_next_weight[:, :old_width] = old_next_weight

                # Count how many times each neuron was cloned
                clone_count = {}
                for src_idx in clone_indices:
                    clone_count[src_idx] = clone_count.get(src_idx, 0) + 1

                # For each cloned neuron, split weight equally across
                # the original and all its clones: divide by (1 + n_clones)
                for src_idx, n_clones in clone_count.items():
                    divisor = float(1 + n_clones)
                    split_weight = old_next_weight[:, src_idx] / divisor
                    new_next_weight[:, src_idx] = split_weight

                for i, src_idx in enumerate(clone_indices):
                    n_clones = clone_count[src_idx]
                    divisor = float(1 + n_clones)
                    new_next_weight[:, old_width + i] = (
                        old_next_weight[:, src_idx] / divisor
                    )

                new_next_layer = nn.Linear(
                    new_width, next_layer.out_features,
                    bias=next_layer.bias is not None,
                )
                new_next_layer.weight.data = new_next_weight
                if next_layer.bias is not None:
                    new_next_layer.bias.data = next_layer.bias.data.clone()

                _replace_module(model, next_name, new_next_layer)

                # Refresh layer list
                linear_layers = []
                for name, module in model.named_modules():
                    if isinstance(module, nn.Linear):
                        linear_layers.append((name, module))

        return model

    def _apply_deeper(
        self, model: Any, plan: GrowthPlan, config: OsmosisConfig
    ) -> Any:
        """Apply Net2DeeperNet — insert identity-initialized layers.

        New layers are initialized as identity transforms so the model's
        output is preserved.
        """
        import torch
        import torch.nn as nn

        for after_name, width in plan.deeper_insertions:
            # Create an identity-initialized layer + ReLU
            new_layer = nn.Linear(width, width)
            # Initialize as identity matrix
            nn.init.eye_(new_layer.weight)
            nn.init.zeros_(new_layer.bias)

            # Insert after the specified layer in a Sequential-like structure
            _insert_after(model, after_name, new_layer, config)

        return model


def _replace_module(model: Any, target_name: str, new_module: Any) -> None:
    """Replace a named module within a model."""
    import torch.nn as nn

    parts = target_name.split(".")
    parent = model
    for part in parts[:-1]:
        if part.isdigit():
            parent = parent[int(part)]
        else:
            parent = getattr(parent, part)

    last = parts[-1]
    if last.isdigit():
        parent[int(last)] = new_module
    else:
        setattr(parent, last, new_module)


def _insert_after(
    model: Any, after_name: str, new_layer: Any, config: OsmosisConfig
) -> None:
    """Insert a new layer after a named layer in a Sequential model.

    For non-Sequential models, wraps the target and new layer together.
    """
    import torch.nn as nn

    parts = after_name.split(".")
    parent = model
    for part in parts[:-1]:
        if part.isdigit():
            parent = parent[int(part)]
        else:
            parent = getattr(parent, part)

    last_part = parts[-1]

    if isinstance(parent, nn.Sequential):
        idx = int(last_part)
        modules = list(parent.children())

        new_modules = []
        for i, mod in enumerate(modules):
            new_modules.append(mod)
            if i == idx:
                new_modules.append(new_layer)
                new_modules.append(nn.ReLU())

        # Rebuild Sequential
        new_sequential = nn.Sequential(*new_modules)

        if len(parts) > 1:
            grandparent = model
            for part in parts[:-2]:
                if part.isdigit():
                    grandparent = grandparent[int(part)]
                else:
                    grandparent = getattr(grandparent, part)
            setattr(grandparent, parts[-2], new_sequential)
        else:
            # The parent IS the model (top-level Sequential)
            for i, mod in enumerate(new_modules):
                if i < len(parent):
                    parent[i] = mod  # type: ignore
            # If new modules are longer, we need to rebuild
            new_seq = nn.Sequential(*new_modules)
            # Copy attributes
            for k, v in new_seq._modules.items():
                parent._modules[k] = v
            # Clean up extra keys
            parent._modules = new_seq._modules

    else:
        # For non-Sequential: wrap the existing module with new_layer
        existing = getattr(parent, last_part)
        wrapper = nn.Sequential(existing, new_layer, nn.ReLU())
        setattr(parent, last_part, wrapper)

    if config.verbose:
        print(
            f"[Grower] Inserted identity layer ({new_layer.in_features} units) "
            f"after '{after_name}'"
        )
