path: root/sdil/core.py

"""
SDIL -- Somato-Dendritic Innovation Learning.

A non-backprop learning rule inspired by Harnett et al. 2026 (Nature),
"Vectorized instructive signals in cortical dendrites".

The Harnett-specific object studied here is the *somato-dendritic residual*:
the teaching signal is the apical feedback MINUS the part of it that is
predictable from the cell's own somatic activity. Only the innovation drives
plasticity. Apical credit signals and learned feedback from node perturbations
both have substantial prior art; see NOVELTY.md for the exact claim boundary.

For hidden layer l:
    u_l = W_l h_{l-1},  h_l = phi(u_l)              (somatic / basal forward)
    a_l = A_l c_l                                   (apical feedback, c_l = top-down signal)
    ahat_l = P_l h_l                                (predictable baseline: soma -> dendrite regression)
    r_l = a_l - ahat_l                              (INNOVATION = teaching signal)

Three learning rules, all local (no weight transport, no reverse-mode graph):
    dW_l = eta   * (r_l ⊙ phi'(u_l)) h_{l-1}^T      (three-factor: pre * gain * innovation)
    dP_l = eta_P * (a_l - P_l h_l) h_l^T            (identified only on neutral periods)
    dA_l = eta_A * (q_l - r_l) c_l^T                (on perturbation trials; q_l = causal node-pert estimate)

q_l is a node-perturbation estimate of the causal descent direction
-grad_{h_l} L, obtained by occasionally perturbing h_l and watching the loss.
Learning A_l this way follows the learned synthetic-feedback approach of
Lansdell, Prakash & Kording (ICLR 2020). Relative to fixed DFA it supplies an
adaptive feedback backbone; it is not itself claimed as an SDIL novelty.

This module deliberately does NOT use autograd for learning. Autograd is used
ONLY inside probes.py to MEASURE the true gradient for alignment diagnostics.
"""

import math
import torch


# --------------------------------------------------------------------------
# activations
# --------------------------------------------------------------------------
def tanh(u):
    return torch.tanh(u)


def tanh_prime(u):
    t = torch.tanh(u)
    return 1.0 - t * t


_SQRT2 = 2.0 ** 0.5
_INV_SQRT_2PI = 0.3989422804014327


def gelu(u):
    return 0.5 * u * (1.0 + torch.erf(u / _SQRT2))


def gelu_prime(u):
    return 0.5 * (1.0 + torch.erf(u / _SQRT2)) + u * _INV_SQRT_2PI * torch.exp(-0.5 * u * u)


def silu(u):
    s = torch.sigmoid(u)
    return u * s


def silu_prime(u):
    s = torch.sigmoid(u)
    return s * (1.0 + u * (1.0 - s))


def relu(u):
    return torch.relu(u)


def relu_prime(u):
    return (u > 0).to(u.dtype)


ACTS = {"tanh": (tanh, tanh_prime), "gelu": (gelu, gelu_prime),
        "silu": (silu, silu_prime), "relu": (relu, relu_prime)}


