"""Original TensorViz teaching implementation of SwiGLU (Shazeer, 2020).

The bias-free ordinary FFN uses width 24; the gated FFN uses width 16.
Both contain 384 parameters at an input/output width of 8.
"""
import torch
import torch.nn as nn


class PlainFeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.up = nn.Linear(8, 24, bias=False)
        self.relu = nn.ReLU()
        self.down = nn.Linear(24, 8, bias=False)

    def forward(self, x):
        return self.down(self.relu(self.up(x)))


class SwiGLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = nn.Linear(8, 16, bias=False)
        self.silu = nn.SiLU()
        self.up = nn.Linear(8, 16, bias=False)
        self.down = nn.Linear(16, 8, bias=False)

    def forward(self, x):
        gate = self.silu(self.gate(x))
        values = self.up(x)
        gated = gate * values
        return self.down(gated)


# tensorviz: input=1,4,8
class GatedFeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.plain = PlainFeedForward()
        self.swiglu = SwiGLU()

    def forward(self, x):
        plain = self.plain(x)
        gated = self.swiglu(x)
        return torch.stack([plain, gated])


def verify():
    torch.manual_seed(0)
    model = GatedFeedForward().double()
    assert sum(p.numel() for p in model.plain.parameters()) == 384
    assert sum(p.numel() for p in model.swiglu.parameters()) == 384
    x = torch.randn(1, 4, 8, dtype=torch.float64, requires_grad=True)
    model.swiglu(x).square().sum().backward()
    for branch in [model.swiglu.gate, model.swiglu.up, model.swiglu.down]:
        assert torch.isfinite(branch.weight.grad).all() and branch.weight.grad.abs().sum() > 0
    with torch.no_grad():
        before = model.plain(x)
        model.swiglu.gate.weight.zero_()
        torch.testing.assert_close(model.swiglu(x), torch.zeros_like(x))
        torch.testing.assert_close(model.plain(x), before)
    return ["Both comparison branches contain exactly 384 parameters", "Gradients reach the gate, value, and output projections", "A zero gate suppresses SwiGLU without changing the ordinary FFN"]


def experiment():
    torch.manual_seed(0)
    model = GatedFeedForward().double().eval()
    x = torch.tensor([1., -2., .5, 3., -1., .25, 2., -3.], dtype=torch.float64)
    cases = []
    original = model.swiglu.gate.weight.detach().clone()
    for identity, label, scale in [("learned", "Initialized gate", 1), ("zero", "Zero gate weights", 0), ("double", "Double gate weights", 2)]:
        with torch.no_grad():
            model.swiglu.gate.weight.copy_(original * scale)
            gate = model.swiglu.silu(model.swiglu.gate(x))
            values = model.swiglu.up(x)
            product = gate * values
            output = model.swiglu(x)
        cases.append({"id": identity, "label": label, "target": "swiglu", "note": "Only the gate weights change. SiLU gates can be negative and are not probabilities. These initialized weights illustrate arithmetic, not learned language behavior.", "vectors": [
            {"label": "SiLU gate · 16 features", "values": gate.tolist()},
            {"label": "Value branch · 16 features", "values": values.tolist()},
            {"label": "Elementwise product · 16 features", "values": product.tolist()},
            {"label": "Output projection · 8 features", "values": output.tolist()},
        ], "metrics": [{"label": "Ordinary FFN parameters", "value": 384}, {"label": "SwiGLU parameters", "value": 384}]})
    return {"kind": "vectors", "title": "One branch controls another.", "description": "Recorded CPU calculations using the same input and seed-0 initialization. The gate is the only changed branch.", "controlLabel": "Gate configuration", "cases": cases}