"""
Tests for Model Osmosis v0.6.0 features:
- New merge strategies (DARE, Breadcrumbs, Consensus)
- Obfuscation quality preservation
- New adapters (TensorFlow, JAX)
- Pipeline API (compare, benchmark)
- Performance improvements (inplace merge)
"""

import copy

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

from model_osmosis.merging.strategies import (
    DAREStrategy,
    BreadcrumbsStrategy,
    ConsensusStrategy,
    get_strategy,
    STRATEGY_REGISTRY,
)
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():
    return [_make_model(seed=i) for i in range(3)]


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


# ---------------------------------------------------------------------------
# New Merge Strategies
# ---------------------------------------------------------------------------


class TestDAREStrategy:
    def test_basic_merge(self):
        strat = DAREStrategy(drop_rate=0.5, seed=42)
        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,)
        assert result.dtype == a.dtype

    def test_single_tensor(self):
        strat = DAREStrategy()
        a = np.array([1.0, 2.0, 3.0])
        result = strat.merge_tensors([a])
        np.testing.assert_array_equal(result, a)

    def test_drop_rate_effect(self):
        """Higher drop rate should produce different results."""
        a = np.random.rand(20)
        b = np.random.rand(20)
        strat_low = DAREStrategy(drop_rate=0.1, seed=42)
        strat_high = DAREStrategy(drop_rate=0.9, seed=42)
        r_low = strat_low.merge_tensors([a, b])
        r_high = strat_high.merge_tensors([a, b])
        assert r_low.shape == (20,)
        assert r_high.shape == (20,)

    def test_2d_tensors(self):
        strat = DAREStrategy(seed=42)
        a = np.random.rand(4, 3).astype(np.float32)
        b = np.random.rand(4, 3).astype(np.float32)
        result = strat.merge_tensors([a, b])
        assert result.shape == (4, 3)
        assert result.dtype == np.float32

    def test_registry(self):
        assert "dare" in STRATEGY_REGISTRY
        s = get_strategy("dare", drop_rate=0.3)
        assert isinstance(s, DAREStrategy)
        assert s.drop_rate == 0.3


class TestBreadcrumbsStrategy:
    def test_basic_merge(self):
        strat = BreadcrumbsStrategy(density=0.5)
        a = np.array([1.0, -2.0, 3.0, -4.0, 0.1])
        b = np.array([1.5, -1.5, 2.5, -3.5, 0.2])
        result = strat.merge_tensors([a, b])
        assert result.shape == (5,)

    def test_single_tensor(self):
        strat = BreadcrumbsStrategy()
        a = np.array([1.0, 2.0])
        result = strat.merge_tensors([a])
        np.testing.assert_array_equal(result, a)

    def test_density_parameter(self):
        strat_low = BreadcrumbsStrategy(density=0.1)
        strat_high = BreadcrumbsStrategy(density=0.9)
        tensors = [np.random.rand(20) for _ in range(3)]
        r_low = strat_low.merge_tensors(tensors)
        r_high = strat_high.merge_tensors(tensors)
        assert r_low.shape == (20,)
        assert r_high.shape == (20,)

    def test_registry(self):
        assert "breadcrumbs" in STRATEGY_REGISTRY
        s = get_strategy("breadcrumbs", density=0.4)
        assert isinstance(s, BreadcrumbsStrategy)


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

    def test_single_tensor(self):
        strat = ConsensusStrategy()
        a = np.array([1.0, 2.0, 3.0])
        result = strat.merge_tensors([a])
        np.testing.assert_array_equal(result, a)

    def test_threshold_clamping(self):
        """Threshold should be clamped to [0.5, 1.0]."""
        strat = ConsensusStrategy(consensus_threshold=0.2)
        assert strat.consensus_threshold == 0.5
        strat = ConsensusStrategy(consensus_threshold=1.5)
        assert strat.consensus_threshold == 1.0

    def test_registry(self):
        assert "consensus" in STRATEGY_REGISTRY
        s = get_strategy("consensus", consensus_threshold=0.8)
        assert isinstance(s, ConsensusStrategy)