# --------------------------------------------------------------------------
# model
# --------------------------------------------------------------------------
class SDILNet:
    """Fully-connected SDIL network.

    Weights/parameters are plain tensors (requires_grad=False). Learning is
    manual and local. The output layer is trained with the exact local delta
    rule (softmax-CE gradient wrt logits IS the local output error), which is
    biologically unproblematic; only hidden layers use the SDIL innovation.

    feedback signal c: by default the output error e = softmax(logits) - onehot,
    broadcast to every hidden layer (dim = n_classes). With P=0 and no traffic,
    learning A_l from node perturbations is the direct learned-feedback method
    of Lansdell et al.; with A_l fixed random it reduces to DFA. SDIL's candidate
    addition is the per-cell innovation under mixed apical traffic.
    """

    def __init__(self, sizes, act="tanh", device="cpu", seed=0,
                 w_scale=1.0, a_scale=1.0, feedback="error", dtype=torch.float32,
                 nuis_rho=0.0, nuis_seed=1234, residual=False,
                 predictor_mode="diagonal", traffic_mode="soma",
                 vectorizer_mode="linear"):
        # sizes = [n_in, n_h1, ..., n_hk, n_out]
        self.sizes = list(sizes)
        self.L = len(sizes) - 1              # number of weight layers
        self.n_classes = sizes[-1]
        self.device = device
        self.dtype = dtype
        self.act, self.act_prime = ACTS[act]
        self.feedback = feedback
        if traffic_mode not in ("none", "soma", "topdown", "mixed"):
            raise ValueError(f"unknown traffic_mode: {traffic_mode}")
        self.traffic_mode = traffic_mode
        # residual (skip) connections on the constant-width hidden blocks: the
        # modern way to train deep nets with NO normalization.
        # h_{l+1}=h_l+alpha*phi(u_l) for the interior blocks (input projection and
        # readout stay plain). alpha=1/sqrt(#hidden) (SkipInit/Fixup/LayerScale)
        # keeps activation variance bounded across depth without any BN.
        self.residual = residual
        n_hidden = self.L - 1
        self.res_alpha = (1.0 / math.sqrt(max(1, n_hidden))) if residual else 1.0
        g = torch.Generator(device="cpu").manual_seed(seed)

        def he(shape, scale, gen=g):
            fan_in = shape[1]
            return (torch.randn(*shape, generator=gen) * (scale / math.sqrt(fan_in))).to(device=device, dtype=dtype)

        # forward weights + bias
        self.W = [he((sizes[i + 1], sizes[i]), w_scale) for i in range(self.L)]
        self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)]

        # Apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l),
        # hidden layers only.  A fixed linear map cannot represent credit that
        # changes across activation regions. Optional zero-initialized gated
        # terms preserve the initial linear model while letting perturbation
        # calibration learn state-dependent feedback locally.
        if vectorizer_mode not in ("linear", "soma_gated", "context_gated"):
            raise ValueError(f"unknown vectorizer_mode: {vectorizer_mode}")
        self.vectorizer_mode = vectorizer_mode
        self.A = [he((sizes[i + 1], self.n_classes), a_scale) for i in range(self.L - 1)]
        if vectorizer_mode == "linear":
            self.A_gate = None
        elif vectorizer_mode == "soma_gated":
            self.A_gate = [torch.zeros_like(weight) for weight in self.A]
        else:
            context_dim = sizes[-2]
            self.A_gate = [torch.zeros((sizes[i + 1], self.n_classes * context_dim),
                                       device=device, dtype=dtype)
                           for i in range(self.L - 1)]
        if predictor_mode not in ("diagonal", "full"):
            raise ValueError(f"unknown predictor_mode: {predictor_mode}")
        self.predictor_mode = predictor_mode

        # ---- soma-predictable nuisance on the apical compartment ----
        # Harnett et al. define each cell's SD residual using a least-squares line
        # between that same cell's somatic and dendritic event magnitudes.  The
        # faithful default is therefore a per-neuron affine predictor, not a dense
        # population map.  "full" preserves the original experimental ablation.
        self.nuis_rho = nuis_rho
        gn = torch.Generator(device="cpu").manual_seed(nuis_seed)
        gc = torch.Generator(device="cpu").manual_seed(nuis_seed + 1)
        if predictor_mode == "diagonal":
            # Positive log-normal slopes model ordinary soma-dendrite coupling.
            self.Bnuis = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gn)).to(
                device=device, dtype=dtype) for i in range(self.L - 1)]
            # A second, independently generated traffic family carries the
            # network's own high-level state back along corresponding feature
            # channels. In constant-width residual networks this models normal
            # contextual/top-down traffic rather than a target constructed from
            # the local predictor. P still sees only the same cell's soma.
            self.Bcontext = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gc)).to(
                device=device, dtype=dtype) for i in range(self.L - 1)]
            self.P = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
                      for i in range(self.L - 1)]
            self.P_bias = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
                           for i in range(self.L - 1)]
        else:
            self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn)
                          for i in range(self.L - 1)]
            context_dim = sizes[-2]
            self.Bcontext = [he((sizes[i + 1], context_dim), 1.0, gen=gc)
                             for i in range(self.L - 1)]
            self.P = [torch.zeros((sizes[i + 1], sizes[i + 1]), device=device, dtype=dtype)
                      for i in range(self.L - 1)]
            self.P_bias = None

        # velocity buffers for momentum on forward weights (optional)
        self.mW = [torch.zeros_like(w) for w in self.W]
        self.mb = [torch.zeros_like(bb) for bb in self.b]

    # ------- forward -------
    def forward(self, x, settle_steps=0, kappa=0.0, c_override=None):
        """Return dict of activations.

        h[0]=x; for l in 1..L-1 hidden tanh; logits = W_L h_{L-1} + b_L (linear).
        If settle_steps>0, run online apical control: after the feedforward pass,
        iterate h_l <- h_l + (1/tau)(-h_l + phi(u_l) + kappa r_l) using the current
        innovation r_l (desired-velocity interpretation). Requires labels via
        c_override or is applied post-hoc by the caller.
        """
        h = [x]
        u = []
        for l in range(self.L):
            pre = h[-1]
            ul = pre @ self.W[l].t() + self.b[l]
            u.append(ul)
            if l < self.L - 1:
                act = self.act(ul)
                # skip on interior blocks only (l>=1 => width->width, dims match)
                if self.residual and l >= 1:
                    h.append(pre + self.res_alpha * act)
                else:
                    h.append(act)
            else:
                h.append(ul)                 # logits (linear output)
        return {"h": h, "u": u}

    def logits(self, x):
        return self.forward(x)["h"][-1]

    # ------- feedback signal -------
    def output_error(self, logits, y_onehot):
        # softmax-CE gradient wrt logits
        p = torch.softmax(logits, dim=1)
        return p - y_onehot

    def apical_traffic(self, l, h_l, context=None):
        """Return ordinary non-teaching traffic in the apical compartment.

        ``soma`` models back-propagating somatic activity / normal
        soma-dendrite coupling. ``topdown`` uses the network's own final hidden
        state, generated independently of P, as contextual feedback. ``mixed``
        combines both at matched total scale. The legacy ``nuis_rho`` argument
        is retained as the traffic amplitude so existing result protocols are
        exactly backward compatible.
        """
        if not self.nuis_rho or self.traffic_mode == "none":
            return torch.zeros_like(h_l)

        if self.predictor_mode == "diagonal":
            soma = self.Bnuis[l] * h_l
        else:
            soma = h_l @ self.Bnuis[l].t()

        if self.traffic_mode == "soma":
            traffic = soma
        else:
            if context is None:
                raise ValueError("topdown apical traffic requires a context state")
            if self.predictor_mode == "diagonal":
                if context.shape[1] != h_l.shape[1]:
                    raise ValueError(
                        "diagonal topdown traffic requires constant hidden width; "
                        f"got context={context.shape[1]} and layer={h_l.shape[1]}")
                topdown = self.Bcontext[l] * context
            else:
                topdown = context @ self.Bcontext[l].t()
            traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2
        return self.nuis_rho * traffic

    def vectorizer(self, l, c, h_l=None, context=None):
        """Return the learned instructional component of apical activity.

        ``soma_gated`` adds ``tanh(h_i) G_i c`` independently in every cell.
        ``context_gated`` adds a linear readout of ``c outer tanh(context)``.
        Bounding the state feature prevents a bilinear positive-feedback loop;
        both terms vanish when c=0, so neutral-period predictor identification
        remains unchanged.
        """
        value = c @ self.A[l].t()
        if self.vectorizer_mode == "linear":
            return value
        if self.vectorizer_mode == "soma_gated":
            if h_l is None:
                raise ValueError("soma-gated vectorizer requires somatic activity")
            return value + torch.tanh(h_l) * (c @ self.A_gate[l].t())
        if context is None:
            raise ValueError("context-gated vectorizer requires a context state")
        features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1)
        return value + features @ self.A_gate[l].t()

    def apical(self, l, c, h_l=None, context=None):
        """a_l = A_l c + ordinary non-teaching apical traffic."""
        a = self.vectorizer(l, c, h_l, context)
        if h_l is not None:
            a = a + self.apical_traffic(l, h_l, context)
        return a

    def baseline(self, l, h_l):
        """ahat_l = P_l h_l."""
        if self.predictor_mode == "diagonal":
            return self.P[l] * h_l + self.P_bias[l]
        return h_l @ self.P[l].t()

    def innovation(self, l, c, h_l, use_residual=True, context=None):
        """r_l = a_l - ahat_l  (or raw a_l if use_residual=False -> ablation)."""
        a = self.apical(l, c, h_l, context)
        if use_residual:
            return a - self.baseline(l, h_l), a
        return a, a


