"""Inference-only unified-attention Forge2Vec architecture."""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoConfig, AutoModel, SiglipVisionConfig, SiglipVisionModel


def _mean_pooling(tokens: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    expanded = mask.unsqueeze(-1).to(tokens.dtype)
    return (tokens * expanded).sum(1) / expanded.sum(1).clamp(min=1e-9)


class UnifiedAttentionForge2Vec(nn.Module):
    """One 256-dimensional embedding model for text, metadata, and poster style."""

    def __init__(self, text_config_path: str, max_style_weight: float = 0.10) -> None:
        super().__init__()
        self.transformer = AutoModel.from_config(AutoConfig.from_pretrained(text_config_path))
        self.vision_encoder = SiglipVisionModel(SiglipVisionConfig())
        self.genre_embedding = nn.Embedding(100, 64, padding_idx=0)
        self.style_bottleneck = nn.Sequential(
            nn.Linear(4608, 256), nn.LayerNorm(256), nn.GELU(),
            nn.Dropout(0.20), nn.Linear(256, 64), nn.LayerNorm(64),
        )
        self.unified_description = nn.Sequential(
            nn.Linear(768, 384), nn.LayerNorm(384), nn.GELU(),
            nn.Dropout(0.10), nn.Linear(384, 128),
        )
        self.unified_title = nn.Sequential(nn.Linear(384, 128), nn.LayerNorm(128), nn.GELU())
        self.unified_genre = nn.Sequential(nn.Linear(64, 128), nn.LayerNorm(128), nn.GELU())
        self.unified_metadata = nn.Sequential(
            nn.Linear(4, 64), nn.LayerNorm(64), nn.GELU(), nn.Linear(64, 128),
        )
        self.unified_style = nn.Sequential(nn.Linear(64, 128), nn.LayerNorm(128), nn.GELU())
        self.unified_cls = nn.Parameter(torch.empty(1, 1, 128))
        self.unified_missing = nn.Parameter(torch.empty(5, 128))
        self.unified_type = nn.Parameter(torch.empty(6, 128))
        semantic_layer = nn.TransformerEncoderLayer(
            d_model=128, nhead=8, dim_feedforward=512, dropout=0.10,
            activation="gelu", batch_first=True, norm_first=True,
        )
        self.semantic_attention = nn.TransformerEncoder(
            semantic_layer, num_layers=2, norm=nn.LayerNorm(128)
        )
        fusion_layer = nn.TransformerEncoderLayer(
            d_model=128, nhead=8, dim_feedforward=512, dropout=0.10,
            activation="gelu", batch_first=True, norm_first=True,
        )
        self.fusion_attention = nn.TransformerEncoder(
            fusion_layer, num_layers=2, norm=nn.LayerNorm(128)
        )
        self.image_gate = nn.Sequential(
            nn.Linear(512, 128), nn.LayerNorm(128), nn.GELU(),
            nn.Dropout(0.10), nn.Linear(128, 1),
        )
        self.unified_output = nn.Sequential(nn.Linear(128, 256), nn.LayerNorm(256))
        self.register_buffer("max_style_weight", torch.tensor(float(max_style_weight)))

    def encode_text(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        output = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
        # Training and production inference use SentenceTransformer's raw
        # mean-pooled vectors here. Normalization happens only at the final
        # Forge2Vec output.
        return _mean_pooling(output.last_hidden_state, attention_mask)

    @staticmethod
    def _present_or_missing(token: torch.Tensor, present: torch.Tensor, missing: torch.Tensor) -> torch.Tensor:
        present = present.reshape(-1, 1, 1)
        return token * present + missing.reshape(1, 1, -1) * (1.0 - present)

    def _semantic_tokens(self, ua, en, titles, genres, metadata, mask):
        batch_size = ua.size(0)
        description = self.unified_description(torch.cat((ua, en), -1)).unsqueeze(1)
        title = self.unified_title(titles).unsqueeze(1)
        genre_vectors = self.genre_embedding(genres.long())
        genre_mask = (genres != 0).unsqueeze(-1).to(genre_vectors.dtype)
        genre_pooled = (genre_vectors * genre_mask).sum(1) / genre_mask.sum(1).clamp(min=1.0)
        genre = self.unified_genre(genre_pooled).unsqueeze(1)
        metadata_token = self.unified_metadata(metadata).unsqueeze(1)
        raw = (description, title, genre, metadata_token)
        tokens = [
            self._present_or_missing(token, mask[:, index], self.unified_missing[index])
            for index, token in enumerate(raw)
        ]
        sequence = torch.cat((self.unified_cls.expand(batch_size, -1, -1), *tokens), dim=1)
        sequence = sequence + self.unified_type[:5].unsqueeze(0)
        encoded = self.semantic_attention(sequence)
        return encoded, encoded[:, 0]

    def encode_style(self, pixels: torch.Tensor) -> torch.Tensor:
        output = self.vision_encoder(pixel_values=pixels, output_hidden_states=True)
        statistics = []
        for tokens in (output.hidden_states[4], output.hidden_states[8], output.hidden_states[-1]):
            statistics.extend((tokens.mean(1), tokens.std(1, unbiased=False)))
        return F.normalize(self.style_bottleneck(torch.cat(statistics, -1)), dim=-1)

    def forward(self, ua, en, titles, genres, metadata, modality_mask, poster_pixels=None):
        semantic_sequence, semantic_summary = self._semantic_tokens(
            ua, en, titles, genres, metadata, modality_mask
        )
        available = modality_mask[:, 4:5]
        image_token = self.unified_missing[4].reshape(1, 128).expand(ua.size(0), -1)
        if poster_pixels is not None and bool((available > 0).any()):
            active = available.squeeze(1) > 0
            image_token = image_token.clone()
            image_token[active] = self.unified_style(self.encode_style(poster_pixels[active]))
        gate_features = torch.cat((
            semantic_summary, image_token, torch.abs(semantic_summary - image_token),
            semantic_summary * image_token,
        ), -1)
        gate = torch.sigmoid(self.image_gate(gate_features)) * available * self.max_style_weight
        sequence = torch.cat((semantic_sequence, (image_token * gate).unsqueeze(1)), 1)
        fused = self.fusion_attention(sequence + self.unified_type.unsqueeze(0))
        return F.normalize(self.unified_output(fused[:, 0]), dim=-1)