"""Projecteurs multimodaux : alignent image/audio sur l'espace texte (Phases 4-8).

Idée : les encodeurs image/audio produisent des features de dimension propre. Pour
que le décodeur texte puisse les lire, un petit **MLP** projette ces features dans
la dimension du texte (`text_dim`). On insère alors ces vecteurs dans la séquence
de tokens texte, encadrés par des tokens de modalité (`<|image_start|>` …).

Choix d'un MLP (2 couches + GELU) plutôt qu'une projection linéaire : il apprend
une transformation non linéaire, ce qui aide l'alignement, pour un coût modeste.

Justification complète : docs/architecture.md §5.
"""

from __future__ import annotations

import torch
import torch.nn as nn


class MLPProjector(nn.Module):
    """MLP ``in_dim -> hidden -> out_dim`` avec GELU et LayerNorm optionnelle."""

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int | None = None,
        dropout: float = 0.0,
        use_norm: bool = True,
    ):
        super().__init__()
        hidden_dim = hidden_dim or out_dim
        layers: list[nn.Module] = []
        if use_norm:
            layers.append(nn.LayerNorm(in_dim))
        layers += [
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, out_dim),
        ]
        if dropout > 0:
            layers.append(nn.Dropout(dropout))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def build_projector(in_dim: int, text_dim: int, hidden_dim: int | None = None,
                    dropout: float = 0.0) -> MLPProjector:
    """Fabrique un projecteur vers l'espace texte (helper)."""
    return MLPProjector(in_dim, text_dim, hidden_dim=hidden_dim, dropout=dropout)