def teaching_signal(net, l, c, h_l, cfg, context=None):
    """Select the signal that drives a hidden layer's local plasticity.

    ``match_innovation_norm`` is a deliberately strong raw-apical control: it
    preserves each sample's raw apical direction exactly while rescaling its
    Euclidean norm to the norm of the corresponding innovation.  Therefore a
    raw-vs-residual difference under this control cannot be attributed merely
    to the nuisance making the teaching vector larger.

    Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic
    probes use one definition of every signal.
    """
    a = net.apical(l, c, h_l, context)
    innovation = a - net.baseline(l, h_l)
    if cfg.use_residual:
        teaching = innovation
    elif cfg.raw_scale_control == "match_innovation_norm":
        raw_norm = a.norm(dim=1, keepdim=True)
        innovation_norm = innovation.norm(dim=1, keepdim=True)
        teaching = a * (innovation_norm / raw_norm.clamp_min(1e-12))
    else:
        teaching = a
    return teaching, a, innovation


# --------------------------------------------------------------------------
# node perturbation: causal estimate of -grad_{h_l} L used to TRAIN A_l
# --------------------------------------------------------------------------
def loss_ce(logits, y):
    return torch.nn.functional.cross_entropy(logits, y, reduction="none")


def node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1):
    """Estimate q_l ~ -grad_{h_l}L for every hidden layer by forward-only
    perturbation of the somatic activity. Returns list over hidden layers of
    (B, n_l) tensors. No autograd.

    We perturb h_l -> h_l + sigma*xi, propagate forward THROUGH THE REST of the
    net (weights only), read the loss change dL, and form the antithetic estimate
        q_l = -((L(+sigma*xi) - L(-sigma*xi)) / (2*sigma)) * xi
    Antithetic pairs (+xi, -xi) cancel the O(sigma^2) even term. Averaging over
    n_dirs independent directions cuts the estimator variance ~1/n_dirs, which
    the predecessor project's failures flagged as the main risk of learning a
    feedback pathway from causal probes.
    """
    with torch.no_grad():
        base = net.forward(x)
        h = base["h"]
        L0 = loss_ce(h[-1], y)                     # (B,)
        qs = []
        for l in range(net.L - 1):                 # hidden layers
            hl = h[l + 1]
            acc = torch.zeros_like(hl)
            for _ in range(n_dirs):
                xi = torch.randn_like(hl)
                dL = _relayer_loss(net, h, l + 1, hl + sigma * xi, y) - L0
                # Correctly-scaled zeroth-order estimator of grad_{h_l}L: with
                # perturbation delta=sigma*xi, ghat = delta*dL/sigma^2 = xi*dL/sigma,
                # so E[ghat]=grad. Dividing by sigma (NOT sigma^2) makes ||q||~||grad||
                # -> the three-factor update is genuine SGD at learning rate eta.
                if antithetic:
                    dL2 = _relayer_loss(net, h, l + 1, hl - sigma * xi, y) - L0
                    coeff = (dL - dL2) / (2.0 * sigma)
                else:
                    coeff = dL / sigma
                acc += -(coeff.unsqueeze(1)) * xi
            qs.append(acc / n_dirs)                 # (B, n_l)
        return qs


def simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2,
                                           antithetic=True, n_dirs=1):
    """Estimate every hidden-layer target with shared forward evaluations.

    Each direction injects an independent Rademacher perturbation at *every*
    hidden activation during one forward pass.  The scalar loss difference is
    demultiplexed locally as q_l = -dL/dsigma * xi_l.  Cross-layer terms add
    variance but vanish in expectation because the xi_l are independent.  This
    changes causal-calibration cost from layerwise O(n_dirs * depth^2) tail
    replays to O(n_dirs * depth) full-network work.
    """
    with torch.no_grad():
        base = net.forward(x)
        h = base["h"]
        batch = x.shape[0]
        xis = [torch.empty((n_dirs,) + hl.shape, device=hl.device, dtype=hl.dtype)
                .bernoulli_(0.5).mul_(2).sub_(1) for hl in h[1:-1]]

        # Fold direction and antithetic sign into the batch dimension.  This is
        # mathematically identical to 2*n_dirs small forwards but turns dozens
        # of tiny Pascal-GPU launches into one efficient matrix multiplication
        # per layer.
        copies = 2 * n_dirs if antithetic else n_dirs
        x_batch = x.repeat(copies, 1)
        y_batch = y.repeat(copies)
        if antithetic:
            noise_batch = [torch.cat((xi, -xi), dim=0).flatten(0, 1) for xi in xis]
        else:
            noise_batch = [xi.flatten(0, 1) for xi in xis]
        losses = _loss_with_hidden_noise(net, x_batch, y_batch, noise_batch, sigma)
        losses = losses.view(copies, batch)
        if antithetic:
            coeff = (losses[:n_dirs] - losses[n_dirs:]) / (2.0 * sigma)
        else:
            coeff = (losses - loss_ce(h[-1], y).unsqueeze(0)) / sigma
        return [-(coeff.unsqueeze(2) * xi).mean(0) for xi in xis]


