# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
# Licensed under Creative Commons Attribution-NonCommercial 4.0 International: free for non-commercial use with attribution to the author.
"""
TetherGate — relative-anchor spring gating that keeps learned gates from collapsing.

A learnable gate scales a residual branch: out = x + g * branch(x). The branch can be an expert (MoE), an adapter
(LoRA, domain adapters), a newly added layer or block, a cross-attention / memory / retrieval path, or any scalar of
the ReZero / LayerScale kind. The gate is driven by a control signal `s` that you choose (any statistic of how much
the branch is needed — or a constant, which turns the gate into an anchored learnable scalar).
Status: experimental — tested in our own training runs, early observations promising; not a validated method.

THE PROBLEM
    When the loss does not "pay" for keeping a branch open — the branch hurts at initialization, or its benefit is
    delayed — the optimizer slowly pushes the gate to zero, and the branch stops contributing to the model.
    Two independent drifts cause it:
      1. early in training, the control signal drifts below a threshold that has not moved yet;
      2. later, the threshold itself drifts upward, because the loss-only equilibrium moves towards "closed".

THE MECHANISM
    g_raw = sigmoid( clamp( beta * (s - theta), arg_min, +8 ) )
    g     = F(t) + (1 - F(t)) * g_raw                  # affine floor, scheduled by the training STEP
    after every optimizer step (outside autograd):
        theta <- theta - k(t) * (theta - theta_star)
        theta_star = EMA(s) - logit(g_target) / beta    # RELATIVE anchor: follows the signal's own drift
    The spring removes the common one-sided drift; the gate's mean openness settles near g_target independently of
    the learning rate (the LR only sets the time constant), while the per-token selectivity driven by `s` stays.
    Gradients through the gate are real in both directions (forward == backward; the wide clamp keeps a live
    gradient even near the floor). Outside the clamp range — beta * (s - theta) below arg_min (default -8) or
    above +8 — the gradient to theta (and to s) is exactly 0; theta then moves only through the spring.

PRACTICAL NOTES
    * g_target: 0.5 is a safe default; use 0.3–0.7 to bias a branch towards closed / open.
    * beta (fixed): 4–8 for a signal of unit scale; normalise `s` (e.g. divide by its running std) if its scale is
      very different — beta * std(s) of about 1–3 keeps the gate selective without saturating.
    * Floor schedule F(t): keep 1.0 (gate fully open, "asleep") until the branch has warmed up — a few hundred to a
      few thousand steps — then let it decay; 'exp' (default) never reaches f_min, 'cos' reaches it at ramp_end.
      Never drive the floor by the loss value.
    * Spring strength: kappa_sleep while the floor is 1.0 centres theta fast; afterwards k = lr * wd removes only the
      common drift. If the signal drifts quickly, schedule k up smoothly — do not jump it.
    * Choosing wd: under a steady closing push the optimizer moves theta by about lr per step and the spring pulls it
      back by lr * wd * (theta - theta_star), so theta settles about 1/wd away from its anchor, whatever the LR. The
      gate's argument then shifts by about beta / wd: keep beta / wd well below 1 (the default wd = 20 with beta = 6
      gives 0.3 and the gate stays near g_target; wd = 0.5 would give 12 and the gate would close to the floor).
      wd = 20 is this file's default (as in the demo); our long production run used wd = 0.5, which leaves
      the loss much more freedom — with a small wd, watch the mean gate opening.
    * Reopening a gate that has closed: use a ramp (raise the floor gradually); a step change on a mature model is a
      loss shock.
    * anchor='global' (one target for all gates) lets layers / experts differentiate through their own signal;
      anchor='local' forces the same openness everywhere.
    * Keep theta in its own parameter group without weight decay (as in the demo below).
    * Monitor the variance of `s`: the anchor keeps the gate open, it cannot restore a signal that lost its
      information.
    * torch.compile: the forward pass has no Python branching on tensors; tethergate_step runs outside the graph.
"""
import math
import torch
import torch.nn as nn


