"""Original TensorViz teaching implementation of RMSNorm (Zhang & Sennrich, 2019).

Compare normalization formulas on eight features. No trained weights are used.
"""
import torch
import torch.nn as nn


class RootMeanSquare(nn.Module):
    def forward(self, x):
        return torch.sqrt(x.square().mean(dim=-1, keepdim=True) + 1e-6)


class RMSNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.rms = RootMeanSquare()
        self.gain = nn.Parameter(torch.ones(8))

    def forward(self, x):
        scale = self.rms(x)
        return (x / scale) * self.gain


# tensorviz: input=1,4,8
class NormalizationComparison(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_norm = nn.LayerNorm(8, eps=1e-6)
        self.rms_norm = RMSNorm()

    def forward(self, x):
        centered = self.layer_norm(x)
        rescaled = self.rms_norm(x)
        return torch.stack([centered, rescaled])


def verify():
    model = NormalizationComparison().double().eval()
    x = torch.tensor([[1., 2., 3., 4., -1., -2., -3., -4.]], dtype=torch.float64)
    with torch.no_grad():
        ln, rms = model(x)
        shifted_ln, shifted_rms = model(x + 5)
        torch.testing.assert_close(ln, shifted_ln)
        assert not torch.allclose(rms, shifted_rms)
        torch.testing.assert_close(rms.square().mean(-1), torch.ones(1, dtype=x.dtype), atol=1e-6, rtol=0)
        torch.testing.assert_close(model.rms_norm(x * 10), rms, atol=1e-6, rtol=0)
        torch.testing.assert_close(model.rms_norm(torch.zeros_like(x)), torch.zeros_like(x))
        model.rms_norm.gain.fill_(2)
        torch.testing.assert_close(model.rms_norm(x), rms * 2)
    return ["LayerNorm removes a common offset; RMSNorm preserves its effect", "RMS normalization and positive rescaling agree within epsilon tolerance", "Zero input stays finite and learned gain rescales the result"]


def experiment():
    model = NormalizationComparison().double().eval()
    base = torch.tensor([1., 2., 3., 4., -1., -2., -3., -4.], dtype=torch.float64)
    cases = []
    for identity, label, x, note in [
        ("balanced", "Balanced vector", base, "Both formulas agree for this zero-mean vector when their gains and epsilon match."),
        ("offset", "Add 5 to every feature", base + 5, "LayerNorm removes the common offset. RMSNorm retains a positive mean while rescaling the vector."),
        ("scaled", "Multiply every feature by 10", base * 10, "Both outputs nearly match the balanced case; epsilon introduces a small numerical difference."),
    ]:
        with torch.no_grad():
            ln, rms = model(x)
        cases.append({"id": identity, "label": label, "note": note, "target": "rms_norm", "vectors": [
            {"label": "Input · 8 features", "values": x.tolist()},
            {"label": "LayerNorm", "values": ln.tolist()},
            {"label": "RMSNorm", "values": rms.tolist()},
        ], "metrics": [{"label": "LayerNorm output mean", "value": ln.mean().item()}, {"label": "RMSNorm output mean", "value": rms.mean().item()}]})
    return {"kind": "vectors", "title": "Keep the vector. Change the normalization.", "description": "Recorded CPU calculations with unit gains, zero LayerNorm bias, and epsilon 0.000001. Choose an input transformation.", "controlLabel": "Input transformation", "cases": cases}