"""
Tests for optimized stacking (Top-K routing, pruning, confidence, quantization).
"""

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

from model_osmosis.stacking import StackedModel, ModelStacker, StackOptimizer
from model_osmosis import OsmosisPipeline, OsmosisConfig


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_model(in_dim: int, hidden: int, out_dim: int, seed: int = 0):
    """Create a simple model with deterministic weights."""
    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 five_models():
    """Five models with varying hidden sizes."""
    return [_make_model(6, h, 3, seed=i) for i, h in enumerate([8, 16, 32, 16, 8])]


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


# ---------------------------------------------------------------------------
# Top-K Sparse Gating Tests
# ---------------------------------------------------------------------------


class TestTopKGating:
    def test_top_k_2_of_5(self, five_models, sample_data):
        """Only 2 of 5 experts should be activated per input."""
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            five_models, sample_data=sample_data, top_k=2, train_gate=False
        )
        assert result.top_k == 2
        assert result.num_experts == 5

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

    def test_top_k_1(self, five_models, sample_data):
        """Top-1: only 1 expert per input (maximum sparsity)."""
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            five_models, sample_data=sample_data, top_k=1, train_gate=False
        )
        assert result.top_k == 1

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

    def test_top_k_equals_num_experts(self, five_models, sample_data):
        """Top-K = num_experts = dense routing (all experts)."""
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            five_models, sample_data=sample_data, top_k=5, train_gate=False
        )
        assert result.top_k == 5

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

    def test_top_k_clamped(self):
        """Top-K should be clamped to num_experts if too high."""
        models = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(models, top_k=10, train_gate=False)
        assert result.top_k == 2  # clamped to num_experts


# ---------------------------------------------------------------------------
# Confidence Tests
# ---------------------------------------------------------------------------


class TestConfidence:
    def test_confidence_heads_created(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(five_models, sample_data=sample_data, train_gate=False)
        assert result.confidence_heads is not None
        assert len(result.confidence_heads) == 5

    def test_no_confidence(self):
        """Can create without confidence heads."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacked = StackedModel(
            experts=experts, input_dim=6, output_dim=3, use_confidence=False
        )
        assert stacked.confidence_heads is None
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = stacked(x)
        assert out.shape == (5, 3)


# ---------------------------------------------------------------------------
# StackOptimizer Tests
# ---------------------------------------------------------------------------


class TestStackOptimizer:
    def test_score_experts(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        scores = opt.score_experts(stacked, sample_data)

        assert len(scores) == 5
        assert all(0.0 <= s <= 1.0 for s in scores.values())
        # At least one expert should have score 1.0 (max)
        assert max(scores.values()) == 1.0

    def test_prune_weak_experts(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        scores = opt.score_experts(stacked, sample_data)
        pruned = opt.prune_experts(stacked, scores, threshold=0.3)

        # Should have fewer (or equal) experts
        assert pruned.num_experts <= stacked.num_experts
        assert pruned.num_experts >= 1  # always keep at least 1

    def test_prune_keeps_at_least_one(self, five_models, sample_data):
        """Even with threshold=1.0, at least 1 expert survives."""
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        scores = opt.score_experts(stacked, sample_data)
        pruned = opt.prune_experts(stacked, scores, threshold=1.0)

        assert pruned.num_experts >= 1

    def test_train_confidence_heads(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        # Should not raise
        opt.train_confidence_heads(stacked, sample_data, epochs=5)

    def test_full_optimize(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        optimized = opt.optimize(
            stacked, sample_data, prune_threshold=0.05, confidence_epochs=5
        )

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

    def test_quantize_experts(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        stacked = stacker.stack(five_models, sample_data=sample_data, train_gate=False)

        opt = StackOptimizer(verbose=False)
        quantized = opt.quantize_experts(stacked)

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


# ---------------------------------------------------------------------------
# Stacker with optimize=True Tests
# ---------------------------------------------------------------------------


class TestStackerOptimize:
    def test_stack_with_optimize(self, five_models, sample_data):
        """Full pipeline: stack + auto-optimize."""
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            five_models,
            sample_data=sample_data,
            optimize=True,
            prune_threshold=0.05,
        )

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

    def test_usage_stats(self, five_models, sample_data):
        """Test usage statistics tracking."""
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            five_models, sample_data=sample_data, top_k=2, train_gate=False
        )

        # Run some inference
        x = torch.rand(100, 6)
        with torch.no_grad():
            result(x)

        stats = result.usage_stats()
        assert len(stats) == 5
        # With top-2 routing, at least 2 experts should have been used
        used = sum(1 for v in stats.values() if v > 0)
        assert used >= 2

    def test_summary(self, five_models, sample_data):
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(five_models, sample_data=sample_data, train_gate=False)
        summary = result.summary()
        assert "Sparse MoE" in summary
        assert "top-2" in summary
        assert "Active per input" in summary


# ---------------------------------------------------------------------------
# Existing Stacker Tests (backward compatibility)
# ---------------------------------------------------------------------------


class TestBackwardCompatibility:
    def test_original_stack_still_works(self):
        models = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        sample = np.random.rand(50, 6).astype(np.float32)
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(models, sample_data=sample)
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)

    def test_pipeline_stack_still_works(self):
        pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
        models = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        sample = np.random.rand(30, 6).astype(np.float32)
        result = pipeline.stack(models, sample_data=sample)
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 3)
