"""
Tests for SafeStacker — guaranteed-better model additions.
"""

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

from model_osmosis.stacking import SafeStacker, StackedModel


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 eval_data():
    return np.random.rand(50, 6).astype(np.float32)


class TestSafeStacker:
    def test_first_model_always_accepted(self, eval_data):
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        model = _make_model(6, 16, 3, seed=0)
        accepted = safe.add(model, name="first")
        assert accepted is True
        assert len(safe.accepted) == 1

    def test_add_multiple_models(self, eval_data):
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        for i in range(3):
            safe.add(_make_model(6, 16, 3, seed=i), name=f"model_{i}")

        # At least 1 accepted (the first one always is)
        assert len(safe.accepted) >= 1
        # Total tried = accepted + rejected
        assert len(safe.accepted) + len(safe.rejected) == 3

    def test_build_returns_stacked_model(self, eval_data):
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        safe.add(_make_model(6, 16, 3, seed=0))
        safe.add(_make_model(6, 32, 3, seed=1))

        result = safe.build()
        assert isinstance(result, StackedModel)

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

    def test_build_without_add_raises(self, eval_data):
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        with pytest.raises(ValueError):
            safe.build()

    def test_score_never_decreases(self, eval_data):
        """The core guarantee: quality never goes down."""
        safe = SafeStacker(eval_data=eval_data, verbose=False)

        # Add multiple models
        for i in range(5):
            safe.add(_make_model(6, 16, 3, seed=i), name=f"m{i}")

        # Scores should be monotonically non-decreasing
        for i in range(1, len(safe._scores)):
            assert safe._scores[i] >= safe._scores[i - 1] - 1e-6, (
                f"Score decreased at step {i}: "
                f"{safe._scores[i-1]:.6f} → {safe._scores[i]:.6f}"
            )

    def test_min_improvement_threshold(self, eval_data):
        """With a high threshold, more models should be rejected."""
        safe_strict = SafeStacker(
            eval_data=eval_data, min_improvement=100.0, verbose=False
        )
        # First is always accepted
        safe_strict.add(_make_model(6, 16, 3, seed=0))
        # With impossibly high threshold, second should be rejected
        safe_strict.add(_make_model(6, 16, 3, seed=1))

        assert len(safe_strict.rejected) >= 1

    def test_summary(self, eval_data):
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        safe.add(_make_model(6, 16, 3, seed=0), name="alpha")
        safe.add(_make_model(6, 32, 3, seed=1), name="beta")

        summary = safe.summary()
        assert "SafeStacker" in summary
        assert "Accepted" in summary

    def test_custom_eval_fn(self, eval_data):
        """User can provide custom evaluation function."""
        def my_scorer(model, data):
            x = torch.from_numpy(data).float()
            with torch.no_grad():
                out = model(x)
            return out.mean().item()

        safe = SafeStacker(
            eval_data=eval_data, eval_fn=my_scorer, verbose=False
        )
        safe.add(_make_model(6, 16, 3, seed=0))
        result = safe.build()
        assert isinstance(result, StackedModel)
