"""
HOW DEEP SHOULD A STACK BE FOR AN n x n IMAGE?

A three-layer stack beat every five-layer arrangement on Fashion at 14x14,
and beat them at half the storage. That says the optimum is below five for
that grid, and says nothing about why or about any other grid.

The obvious candidate rule is already refuted. A stack of L 3x3 layers sees
(2L+1)^2 positions, so covering n x n by reach alone needs L >= (n-1)/2 —
seven layers at n=14, more than double the measured optimum. And the
butterfly comparison went the same way: three layers reaching 16 of 256
positions beat a butterfly reaching all 256. COVERAGE IS NOT THE CRITERION.

So this measures the optimum directly at three grid sizes and asks whether
it moves with n.

WHAT IS HELD FIXED AND WHY. The hidden width is held near constant across
grids by scaling channels inversely with area — 64 channels at 7x7, 16 at
14x14, 4 at 28x28 — so every arm computes the same size of matrix product
and the comparison is not a compute sweep in disguise. Storage still varies
enormously across grids (a 7x7 arm holds hundreds of times more than a
28x28 one), so ACROSS-GRID accuracies are not comparable. WITHIN a grid
they are, and where the peak falls is the whole question.

The reported quantities are accuracy, stored values and multiplies, so the
peak can be read three ways: best accuracy, best accuracy per value, and
best accuracy per multiply. Those need not agree, and if they do not, that
is the more useful finding — it says which currency the answer depends on.
"""

import numpy as np
import time
import json
import os

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 windowed(g, c_in, k, c_out):
    ni, no = c_in*g*g, c_out*g*g
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g*g), ii % (g*g)
    co, po = jj // (g*g), jj % (g*g)
    dr = pi // g - (po // g - k//2)
    dc = pi % g - (po % g - 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, no*(c_in*k*k)


DETERMINISTIC = True
_FIXED = {}


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

    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):
    if not DETERMINISTIC:
        g = xp.zeros(K, DT)
        if _GPU:
            import cupyx
            cupyx.scatter_add(g, idx, dW.reshape(-1))
        else:
            np.add.at(g, idx, dW.reshape(-1))
        return g
    key = (id(idx), K)
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


def train(Xtr, Ytr, Xte, yte, g, chan, depth, cfg, seed):
    D = Xtr.shape[1]
    rg = np.random.default_rng(seed)
    layers, cin = [], 1
    for l in range(depth):
        idx, K, no, macs = windowed(g, cin, 3, chan)
        layers.append(dict(idx=to_dev(idx, np.int32) if _GPU else idx,
                           K=K, out=no, macs=macs,
                           ins=D if l == 0 else layers[-1]["out"],
                           taps=cin*9))
        cin = chan
    L = depth
    P = []
    for l in layers:
        v = rg.normal(0, np.sqrt(2.0/l["taps"]), l["K"]).astype(np.float32)
        v[-1] = 0.0
        P.append(to_dev(v))
    P += [xp.ones(l["out"], DT) for l in layers]
    P += [xp.zeros(l["out"], DT) for l in layers]
    P += [to_dev(rg.normal(0, np.sqrt(2.0/layers[-1]["out"]),
                           (layers[-1]["out"], 10))), xp.zeros(10, DT)]
    HEAD, OB = 3*L, 3*L+1
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = Xtr.shape[0]; t = 0
    ag = np.random.default_rng(seed + 991)

    def fwd(x, keep=False):
        cache = []; h = x
        for li, l in enumerate(layers):
            W = P[li][l["idx"]].reshape(l["ins"], l["out"])
            z = h @ W
            var = z.var(1, keepdims=True) + 1e-5
            zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
            zs = zn*P[L+li] + P[2*L+li]
            a = xp.maximum(zs, 0)
            if keep:
                cache.append((h, W, var, zn, zs))
            h = a
        return h, cache

    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]; y = Ytr[b]
            h, cache = fwd(x, keep=True)
            lg = h @ P[HEAD] + P[OB]
            e = xp.exp(lg - lg.max(1, keepdims=True))
            d = (e/e.sum(1, keepdims=True) - y)/len(b)
            G = [xp.zeros_like(p) for p in P]
            G[HEAD] = h.T @ d; G[OB] = d.sum(0)
            dh = d @ P[HEAD].T
            for li in range(L-1, -1, -1):
                hin, W, var, zn, zs = 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)
                G[li] = scatter(hin.T @ dz, layers[li]["idx"], layers[li]["K"])
                if li > 0:
                    dh = dz @ W.T
            t += 1
            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_ - cfg["lr"]*(M[i]/(1-0.9**t)) \
                    / (xp.sqrt(V[i]/(1-0.999**t))+1e-8)
    out = []
    for s in range(0, Xte.shape[0], 4096):
        h, _ = fwd(Xte[s:s+4096])
        out.append(to_host(h @ P[HEAD] + P[OB]))
    acc = float((np.concatenate(out).argmax(1) == yte).mean())
    # COUNT THE HEAD. It is dense and unfolded, and on Fashion at 14x14 it
    # is 99.5% of a one-layer model's parameters — reporting the folded
    # body alone made a 15% storage difference look like 65x.
    head = layers[-1]["out"]*10 + 10
    return (acc, sum(l["K"] for l in layers) + head,
            sum(l["macs"] for l in layers) + layers[-1]["out"]*10, head)