def _loss_with_hidden_noise(net, x, y, noises, scale):
    """Forward loss with additive interventions after every hidden layer."""
    cur = x
    for l in range(net.L):
        pre = cur
        u = pre @ net.W[l].t() + net.b[l]
        if l < net.L - 1:
            act = net.act(u)
            cur = pre + net.res_alpha * act if net.residual and l >= 1 else act
            cur = cur + scale * noises[l]
        else:
            cur = u
    return loss_ce(cur, y)


def _relayer_loss(net, h, start_idx, h_start_new, y):
    """Recompute loss when hidden activation at layer index `start_idx`
    (1-based into h) is replaced by h_start_new, propagating forward.
    start_idx>=1 so every layer in the tail is an interior/readout layer."""
    cur = h_start_new
    for l in range(start_idx, net.L):
        pre = cur
        ul = pre @ net.W[l].t() + net.b[l]
        if l < net.L - 1:
            act = net.act(ul)
            cur = pre + net.res_alpha * act if net.residual else act    # interior blocks skip (l>=1)
        else:
            cur = ul
    return loss_ce(cur, y)


# --------------------------------------------------------------------------
# the SDIL update
# --------------------------------------------------------------------------
class SDILConfig:
    def __init__(self, eta=0.05, eta_A=0.02, eta_P=0.002,
                 use_residual=True, learn_A=True, learn_P=True,
                 pert_sigma=1e-2, pert_every=5, pert_ndirs=1, momentum=0.0, wd=0.0,
                 settle_steps=0, kappa=0.0, feedback="error", p_update_on_neutral=True,
                 normalize_delta=False, pert_mode="layerwise",
                 raw_scale_control="none", direct_node_pert=False,
                 vectorizer_optimizer="sgd", vectorizer_eps=1e-6):
        self.eta = eta
        self.eta_A = eta_A
        self.eta_P = eta_P
        self.use_residual = use_residual      # ablation: residual vs raw apical
        if raw_scale_control not in ("none", "match_innovation_norm"):
            raise ValueError(f"unknown raw_scale_control: {raw_scale_control}")
        self.raw_scale_control = raw_scale_control
        self.learn_A = learn_A                # ablation: perturbation-trained vs fixed-random A
        self.learn_P = learn_P                # ablation: predictor on/off
        self.pert_sigma = pert_sigma
        self.pert_every = pert_every          # run node-perturbation every k steps
        self.pert_ndirs = pert_ndirs          # directions averaged per perturbation
        if pert_mode not in ("layerwise", "simultaneous"):
            raise ValueError(f"unknown pert_mode: {pert_mode}")
        self.pert_mode = pert_mode
        self.momentum = momentum
        self.wd = wd
        self.settle_steps = settle_steps      # online apical control iterations
        self.kappa = kappa
        self.feedback = feedback              # "error" | "error_deriv"
        self.p_update_on_neutral = p_update_on_neutral
        # node perturbation estimates the descent DIRECTION well but its
        # magnitude is noisy, so A can learn a mis-scaled per-layer step.
        # Normalising each layer's delta to unit RMS decouples step size (set by
        # eta) from that noise -- like normalized SGD.
        self.normalize_delta = normalize_delta
        self.direct_node_pert = direct_node_pert
        if vectorizer_optimizer not in ("sgd", "nlms"):
            raise ValueError(f"unknown vectorizer optimizer: {vectorizer_optimizer}")
        self.vectorizer_optimizer = vectorizer_optimizer
        self.vectorizer_eps = vectorizer_eps


