"""Xavante - fusion_layer.py - Fusao multimodal via projecao isometrica."""
from __future__ import annotations

import logging

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


class FusionLayer(nn.Module):
    """
    Fusao de N modalidades via:
    1. Projecao isometrica (preserva norma)
    2. Gated combination
    """

    def __init__(self, d_model: int, n_modalities: int = 4):
        super().__init__()
        self.n_modalities = n_modalities
        self.gates = nn.ModuleList(
            [nn.Linear(d_model * n_modalities, 1) for _ in range(n_modalities)]
        )
        self.proj = nn.Linear(d_model, d_model, bias=False)
        # Isometria: inicializar proj como ortogonal
        nn.init.orthogonal_(self.proj.weight)

    def forward(self, modality_embs: list) -> torch.Tensor:
        """modality_embs: lista de [B, d_model] (n_modalidades deles)."""
        if len(modality_embs) != self.n_modalities:
            # Pad with zeros if missing
            while len(modality_embs) < self.n_modalities:
                modality_embs.append(torch.zeros_like(modality_embs[0]))
            modality_embs = modality_embs[: self.n_modalities]
        cat = torch.cat(modality_embs, dim=-1)  # [B, n*d]
        # Gates
        gate_scores = [torch.sigmoid(g(cat)) for g in self.gates]
        # Stack: [n, B, 1] -> [B, n, 1]
        gates_stacked = torch.cat(gate_scores, dim=-1)  # [B, n]
        # Normalize via softmax
        gates_norm = torch.softmax(gates_stacked, dim=-1)
        # Weighted sum
        stacked = torch.stack(modality_embs, dim=1)  # [B, n, d]
        fused = (stacked * gates_norm.unsqueeze(-1)).sum(dim=1)  # [B, d]
        # Isometric projection
        return self.proj(fused)


__all__ = ["FusionLayer"]