class TestNewStrategyMerger:
    """Test new strategies through the full MultiModelMerger flow."""

    def test_merge_dare(self, three_models):
        merger = MultiModelMerger(strategy="dare", 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_breadcrumbs(self, three_models):
        merger = MultiModelMerger(strategy="breadcrumbs", 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_consensus(self, three_models):
        merger = MultiModelMerger(strategy="consensus", 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_all_strategies_in_registry(self):
        expected = {
            "weighted_average", "slerp", "ties",
            "dare", "breadcrumbs", "consensus",
        }
        assert expected == set(STRATEGY_REGISTRY.keys())


# ---------------------------------------------------------------------------
# Obfuscation Quality Preservation
# ---------------------------------------------------------------------------


class TestObfuscationQuality:
    def test_preserve_quality_default(self):
        """preserve_quality should be True by default."""
        from model_osmosis.obfuscation import ObfuscationConfig
        config = ObfuscationConfig()
        assert config.preserve_quality is True
        assert config.knowledge_noise_scale == 0.01
        assert config.verification_tolerance == 1e-6

    def test_quality_preserving_obfuscation(self):
        """With preserve_quality=True, obfuscated model should produce
        nearly identical outputs (within float precision)."""
        from model_osmosis.obfuscation import WeightObfuscator, ObfuscationConfig

        model = _make_model(seed=42)
        sample = np.random.rand(20, 6).astype(np.float32)

        # Get original outputs
        x = torch.from_numpy(sample).float()
        model.eval()
        with torch.no_grad():
            original_out = model(x).numpy()

        # Obfuscate with quality preservation
        config = ObfuscationConfig(
            preserve_quality=True,
            verify=False,  # We'll check manually
            seed=42,
        )
        obfuscator = WeightObfuscator(config)
        obfuscated_model = obfuscator.obfuscate_torch_model(
            copy.deepcopy(model), sample
        )

        # Get obfuscated outputs
        obfuscated_model.eval()
        with torch.no_grad():
            obfuscated_out = obfuscated_model(x).numpy()

        # Should be EXACT (permutation + scaling are mathematically exact)
        max_error = np.max(np.abs(original_out - obfuscated_out))
        assert max_error < 1e-5, (
            f"Quality-preserving obfuscation error too high: {max_error:.2e}"
        )

    def test_knowledge_no_noise_in_quality_mode(self):
        """With noise_scale=0, knowledge should only be permuted, not corrupted."""
        from model_osmosis.obfuscation.transforms import obfuscate_knowledge

        rng = np.random.default_rng(42)
        knowledge = {
            "weights": np.random.rand(10, 5).astype(np.float32),
            "activations": np.random.rand(8, 3).astype(np.float32),
        }

        obfuscated = obfuscate_knowledge(knowledge, rng, noise_scale=0.0)

        # Values should be unchanged (just permuted along axis 0)
        for key in knowledge:
            orig = knowledge[key]
            obf = obfuscated[key]
            assert orig.shape == obf.shape
            # Each row in obfuscated should exist in original (permutation only)
            for row in obf:
                found = any(np.allclose(row, orig_row) for orig_row in orig)
                assert found, f"Row not found in original - noise was added!"


# ---------------------------------------------------------------------------
# New Adapters
# ---------------------------------------------------------------------------


class TestTensorFlowAdapter:
    def test_import(self):
        """TensorFlowAdapter should be importable from adapters."""
        from model_osmosis.adapters.tensorflow import TensorFlowAdapter
        adapter = TensorFlowAdapter()
        assert adapter is not None

    def test_expert_from_weights(self):
        """TensorFlowExpert should work with raw weight dicts."""
        from model_osmosis.adapters.tensorflow import TensorFlowExpert

        class DummyModel:
            name = "test"
            layers = []
            input_shape = (None, 10)
            output_shape = (None, 3)
            def predict(self, x, verbose=0):
                return x[:, :3]

        expert = TensorFlowExpert(DummyModel(), name="dummy")
        assert expert.name == "dummy"


class TestJAXAdapter:
    def test_import(self):
        """JAXAdapter should be importable from adapters."""
        from model_osmosis.adapters.jax import JAXAdapter
        adapter = JAXAdapter()
        assert adapter is not None

    def test_flatten_unflatten(self):
        """Pytree flatten/unflatten should be invertible."""
        from model_osmosis.adapters.jax import _flatten_pytree, _unflatten_pytree

        tree = {
            "encoder": {
                "dense_0": {"kernel": np.ones((3, 4)), "bias": np.zeros(4)},
                "dense_1": {"kernel": np.ones((4, 2)), "bias": np.zeros(2)},
            }
        }

        flat = _flatten_pytree(tree)
        assert "encoder.dense_0.kernel" in flat
        assert "encoder.dense_1.bias" in flat
        assert flat["encoder.dense_0.kernel"].shape == (3, 4)

        # Unflatten should reconstruct
        restored = _unflatten_pytree(flat)
        assert "encoder" in restored
        assert "dense_0" in restored["encoder"]
        np.testing.assert_array_equal(
            restored["encoder"]["dense_0"]["kernel"],
            np.ones((3, 4)),
        )

    def test_expert_flat_params(self):
        """JAXExpert should flatten nested params."""
        from model_osmosis.adapters.jax import JAXExpert

        params = {
            "layer_0": {"weight": np.ones((3, 4))},
            "layer_1": {"weight": np.ones((4, 2))},
        }
        expert = JAXExpert(params)
        flat = expert.flat_params
        assert "layer_0.weight" in flat
        assert "layer_1.weight" in flat

    def test_adapter_get(self):
        """get_adapter should return JAXAdapter and TensorFlowAdapter."""
        from model_osmosis import get_adapter
        tf_cls = get_adapter("tensorflow")
        jax_cls = get_adapter("jax")
        assert tf_cls.__name__ == "TensorFlowAdapter"
        assert jax_cls.__name__ == "JAXAdapter"

        # Aliases
        assert get_adapter("tf").__name__ == "TensorFlowAdapter"
        assert get_adapter("flax").__name__ == "JAXAdapter"


# ---------------------------------------------------------------------------
# Pipeline API: compare() and benchmark()
# ---------------------------------------------------------------------------


class TestPipelineCompare:
    def test_compare_identical_models(self):
        model = _make_model(seed=0)
        sample = np.random.rand(20, 6).astype(np.float32)
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        result = pipeline.compare(model, model, sample)

        assert result["cosine_similarity"] > 0.999
        assert result["mse"] < 1e-10
        assert result["correlation"] > 0.999
        assert result["max_error"] < 1e-6
        assert result["num_samples"] == 20

    def test_compare_different_models(self):
        model_a = _make_model(seed=0)
        model_b = _make_model(seed=99)
        sample = np.random.rand(20, 6).astype(np.float32)
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        result = pipeline.compare(model_a, model_b, sample)

        # Different models should have measurably different outputs
        assert result["cosine_similarity"] < 1.0
        assert result["mse"] > 0


class TestPipelineBenchmark:
    def test_benchmark_returns_ranked_results(self, three_models, sample_data):
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        results = pipeline.benchmark(three_models, sample_data)

        assert len(results) > 0
        assert all("strategy" in r for r in results)
        assert all("score" in r for r in results)
        assert all("weights" in r for r in results)

        # Should be sorted by score descending
        scores = [r["score"] for r in results]
        assert scores == sorted(scores, reverse=True)

    def test_benchmark_specific_strategies(self, three_models, sample_data):
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        results = pipeline.benchmark(
            three_models, sample_data,
            strategies=["weighted_average", "slerp"],
        )
        names = [r["strategy"] for r in results]
        assert set(names) == {"weighted_average", "slerp"}


# ---------------------------------------------------------------------------
# Performance: Inplace Merge
# ---------------------------------------------------------------------------


class TestInplaceMerge:
    def test_inplace_merge(self, three_models):
        """Inplace merge should modify the first model."""
        first_id = id(three_models[0])
        merger = MultiModelMerger(strategy="weighted_average", verbose=False)
        result = merger.merge(three_models, inplace=True)

        # Result should BE the first model (same object)
        assert id(result) == first_id

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


# ---------------------------------------------------------------------------
# Version
# ---------------------------------------------------------------------------


class TestVersion:
    def test_version_bump(self):
        import model_osmosis
        assert model_osmosis.__version__ == "0.8.0"
