"""
Tests for multi-model merging.
"""

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

from model_osmosis.merging.strategies import (
    WeightedAverageStrategy,
    SLERPStrategy,
    TIESStrategy,
    get_strategy,
)
from model_osmosis.merging.streaming import StreamingParameterLoader
from model_osmosis.merging.evaluator import MergeEvaluator
from model_osmosis.merging.merger import MultiModelMerger
from model_osmosis import OsmosisPipeline, OsmosisConfig


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _make_model(seed: int = 0):
    """Create a simple model with deterministic weights."""
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(6, 16),
        nn.ReLU(),
        nn.Linear(16, 8),
        nn.ReLU(),
        nn.Linear(8, 3),
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_normal_(p)
    model.eval()
    return model


@pytest.fixture
def three_models():
    """Three models with different weights."""
    return [_make_model(seed=i) for i in range(3)]


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


# ---------------------------------------------------------------------------
# Strategy Tests
# ---------------------------------------------------------------------------


class TestWeightedAverage:
    def test_uniform_average(self):
        strat = WeightedAverageStrategy()
        a = np.array([1.0, 2.0, 3.0])
        b = np.array([5.0, 6.0, 7.0])
        result = strat.merge_tensors([a, b])
        np.testing.assert_allclose(result, [3.0, 4.0, 5.0])

    def test_weighted_average(self):
        strat = WeightedAverageStrategy()
        a = np.array([0.0, 0.0])
        b = np.array([10.0, 10.0])
        result = strat.merge_tensors([a, b], weights=[0.8, 0.2])
        np.testing.assert_allclose(result, [2.0, 2.0])

    def test_three_tensors(self):
        strat = WeightedAverageStrategy()
        a = np.ones(5) * 1
        b = np.ones(5) * 2
        c = np.ones(5) * 3
        result = strat.merge_tensors([a, b, c])
        np.testing.assert_allclose(result, [2.0] * 5)


class TestSLERP:
    def test_two_tensors(self):
        strat = SLERPStrategy()
        a = np.array([1.0, 0.0, 0.0])
        b = np.array([0.0, 1.0, 0.0])
        result = strat.merge_tensors([a, b], weights=[0.5, 0.5])
        # SLERP at t=0.5 should be on the arc between a and b
        assert result.shape == (3,)
        # Both components should be positive and similar
        assert result[0] > 0 and result[1] > 0

    def test_identical_tensors(self):
        strat = SLERPStrategy()
        a = np.array([1.0, 2.0, 3.0])
        result = strat.merge_tensors([a, a])
        np.testing.assert_allclose(result, a, atol=1e-5)

    def test_three_tensors(self):
        strat = SLERPStrategy()
        tensors = [np.random.rand(10) for _ in range(3)]
        result = strat.merge_tensors(tensors)
        assert result.shape == (10,)


class TestTIES:
    def test_basic_merge(self):
        strat = TIESStrategy(density=0.5)
        a = np.array([1.0, -2.0, 3.0, -4.0])
        b = np.array([1.5, -1.5, 2.5, -3.5])
        c = np.array([0.5, -2.5, 3.5, -4.5])
        result = strat.merge_tensors([a, b, c])
        assert result.shape == (4,)

    def test_density_parameter(self):
        strat_high = TIESStrategy(density=0.9)
        strat_low = TIESStrategy(density=0.1)
        tensors = [np.random.rand(20) for _ in range(3)]
        r_high = strat_high.merge_tensors(tensors)
        r_low = strat_low.merge_tensors(tensors)
        # Both should produce valid outputs
        assert r_high.shape == (20,)
        assert r_low.shape == (20,)


class TestStrategyRegistry:
    def test_get_valid_strategy(self):
        s = get_strategy("weighted_average")
        assert isinstance(s, WeightedAverageStrategy)

    def test_get_invalid_raises(self):
        with pytest.raises(ValueError):
            get_strategy("nonexistent")


# ---------------------------------------------------------------------------
# Streaming Tests
# ---------------------------------------------------------------------------


class TestStreamingParameterLoader:
    def test_stream_from_module(self):
        model = _make_model()
        loader = StreamingParameterLoader()
        params = list(loader.stream_parameters(model))

        assert len(params) > 0
        for name, arr in params:
            assert isinstance(name, str)
            assert isinstance(arr, np.ndarray)

    def test_stream_from_dict(self):
        model = _make_model()
        state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
        loader = StreamingParameterLoader()
        params = list(loader.stream_parameters(state_dict))
        assert len(params) > 0

    def test_get_param_names(self):
        model = _make_model()
        loader = StreamingParameterLoader()
        names = loader.get_all_param_names(model)
        assert len(names) > 0
        assert all(isinstance(n, str) for n in names)


# ---------------------------------------------------------------------------
# Evaluator Tests
# ---------------------------------------------------------------------------


class TestMergeEvaluator:
    def test_find_best_weights(self, three_models, sample_data):
        evaluator = MergeEvaluator()
        strat = WeightedAverageStrategy()
        weights, score = evaluator.find_best_weights(
            three_models, sample_data, strat, n_trials=5
        )
        assert len(weights) == 3
        assert all(w >= 0 for w in weights)
        assert abs(sum(weights) - 1.0) < 1e-5

    def test_single_model_returns_default(self):
        evaluator = MergeEvaluator()
        model = _make_model()
        strat = WeightedAverageStrategy()
        weights, score = evaluator.find_best_weights(
            [model], np.random.rand(10, 6).astype(np.float32), strat
        )
        assert weights == [1.0]


# ---------------------------------------------------------------------------
# MultiModelMerger Tests
# ---------------------------------------------------------------------------


class TestMultiModelMerger:
    def test_merge_weighted_average(self, three_models):
        merger = MultiModelMerger(strategy="weighted_average", verbose=False)
        result = merger.merge(three_models)

        # Should produce a valid model
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_merge_slerp(self, three_models):
        merger = MultiModelMerger(strategy="slerp", verbose=False)
        result = merger.merge(three_models)

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_merge_ties(self, three_models):
        merger = MultiModelMerger(strategy="ties", verbose=False)
        result = merger.merge(three_models)

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_merge_auto(self, three_models, sample_data):
        merger = MultiModelMerger(strategy="auto", verbose=False)
        result = merger.merge(three_models, eval_data=sample_data)

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_merge_custom_weights(self, three_models):
        merger = MultiModelMerger(
            strategy="weighted_average",
            model_weights=[0.6, 0.3, 0.1],
            verbose=False,
        )
        result = merger.merge(three_models)

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_merged_differs_from_originals(self, three_models):
        """Merged model should produce different outputs than any single model."""
        merger = MultiModelMerger(strategy="weighted_average", verbose=False)
        result = merger.merge(three_models)

        x = torch.rand(10, 6)
        with torch.no_grad():
            merged_out = result(x).numpy()
            original_outs = [m(x).numpy() for m in three_models]

        # Merged should not be identical to any single model
        for orig in original_outs:
            assert not np.allclose(merged_out, orig, atol=1e-6)


# ---------------------------------------------------------------------------
# Pipeline.merge() Tests
# ---------------------------------------------------------------------------


class TestPipelineMerge:
    def test_pipeline_merge(self, three_models):
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        result = pipeline.merge(
            three_models,
            strategy="weighted_average",
        )

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_pipeline_merge_auto(self, three_models, sample_data):
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        result = pipeline.merge(
            three_models,
            strategy="auto",
            eval_data=sample_data,
        )

        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)