class TetherGate(nn.Module):
    """Gate for one branch (or one layer). Call `gate = g(s)`, then `out = x + gate * branch(x)`.

    s: tensor of any shape (per token / per sample / per layer) or a python float; the gate
       broadcasts like `s`. With a constant s the module reduces to an anchored learnable scalar.
    detach_signal: True (default) treats `s` as data — the branch cannot learn to manipulate
       its own gate signal. Set False only if you deliberately want gradients to flow into `s`.
    """

    def __init__(self, beta=6.0, g_target=0.5, theta_init=0.0, arg_min=-8.0, ema_momentum=0.99,
                 detach_signal=True):
        super().__init__()
        self.detach_signal = bool(detach_signal)
        self.beta = float(beta)                      # fixed
        self.g_target = float(g_target)
        self.arg_min = float(arg_min)
        self.ema_momentum = float(ema_momentum)
        self.theta = nn.Parameter(torch.tensor(float(theta_init)))
        self.register_buffer('floor', torch.tensor(1.0))     # F(t); 1.0 = gate "asleep" (full pass)
        self.register_buffer('ema_s', torch.tensor(float('nan')))

    def forward(self, s):
        s = torch.as_tensor(s, dtype=self.theta.dtype, device=self.theta.device)
        if self.detach_signal:
            s = s.detach()
        if self.training:
            with torch.no_grad():                    # statistics only — no Python branching (torch.compile-friendly)
                m = s.detach().float().mean()
                ema = self.ema_momentum * self.ema_s + (1.0 - self.ema_momentum) * m
                self.ema_s.copy_(torch.where(torch.isnan(self.ema_s), m, ema))
        g_raw = torch.sigmoid((self.beta * (s - self.theta)).clamp(self.arg_min, 8.0))
        return self.floor + (1.0 - self.floor) * g_raw

    def theta_star(self):
        return float(self.ema_s) - math.log(self.g_target / (1.0 - self.g_target)) / self.beta


def floor_schedule(step, ramp_start=3000, f_min=0.05, tau=1800.0, mode='exp', ramp_end=4500):
    """F(t): 1.0 until ramp_start, then decays to f_min. 'exp' approaches f_min asymptotically
    (never reaches it); 'cos' reaches f_min at ramp_end. Driven by the STEP, never by the loss."""
    if step < ramp_start:
        return 1.0
    if mode == 'exp':
        return f_min + (1.0 - f_min) * math.exp(-(step - ramp_start) / tau)
    x = min(1.0, (step - ramp_start) / max(1, ramp_end - ramp_start))
    return f_min + (1.0 - f_min) * 0.5 * (1.0 + math.cos(math.pi * x))


@torch.no_grad()
def tethergate_step(gates, step, lr, wd=20.0, kappa_sleep=5e-3, anchor='global', **floor_kw):
    """Call once per training step, AFTER optimizer.step().
    Sets the floor F(t) and pulls every theta towards its relative anchor.
    k = kappa_sleep while the floor is 1.0 (theta gets no loss gradient then — center it fast),
    afterwards k = lr * wd (slow mode: only the common drift is removed).
    anchor='global': one target from the mean EMA over all gates (keeps differentiation);
    anchor='local' : one target per gate."""
    f = floor_schedule(step, **floor_kw)
    gates = [g for g in gates if not torch.isnan(g.ema_s)]
    for g in gates:
        g.floor.fill_(f)
    if not gates:
        return f
    k = kappa_sleep if f >= 1.0 else lr * wd
    if anchor == 'global':
        t_star = sum(g.theta_star() for g in gates) / len(gates)
        for g in gates:
            g.theta.sub_(k * (g.theta - t_star))
    else:
        for g in gates:
            g.theta.sub_(k * (g.theta - g.theta_star()))
    return f


if __name__ == '__main__':
    # Minimal demo: two layers; each new branch is gated. `s` can be any statistic you choose.
    torch.manual_seed(0)
    d = 16
    blocks = nn.ModuleList(nn.Linear(d, d) for _ in range(2))
    gates = nn.ModuleList(TetherGate(g_target=0.5) for _ in range(2))
    opt = torch.optim.AdamW([{'params': blocks.parameters(), 'weight_decay': 0.1},
                             {'params': gates.parameters(), 'weight_decay': 0.0}], lr=1e-3)
    for step in range(400):
        x = torch.randn(32, d); h = x; opened = []
        for blk, gt in zip(blocks, gates):
            u = blk(h)
            s = u.abs().mean(-1, keepdim=True)        # example signal: branch activity per token
            g = gt(s); opened.append(float(g.detach().mean()))
            h = h + g * u
        loss = (h - x.flip(-1)).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        tethergate_step(gates, step, lr=1e-3, wd=20.0, kappa_sleep=0.05, ramp_start=100, tau=100.0)
        if step % 100 == 99:
            print(f'step {step+1}: floor {float(gates[0].floor):.3f} | mean gate {[round(o, 3) for o in opened]} '
                  f'| theta {[round(g.theta.item(), 3) for g in gates]} | anchor {[round(g.theta_star(), 3) for g in gates]}')