"""Compatibility-based routing model.

This module provides a learned compatibility scorer that evaluates
how well a query matches each candidate LLM based on their respective
embeddings.

Architecture Overview:

    Query Text ──→ Embedder ──→ Query Embedding
                                      │
                                      ▼
                              CompatibilityScorer
                                      │
    LLM Descriptions ──→ Embedder ──→ LLM Embeddings
                                      │
                                      ▼
                              Affinity Scores
                                      │
                                      ▼
                              Softmax → Routing

The model learns to predict which LLM will perform best for a given
query by comparing learned representations of the query and each
candidate model's capabilities.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class CompatibilityConfig:
    """Configuration for the compatibility scorer.

    Attributes:
        query_dim: Dimension of query embeddings.
        llm_dim: Dimension of LLM description embeddings.
        hidden_dim: Internal hidden dimension.
        num_layers: Number of hidden layers in the scorer.
        num_llms: Number of target LLMs.
        dropout: Dropout probability.
    """

    query_dim: int = 384
    llm_dim: int = 384
    hidden_dim: int = 256
    num_layers: int = 3
    num_llms: int = 7
    dropout: float = 0.1


class CompatibilityScorer(nn.Module):
    """Learns query-LLM compatibility scores.

    This model takes query embeddings and LLM description embeddings,
    projects them to a common space, and computes compatibility scores
    for routing decisions.

    The scoring mechanism learns which query characteristics match
    which LLM capabilities through supervised training on routing
    performance data.
    """

    def __init__(self, config: CompatibilityConfig | None = None) -> None:
        """Initialize the compatibility scorer.

        Args:
            config: Model configuration. Defaults to CompatibilityConfig().
        """
        super().__init__()
        self.config = config or CompatibilityConfig()

        # Project query to hidden space
        self.query_encoder = self._build_encoder(
            self.config.query_dim,
            self.config.hidden_dim,
            self.config.num_layers,
        )

        # Project LLM descriptions to hidden space
        self.llm_encoder = self._build_encoder(
            self.config.llm_dim,
            self.config.hidden_dim,
            self.config.num_layers,
        )

        # Compatibility scoring head
        self.scorer = nn.Sequential(
            nn.Linear(self.config.hidden_dim * 2, self.config.hidden_dim),
            nn.GELU(),
            nn.Dropout(self.config.dropout),
            nn.Linear(self.config.hidden_dim, 1),
        )

    def _build_encoder(
        self,
        in_dim: int,
        out_dim: int,
        num_layers: int,
    ) -> nn.Module:
        """Build an encoder network."""
        if num_layers <= 1:
            return nn.Linear(in_dim, out_dim)

        layers: list[nn.Module] = []
        curr_dim = in_dim
        hidden = (in_dim + out_dim) // 2

        for _ in range(num_layers - 1):
            layers.extend([
                nn.Linear(curr_dim, hidden),
                nn.LayerNorm(hidden),
                nn.GELU(),
                nn.Dropout(self.config.dropout),
            ])
            curr_dim = hidden

        layers.append(nn.Linear(curr_dim, out_dim))
        return nn.Sequential(*layers)

    def forward(
        self,
        query_embeddings: torch.Tensor,
        llm_embeddings: torch.Tensor,
    ) -> torch.Tensor:
        """Compute compatibility scores.

        Args:
            query_embeddings: Query embeddings (batch_size, query_dim).
            llm_embeddings: LLM embeddings (num_llms, llm_dim).

        Returns:
            Compatibility logits (batch_size, num_llms).
        """
        batch_size = query_embeddings.size(0)
        num_llms = llm_embeddings.size(0)

        # Encode queries and LLMs
        query_h = self.query_encoder(query_embeddings)  # (batch, hidden)
        llm_h = self.llm_encoder(llm_embeddings)  # (num_llms, hidden)

        # Compute pairwise compatibility scores
        query_expanded = query_h.unsqueeze(1).expand(-1, num_llms, -1)
        llm_expanded = llm_h.unsqueeze(0).expand(batch_size, -1, -1)

        # Concatenate and score
        pair_features = torch.cat([query_expanded, llm_expanded], dim=-1)
        logits = self.scorer(pair_features).squeeze(-1)

        return logits

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

        Args:
            query_embeddings: Query embeddings (batch_size, query_dim).
            llm_embeddings: LLM embeddings (num_llms, llm_dim).

        Returns:
            Tuple of:
                - Selected LLM indices (batch_size,)
                - Confidence scores (batch_size,)
        """
        logits = self.forward(query_embeddings, llm_embeddings)
        probabilities = F.softmax(logits, dim=-1)
        confidences, selected = torch.max(probabilities, dim=-1)
        return selected, confidences


def save_model(model: CompatibilityScorer, path: str) -> None:
    """Save model checkpoint."""
    from pathlib import Path

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

    torch.save(
        {
            "state_dict": model.state_dict(),
            "config": {
                "query_dim": model.config.query_dim,
                "llm_dim": model.config.llm_dim,
                "hidden_dim": model.config.hidden_dim,
                "num_layers": model.config.num_layers,
                "num_llms": model.config.num_llms,
                "dropout": model.config.dropout,
            },
        },
        path_obj,
    )


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

    checkpoint = torch.load(Path(path), map_location=device, weights_only=True)
    config = CompatibilityConfig(**checkpoint["config"])
    model = CompatibilityScorer(config)
    model.load_state_dict(checkpoint["state_dict"])
    model.to(device)
    model.eval()
    return model


# Backwards compatibility aliases
HeteroGraphConfig = CompatibilityConfig
HeteroGraphRouter = CompatibilityScorer
