"""Original LoRA teaching layer, Hu et al. (2021). No pretrained weights.

W is frozen. The rank-two branch computes (alpha/r) B A x. Standard zero-B
initialization starts as the base layer; recorded cases then perform toy fitting.
"""
import torch
import torch.nn as nn


# tensorviz: input=1,4,8
class LowRankAdapter(nn.Module):
    def __init__(self):
        super().__init__()
        self.base = nn.Linear(8, 8, bias=False)
        self.base.weight.requires_grad_(False)
        self.down = nn.Linear(8, 2, bias=False)
        self.up = nn.Linear(2, 8, bias=False)
        nn.init.zeros_(self.up.weight)
        self.scale = 2.0 / 2

    def forward(self, x):
        frozen = self.base(x)
        update = self.up(self.down(x)) * self.scale
        return frozen + update


def fitted(steps):
    torch.manual_seed(5)
    model = LowRankAdapter().double()
    x = torch.eye(8, dtype=torch.float64)
    teacher_update = torch.zeros(8, 8, dtype=torch.float64)
    teacher_update[0, 0] = .8
    teacher_update[1, 1] = -.6
    target = x @ (model.base.weight.detach() + teacher_update).T
    optimizer = torch.optim.SGD([model.down.weight, model.up.weight], lr=1.)
    losses = []
    for _ in range(steps):
        optimizer.zero_grad()
        loss = (model(x) - target).square().mean()
        losses.append(loss.item())
        loss.backward()
        optimizer.step()
    return model, x, target, losses


def verify():
    model, x, target, _ = fitted(0)
    frozen = model.base.weight.detach().clone()
    torch.testing.assert_close(model(x), model.base(x))
    (model(x)-target).square().mean().backward()
    assert model.base.weight.grad is None
    assert model.up.weight.grad.norm() > 0
    assert model.down.weight.grad.norm() == 0  # Zero B blocks A's initial gradient.
    model, x, target, losses = fitted(320)
    assert losses[-1] < losses[0] / 10
    torch.testing.assert_close(model.base.weight, frozen, rtol=0, atol=0)
    merged = model.base.weight + model.scale * model.up.weight @ model.down.weight
    torch.testing.assert_close(model(x), x @ merged.T)
    assert torch.linalg.matrix_rank(model.up.weight @ model.down.weight) <= 2
    assert sum(p.numel() for p in model.parameters() if p.requires_grad) == 32
    return ['Zero-B initialization exactly preserves the frozen base output', 'Only B has a nonzero adapter gradient at the first step; the base receives no gradient', 'Toy training reduces fitting loss while base weights remain byte-identical', 'Merged weights reproduce the two-branch output and the update rank is at most two']


def experiment():
    cases = []
    for steps in [0, 1, 40, 320]:
        model, x, target, losses = fitted(steps)
        with torch.no_grad():
            merged = model.base.weight + model.scale * model.up.weight @ model.down.weight
            out = model(x)
            cases.append({'id':f'step-{steps}', 'label':f'After {steps} SGD steps', 'target':'up',
                'note':'A fixed eight-example toy task asks for a rank-two weight change. The base is frozen throughout; only A (8 → 2) and B (2 → 8) are fitted. At initialization B is zero. This records adapter arithmetic, not fine-tuning quality on language tasks.',
                'vectors':[{'label':'Frozen base · first example', 'values':model.base(x)[0].tolist()}, {'label':'Adapter update', 'values':(out-model.base(x))[0].tolist()}, {'label':'Adapted output', 'values':out[0].tolist()}, {'label':'Toy target', 'values':target[0].tolist()}],
                'metrics':[{'label':'Mean squared error · all eight examples', 'value':(out-target).square().mean().item()}, {'label':'Maximum merge error', 'value':(out-x @ merged.T).abs().max().item()}, {'label':'Trainable adapter parameters', 'value':32}, {'label':'Frozen base parameters', 'value':64}]})
    return {'kind':'vectors', 'title':'Fit an update, then merge the branches.', 'description':'Recorded deterministic SGD on a synthetic rank-two task. W + (alpha/r) B A gives the same output after merging, up to floating-point rounding.', 'controlLabel':'Recorded training checkpoint', 'cases':cases}