"""Two Tower Router baseline.

This architecture follows the industry-standard two-tower pattern used by
TikTok, YouTube, Pinterest, and other large-scale recommendation systems.

Architecture:
    Query  -> [Query Tower]  -> query_repr (384-dim)
    Model  -> [Model Tower]  -> model_repr (384-dim)

    Score = dot(query_repr, model_repr) / temperature

Key Features:
    - Separate encoders enable pre-computation of model embeddings
    - Dot product scoring allows ANN retrieval for large catalogs
    - Light-SE blocks learn feature importance
    - Optional fine-grained interaction (FE-Block inspired by ColBERT)
    - Contrastive learning objective for embedding space structure

Advantages:
    - Pre-computable model embeddings (store in Faiss/ScaNN)
    - Sub-millisecond inference via dot product
    - Zero-shot routing to new models (just compute embedding)
    - Interpretable embedding space (visualize with t-SNE/UMAP)
    - Industry-proven at billion-scale

Disadvantages:
    - Dot product limits expressiveness vs cross-attention
    - May underfit for complex query-model interactions
    - Requires careful temperature tuning

When to use:
    - Large model catalog (20+ models)
    - Frequent model additions
    - Latency-critical applications
    - Need for interpretable embeddings

References:
    - Yi et al., "Sampling-Bias-Corrected Neural Modeling for Large Corpus
      Item Recommendations" (RecSys 2019) - YouTube two-tower
    - Huang et al., "Embedding-based Retrieval in Facebook Search" (KDD 2020)
    - IntTower: "IntTower: The Next Generation of Two-Tower Model for
      Pre-Ranking System" (CIKM 2022)
"""

from __future__ import annotations

from typing import Optional

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

from .base import BaseRouter, ResidualBlock, RouterConfig, SEBlock


