"""
Tests for CrossMerger — cross-architecture model merging.
"""

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

from model_osmosis.merging.cross_merger import CrossMerger, MergedStudent


def _make_model(in_dim, hidden, out_dim, seed=0):
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(in_dim, hidden),
        nn.ReLU(),
        nn.Linear(hidden, out_dim),
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_normal_(p)
    model.eval()
    return model


@pytest.fixture
def sample_data():
    return np.random.rand(100, 10).astype(np.float32)


class TestCrossMerger:
    def test_merge_same_architecture(self, sample_data):
        """Merge models with same arch (basic sanity check)."""
        models = [_make_model(10, 32, 5, seed=i) for i in range(3)]
        merger = CrossMerger(verbose=False)
        merged = merger.merge(models, sample_data, epochs=30)

        assert isinstance(merged, MergedStudent)
        x = torch.rand(5, 10)
        with torch.no_grad():
            out = merged(x)
        assert out.shape == (5, 5)

    def test_merge_different_architectures(self, sample_data):
        """Merge models with DIFFERENT architectures — the main use case."""
        model_a = _make_model(10, 64, 8, seed=0)   # big, 8-dim output
        model_b = _make_model(10, 16, 3, seed=1)   # small, 3-dim output
        model_c = _make_model(10, 32, 12, seed=2)  # medium, 12-dim output

        merger = CrossMerger(verbose=False)
        merged = merger.merge(
            [model_a, model_b, model_c],
            sample_data,
            epochs=30,
        )

        x = torch.rand(5, 10)
        with torch.no_grad():
            out = merged(x)
        # Output should match the max output dim (12)
        assert out.shape == (5, 12)

    def test_merge_two_models(self, sample_data):
        """Simplest case: merge 2 models."""
        model_a = _make_model(10, 32, 5, seed=0)
        model_b = _make_model(10, 64, 8, seed=1)

        merger = CrossMerger(verbose=False)
        merged = merger.merge([model_a, model_b], sample_data, epochs=20)

        x = torch.rand(3, 10)
        with torch.no_grad():
            out = merged(x)
        assert out.shape == (3, 8)

    def test_custom_output_dim(self, sample_data):
        """User can specify output dimension."""
        models = [_make_model(10, 32, 5, seed=i) for i in range(2)]
        merger = CrossMerger(verbose=False)
        merged = merger.merge(
            models, sample_data, output_dim=20, epochs=20
        )

        x = torch.rand(5, 10)
        with torch.no_grad():
            out = merged(x)
        assert out.shape == (5, 20)

    def test_custom_weights(self, sample_data):
        """Can weight models differently."""
        models = [_make_model(10, 32, 5, seed=i) for i in range(3)]
        merger = CrossMerger(verbose=False)

        # Give more weight to first model
        merged = merger.merge(
            models, sample_data, weights=[0.8, 0.1, 0.1], epochs=20
        )

        x = torch.rand(5, 10)
        with torch.no_grad():
            out = merged(x)
        assert out.shape == (5, 5)

    def test_single_model(self, sample_data):
        """Single model should work (essentially distillation)."""
        model = _make_model(10, 64, 5, seed=0)
        merger = CrossMerger(verbose=False)
        merged = merger.merge([model], sample_data, epochs=30)

        assert isinstance(merged, MergedStudent)

    def test_merged_is_clean_module(self, sample_data):
        """Result should be a clean nn.Module with no dependencies."""
        models = [_make_model(10, 32, 5, seed=i) for i in range(2)]
        merger = CrossMerger(verbose=False)
        merged = merger.merge(models, sample_data, epochs=20)

        # Should be a standalone nn.Module
        assert isinstance(merged, nn.Module)
        # Should have trainable parameters
        params = sum(p.numel() for p in merged.parameters())
        assert params > 0

    def test_wildly_different_architectures(self, sample_data):
        """Merge models with very different sizes and outputs."""
        tiny = _make_model(10, 8, 2, seed=0)     # tiny: 2-dim output
        huge = _make_model(10, 256, 50, seed=1)   # huge: 50-dim output

        merger = CrossMerger(verbose=False)
        merged = merger.merge([tiny, huge], sample_data, epochs=20)

        x = torch.rand(5, 10)
        with torch.no_grad():
            out = merged(x)
        assert out.shape == (5, 50)