def _update_apical_vectorizer(net, h, c, r_list, qs, cfg):
    """Apply the local causal-regression update to all apical pathways."""
    B = c.shape[0]
    context = h[-2]
    context_features = None
    if net.vectorizer_mode == "context_gated":
        context_features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1)
    for l in range(net.L - 1):
        calibration_error = qs[l] - r_list[l]
        scaled_error = calibration_error
        if cfg.vectorizer_optimizer == "nlms":
            if net.vectorizer_mode == "linear":
                feature_power = c.square().sum(1, keepdim=True)
            elif net.vectorizer_mode == "soma_gated":
                # Each cell i sees [c, tanh(h_i)c], so its effective feature
                # power is ||c||^2 (1+tanh(h_i)^2).
                feature_power = (c.square().sum(1, keepdim=True)
                                 * (1.0 + torch.tanh(h[l + 1]).square()))
            else:
                feature_power = (c.square().sum(1, keepdim=True)
                                 + context_features.square().sum(1, keepdim=True))
            scaled_error = calibration_error / feature_power.clamp_min(cfg.vectorizer_eps)
        dA = scaled_error.t() @ c / B
        net.A[l] += cfg.eta_A * dA
        if net.vectorizer_mode == "soma_gated":
            dgate = (scaled_error * torch.tanh(h[l + 1])).t() @ c / B
            net.A_gate[l] += cfg.eta_A * dgate
        elif net.vectorizer_mode == "context_gated":
            dgate = scaled_error.t() @ context_features / B
            net.A_gate[l] += cfg.eta_A * dgate


