# coding=utf-8
"""Dual-Stream Feed-Forward (DSFF).

Implements Section IX of the Wiola paper. Two parallel *dense* streams of
different widths and activations are fused by a learned per-dimension gate:

    Stream A (narrow, SwiGLU):  a = D_A( SiLU(G_A x) * (U_A x) )
    Stream B (wide,   GELU):    b = D_B( GELU(U_B x) )
    gate:                       alpha = sigmoid(W_f [a; b])  in (0,1)^d
    output:                     alpha * a + (1 - alpha) * b

Setting ``W_f = 0`` yields ``alpha = 0.5``, reducing DSFF to a simple ensemble
average; DSFF therefore strictly generalises a two-stream ensemble.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class DualStreamFeedForward(nn.Module):
    def __init__(self, hidden_size: int, narrow_size: int, wide_size: int):
        super().__init__()
        # Stream A: narrow SwiGLU.
        self.gate_a = nn.Linear(hidden_size, narrow_size, bias=False)  # G_A
        self.up_a = nn.Linear(hidden_size, narrow_size, bias=False)  # U_A
        self.down_a = nn.Linear(narrow_size, hidden_size, bias=False)  # D_A
        # Stream B: wide GELU.
        self.up_b = nn.Linear(hidden_size, wide_size, bias=False)  # U_B
        self.down_b = nn.Linear(wide_size, hidden_size, bias=False)  # D_B
        # Per-dimension fusion gate from concatenated stream outputs.
        self.fusion = nn.Linear(2 * hidden_size, hidden_size, bias=False)  # W_f

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a = self.down_a(F.silu(self.gate_a(x)) * self.up_a(x))
        b = self.down_b(F.gelu(self.up_b(x)))
        alpha = torch.sigmoid(self.fusion(torch.cat((a, b), dim=-1)))
        return alpha * a + (1.0 - alpha) * b