"""
JEPA Predictor — the core of VL-JEPA architecture.

Takes visual tokens (from V-JEPA encoder) and query tokens (from BPE tokenizer),
and predicts a dense embedding in the same space as the Y-encoder output.

Key design choices:
- Bidirectional attention (NOT causal) — visual and query tokens attend to each other freely
- Non-autoregressive — single forward pass produces the embedding (no sequential generation)
- Output is a 1536-D embedding vector, not text tokens
"""

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


class JEPAPredictor(nn.Module):
    """
    JEPA Predictor — predicts text embeddings from visual + query tokens.

    Architecture:
        [visual_tokens (576×768), query_tokens (≤512×768)] → N × TransformerBlock(bidirectional) → AvgPool → Linear → L2 normalize

    This is the trainable core. During JEPA pretraining, the X-Encoder may be frozen
    while the predictor learns to map visual+query → text embedding space.

    Args:
        hidden_dim: Transformer dimension (768)
        embed_dim: Output embedding dimension (1536)
        num_heads: Number of attention heads (12)
        num_blocks: Number of transformer blocks (8)
        vocab_size: BPE vocabulary size (for query token embedding)
        max_query_len: Maximum query token length (512)
        dropout: Dropout rate
    """

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

        # Query token embedding (for text queries)
        self.query_embed = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
        self.query_pos = nn.Parameter(torch.randn(1, max_query_len, hidden_dim) * 0.02)

        # Modality type embeddings to distinguish visual vs query tokens
        self.visual_type_embed = nn.Parameter(torch.zeros(1, 1, hidden_dim))
        self.query_type_embed = nn.Parameter(torch.zeros(1, 1, hidden_dim))

        self.embed_dropout = nn.Dropout(dropout)

        # Bidirectional transformer blocks — visual and query attend to each other
        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,
        visual_tokens: torch.Tensor,
        query_ids: torch.Tensor | None = None,
        query_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Args:
            visual_tokens: [batch, num_patches, hidden_dim] — from V-JEPA encoder (576×768)
            query_ids: [batch, query_len] — BPE token IDs for the query (optional for captioning)
            query_padding_mask: [batch, query_len] — True for non-pad positions
        Returns:
            [batch, embed_dim] — L2-normalized predicted embedding (1536-D)
        """
        B = visual_tokens.shape[0]

        # Add visual type embedding
        visual = visual_tokens + self.visual_type_embed  # [B, 576, 768]

        if query_ids is not None:
            Q = query_ids.shape[1]
            # Embed query tokens + positional + type
            query = self.query_embed(query_ids) + self.query_pos[:, :Q, :] + self.query_type_embed
            # Concatenate: [visual_tokens, query_tokens]
            x = torch.cat([visual, query], dim=1)  # [B, 576+Q, 768]
        else:
            x = visual  # [B, 576, 768] — captioning mode (no query)

        x = self.embed_dropout(x)

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

        x = self.norm(x)

        # Average pooling over all non-padding positions
        if query_ids is not None and query_padding_mask is not None:
            # Build combined mask: all visual tokens are valid + query padding mask
            visual_mask = torch.ones(B, visual_tokens.shape[1], device=visual_tokens.device, dtype=torch.bool)
            combined_mask = torch.cat([visual_mask, query_padding_mask], dim=1)  # [B, 576+Q]
            mask = combined_mask.unsqueeze(-1).float()  # [B, 576+Q, 1]
            x = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        else:
            x = x.mean(dim=1)  # [B, 768]

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

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

        return x