def apical_calibration_step(net, x, y, y_onehot, cfg,
                            pert_mode=None, pert_ndirs=None):
    """Calibrate the apical vectorizer while keeping all forward weights fixed.

    This is the feedback-first half of a two-timescale protocol. It uses labels
    and forward-only causal interventions exactly like an ordinary perturbation
    event, but performs no hidden or output weight update. The caller controls
    and accounts for the number of prefix steps and perturbation work.
    """
    if not cfg.learn_A:
        raise ValueError("apical calibration requires learn_A=True")
    if cfg.feedback != "error":
        raise ValueError("feedback-first calibration currently requires error feedback")
    mode = cfg.pert_mode if pert_mode is None else pert_mode
    directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs
    if mode not in ("layerwise", "simultaneous"):
        raise ValueError(f"unknown perturbation mode: {mode}")
    if directions < 1:
        raise ValueError("perturbation directions must be positive")

    with torch.no_grad():
        fwd = net.forward(x)
        h = fwd["h"]
        logits = h[-1]
        loss = loss_ce(logits, y).mean().item()
        c = net.output_error(logits, y_onehot)
        context = h[-2]
        r_list = [teaching_signal(net, l, c, h[l + 1], cfg, context)[0]
                  for l in range(net.L - 1)]
        estimator = (simultaneous_node_perturbation_targets
                     if mode == "simultaneous" else node_perturbation_targets)
        qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=directions)
        _update_apical_vectorizer(net, h, c, r_list, qs, cfg)
    return loss, {"error": c, "did_pert": True}


def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
    """One SDIL training step (minibatch). Returns (loss, aux dict). No autograd."""
    with torch.no_grad():
        B = x.shape[0]
        fwd = net.forward(x)
        h, u = fwd["h"], fwd["u"]
        logits = h[-1]
        loss = loss_ce(logits, y).mean().item()

        e = net.output_error(logits, y_onehot)        # (B, n_classes)
        # feedback signal c
        if cfg.feedback == "error_deriv" and prev_error is not None:
            c = e - prev_error                         # temporal error difference
        else:
            c = e

        # ---- optional online apical control (desired-velocity dynamics) ----
        if cfg.settle_steps > 0:
            h = _settle(net, h, u, c, cfg)

        context = h[-2]

        # ---- compute innovations for hidden layers ----
        r_list, a_list = [], []
        for l in range(net.L - 1):
            r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg, context)
            r_list.append(r)
            a_list.append(a)

        # Calibrate A against the same network state that produced r and c.
        # Measuring q after mutating W would pair a post-update causal target
        # with a pre-update prediction, introducing an avoidable stale-target
        # error (especially at large learning rates).
        did_pert = ((cfg.learn_A or cfg.direct_node_pert)
                    and step % cfg.pert_every == 0)
        qs = None
        if did_pert:
            estimator = (simultaneous_node_perturbation_targets
                         if cfg.pert_mode == "simultaneous"
                         else node_perturbation_targets)
            qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)

        # ================= forward weight updates =================
        # hidden layers: three-factor local rule
        teaching_list = qs if cfg.direct_node_pert else r_list
        if cfg.direct_node_pert and qs is None:
            raise RuntimeError("direct node perturbation requires a target on every update step")
        for l in range(net.L - 1):
            gain = net.act_prime(u[l])                 # phi'(u_l)   (B, n_l)
            delta = teaching_list[l] * gain            # (B, n_l)
            if cfg.normalize_delta:
                delta = delta / (delta.pow(2).mean().sqrt() + 1e-8)
            # For an interior residual block h'=h+alpha*phi(Wh), the local
            # Jacobian dh'/dW contains alpha.  Omitting it preserves direction
            # but makes the effective step grow as 1/alpha=sqrt(depth), which
            # confounds depth-scaling comparisons.  The input projection (l=0)
            # is plain rather than residual and therefore has scale 1.
            if net.residual and l >= 1:
                delta = net.res_alpha * delta
            dW = delta.t() @ h[l] / B                  # (n_l, n_{l-1})
            db = delta.mean(0)
            _apply(net, l, dW, db, cfg)
        # output layer: exact local delta rule (output error is locally available)
        gL = e                                         # grad wrt logits
        dWL = -(gL.t() @ h[net.L - 1]) / B             # descent
        dbL = -gL.mean(0)
        _apply(net, net.L - 1, dWL, dbL, cfg)

        # ================= apical vectorizer A via node perturbation =======
        if did_pert and cfg.learn_A:
            _update_apical_vectorizer(net, h, c, r_list, qs, cfg)

        # ================= predictor P (neutral) ==========================
        # KEY identification condition. P must learn the soma->apical coupling
        # WITHOUT absorbing the teaching signal. On "neutral periods"
        # the apical compartment carries only the soma-predictable nuisance drive
        # (top-down teaching c is absent), so we fit P to the c=0 apical. This
        # makes P -> nuisance map (rho*Bnuis) and leaves the innovation intact.
        # With p_update_on_neutral=False we fit P to the full task-period apical
        # (ablation) -- that path demonstrably eats the teaching signal.
        if cfg.learn_P:
            zero_c = torch.zeros_like(c)
            for l in range(net.L - 1):
                if cfg.p_update_on_neutral:
                    target = net.apical(l, zero_c, h[l + 1], context)  # traffic only
                else:
                    target = a_list[l]                         # full apical (eats signal)
                resid = target - net.baseline(l, h[l + 1])
                _update_predictor(net, l, resid, h[l + 1], cfg.eta_P)

    return loss, {"error": e, "did_pert": did_pert}


