"""
Tests for model stacking (MoE).
"""

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

from model_osmosis.stacking import StackedModel, ModelStacker
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


# ---------------------------------------------------------------------------
# StackedModel Tests
# ---------------------------------------------------------------------------


class TestStackedModel:
    def test_forward_same_arch(self):
        """Stack 3 same-architecture models."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        stacked = StackedModel(
            experts=experts,
            input_dim=6,
            output_dim=3,
        )
        x = torch.rand(10, 6)
        out = stacked(x)
        assert out.shape == (10, 3)

    def test_forward_different_hidden(self):
        """Stack models with different hidden sizes but same I/O."""
        experts = [
            _make_model(6, 8, 3, seed=0),
            _make_model(6, 32, 3, seed=1),
            _make_model(6, 64, 3, seed=2),
        ]
        stacked = StackedModel(
            experts=experts,
            input_dim=6,
            output_dim=3,
        )
        x = torch.rand(10, 6)
        out = stacked(x)
        assert out.shape == (10, 3)

    def test_forward_different_output_dims(self):
        """Stack models with different output sizes — projectors adjust."""
        experts = [
            _make_model(6, 16, 2, seed=0),
            _make_model(6, 16, 5, seed=1),
            _make_model(6, 16, 3, seed=2),
        ]
        stacked = StackedModel(
            experts=experts,
            input_dim=6,
            output_dim=4,
            expert_output_dims=[2, 5, 3],
        )
        x = torch.rand(10, 6)
        out = stacked(x)
        assert out.shape == (10, 4)

    def test_forward_different_input_dims(self):
        """Stack models with different input sizes — projectors adjust."""
        experts = [
            _make_model(4, 16, 3, seed=0),
            _make_model(8, 16, 3, seed=1),
        ]
        stacked = StackedModel(
            experts=experts,
            input_dim=6,
            output_dim=3,
            expert_input_dims=[4, 8],
        )
        x = torch.rand(10, 6)
        out = stacked(x)
        assert out.shape == (10, 3)

    def test_preserves_all_params(self):
        """Total params should be sum of all experts + gate + projectors."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        expert_total = sum(sum(p.numel() for p in e.parameters()) for e in experts)

        stacked = StackedModel(experts=experts, input_dim=6, output_dim=3)
        # Stacked should have MORE params (experts + gate)
        assert stacked.total_params() > expert_total

    def test_summary(self):
        experts = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacked = StackedModel(experts=experts, input_dim=6, output_dim=3)
        summary = stacked.summary()
        assert "expert_0" in summary
        assert "expert_1" in summary
        assert "Total" in summary

    def test_expert_params(self):
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        stacked = StackedModel(
            experts=experts,
            input_dim=6,
            output_dim=3,
            expert_names=["alpha", "beta", "gamma"],
        )
        ep = stacked.expert_params()
        assert len(ep) == 3
        assert ep[0][0] == "alpha"
        assert all(count > 0 for _, count in ep)


# ---------------------------------------------------------------------------
# ModelStacker Tests
# ---------------------------------------------------------------------------


class TestModelStacker:
    def test_stack_same_arch(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)

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

    def test_stack_different_architectures(self):
        """Stack models with different sizes — the key use case."""
        models = [
            _make_model(6, 8, 3, seed=0),
            _make_model(6, 32, 3, seed=1),
            _make_model(6, 64, 3, seed=2),
        ]
        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_stack_with_names(self):
        models = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        sample = np.random.rand(30, 6).astype(np.float32)

        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            models,
            sample_data=sample,
            expert_names=["code_model", "math_model"],
        )
        ep = result.expert_params()
        assert ep[0][0] == "code_model"
        assert ep[1][0] == "math_model"

    def test_stack_without_gate_training(self):
        models = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacker = ModelStacker(verbose=False)
        result = stacker.stack(models, train_gate=False)

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

    def test_originals_not_mutated(self):
        """Stacking should deep-copy experts, not mutate originals."""
        models = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        original_params = [
            {n: p.clone() for n, p in m.named_parameters()} for m in models
        ]

        stacker = ModelStacker(verbose=False)
        result = stacker.stack(models)

        # Original models should be unchanged
        for model, orig in zip(models, original_params):
            for name, param in model.named_parameters():
                torch.testing.assert_close(param, orig[name])


# ---------------------------------------------------------------------------
# Pipeline.stack() Tests
# ---------------------------------------------------------------------------


class TestPipelineStack:
    def test_pipeline_stack(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)
        assert result.total_params() > sum(
            sum(p.numel() for p in m.parameters()) for m in models
        )