class TwoTowerRouter(BaseRouter):
    """Two-tower router with separate query and model encoders.

    This is the recommended baseline for most LLM routing scenarios,
    balancing expressiveness with inference efficiency.
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize Two Tower router.

        Args:
            config: RouterConfig with architecture parameters.
        """
        super().__init__(config)

        # Query Tower
        self.query_tower = self._build_tower(
            input_dim=config.query_dim,
            hidden_dim=config.hidden_dim,
            num_layers=config.num_layers,
            dropout=config.dropout,
            use_layer_norm=config.use_layer_norm,
        )

        # Model Tower
        self.model_tower = self._build_tower(
            input_dim=config.model_dim,
            hidden_dim=config.hidden_dim,
            num_layers=config.num_layers,
            dropout=config.dropout,
            use_layer_norm=config.use_layer_norm,
        )

        # Light-SE blocks for feature importance
        self.query_se = None
        self.model_se = None
        if config.use_se_block:
            self.query_se = SEBlock(config.hidden_dim, reduction=4)
            self.model_se = SEBlock(config.hidden_dim, reduction=4)

        # Learnable temperature for calibration
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(config.temperature))
        )

        # Cache for pre-computed model embeddings
        self._cached_model_repr: Optional[torch.Tensor] = None

    def _build_tower(
        self,
        input_dim: int,
        hidden_dim: int,
        num_layers: int,
        dropout: float,
        use_layer_norm: bool,
    ) -> nn.Module:
        """Build an encoder tower.

        Args:
            input_dim: Input embedding dimension.
            hidden_dim: Hidden/output dimension.
            num_layers: Number of layers.
            dropout: Dropout probability.
            use_layer_norm: Whether to use layer normalization.

        Returns:
            Tower module.
        """
        layers = []

        # Input projection
        layers.append(nn.Linear(input_dim, hidden_dim))
        if use_layer_norm:
            layers.append(nn.LayerNorm(hidden_dim))
        layers.append(nn.GELU())
        layers.append(nn.Dropout(dropout))

        # Additional layers with residual connections
        for _ in range(num_layers - 1):
            layers.append(
                ResidualBlock(
                    dim=hidden_dim,
                    hidden_mult=2.0,
                    dropout=dropout,
                    use_layer_norm=use_layer_norm,
                )
            )

        # Final projection with normalization
        layers.append(nn.Linear(hidden_dim, hidden_dim))
        if use_layer_norm:
            layers.append(nn.LayerNorm(hidden_dim))

        return nn.Sequential(*layers)

    def encode_query(self, query_emb: torch.Tensor) -> torch.Tensor:
        """Encode query embeddings through query tower.

        Args:
            query_emb: Query embeddings, shape (batch, query_dim).

        Returns:
            Encoded query representation, shape (batch, hidden_dim).
        """
        query_repr = self.query_tower(query_emb)

        if self.query_se is not None:
            query_repr = self.query_se(query_repr)

        # L2 normalize for cosine similarity
        query_repr = F.normalize(query_repr, p=2, dim=-1)

        return query_repr

    def encode_models(self, model_embs: torch.Tensor) -> torch.Tensor:
        """Encode model embeddings through model tower.

        This can be pre-computed and cached for inference efficiency.

        Args:
            model_embs: Model embeddings, shape (num_models, model_dim).

        Returns:
            Encoded model representations, shape (num_models, hidden_dim).
        """
        model_repr = self.model_tower(model_embs)

        if self.model_se is not None:
            model_repr = self.model_se(model_repr)

        # L2 normalize for cosine similarity
        model_repr = F.normalize(model_repr, p=2, dim=-1)

        # Cache for potential reuse
        self._cached_model_repr = model_repr

        return model_repr

    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute routing scores via dot product.

        Since representations are L2-normalized, this is cosine similarity
        scaled by learned temperature.

        Args:
            query_repr: Encoded queries, shape (batch, hidden_dim).
            model_repr: Encoded models, shape (num_models, hidden_dim).

        Returns:
            Routing scores, shape (batch, num_models).
        """
        # Dot product: (batch, hidden) @ (hidden, num_models) -> (batch, num_models)
        scores = torch.matmul(query_repr, model_repr.T)

        # Apply learned temperature
        temperature = torch.exp(self.log_temperature)
        scores = scores / temperature

        return scores

    def get_cached_model_embeddings(self) -> Optional[torch.Tensor]:
        """Get cached model embeddings for analysis.

        Returns:
            Cached model representations or None if not computed.
        """
        return self._cached_model_repr

    def precompute_model_embeddings(
        self, model_embs: torch.Tensor
    ) -> torch.Tensor:
        """Pre-compute and cache model embeddings.

        Call this once at initialization or when model catalog changes.

        Args:
            model_embs: Model embeddings, shape (num_models, model_dim).

        Returns:
            Pre-computed model representations.
        """
        with torch.no_grad():
            model_repr = self.encode_models(model_embs)
            self._cached_model_repr = model_repr.detach()
        return self._cached_model_repr

    def forward_with_cached_models(
        self, query_emb: torch.Tensor
    ) -> torch.Tensor:
        """Forward pass using cached model embeddings.

        For inference efficiency when model catalog is static.

        Args:
            query_emb: Query embeddings, shape (batch, query_dim).

        Returns:
            Routing scores, shape (batch, num_models).

        Raises:
            ValueError: If model embeddings have not been pre-computed.
        """
        if self._cached_model_repr is None:
            raise ValueError(
                "Model embeddings not cached. Call precompute_model_embeddings first."
            )

        query_repr = self.encode_query(query_emb)
        return self.compute_scores(query_repr, self._cached_model_repr)


class TwoTowerWithInteraction(BaseRouter):
    """Two Tower variant with fine-grained interaction (IntTower-inspired).

    Adds a lightweight cross-attention mechanism between query and model
    representations before the final dot product. This increases expressiveness
    while maintaining most of the efficiency benefits.

    Architecture:
        Query -> [Query Tower] -> query_repr
        Model -> [Model Tower] -> model_repr

        # Fine-grained interaction (FE-Block)
        query_repr' = query_repr + CrossAttn(query_repr, model_repr)

        Score = dot(query_repr', model_repr) / temperature

    This allows the query representation to "see" the model representations
    before scoring, capturing conditional preferences like "for coding tasks,
    I prefer models with long context windows."
    """

    def __init__(self, config: RouterConfig) -> None:
        """Initialize Two Tower with interaction.

        Args:
            config: RouterConfig with architecture parameters.
        """
        super().__init__(config)

        # Base towers
        self.query_tower = self._build_tower(config)
        self.model_tower = self._build_tower(config)

        # SE blocks
        if config.use_se_block:
            self.query_se = SEBlock(config.hidden_dim, reduction=4)
            self.model_se = SEBlock(config.hidden_dim, reduction=4)
        else:
            self.query_se = nn.Identity()
            self.model_se = nn.Identity()

        # Cross-attention for fine-grained interaction
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=config.hidden_dim,
            num_heads=config.num_heads,
            dropout=config.dropout,
            batch_first=True,
        )
        self.attn_norm = nn.LayerNorm(config.hidden_dim)

        # Gating mechanism for residual connection
        self.gate = nn.Sequential(
            nn.Linear(config.hidden_dim * 2, config.hidden_dim),
            nn.Sigmoid(),
        )

        # Temperature
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(config.temperature))
        )

    def _build_tower(self, config: RouterConfig) -> nn.Module:
        """Build encoder tower."""
        return nn.Sequential(
            nn.Linear(config.query_dim, config.hidden_dim),
            nn.LayerNorm(config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            ResidualBlock(config.hidden_dim, dropout=config.dropout),
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.LayerNorm(config.hidden_dim),
        )

    def encode_query(self, query_emb: torch.Tensor) -> torch.Tensor:
        """Encode query through tower and SE block."""
        query_repr = self.query_tower(query_emb)
        query_repr = self.query_se(query_repr)
        return query_repr

    def encode_models(self, model_embs: torch.Tensor) -> torch.Tensor:
        """Encode models through tower and SE block."""
        model_repr = self.model_tower(model_embs)
        model_repr = self.model_se(model_repr)
        return model_repr

    def compute_scores(
        self, query_repr: torch.Tensor, model_repr: torch.Tensor
    ) -> torch.Tensor:
        """Compute scores with fine-grained interaction.

        Args:
            query_repr: Encoded queries, shape (batch, hidden_dim).
            model_repr: Encoded models, shape (num_models, hidden_dim).

        Returns:
            Routing scores, shape (batch, num_models).
        """
        batch_size = query_repr.shape[0]
        num_models = model_repr.shape[0]

        # Expand for cross-attention
        # query: (batch, 1, hidden) attends to models: (batch, num_models, hidden)
        query_for_attn = query_repr.unsqueeze(1)  # (batch, 1, hidden)
        model_for_attn = model_repr.unsqueeze(0).expand(
            batch_size, -1, -1
        )  # (batch, num_models, hidden)

        # Cross-attention: query attends to all models
        attn_output, _ = self.cross_attention(
            query=query_for_attn,
            key=model_for_attn,
            value=model_for_attn,
        )  # (batch, 1, hidden)
        attn_output = attn_output.squeeze(1)  # (batch, hidden)

        # Gated residual connection
        gate_input = torch.cat([query_repr, attn_output], dim=-1)
        gate_weight = self.gate(gate_input)
        query_enhanced = query_repr + gate_weight * attn_output
        query_enhanced = self.attn_norm(query_enhanced)

        # L2 normalize for dot product
        query_enhanced = F.normalize(query_enhanced, p=2, dim=-1)
        model_repr_norm = F.normalize(model_repr, p=2, dim=-1)

        # Dot product scoring
        scores = torch.matmul(query_enhanced, model_repr_norm.T)

        # Temperature
        temperature = torch.exp(self.log_temperature)
        scores = scores / temperature

        return scores


class TwoTowerContrastiveLoss(nn.Module):
    """Contrastive loss for Two Tower training.

    Implements InfoNCE loss with in-batch negatives and optional
    hard negative mining.

    For each query:
        - Positive: The best model (highest ground truth score)
        - Negatives: All other models (in-batch)

    Loss = -log(exp(sim(q, m+) / tau) / sum_m(exp(sim(q, m) / tau)))
    """

    def __init__(
        self,
        temperature: float = 0.07,
        use_hard_negatives: bool = True,
        hard_negative_weight: float = 0.5,
    ) -> None:
        """Initialize contrastive loss.

        Args:
            temperature: Softmax temperature (lower = harder negatives).
            use_hard_negatives: Whether to weight hard negatives more.
            hard_negative_weight: Weight for hard negative term.
        """
        super().__init__()
        self.temperature = temperature
        self.use_hard_negatives = use_hard_negatives
        self.hard_negative_weight = hard_negative_weight

    def forward(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> torch.Tensor:
        """Compute contrastive loss.

        Args:
            scores: Router scores, shape (batch, num_models).
            targets: Ground truth effects, shape (batch, num_models).

        Returns:
            Scalar loss value.
        """
        # Identify positive (best model per query)
        best_idx = targets.argmax(dim=-1)  # (batch,)

        # Scale scores by temperature
        scaled_scores = scores / self.temperature

        # InfoNCE: log_softmax then gather positive
        log_probs = F.log_softmax(scaled_scores, dim=-1)
        positive_log_prob = log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1)

        loss = -positive_log_prob.mean()

        # Optional: Hard negative mining
        if self.use_hard_negatives:
            # Hard negatives: high-scoring models that aren't the best
            probs = F.softmax(scaled_scores, dim=-1)

            # Create mask for non-positive models
            mask = torch.ones_like(probs, dtype=torch.bool)
            mask.scatter_(1, best_idx.unsqueeze(1), False)

            # Weight hard negatives (high prob but wrong)
            hard_neg_probs = probs * mask.float()
            hard_neg_loss = hard_neg_probs.sum(dim=-1).mean()

            loss = loss + self.hard_negative_weight * hard_neg_loss

        return loss


class TwoTowerHybridLoss(nn.Module):
    """Hybrid loss combining contrastive and listwise objectives.

    Combines:
        1. Contrastive (InfoNCE): Pulls query toward best model
        2. Listwise KL: Matches full distribution of model scores
        3. Uniformity: Prevents embedding collapse

    This is the recommended loss for Two Tower training.
    """

    def __init__(
        self,
        contrastive_weight: float = 0.5,
        kl_weight: float = 0.3,
        uniformity_weight: float = 0.2,
        temperature: float = 0.07,
        kl_temperature: float = 1.5,
    ) -> None:
        """Initialize hybrid loss.

        Args:
            contrastive_weight: Weight for InfoNCE term.
            kl_weight: Weight for KL divergence term.
            uniformity_weight: Weight for uniformity regularizer.
            temperature: Temperature for contrastive loss.
            kl_temperature: Temperature for KL soft targets.
        """
        super().__init__()
        self.contrastive_weight = contrastive_weight
        self.kl_weight = kl_weight
        self.uniformity_weight = uniformity_weight
        self.temperature = temperature
        self.kl_temperature = kl_temperature

    def forward(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
        query_repr: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Compute hybrid loss.

        Args:
            scores: Router scores, shape (batch, num_models).
            targets: Ground truth effects, shape (batch, num_models).
            query_repr: Query representations for uniformity (optional).

        Returns:
            Scalar loss value.
        """
        loss = torch.tensor(0.0, device=scores.device)

        # 1. Contrastive (InfoNCE)
        if self.contrastive_weight > 0:
            best_idx = targets.argmax(dim=-1)
            scaled_scores = scores / self.temperature
            log_probs = F.log_softmax(scaled_scores, dim=-1)
            positive_log_prob = log_probs.gather(
                1, best_idx.unsqueeze(1)
            ).squeeze(1)
            contrastive_loss = -positive_log_prob.mean()
            loss = loss + self.contrastive_weight * contrastive_loss

        # 2. Listwise KL divergence
        if self.kl_weight > 0:
            soft_targets = F.softmax(
                targets / self.kl_temperature, dim=-1
            ).clamp_min(1e-8)
            log_probs = F.log_softmax(scores, dim=-1)
            kl_loss = F.kl_div(log_probs, soft_targets, reduction="batchmean")
            loss = loss + self.kl_weight * kl_loss

        # 3. Uniformity regularizer (prevents collapse)
        if self.uniformity_weight > 0 and query_repr is not None:
            # Pairwise squared distances on unit sphere
            query_repr_norm = F.normalize(query_repr, p=2, dim=-1)
            sq_pdist = torch.cdist(query_repr_norm, query_repr_norm, p=2).pow(2)

            # Uniformity: -log(mean(exp(-sq_dist / t)))
            # Lower is more uniform (spread out on hypersphere)
            uniformity = sq_pdist.mul(-2).exp().mean().log()
            loss = loss + self.uniformity_weight * uniformity

        return loss
