"""Tests for the router MLP model."""

import pytest
import torch

from partner_inference_harness.model.router_mlp import (
    RouterConfig,
    RouterMLP,
    load_model,
    save_model,
)


class TestRouterConfig:
    """Tests for RouterConfig."""

    def test_default_values(self) -> None:
        """Default config has expected values."""
        config = RouterConfig()
        assert config.input_dim == 256
        assert config.hidden_dim == 2048
        assert config.num_layers == 4
        assert config.num_models == 7
        assert config.dropout == 0.1

    def test_custom_values(self) -> None:
        """Custom config values are preserved."""
        config = RouterConfig(input_dim=512, num_models=10)
        assert config.input_dim == 512
        assert config.num_models == 10


class TestRouterMLP:
    """Tests for RouterMLP model."""

    def test_forward_shape(self) -> None:
        """Forward pass produces correct output shape."""
        config = RouterConfig(input_dim=256, num_models=7)
        model = RouterMLP(config)

        batch_size = 4
        x = torch.randn(batch_size, config.input_dim)
        output = model(x)

        assert output.shape == (batch_size, config.num_models)

    def test_predict_returns_indices_and_confidences(self) -> None:
        """Predict returns valid indices and confidence scores."""
        config = RouterConfig(num_models=5)
        model = RouterMLP(config)

        x = torch.randn(3, config.input_dim)
        indices, confidences = model.predict(x)

        assert indices.shape == (3,)
        assert confidences.shape == (3,)
        assert torch.all(indices >= 0)
        assert torch.all(indices < config.num_models)
        assert torch.all(confidences >= 0)
        assert torch.all(confidences <= 1)

    def test_parameter_count(self) -> None:
        """Model has expected parameter count (roughly 8M)."""
        config = RouterConfig()
        model = RouterMLP(config)
        param_count = model.count_parameters()

        # Should be roughly 8M parameters
        assert 5_000_000 < param_count < 20_000_000

    def test_save_and_load(self, tmp_path) -> None:
        """Model can be saved and loaded."""
        config = RouterConfig(input_dim=128, num_models=3)
        model = RouterMLP(config)
        model.eval()  # Disable dropout for deterministic output

        # Get output before save
        x = torch.randn(2, config.input_dim)
        with torch.no_grad():
            original_output = model(x)

        # Save and load
        model_path = tmp_path / "model.pth"
        save_model(model, model_path)
        loaded_model = load_model(model_path)

        # Output should match
        with torch.no_grad():
            loaded_output = loaded_model(x)
        assert torch.allclose(original_output, loaded_output)

    def test_eval_mode_after_load(self, tmp_path) -> None:
        """Loaded model is in eval mode."""
        model = RouterMLP()
        model_path = tmp_path / "model.pth"
        save_model(model, model_path)

        loaded_model = load_model(model_path)
        assert not loaded_model.training
