"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only."""
from __future__ import annotations

import torch
import torch.nn as nn

from v4.classes.accessory.se_block import SEBlock


class FusionBridge(nn.Module):
    """Project N input embeddings to a shared dim, fuse via element-wise product.

    Pure embedding producer — no classification head.  Attach a head stage in
    the pipeline config to produce logits.

    Parameters
    ----------
    input_dims   : ordered list of input embedding dims
    fusion_dim   : projection / output dimension
    use_se       : SE gate on the fused vector
    se_reduction : SE reduction factor
    se_pre_norm  : LayerNorm before each projection; else Identity
    """

    def __init__(
        self,
        input_dims:   list[int],
        fusion_dim:   int  = 256,
        use_se:       bool = True,
        se_reduction: int  = 16,
        se_pre_norm:  bool = True,
    ):
        super().__init__()
        self.out_dim = fusion_dim
        self.W  = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
        self.ln = nn.ModuleList(
            [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
             for _ in input_dims]
        )
        self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None

    def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
        assert len(embeddings) == len(self.W), (
            f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
        )
        h = self.ln[0](self.W[0](embeddings[0]))
        for i in range(1, len(embeddings)):
            h = h * self.ln[i](self.W[i](embeddings[i]))
        if self.se is not None:
            h, _ = self.se(h)
        return h

    def set_phase(self, phase: str) -> None:
        """Freeze bridge during tower_warmup; trainable otherwise."""
        enabled = phase not in ("tower_warmup", "cd_warmup")
        for p in self.parameters():
            p.requires_grad_(enabled)