"""
THE TESTBED.

Everything measured so far has run on models small enough to hold in a
sentence, and the findings that matter most are corrections to how those
models were built rather than to the framework itself. This trains the
model those corrections imply, once, and saves it so that every member,
borrowing and marketplace question afterwards costs a ridge solve rather
than a training run.

SIX MEASUREMENTS DECIDE THE SPECIFICATION.

  POOL BEFORE THE HEAD. Without it a head reads every hidden unit and is
  larger than the entire folded body, which inflated every storage and
  arithmetic ratio in the programme and made depth look useless because
  the readout was doing depth's work. Here the head reads 192 numbers
  rather than 3,072.

  NARROW FINAL FEATURES, and this is the real driver rather than a storage
  optimisation. A member costs (features x classes) and a TRANSLATION
  between two bases costs (features x features). At 3,136 features a
  translation was 9.8 million unknowns fitted on fifty examples, which is
  why borrowing behaved erratically. At 192 it is 37 thousand, and a few
  hundred examples determine it.

  SPATIAL DOWNSAMPLING, which is the piece that does not exist yet. A fold
  as written maps a g x g grid to a g x g grid, so hidden width stays at
  channels x 1024 and the dense weight matrix is impossible at 32 x 32.
  A strided fold is the same construction with the output grid coarser
  than the input: an output position reads a window centred at TWICE its
  own coordinates. The value index is unchanged in meaning, so everything
  downstream still applies.

  CLASSES HELD OUT FROM THE START. Every member result needs a base that
  has genuinely never seen something. Twenty of CIFAR-100's hundred are
  reserved and no base ever touches them, which also makes confusable
  groups natural rather than contrived.

  TWO BASES, trained on DISJOINT halves with their own initialisations.
  Every borrowing and marketplace question needs a pair, and training them
  together now costs one run rather than two later.

  DETERMINISTIC THROUGHOUT, saved with the index specification and a hash
  of the values, so every future comparison inherits reproducibility rather
  than the 0.002 floor that silently swallowed three earlier results.

WHAT IS SAVED. Both bases' values, scales and shifts; the specification
needed to rebuild their indices exactly; cached features for the held-out
classes on both bases, so member work needs no forward pass at all; and a
hash of each base for content addressing. Nothing downstream needs this
script again.
"""

