# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Toy collapse benchmark (CPU, seconds). NOT evidence about large models — a reproducible
illustration of the failure mode and of what the spring does about it.

Task: h = x + g * branch(x), target y = x + A x. The branch is needed, but its random init
is harmful, so early on the loss pushes a free gate towards zero; with plain SGD a closed
gate also starves the branch of gradient. Variants:
  open   : g = 1 (no gate)                      — reference
  free   : g = sigmoid(beta * (s - theta)), theta learned, no floor, no spring
  tether : TetherGate (floor 1.0 for `ramp_start` steps, then decay; spring towards g_target)
s = mean |branch(x)| per token (detached), the same for both gated variants.

    python demo_collapse.py [--steps 3000] [--seeds 0 1 2]
"""
import argparse

import torch
import torch.nn as nn

from tethergate import TetherGate, tethergate_step


def run(variant, seed, steps, d=32, lr=0.05, beta=6.0, init_scale=3.0, ramp_start=300, tau=300.0,
        g_target=0.5, wd=0.5):
    torch.manual_seed(seed)
    A = torch.randn(d, d) * 0.5 / d ** 0.5
    branch = nn.Linear(d, d, bias=False)
    with torch.no_grad():
        branch.weight.mul_(init_scale)
    theta = nn.Parameter(torch.tensor(0.0))
    gate = TetherGate(beta=beta, g_target=g_target) if variant == 'tether' else None
    params = list(branch.parameters()) + ([theta] if variant == 'free' else []) + \
        (list(gate.parameters()) if gate is not None else [])
    opt = torch.optim.SGD(params, lr=lr)
    gen = torch.Generator().manual_seed(1000 + seed)
    log = []
    for step in range(steps):
        x = torch.randn(256, d, generator=gen)
        u = branch(x)
        s = u.detach().abs().mean(-1, keepdim=True)
        if variant == 'open':
            g = torch.ones_like(s)
        elif variant == 'free':
            g = torch.sigmoid((beta * (s - theta)).clamp(-8.0, 8.0))
        else:
            g = gate(s)
        loss = ((x + g * u) - (x + x @ A.t())).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if gate is not None:
            tethergate_step([gate], step, lr=lr, wd=wd, kappa_sleep=0.05, ramp_start=ramp_start, tau=tau)
        if step % max(1, steps // 6) == 0 or step == steps - 1:
            log.append((step, loss.item(), g.mean().item()))
    with torch.no_grad():                       # held-out loss on fresh data
        x = torch.randn(4096, d, generator=torch.Generator().manual_seed(77))
        u = branch(x)
        s = u.abs().mean(-1, keepdim=True)
        if variant == 'open':
            g = torch.ones_like(s)
        elif variant == 'free':
            g = torch.sigmoid((beta * (s - theta)).clamp(-8.0, 8.0))
        else:
            gate.eval(); g = gate(s)
        held = float(((x + g * u) - (x + x @ A.t())).pow(2).mean())
    return held, float(g.mean()), log


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--steps', type=int, default=3000)
    ap.add_argument('--seeds', type=int, nargs='+', default=[0, 1, 2])
    ap.add_argument('--verbose', action='store_true')
    a = ap.parse_args()
    print(f'steps={a.steps} seeds={a.seeds}  (held-out MSE; mean gate on held-out data)')
    for variant in ('open', 'free', 'tether'):
        for seed in a.seeds:
            held, gm, log = run(variant, seed, a.steps)
            print(f'  {variant:6s} seed {seed}: held-out MSE {held:.5f} | mean gate {gm:.3f}')
            if a.verbose:
                for st, l, g in log:
                    print(f'           step {st:5d} loss {l:.5f} gate {g:.3f}')