"""Original TensorViz comparison of normalization placement (Xiong et al., 2020).

F is one bias-free linear layer so the identity path is easy to isolate. Both
branches start with identical F weights; this is not a training benchmark.
"""
import torch
import torch.nn as nn


class PreNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = nn.LayerNorm(8, eps=1e-6)
        self.branch = nn.Linear(8, 8, bias=False)

    def forward(self, x):
        normalized = self.norm(x)
        learned = self.branch(normalized)
        return x + learned


class PostNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.branch = nn.Linear(8, 8, bias=False)
        self.norm = nn.LayerNorm(8, eps=1e-6)

    def forward(self, x):
        learned = self.branch(x)
        residual = x + learned
        return self.norm(residual)


# tensorviz: input=1,4,8
class NormPlacement(nn.Module):
    def __init__(self):
        super().__init__()
        self.pre = PreNorm()
        self.post = PostNorm()
        self.post.branch.load_state_dict(self.pre.branch.state_dict())

    def forward(self, x):
        pre = self.pre(x)
        post = self.post(x)
        return torch.stack([pre, post])


def verify():
    torch.manual_seed(0)
    model = NormPlacement().double()
    torch.testing.assert_close(model.pre.branch.weight, model.post.branch.weight)
    x = torch.arange(1, 9, dtype=torch.float64).requires_grad_()
    with torch.no_grad():
        model.pre.branch.weight.zero_()
        model.post.branch.weight.zero_()
    pre, post = model(x)
    torch.testing.assert_close(pre, x)
    torch.testing.assert_close(post, model.post.norm(x))
    assert not torch.allclose(pre, post)
    gradient = torch.autograd.grad(pre.sum(), x)[0]
    torch.testing.assert_close(gradient, torch.ones_like(x))
    return ["Both placements use identical branch weights", "Zeroing F leaves pre-norm's identity path intact but post-norm still normalizes", "The zero-branch pre-norm path has the identity derivative"]


def experiment():
    torch.manual_seed(0)
    model = NormPlacement().double().eval()
    original = model.pre.branch.weight.detach().clone()
    x = torch.arange(1, 9, dtype=torch.float64)
    cases = []
    for identity, label, strength, note in [
        ("same-weights", "Same initialized F", 1, "Both branches have the same F weights and LayerNorm parameters. Only the placement changes: x + F(Norm(x)) versus Norm(x + F(x))."),
        ("zero-branch", "Set F to zero", 0, "The pre-norm residual is exactly x. Post-norm still transforms x through LayerNorm. This isolates the identity path without making a training-stability claim."),
    ]:
        with torch.no_grad():
            model.pre.branch.weight.copy_(original * strength)
            model.post.branch.weight.copy_(original * strength)
            pre, post = model(x)
        cases.append({"id": identity, "label": label, "note": note, "target": "pre", "vectors": [
            {"label": "Input", "values": x.tolist()},
            {"label": "Pre-norm output", "values": pre.tolist()},
            {"label": "Post-norm output", "values": post.tolist()},
        ], "metrics": [{"label": "Pre-norm distance from input", "value": (pre-x).norm().item()}, {"label": "Post-norm distance from input", "value": (post-x).norm().item()}]})
    return {"kind": "vectors", "title": "Move the normalization. Keep F fixed.", "description": "Recorded CPU calculations with the same eight-feature input, matching branch weights, and matching LayerNorm parameters.", "controlLabel": "Residual branch", "cases": cases}