"""Router MLP model for model selection.

This module provides a simple multi-layer perceptron for routing queries to
appropriate language models. The architecture is designed for efficient GPU
inference with a compute profile suitable for load testing.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch
import torch.nn as nn

if TYPE_CHECKING:
    from pathlib import Path


@dataclass
class RouterConfig:
    """Configuration for the router MLP.

    Attributes:
        input_dim: Input embedding dimension.
        hidden_dim: Hidden layer dimension.
        num_layers: Number of hidden layers.
        num_models: Number of output models to route to.
        dropout: Dropout probability during training.
    """

    input_dim: int = 256
    hidden_dim: int = 2048
    num_layers: int = 4
    num_models: int = 7
    dropout: float = 0.1


class RouterMLP(nn.Module):
    """Multi-layer perceptron for query-to-model routing.

    A simple feedforward network that takes query embeddings and produces
    routing scores over a set of target models.

    Architecture:
        Linear(input_dim, hidden_dim) -> ReLU -> LayerNorm -> Dropout
        -> [Linear(hidden_dim, hidden_dim) -> ReLU -> LayerNorm -> Dropout] x (num_layers-1)
        -> Linear(hidden_dim, num_models)

    Attributes:
        config: Router configuration parameters.
    """

    def __init__(self, config: RouterConfig | None = None) -> None:
        """Initialize the router MLP.

        Args:
            config: Configuration for the model architecture. Defaults to
                RouterConfig() if not provided.
        """
        super().__init__()
        self.config = config or RouterConfig()

        # Build layers
        layers: list[nn.Module] = []

        # Input projection
        layers.append(nn.Linear(self.config.input_dim, self.config.hidden_dim))
        layers.append(nn.ReLU())
        layers.append(nn.LayerNorm(self.config.hidden_dim))
        layers.append(nn.Dropout(self.config.dropout))

        # Hidden layers
        for _ in range(self.config.num_layers - 1):
            layers.append(nn.Linear(self.config.hidden_dim, self.config.hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.LayerNorm(self.config.hidden_dim))
            layers.append(nn.Dropout(self.config.dropout))

        self.backbone = nn.Sequential(*layers)

        # Output head
        self.head = nn.Linear(self.config.hidden_dim, self.config.num_models)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the router.

        Args:
            x: Input tensor of shape (batch_size, input_dim).

        Returns:
            Logits tensor of shape (batch_size, num_models).
        """
        features = self.backbone(x)
        logits = self.head(features)
        return logits

    def predict(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Predict routing with confidence scores.

        Args:
            x: Input tensor of shape (batch_size, input_dim).

        Returns:
            Tuple of:
                - selected_indices: Tensor of shape (batch_size,) with selected model indices.
                - confidences: Tensor of shape (batch_size,) with confidence scores.
        """
        logits = self.forward(x)
        probabilities = torch.softmax(logits, dim=-1)
        confidences, selected_indices = torch.max(probabilities, dim=-1)
        return selected_indices, confidences

    def count_parameters(self) -> int:
        """Count total trainable parameters.

        Returns:
            Number of trainable parameters in the model.
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def get_memory_footprint_mb(self) -> float:
        """Estimate model memory footprint in MB.

        Returns:
            Estimated memory in megabytes.
        """
        param_bytes = sum(
            p.numel() * p.element_size() for p in self.parameters()
        )
        return param_bytes / (1024 * 1024)


def save_model(model: RouterMLP, path: Path | str) -> None:
    """Save model weights and configuration.

    Args:
        model: RouterMLP instance to save.
        path: File path for the saved model.
    """
    import yaml
    from pathlib import Path

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    # Save weights
    torch.save(
        {
            "state_dict": model.state_dict(),
            "config": {
                "input_dim": model.config.input_dim,
                "hidden_dim": model.config.hidden_dim,
                "num_layers": model.config.num_layers,
                "num_models": model.config.num_models,
                "dropout": model.config.dropout,
            },
        },
        path,
    )

    # Save config as YAML for readability
    config_path = path.with_suffix(".yaml")
    with open(config_path, "w") as f:
        yaml.dump(
            {
                "input_dim": model.config.input_dim,
                "hidden_dim": model.config.hidden_dim,
                "num_layers": model.config.num_layers,
                "num_models": model.config.num_models,
                "dropout": model.config.dropout,
                "parameter_count": model.count_parameters(),
            },
            f,
            default_flow_style=False,
        )


def load_model(path: Path | str, device: str = "cpu") -> RouterMLP:
    """Load model from checkpoint.

    Args:
        path: Path to the saved model checkpoint.
        device: Device to load the model onto.

    Returns:
        Loaded RouterMLP instance.
    """
    from pathlib import Path

    path = Path(path)
    checkpoint = torch.load(path, map_location=device, weights_only=True)

    config = RouterConfig(**checkpoint["config"])
    model = RouterMLP(config)
    model.load_state_dict(checkpoint["state_dict"])
    model.to(device)
    model.eval()

    return model