def load(grid, cfg):
    from tensorflow import keras
    (a, b), (c, d) = keras.datasets.fashion_mnist.load_data()
    X = np.concatenate([a, c]).astype(np.float32)/255.0
    y = np.concatenate([b, d]).ravel().astype(np.int64)
    if grid != 28:
        s = 28//grid
        X = X.reshape(-1, grid, s, grid, s).mean(axis=(2, 4))
    rg = np.random.default_rng(0); p = rg.permutation(len(X))
    tr, te = p[:cfg["n_train"]], p[cfg["n_train"]:cfg["n_train"]+10000]
    mu, sd = X[tr].mean(), X[tr].std()+1e-8
    f = lambda Z: ((Z-mu)/sd).reshape(len(Z), -1)
    Y = np.zeros((len(tr), 10), np.float32); Y[np.arange(len(tr)), y[tr]] = 1
    return f(X[tr]), Y, f(X[te]), y[te]


CFG = dict(n_train=20000, batch=128, lr=1e-3, epochs=40, seeds=(0, 1),
           hidden_target=3136, grids=(7, 14, 28), depths=(1, 2, 3, 4, 5))


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("HOW DEEP SHOULD A STACK BE FOR AN n x n IMAGE?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:14s} = {v}")
    print(f"\n  channels scale inversely with area to hold the hidden width")
    print(f"  near {CFG['hidden_target']}, so no arm is favoured by having a")
    print(f"  bigger matrix product:")
    for g in CFG["grids"]:
        c = max(1, CFG["hidden_target"]//(g*g))
        print(f"    {g:2d}x{g:<2d} -> {c:3d} channels, hidden {c*g*g:,}")
    print(f"\n  reach says L >= (n-1)/2 to cover the grid: "
          + ", ".join(f"{g}x{g} needs {int(np.ceil((g-1)/2))}"
                      for g in CFG["grids"]))
    print("=" * 78, flush=True)

    res = {}
    for g in CFG["grids"]:
        chan = max(1, CFG["hidden_target"]//(g*g))
        Xtr, Ytr, Xte, yte = load(g, CFG)
        Xtr, Ytr, Xte = to_dev(Xtr), to_dev(Ytr), to_dev(Xte)
        print(f"\n  {g}x{g}, {chan} channels")
        print(f"  (the head alone is "
              f"{chan*g*g*10 + 10:,} parameters, counted in every row)")
        print(f"  {'depth':>6s} {'TOTAL vals':>11s} {'multiplies':>12s} "
              f"{'accuracy':>9s} {'sd':>7s} {'reach':>7s}")
        for L in CFG["depths"]:
            # the fixed-order tables are keyed by the index array's identity
            # and each arm builds new ones, so the cache would grow across
            # fifteen arms until the device ran out
            _FIXED.clear()
            if _GPU:
                _cp.get_default_memory_pool().free_all_blocks()
            out = [train(Xtr, Ytr, Xte, yte, g, chan, L, CFG, s)
                   for s in CFG["seeds"]]
            acc = [o[0] for o in out]
            K, mc, hd = out[0][1], out[0][2], out[0][3]
            reach = min(g, 2*L+1)
            res[f"{g}/{L}"] = dict(acc=float(np.mean(acc)),
                                   sd=float(np.std(acc)), K=K, macs=mc,
                                   reach=reach, chan=chan, head=hd)
            print(f"  {L:6d} {K:11,} {mc:12,} {np.mean(acc):9.4f} "
                  f"{np.std(acc):7.4f} {reach:5d}x{reach:<2d}"
                  f"   [{time.time()-t0:.0f}s]", flush=True)
            json.dump(res, open("depth_law.json", "w"), indent=2)

    print("\n" + "=" * 78)
    print("  WHERE IS THE PEAK, AND DOES IT MOVE WITH n?")
    print("=" * 78)
    print(f"  {'grid':>6s} {'best accuracy':>15s} {'best per value':>16s} "
          f"{'best per multiply':>18s} {'reach would say':>16s}")
    peaks = {}
    for g in CFG["grids"]:
        ds = [L for L in CFG["depths"] if f"{g}/{L}" in res]
        acc = {L: res[f"{g}/{L}"]["acc"] for L in ds}
        perv = {L: res[f"{g}/{L}"]["acc"]/res[f"{g}/{L}"]["K"] for L in ds}
        perm = {L: res[f"{g}/{L}"]["acc"]/res[f"{g}/{L}"]["macs"] for L in ds}
        ba = max(acc, key=acc.get); bv = max(perv, key=perv.get)
        bm = max(perm, key=perm.get)
        peaks[g] = (ba, bv, bm)
        print(f"  {g:6d} {ba:15d} {bv:16d} {bm:18d} "
              f"{int(np.ceil((g-1)/2)):16d}")

    sd = max(r["sd"] for r in res.values())
    print(f"\n  seed spread (worst) {sd:.4f}\n")
    ba = [peaks[g][0] for g in CFG["grids"]]
    print(f"  best-accuracy depth by grid: "
          + ", ".join(f"{g}x{g}={peaks[g][0]}" for g in CFG["grids"]))
    if len(set(ba)) == 1:
        print(f"\n  THE OPTIMUM DOES NOT MOVE WITH n. Depth {ba[0]} is best at")
        print(f"  every grid size tried, so the answer is a constant rather")
        print(f"  than a law in n — and reach, which would have demanded")
        print(f"  {', '.join(str(int(np.ceil((g-1)/2))) for g in CFG['grids'])}, is refuted a second time.")
    elif all(x <= y for x, y in zip(ba, ba[1:])):
        print(f"\n  THE OPTIMUM GROWS WITH n, which is what a law would look")
        print(f"  like. It grows far more slowly than reach demands, so the")
        print(f"  criterion is not covering the image — fit a rule to these")
        print(f"  three points and check it once on CIFAR before trusting it.")
    else:
        print(f"\n  NO CLEAN PATTERN across these three grids. Either the")
        print(f"  optimum depends on something else held fixed here, or two")
        print(f"  seeds are too few to locate a peak this flat.")
    flat = [g for g in CFG["grids"]
            if max(res[f"{g}/{L}"]["acc"] for L in CFG["depths"])
            - min(res[f"{g}/{L}"]["acc"] for L in CFG["depths"]) < 4*sd]
    if flat:
        print(f"\n  CAUTION: at {flat} the whole depth range spans less than")
        print(f"  four seed deviations, so the peak there is not located.")
    print(f"\n  and the three currencies "
          f"{'AGREE' if all(len(set(peaks[g])) == 1 for g in CFG['grids']) else 'DISAGREE'}"
          f" about the best depth, which decides")
    print(f"  whether 'how deep' has one answer or one answer per budget.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote depth_law.json")


if __name__ == "__main__":
    main()