"""Router architecture baselines for comparative evaluation.

This module provides three baseline architectures for LLM routing:

1. **GNNRouter**: Graph Neural Network with message passing (current production)
   - Uses heterogeneous bipartite graph (queries ↔ LLMs)
   - GAT attention for message passing
   - Requires full training graph at inference

2. **MultiHeadMLPRouter**: Multi-headed MLP with shared encoder
   - Separate prediction head per model
   - Simple, fast, interpretable
   - Scales linearly with model count

3. **TwoTowerRouter**: Two-tower architecture (TikTok/industry standard)
   - Separate query and model encoders
   - Dot product scoring in shared embedding space
   - Pre-computable model embeddings for fast inference

All baselines implement the same interface for fair comparison:
    - forward(query_emb, model_embs) -> scores
    - encode_query(query_emb) -> query_repr
    - get_routing_signals(scores) -> DifficultySignals

Example:
    >>> from embergraphrouter.model.baselines import TwoTowerRouter, RouterConfig
    >>> config = RouterConfig(query_dim=512, model_dim=512, hidden_dim=384)
    >>> router = TwoTowerRouter(config)
    >>> scores = router(query_emb, model_embs)
    >>> best_model = scores.argmax(dim=-1)
"""

from .base import BaseRouter, DifficultySignals, RouterConfig, RouterOutput
from .evaluation import BaselineEvaluator, EvaluationResult, run_baseline_comparison
from .gnn_baseline import GNNRouter, SimplifiedGNNRouter, create_routing_graph
from .multi_head_mlp import MultiHeadMLPRouter, MultiHeadMLPWithModelFeatures
from .trainer import BaselineTrainer, RouterDataset, TrainerConfig, TrainingMetrics
from .two_tower import (
    TwoTowerContrastiveLoss,
    TwoTowerHybridLoss,
    TwoTowerRouter,
    TwoTowerWithInteraction,
)

__all__ = [
    # Base classes
    "BaseRouter",
    "RouterConfig",
    "RouterOutput",
    "DifficultySignals",
    # Router implementations
    "GNNRouter",
    "SimplifiedGNNRouter",
    "MultiHeadMLPRouter",
    "MultiHeadMLPWithModelFeatures",
    "TwoTowerRouter",
    "TwoTowerWithInteraction",
    # Loss functions
    "TwoTowerContrastiveLoss",
    "TwoTowerHybridLoss",
    # Training
    "BaselineTrainer",
    "TrainerConfig",
    "TrainingMetrics",
    "RouterDataset",
    # Evaluation
    "BaselineEvaluator",
    "EvaluationResult",
    "run_baseline_comparison",
    # Utilities
    "create_routing_graph",
]