import numpy as np
import time
import json
import os
import hashlib

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def strided_window(g_in, c_in, k, c_out, stride):
    """A 3x3 convolution that also DOWNSAMPLES.

    An output position on the coarser grid reads a window centred at
    stride times its own coordinates on the finer one. The value index is
    still (input channel, output channel, tap), so a value means exactly
    what it meant before and every result about partitions carries over
    unchanged. Only the shape of the map changes.

    Returns the flattened index, the number of distinct values, the output
    width, and the multiplies a materialised forward pass costs."""
    g_out = g_in // stride
    ni, no = c_in*g_in*g_in, c_out*g_out*g_out
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g_in*g_in), ii % (g_in*g_in)
    co, po = jj // (g_out*g_out), jj % (g_out*g_out)
    # the centre of the window this output position reads
    cr, cc = (po // g_out)*stride, (po % g_out)*stride
    dr = pi // g_in - (cr - k//2)
    dc = pi % g_in - (cc - k//2)
    inside = (dr >= 0) & (dr < k) & (dc >= 0) & (dc < k)
    K = c_in*c_out*k*k + 1
    idx = np.where(inside, (ci*c_out + co)*k*k + dr*k + dc, K-1)
    return (idx.ravel().astype(np.int32), K, no,
            int(inside.sum())//1)          # real multiplies, not ni*no


_FIXED = {}


class FixedScatter:
    """Fixed-order accumulation, so a rerun reproduces exactly.

    Rectangle for groups near the median; any oversized group summed over
    its own contiguous slice, because a convolution index is wildly skewed:
    every out-of-window position lands in one padding group."""

    def __init__(self, idx, K, cap=8192):
        h = to_host(idx).astype(np.int64).reshape(-1)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        big = np.where(counts > cap)[0]; small = np.where(counts <= cap)[0]
        self.K = K
        self.order = to_dev(order, np.int64) if _GPU else order
        self.big = [(int(b), int(starts[b]), int(starts[b]+counts[b]))
                    for b in big]
        self.small = to_dev(small, np.int64) if _GPU else small
        self.width = int(counts[small].max()) if len(small) else 0
        if self.width:
            pos = np.concatenate([np.arange(counts[s]) for s in small])
            src = np.concatenate([np.arange(starts[s], starts[s]+counts[s])
                                  for s in small])
            row = np.repeat(np.arange(len(small)), counts[small])
            self.src = to_dev(src, np.int64) if _GPU else src
            sl = row*self.width + pos
            self.slot = to_dev(sl, np.int64) if _GPU else sl
            self.buf = xp.zeros(len(small)*self.width, DT)
        self._keep = idx

    def __call__(self, g):
        gs = g.reshape(-1)[self.order]
        out = xp.zeros(self.K, DT)
        if self.width:
            self.buf[:] = 0
            self.buf[self.slot] = gs[self.src]
            out[self.small] = self.buf.reshape(-1, self.width).sum(1)
        for b, a, z in self.big:
            out[b] = gs[a:z].sum()
        return out


def scatter(dW, idx, K):
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


CFG = dict(
    # 32 -> 16 -> 8 -> 4, channels rising, then pooled to 192
    grid=32, c_in=3, stack=((48, 2), (96, 2), (192, 2)), fast=True,
    n_classes=100, n_held=20,
    batch=128, lr=1e-3, epochs=150, augment=True, aug_pad=4,
    # COSINE DECAY. At a constant rate the first run stopped improving at
    # epoch 100 while train loss kept falling, 0.5047 -> 0.5068 over the
    # last fifty epochs against 1.42 -> 1.16 on train. That is fifty
    # epochs of overfitting, and a decaying rate spends them refining
    # instead. The schedule sweep has been queued since before this
    # existed and still has not run; this is the obvious half of it.
    schedule="cosine",
    # A THIRD BASE ON ALL THE DATA, as a strength reference. Bases 0 and 1
    # stay disjoint, which is what every borrowing question needs. Base 2
    # sees both halves and therefore OVERLAPS both, so it must never be
    # used as a counterparty: it exists only to say how much of the
    # accuracy is the architecture and how much is the split.
    # FOUR BASES, and which pair to use depends on the question.
    #   0 and 1  disjoint halves, same architecture. The CONSERVATIVE
    #            counterparty pair: two parties who never shared data,
    #            which is the realistic marketplace case and which cannot
    #            flatter a transfer result by having learned from the same
    #            examples.
    #   2        all the data, same architecture. A STRENGTH REFERENCE, and
    #            never a counterparty: it overlaps both halves.
    #   3        all the data, HALF THE CHANNELS. A strong but SMALLER
    #            counterparty, which is what a marketplace actually looks
    #            like and which §4.9 of the paper lists as untested: every
    #            transfer measured so far has been between bases of
    #            identical shape.
    # Base 3 shares data with 0, 1 and 2, so a transfer measured against it
    # is an upper bound rather than the conservative figure.
    seeds=(0, 1, 2, 3), report_every=10,
    # SAVE SOMEWHERE THAT SURVIVES. /content is wiped when a session ends,
    # and this is the one artefact everything afterwards reads: four bases,
    # their cached features and the index specification. Everything this
    # programme has produced so far was written to a perishable directory.
    out="/content/drive/MyDrive/pin_testbed",
    handicap={3: ((24, 2), (48, 2), (96, 2))},
)


def spec(cfg):
    """The stack's shape, and what it costs, before anything is built."""
    rows, g, cin = [], cfg["grid"], cfg["c_in"]
    for ch, st in cfg["stack"]:
        g_out = g//st
        ni, no = cin*g*g, ch*g_out*g_out
        K = cin*ch*9 + 1
        rows.append(dict(g_in=g, g_out=g_out, c_in=cin, c_out=ch, stride=st,
                         ins=ni, out=no, K=K, entries=ni*no))
        g, cin = g_out, ch
    return rows, g, cin


def announce(cfg):
    rows, g, cin = spec(cfg)
    n_base = cfg["n_classes"] - cfg["n_held"]
    body = sum(r["K"] for r in rows)
    head = cin*n_base + n_base
    ent = sum(r["entries"] for r in rows)
    print(f"  {'layer':>5s} {'grid':>9s} {'channels':>9s} {'hidden':>8s} "
          f"{'values':>9s} {'index entries':>14s}")
    for i, r in enumerate(rows):
        print(f"  {i+1:5d} {r['g_in']:4d}->{r['g_out']:<4d} "
              f"{r['c_in']:4d}->{r['c_out']:<4d} {r['out']:8,} "
              f"{r['K']:9,} {r['entries']:14,}")
    print(f"  {'pool':>5s} {'':>9s} {'':>9s} {cin:8,} {'':>9s}")
    print(f"  {'head':>5s} {'':>9s} {'':>9s} {n_base:8,} {head:9,}")
    print(f"\n  folded body {body:,} values, head {head:,}, "
          f"so the body is {body/(body+head):.0%} of the model")
    print(f"  (with a flat head it would be "
          f"{cin*rows[-1]['g_out']**2*n_base:,}, and the body {body/(body+cin*rows[-1]['g_out']**2*n_base):.0%})")
    gb = ent*(4 + 8 + 4)/1e9
    print(f"\n  peak device memory for indices, orderings and materialised")
    print(f"  weights: about {gb:.1f} GB. Abort now if that will not fit.")
    return rows, g, cin, body, head


def load(cfg):
    from tensorflow import keras
    (a, b), (c, d) = keras.datasets.cifar100.load_data()
    X = np.concatenate([a, c]).astype(np.float32)/255.0
    y = np.concatenate([b, d]).ravel().astype(np.int64)
    rg = np.random.default_rng(0)
    order = rg.permutation(cfg["n_classes"])
    held = np.sort(order[:cfg["n_held"]])
    base = np.sort(order[cfg["n_held"]:])
    remap = -np.ones(cfg["n_classes"], np.int64)
    remap[base] = np.arange(len(base))
    is_base = np.isin(y, base)
    mu = X[is_base].mean((0, 1, 2)); sd = X[is_base].std((0, 1, 2)) + 1e-8
    X = ((X - mu)/sd).astype(np.float32)
    # CIFAR arrives as (N, 32, 32, 3) and everything here expects a flat
    # (N, channel*row*column) in CHANNEL-FIRST order, because the fold's
    # index is built as channel*g*g + row*g + column. A first version
    # returned the raw array: training still ran, because augment()
    # reshapes to (n, c, g, g) and flattens, so it silently trained on
    # CHANNEL-SCRAMBLED images for a whole epoch. Evaluation skips
    # augmentation and crashed, which is the only reason it was caught.
    X = np.ascontiguousarray(X.transpose(0, 3, 1, 2)).reshape(len(X), -1)
    return X, y, base, held, remap


def augment(x, cfg, g, rg):
    """Random crop with reflection padding, and a horizontal flip."""
    n = x.shape[0]; p = cfg["aug_pad"]
    im = x.reshape(n, cfg["c_in"], g, g)
    pad = xp.zeros((n, cfg["c_in"], g+2*p, g+2*p), DT)
    pad[:, :, p:p+g, p:p+g] = im
    pad[:, :, :p, p:p+g] = im[:, :, p:0:-1, :]
    pad[:, :, p+g:, p:p+g] = im[:, :, -2:-p-2:-1, :]
    pad[:, :, :, :p] = pad[:, :, :, 2*p:p:-1]
    pad[:, :, :, p+g:] = pad[:, :, :, -p-2:-2*p-2:-1]
    oy, ox = rg.integers(0, 2*p+1, 2)
    out = pad[:, :, oy:oy+g, ox:ox+g]
    if rg.random() < 0.5:
        out = out[:, :, :, ::-1]
    return out.reshape(n, -1)


class FastConv:
    """A folded convolution computed without ever building its matrix.

    Two index tables are built once. The FORWARD one says, for each output
    position, which input elements its window covers. The BACKWARD one says,
    for each input element, which window slots read it, which turns the
    gradient back to the input from a scatter into a gather and removes the
    last source of non-determinism from the layer."""

    def __init__(self, g_in, c_in, k, c_out, stride):
        self.g_in, self.c_in, self.k = g_in, c_in, k
        self.c_out, self.stride = c_out, stride
        self.g_out = g_in // stride
        self.ins = c_in*g_in*g_in
        self.out = c_out*self.g_out*self.g_out
        self.K = c_in*c_out*k*k + 1
        go, gi = self.g_out, g_in

        # forward: (positions, c_in*k*k) into the input, with a trailing
        # zero slot standing in for anything outside the image
        oy, ox = np.divmod(np.arange(go*go), go)
        ci = np.arange(c_in)[:, None, None]
        dy = np.arange(k)[None, :, None]
        dx = np.arange(k)[None, None, :]
        iy = oy[:, None, None, None]*stride - k//2 + dy
        ix = ox[:, None, None, None]*stride - k//2 + dx
        ok = (iy >= 0) & (iy < gi) & (ix >= 0) & (ix < gi)
        flat = (ci*gi*gi + np.clip(iy, 0, gi-1)*gi + np.clip(ix, 0, gi-1))
        flat = np.where(ok, flat, self.ins)
        self.fwd_idx = to_dev(flat.reshape(go*go, c_in*k*k), np.int64) \
            if _GPU else flat.reshape(go*go, c_in*k*k).astype(np.int64)
        self.taps = c_in*k*k

        # backward: for each input element, which (position, tap) slots
        # read it. At most (k+stride-1)//stride squared, so four at k=3,
        # stride=2, which keeps the table small and regular.
        col = flat.reshape(go*go, c_in*k*k)
        readers = {}
        for p in range(go*go):
            for tcol in range(c_in*k*k):
                s = int(col[p, tcol])
                if s < self.ins:
                    readers.setdefault(s, []).append(p*(c_in*k*k) + tcol)
        w = max((len(v) for v in readers.values()), default=1)
        tbl = np.full((self.ins, w), go*go*c_in*k*k, np.int64)
        for s, v in readers.items():
            tbl[s, :len(v)] = v
        self.bwd_idx = to_dev(tbl, np.int64) if _GPU else tbl
        self.bwd_width = w

    def gather(self, x):
        """The compact window matrix: (batch*positions, c_in*k*k)."""
        n = x.shape[0]
        xz = xp.concatenate([x, xp.zeros((n, 1), DT)], 1)
        return xz[:, self.fwd_idx].reshape(n*self.g_out**2, self.taps)

    def weights(self, v):
        """The values arranged for the matmul: (c_in*k*k, c_out).

        A value index is (in-channel, out-channel, tap), so this is a
        reshape and a transpose of the value vector and nothing more."""
        return v[:-1].reshape(self.c_in, self.c_out, self.k*self.k) \
            .transpose(0, 2, 1).reshape(self.taps, self.c_out)

    def forward(self, v, x):
        n = x.shape[0]
        col = self.gather(x)
        z = (col @ self.weights(v)).reshape(n, self.g_out**2, self.c_out)
        return z.transpose(0, 2, 1).reshape(n, self.out), col

    def backward(self, v, col, dz, need_input=True):
        """dz has the layer's output shape. Returns the value gradient and,
        if asked, the gradient back to the input."""
        n = dz.shape[0]
        d = dz.reshape(n, self.c_out, self.g_out**2).transpose(0, 2, 1) \
            .reshape(n*self.g_out**2, self.c_out)
        # THE VALUE GRADIENT, straight out of a matmul. No scatter.
        dW = col.T @ d
        gv = xp.zeros(self.K, DT)
        gv[:-1] = dW.reshape(self.c_in, self.k*self.k, self.c_out) \
            .transpose(0, 2, 1).reshape(-1)
        if not need_input:
            return gv, None
        dcol = (d @ self.weights(v).T).reshape(n, -1)
        pad = xp.concatenate([dcol, xp.zeros((n, 1), DT)], 1)
        dx = pad[:, self.bwd_idx].sum(2)
        return gv, dx


_FAST = True


def build(cfg, seed):
    rows, g_fin, c_fin = spec(cfg)
    n_out = cfg["n_classes"] - cfg["n_held"]
    rg = np.random.default_rng(seed)
    layers = []
    for r in rows:
        idx, K, no, macs = strided_window(r["g_in"], r["c_in"], 3,
                                          r["c_out"], r["stride"])
        v = rg.normal(0, np.sqrt(2.0/(r["c_in"]*9)), K).astype(np.float32)
        v[-1] = 0.0
        fc = (FastConv(r["g_in"], r["c_in"], 3, r["c_out"], r["stride"])
              if (_FAST and cfg.get("fast", True)) else None)
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, ins=r["ins"], out=no, macs=macs,
                           c_out=r["c_out"], g_out=r["g_out"], fc=fc,
                           v=to_dev(v), gam=xp.ones(no, DT),
                           bet=xp.zeros(no, DT)))
    head = to_dev(rg.normal(0, np.sqrt(2.0/c_fin), (c_fin, n_out)))
    return layers, head, xp.zeros(n_out, DT), c_fin


def forward(layers, head, ob, x, keep=False):
    cache = []; h = x
    for l in layers:
        if l["fc"] is not None:
            z, col = l["fc"].forward(l["v"], h); W = None
        else:
            W = l["v"][l["idx"]].reshape(l["ins"], l["out"])
            z = h @ W; col = None
        var = z.var(1, keepdims=True) + 1e-5
        zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
        zs = zn*l["gam"] + l["bet"]
        a = xp.maximum(zs, 0)
        if keep:
            cache.append((h, W, var, zn, zs, col))
        h = a
    # POOL over spatial positions, so the head reads one number a channel
    last = layers[-1]
    feat = h.reshape(h.shape[0], last["c_out"], last["g_out"]**2).mean(2)
    return (feat, feat @ head + ob, cache) if keep else (feat,
                                                         feat @ head + ob)


def train_one(X, Yl, Xte, yte, cfg, seed, tag):
    layers, head, ob, c_fin = build(cfg, seed)
    L = len(layers)
    P = [l["v"] for l in layers] + [l["gam"] for l in layers] \
        + [l["bet"] for l in layers] + [head, ob]
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = X.shape[0]; t = 0
    rg = np.random.default_rng(seed + 991)
    npos = layers[-1]["g_out"]**2; cout = layers[-1]["c_out"]
    t0 = time.time()
    steps_total = cfg["epochs"]*int(np.ceil(n/cfg["batch"]))
    for ep in range(cfg["epochs"]):
        perm = rg.permutation(n)
        tot, nb = 0.0, 0
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = X[b]
            if cfg["augment"]:
                x = augment(x, cfg, cfg["grid"], rg)
            y = Yl[b]
            feat, lg, cache = forward(layers, P[3*L], P[3*L+1], x, keep=True)
            mx = lg.max(1, keepdims=True)
            e = xp.exp(lg - mx); se = e.sum(1, keepdims=True)
            tot += float(to_host((-(lg-mx-xp.log(se))*y).sum(1).mean()))
            nb += 1
            d = (e/se - y)/len(b)
            G = [xp.zeros_like(p) for p in P]
            G[3*L] = feat.T @ d; G[3*L+1] = d.sum(0)
            dfeat = d @ P[3*L].T
            # back through the spatial pool
            dh = xp.broadcast_to(dfeat[:, :, None]/npos,
                                 (dfeat.shape[0], cout, npos)
                                 ).reshape(dfeat.shape[0], cout*npos)
            for li in range(L-1, -1, -1):
                hin, W, var, zn, zs, col = cache[li]
                dzs = dh*(zs > 0)
                G[L+li] = (dzs*zn).sum(0); G[2*L+li] = dzs.sum(0)
                dzn = dzs*P[L+li]
                dz = (dzn - dzn.mean(1, keepdims=True)
                      - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var)
                fc = layers[li]["fc"]
                if fc is not None:
                    G[li], dnext = fc.backward(P[li], col, dz,
                                               need_input=li > 0)
                    if li > 0:
                        dh = dnext
                else:
                    G[li] = scatter(hin.T @ dz, layers[li]["idx"],
                                    layers[li]["K"])
                    if li > 0:
                        dh = dz @ W.T
            t += 1
            lr = (cfg["lr"]*0.5*(1 + np.cos(np.pi*t/steps_total))
                  if cfg.get("schedule") == "cosine" else cfg["lr"])
            for i, (p_, gr) in enumerate(zip(P, G)):
                M[i] = 0.9*M[i] + 0.1*gr
                V[i] = 0.999*V[i] + 0.001*gr*gr
                P[i] = p_ - lr*(M[i]/(1-0.9**t)) \
                    / (xp.sqrt(V[i]/(1-0.999**t))+1e-8)
            for li, l in enumerate(layers):
                # PIN THE PADDING VALUE. v[-1] stands for every
                # out-of-window position and starts at zero, but the
                # general path accumulates gradient into it and Adam then
                # moves it, so out-of-window positions acquire a learned
                # shared offset and the layer stops being strictly a
                # convolution. That is true of every earlier run here.
                # Pinned, the model is an actual convolution and the fast
                # path in fastconv.py computes exactly the same thing.
                P[li][-1] = 0.0
                l["v"], l["gam"], l["bet"] = P[li], P[L+li], P[2*L+li]
        if (ep+1) % cfg["report_every"] == 0 or ep == 0:
            acc = evaluate(layers, P[3*L], P[3*L+1], Xte, yte)
            print(f"    {tag} epoch {ep+1:4d}/{cfg['epochs']}   "
                  f"train {tot/max(nb,1):.4f}   test {acc:.4f}   "
                  f"lr {lr:.2e}   [{time.time()-t0:.0f}s]", flush=True)
    return layers, P[3*L], P[3*L+1]


def evaluate(layers, head, ob, Xte, yte):
    out = []
    for s in range(0, Xte.shape[0], 512):
        _, lg = forward(layers, head, ob, Xte[s:s+512])
        out.append(to_host(lg))
    return float((np.concatenate(out).argmax(1) == yte).mean())


def digest(layers, head, ob):
    h = hashlib.sha256()
    for l in layers:
        h.update(to_host(l["v"]).tobytes())
        h.update(to_host(l["gam"]).tobytes())
        h.update(to_host(l["bet"]).tobytes())
    h.update(to_host(head).tobytes()); h.update(to_host(ob).tobytes())
    return h.hexdigest()[:16]


def ensure_output(path):
    """Put the results where they will still exist tomorrow.

    If the target is on Drive and Drive is not mounted, mount it. If that
    fails, fall back to the working directory and SAY SO LOUDLY, because a
    silent fallback is how four bases end up in a folder that is deleted
    when the session closes."""
    if path.startswith("/content/drive"):
        if not os.path.isdir("/content/drive/MyDrive"):
            try:
                from google.colab import drive
                print("  mounting Drive...", flush=True)
                drive.mount("/content/drive")
            except Exception as e:
                alt = os.path.basename(path.rstrip("/")) or "testbed"
                print(f"\n  ** DRIVE IS NOT AVAILABLE ({e}) **")
                print(f"  ** writing to ./{alt} instead, which does NOT")
                print(f"  ** survive the session. Copy it somewhere before")
                print(f"  ** you close this.\n", flush=True)
                os.makedirs(alt, exist_ok=True)
                return alt
    os.makedirs(path, exist_ok=True)
    print(f"  results will be written to {path}", flush=True)
    return path


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("THE TESTBED: CIFAR-100, two bases, twenty classes held back")
    print("  build 2026-08-15f: one file, one paste, saved to Drive.")
    print("  Cosine decay, a strength reference, and a smaller counterparty")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:12s} = {v}")
    print()
    rows, g_fin, c_fin, body, head_n = announce(CFG)
    print("=" * 78, flush=True)
    print(f"\n  checking the fast path against the general one before")
    print(f"  anything is trained on it:")
    ok = True
    for (gg, ci, co, st) in ((16, 3, 8, 2), (32, 3, 6, 2)):
        idx, K, no, _ = strided_window(gg, ci, 3, co, st)
        f = FastConv(gg, ci, 3, co, st)
        r = np.random.default_rng(0)
        v = r.normal(size=K).astype(np.float32); v[-1] = 0.0
        x = r.normal(size=(4, ci*gg*gg)).astype(np.float32)
        vd, xd = to_dev(v), to_dev(x)
        ref = to_host(xd @ vd[to_dev(idx, np.int64)].reshape(ci*gg*gg, no))
        got, col = f.forward(vd, xd)
        dz = to_dev(r.normal(size=(4, no)))
        gv, dx = f.backward(vd, col, dz)
        gvr = scatter(xd.T @ dz, to_dev(idx, np.int32) if _GPU else idx, K)
        e = max(np.abs(to_host(got)-ref).max(),
                np.abs(to_host(gv)[:-1]-to_host(gvr)[:-1]).max()/max(
                    1.0, np.abs(to_host(gvr)[:-1]).max()),
                np.abs(to_host(dx) - to_host(dz @ vd[to_dev(idx, np.int64)]
                                             .reshape(ci*gg*gg, no).T)).max())
        print(f"    {gg}x{gg} {ci}->{co}: worst disagreement {e:.1e}")
        ok = ok and e < 1e-3
    _FIXED.clear()
    if not ok:
        raise SystemExit("  the fast path disagrees; refusing to train")
    print(f"  agreed. Training now.\n", flush=True)

    X, y, base_cls, held_cls, remap = load(CFG)
    print(f"\n  {len(base_cls)} classes for the bases, {len(held_cls)} held")
    print(f"  back entirely: {held_cls.tolist()}")
    rg = np.random.default_rng(7)
    bi = np.where(np.isin(y, base_cls))[0]; rg.shuffle(bi)
    hi = np.where(np.isin(y, held_cls))[0]
    cut = int(0.85*len(bi))
    tr, te = bi[:cut], bi[cut:]
    half = len(tr)//2
    splits = {0: tr[:half], 1: tr[half:], 2: tr, 3: tr}
    print(f"  bases 0 and 1 train on {half:,} images each, DISJOINT;")
    print(f"  base 2 trains on all {len(tr):,} as a STRENGTH REFERENCE and")
    print(f"  overlaps both, so it is never a counterparty in a borrowing")
    print(f"  experiment;")
    print(f"  {len(te):,} held for testing; {len(hi):,} images of the held")
    print(f"  classes are never seen by either", flush=True)

    Xte = to_dev(X[te]); yte = remap[y[te]]
    out = ensure_output(CFG["out"])
    saved = {}
    for s in CFG["seeds"]:
        ix = splits[s]
        if s == 2:
            print(f"\n  (base 2 is the strength reference, not a "
                  f"counterparty)")
        if s in CFG.get("handicap", {}):
            print(f"\n  (base {s} is HANDICAPPED to "
                  f"{CFG['handicap'][s]}, a smaller counterparty)")
        Xtr = to_dev(X[ix])
        Ytr = to_dev(np.eye(len(base_cls), dtype=np.float32)[remap[y[ix]]])
        print(f"\n  base {s}: {len(ix):,} images", flush=True)
        C = dict(CFG)
        if s in CFG.get("handicap", {}):
            C["stack"] = CFG["handicap"][s]
        layers, hd, ob = train_one(Xtr, Ytr, Xte, yte, C, 100*s + 3,
                                   f"base{s}")
        acc = evaluate(layers, hd, ob, Xte, yte)
        dg = digest(layers, hd, ob)
        # cached features for the held classes, so member work needs no
        # forward pass ever again
        feats = []
        for st in range(0, len(hi), 512):
            f, _ = forward(layers, hd, ob, to_dev(X[hi[st:st+512]]))
            feats.append(to_host(f))
        feats = np.concatenate(feats)
        fte = []
        for st in range(0, len(te), 512):
            f, _ = forward(layers, hd, ob, to_dev(X[te[st:st+512]]))
            fte.append(to_host(f))
        np.savez_compressed(
            f"{out}/base{s}.npz",
            values=np.array([to_host(l["v"]) for l in layers], dtype=object),
            gam=np.array([to_host(l["gam"]) for l in layers], dtype=object),
            bet=np.array([to_host(l["bet"]) for l in layers], dtype=object),
            head=to_host(hd), ob=to_host(ob),
            held_feats=feats, held_labels=y[hi], held_index=hi,
            test_feats=np.concatenate(fte), test_labels=yte, test_index=te,
            train_index=ix, digest=dg, accuracy=acc)
        rws, _, cf = spec(C)
        saved[s] = dict(accuracy=acc, digest=dg,
                        values=int(sum(r["K"] for r in rws)),
                        features=int(cf), n_train=int(len(ix)),
                        stack=[list(x) for x in C["stack"]],
                        role=("reference" if s == 2 else
                              "small counterparty" if s in
                              CFG.get("handicap", {}) else "counterparty"))
        print(f"  base {s}: {acc:.4f} on {len(base_cls)} classes, "
              f"digest {dg}   [{time.time()-t0:.0f}s]", flush=True)
        _FIXED.clear()
        if _GPU:
            _cp.get_default_memory_pool().free_all_blocks()

    json.dump(dict(cfg={k: v for k, v in CFG.items()},
                   base_classes=base_cls.tolist(),
                   held_classes=held_cls.tolist(),
                   layers=[{k: int(v) for k, v in r.items()} for r in rows],
                   features=int(c_fin), bases=saved),
              open(f"{out}/spec.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  SAVED")
    print("=" * 78)
    for s, r in saved.items():
        print(f"  base{s}.npz  {r['accuracy']:.4f}  {r['n_train']:6,} imgs  "
              f"{r['values']:7,} values  {r['features']:3d} feat  "
              f"{r['role']}")
    if 2 in saved and 0 in saved:
        d = saved[2]["accuracy"] - saved[0]["accuracy"]
        print(f"\n  doubling the data is worth {d:+.4f}, so "
              f"{'the split costs more than the architecture' if d > 0.04 else 'the split is not the main limit'}")
    if 3 in saved and 2 in saved:
        h = saved[3]["accuracy"] - saved[2]["accuracy"]
        print(f"  halving the channels costs {h:+.4f} at the same data, and")
        print(f"  gives a {saved[3]['features']}-feature counterparty against")
        print(f"  {saved[2]['features']}, so a translation between them is")
        print(f"  RECTANGULAR rather than square for the first time")
    print(f"  spec.json    the index specification, class split and layer")
    print(f"               shapes, so a base can be rebuilt exactly")
    print(f"\n  all of it in {out}\n")
    print(f"  each archive also carries CACHED FEATURES for the twenty held")
    print(f"  classes ({c_fin} numbers an image) and for the test set, so")
    print(f"  every member, borrowing and marketplace question afterwards is")
    print(f"  a ridge solve and needs no forward pass at all.")
    print(f"\n  total {time.time()-t0:.0f}s")


# runs whether this is executed as a script or pasted into a cell
main()