def neutral_p_update(net, x, eta):
    """Fit the predictor P to the neutral-period (c=0) apical drive, i.e. the
    soma-predictable nuisance. Used both as a pre-task warmup and (optionally)
    ongoing. Because the nuisance is linear in h, a converged P cancels it for
    any h, so warming P up before task plasticity prevents the nuisance from
    corrupting the forward weights early. No autograd."""
    with torch.no_grad():
        fwd = net.forward(x)
        h = fwd["h"]
        context = h[-2]
        B = x.shape[0]
        zc = torch.zeros(B, net.n_classes, device=x.device, dtype=x.dtype)
        for l in range(net.L - 1):
            hl = h[l + 1]
            target = net.apical(l, zc, hl, context)  # ordinary traffic only (c=0)
            resid = target - net.baseline(l, hl)
            _update_predictor(net, l, resid, hl, eta)


def _update_predictor(net, l, resid, h_l, eta):
    """Local least-squares update for the selected soma-dendrite predictor."""
    if net.predictor_mode == "diagonal":
        # Normalized LMS makes the per-cell slope learning rate independent of
        # that cell's activity variance. Centering separates slope and intercept
        # updates, matching an online least-squares line fit.
        h_centered = h_l - h_l.mean(0)
        resid_centered = resid - resid.mean(0)
        variance = h_centered.pow(2).mean(0)
        net.P[l] += eta * ((resid_centered * h_centered).mean(0) / (variance + 1e-6))
        net.P_bias[l] += eta * resid.mean(0)
    else:
        net.P[l] += eta * (resid.t() @ h_l / h_l.shape[0])


def _settle(net, h, u, c, cfg):
    """Online apical control: tau dh = -h + phi(u) + kappa r, a few Euler steps.
    Updates hidden states in place-ish (returns new h list)."""
    h = list(h)
    dt_over_tau = 1.0 / max(1, cfg.settle_steps)
    for _ in range(cfg.settle_steps):
        context = h[-2]
        for l in range(net.L - 1):
            r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg, context)
            target = net.act(u[l]) + cfg.kappa * r
            h[l + 1] = h[l + 1] + dt_over_tau * (-h[l + 1] + target)
        # recompute downstream pre-activations from settled hidden states
        for l in range(1, net.L):
            u[l] = h[l] @ net.W[l].t() + net.b[l]
            if l < net.L - 1:
                pass  # h[l+1] already being settled; keep u fresh for gains
    return h


def _apply(net, l, dW, db, cfg):
    if cfg.wd:
        dW = dW - cfg.wd * net.W[l]
    if cfg.momentum:
        net.mW[l].mul_(cfg.momentum).add_(dW)
        net.mb[l].mul_(cfg.momentum).add_(db)
        net.W[l] += cfg.eta * net.mW[l]
        net.b[l] += cfg.eta * net.mb[l]
    else:
        net.W[l] += cfg.eta * dW
        net.b[l] += cfg.eta * db