"""
Y-Encoder — Text Encoder for producing target embeddings during JEPA training.

The Y-Encoder takes tokenized text (answers/captions) and produces a dense
embedding in the same space as the predictor output. During training, the
InfoNCE loss aligns predictor embeddings with Y-encoder embeddings.

Trained with a slow learning rate (0.05x multiplier) to provide stable targets.
"""

import torch
import torch.nn as nn
from model.transformer import TransformerBlock


class YEncoder(nn.Module):
    """
    Text encoder that produces target embeddings for contrastive JEPA training.

    Architecture:
        Token IDs → Embedding → + Positional → N × TransformerBlock → AvgPool → Linear → L2 normalize

    Output: 1536-D normalized embedding vector.

    Args:
        vocab_size: BPE vocabulary size (8192)
        hidden_dim: Transformer dimension (768)
        embed_dim: Output embedding dimension (1536)
        num_heads: Number of attention heads (12)
        num_blocks: Number of transformer blocks (6)
        max_seq_len: Maximum sequence length (512)
        dropout: Dropout rate
    """

    def __init__(
        self,
        vocab_size: int = 8192,
        hidden_dim: int = 768,
        embed_dim: int = 1536,
        num_heads: int = 12,
        num_blocks: int = 6,
        max_seq_len: int = 512,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embed_dim = embed_dim

        # Token and position embeddings
        self.token_embed = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
        self.pos_embed = nn.Parameter(torch.randn(1, max_seq_len, hidden_dim) * 0.02)

        self.embed_dropout = nn.Dropout(dropout)

        # Bidirectional transformer blocks
        self.blocks = nn.ModuleList([
            TransformerBlock(hidden_dim, num_heads, dropout, mode="bidirectional")
            for _ in range(num_blocks)
        ])

        self.norm = nn.LayerNorm(hidden_dim)

        # Project to embedding space
        self.proj = nn.Linear(hidden_dim, embed_dim)

    def forward(self, token_ids: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor:
        """
        Args:
            token_ids: [batch, seq_len] — BPE token IDs
            padding_mask: [batch, seq_len] — True for non-pad positions
        Returns:
            [batch, embed_dim] — L2-normalized text embedding (1536-D)
        """
        B, T = token_ids.shape

        # Token + positional embedding
        x = self.token_embed(token_ids) + self.pos_embed[:, :T, :]
        x = self.embed_dropout(x)

        # Pass through transformer blocks
        for block in self.blocks:
            x = block(x)

        x = self.norm(x)

        # Average pooling over non-padding positions
        if padding_mask is not None:
            mask = padding_mask.unsqueeze(-1).float()  # [B, T, 1]
            x = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)  # [B, hidden_dim]
        else:
            x = x.mean(dim=1)  # [B, hidden_dim]

        # Project to embedding space
        x = self.proj(x)  # [B, embed_dim]

        # L2 normalize
        x = nn.functional.normalize(x, p=2, dim=-1)

        return x