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


def _round_to_multiple(x: int, multiple_of: int) -> int:
    multiple_of = int(multiple_of)
    if multiple_of <= 0:
        return int(x)
    x = int(x)
    return int(((x + multiple_of - 1) // multiple_of) * multiple_of)


class VisualSwiGLUFFN(nn.Module):
    """
    Plain pointwise SwiGLU FFN:

        y = W_o( silu(g) * v ),  where [g, v] = W_proj(x)

    No convolutions, no spatial/temporal mixing, no grid reshapes. Operates on the
    last dim only, so it accepts any token layout ([B, L, C], [B, F, HW, C], ...)
    and returns the same shape. Mixing between tokens is left entirely to attention.
    """

    def __init__(
        self,
        dim: int,
        mlp_ratio: float = 4.0,
        multiple_of: int = 16,
        bias: bool = True,
    ):
        super().__init__()
        base_hidden = int(dim * mlp_ratio)          # e.g. 4d
        hidden = int((2.0 / 3.0) * base_hidden)     # 2/3 trick (SwiGLU compute match)
        hidden = _round_to_multiple(hidden, multiple_of)

        self.hidden_dim = hidden
        self.proj = nn.Linear(dim, 2 * hidden, bias=bias)
        self.out = nn.Linear(hidden, dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g, v = self.proj(x).chunk(2, dim=-1)
        return self.out(F.silu(g) * v)