"""
Tests for dynamic architecture growers.
"""

import numpy as np
import pytest
import torch
import torch.nn as nn

from model_osmosis.core.config import GrowthStrategyType, OsmosisConfig
from model_osmosis.growers.capacity import CapacityEstimator, GrowthPlan
from model_osmosis.growers.neural_grower import NeuralGrower
from model_osmosis.extractors import PyTorchExtractor
from model_osmosis.knowledge.representation import (
    FunctionalKnowledge,
    MetaKnowledge,
    StructuralKnowledge,
    UniversalKnowledge,
)


@pytest.fixture
def small_model():
    """Small target model."""
    return nn.Sequential(
        nn.Linear(8, 16),
        nn.ReLU(),
        nn.Linear(16, 4),
    )


@pytest.fixture
def large_source_knowledge():
    """Knowledge from a larger source model."""
    return UniversalKnowledge(
        structural=StructuralKnowledge(
            model_type="pytorch",
            total_params=5000,  # much larger than small_model
            layer_info=[
                {"name": "0", "type": "Linear", "in_features": 8, "out_features": 64},
                {"name": "2", "type": "Linear", "in_features": 64, "out_features": 32},
                {"name": "4", "type": "Linear", "in_features": 32, "out_features": 4},
            ],
            parameters={},
        ),
        functional=FunctionalKnowledge(
            inputs=np.random.rand(100, 8).astype(np.float32),
            outputs=np.random.rand(100, 4).astype(np.float32),
        ),
        meta=MetaKnowledge(
            source_framework="pytorch",
            input_shape=(8,),
            output_shape=(4,),
        ),
    )


@pytest.fixture
def small_source_knowledge():
    """Knowledge from a model equal/smaller than target."""
    return UniversalKnowledge(
        structural=StructuralKnowledge(
            model_type="pytorch",
            total_params=50,  # smaller than small_model
        ),
        functional=FunctionalKnowledge(
            inputs=np.random.rand(50, 8).astype(np.float32),
            outputs=np.random.rand(50, 4).astype(np.float32),
        ),
        meta=MetaKnowledge(
            source_framework="pytorch",
            input_shape=(8,),
            output_shape=(4,),
        ),
    )


class TestCapacityEstimator:
    def test_no_growth_needed(self, small_model, small_source_knowledge):
        estimator = CapacityEstimator()
        config = OsmosisConfig(verbose=False)
        plan = estimator.estimate(small_model, small_source_knowledge, config)
        assert plan.should_grow is False

    def test_growth_needed(self, small_model, large_source_knowledge):
        estimator = CapacityEstimator()
        config = OsmosisConfig(verbose=False)
        plan = estimator.estimate(small_model, large_source_knowledge, config)
        assert plan.should_grow is True
        assert plan.growth_ratio > 1.0

    def test_wider_targets_generated(self, small_model, large_source_knowledge):
        estimator = CapacityEstimator()
        config = OsmosisConfig(
            growth_strategy=GrowthStrategyType.WIDER, verbose=False
        )
        plan = estimator.estimate(small_model, large_source_knowledge, config)
        assert len(plan.wider_targets) > 0


class TestNeuralGrower:
    def test_can_handle(self, small_model):
        grower = NeuralGrower()
        assert grower.can_handle(small_model) is True

    def test_no_growth_when_sufficient(self, small_model, small_source_knowledge):
        grower = NeuralGrower()
        config = OsmosisConfig(verbose=False)
        result = grower.grow(small_model, small_source_knowledge, config)
        # Should return same model (no growth needed)
        old_params = sum(p.numel() for p in small_model.parameters())
        new_params = sum(p.numel() for p in result.parameters())
        assert new_params == old_params

    def test_wider_growth_increases_params(self, small_model, large_source_knowledge):
        grower = NeuralGrower()
        config = OsmosisConfig(
            growth_strategy=GrowthStrategyType.WIDER,
            verbose=False,
        )
        old_params = sum(p.numel() for p in small_model.parameters())
        grown = grower.grow(small_model, large_source_knowledge, config)
        new_params = sum(p.numel() for p in grown.parameters())
        assert new_params > old_params

    def test_wider_growth_preserves_output(self, large_source_knowledge):
        """After wider growth, model should produce same-shaped outputs."""
        model = nn.Sequential(
            nn.Linear(8, 16),
            nn.ReLU(),
            nn.Linear(16, 4),
        )
        model.eval()

        x = torch.rand(5, 8)
        with torch.no_grad():
            original_out = model(x)

        grower = NeuralGrower()
        config = OsmosisConfig(
            growth_strategy=GrowthStrategyType.WIDER,
            noise_std=0.0,  # no noise → exact function preservation
            verbose=False,
        )
        grown = grower.grow(model, large_source_knowledge, config)
        grown.eval()

        with torch.no_grad():
            grown_out = grown(x)

        # Same output shape
        assert grown_out.shape == original_out.shape

        # With zero noise, outputs should be approximately preserved
        torch.testing.assert_close(grown_out, original_out, atol=1e-5, rtol=1e-5)

    def test_deeper_growth_increases_params(self, small_model, large_source_knowledge):
        grower = NeuralGrower()
        config = OsmosisConfig(
            growth_strategy=GrowthStrategyType.DEEPER,
            verbose=False,
        )
        old_params = sum(p.numel() for p in small_model.parameters())
        grown = grower.grow(small_model, large_source_knowledge, config)
        new_params = sum(p.numel() for p in grown.parameters())
        assert new_params > old_params

    def test_hybrid_growth(self, small_model, large_source_knowledge):
        grower = NeuralGrower()
        config = OsmosisConfig(
            growth_strategy=GrowthStrategyType.HYBRID,
            verbose=False,
        )
        old_params = sum(p.numel() for p in small_model.parameters())
        grown = grower.grow(small_model, large_source_knowledge, config)
        new_params = sum(p.numel() for p in grown.parameters())
        assert new_params > old_params
