"""Perturbation experiments on the levelized host and on the amplified host."""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import subprocess
import sys
import tempfile
import time
from typing import Dict, List, Tuple

from host import (
    encode,
    M_P,
    M_STAR,
    HALT_PC,
    Tape,
    describe,
    inst,
    read_host,
    ser,
    tau_star,
    STATE_LAYOUT,
    IO_CELLS,
)
from paper import levelize, parse_sigma, write_flat


REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def union_bound(s: float, M: int, T: int) -> float:
    return T * M * math.erfc(1.0 / (2.0 * math.sqrt(2.0) * s))


class Devices:
    """The tape device applied to every copy of a batched state."""

    def __init__(self, L, tapes, device):
        import torch
        from host import Tape
        self.torch = torch
        self.L = L
        self.device = device
        self.n = len(tapes)
        self.tapes = [Tape(t) for t in tapes]
        self.done = [False] * self.n
        self.shadow = [[0] * 256 for _ in range(self.n)]
        self.cell_idx = torch.tensor(
            [9 + j * 8 + k for j in L._DEVCELLS for k in range(8)],
            device=device)
        self.pow2 = torch.tensor([1 << (7 - k) for k in range(8)],
                                 device=device)
        self.shift = torch.tensor([7 - k for k in range(8)])
        self.ncell = len(L._DEVCELLS)

    def step(self, v):
        """Apply the device to every copy; return the halt flags."""
        torch = self.torch
        vals = (v[:, self.cell_idx].reshape(self.n, self.ncell, 8)
                * self.pow2).sum(-1).to(torch.int64).cpu().tolist()
        halted = (v[:, 8] >= 0.5).cpu().tolist()
        for c in range(self.n):
            sh = self.shadow[c]
            row = vals[c]
            for ci, j in enumerate(self.L._DEVCELLS):
                sh[j] = row[ci]
            if not self.done[c]:
                self.tapes[c].apply(sh)     # the halting step is a step
            for ci, j in enumerate(self.L._DEVCELLS):
                row[ci] = sh[j]
        for c in range(self.n):
            self.done[c] = self.done[c] or halted[c]
        new = torch.tensor(vals, dtype=torch.int64)
        bits = ((new.unsqueeze(-1) >> self.shift) & 1).float()
        v[:, self.cell_idx] = bits.reshape(self.n, self.ncell * 8).to(self.device)
        return halted


# ---------------------------------------------------------------------------
def run_omega(args) -> int:
    """The self-reproducing instance under read noise, to completion."""
    import torch
    from host import M_STAR, LevEvaluator, read_host, ser, tau_star
    sigma = read_host()
    # the two forms of Lev(N_host) compute the same pre-activations, and this
    # run is long enough that the faster one is worth having
    L = LevEvaluator(sigma, device=args.device, dense=False)
    M = L.info["size"]
    tau = tau_star(sigma, M_STAR)
    target = ser(sigma, bytes(M_STAR), tau)
    print(f"[omega] {L.info['layers']} layers, M = {M:,}, target "
          f"{len(target):,} bytes", flush=True)

    rows = []
    for s, budget in ((0.05, 1 << 21), (0.10, 50000), (0.15, 50000),
                      (0.20, 50000)):
        gen = torch.Generator(device=args.device).manual_seed(20260910)
        v = L._vec(0, M_STAR).unsqueeze(0).to(args.device)
        D = Devices(L, [tau], args.device)
        steps = 0
        wrong = None
        halted = False
        t0 = time.perf_counter()
        while steps < budget:
            v = L.step_noisy(v, s, gen)
            steps += 1
            before = len(D.tapes[0].out)
            halted = D.step(v)[0]
            out = D.tapes[0].out
            if len(out) > before:
                k = len(out) - 1
                if k >= len(target) or out[k] != target[k]:
                    wrong = k
                    break
            if halted:
                break
            if steps % 100000 == 0:
                print(f"    s={s:.2f}: {steps:,} steps, {len(out):,} bytes "
                      f"({steps / (time.perf_counter() - t0):.0f}/s)",
                      flush=True)
        dt = time.perf_counter() - t0
        emitted = len(D.tapes[0].out)
        complete = emitted == len(target) and wrong is None
        outcome = ("emitted the whole instance" if complete else
                   f"a wrong byte after {wrong} correct" if wrong is not None
                   else "the halt bit was set" if halted
                   else "the budget was exhausted")
        rows.append({"sigma": s, "steps": steps, "bytes": emitted,
                     "complete": complete, "first_wrong_byte": wrong,
                     "halted": halted, "seconds": dt,
                     "bound": union_bound(s, M, max(steps, 1))})
        print(f"    s={s:.2f}: {steps:,} steps, {emitted:,} bytes, "
              f"{outcome} ({dt / 60:.1f} min)", flush=True)
    json.dump({"M": M, "target_bytes": len(target), "rows": rows},
              open(runs_path("paper_noise_omega.json"), "w"), indent=1)
    return 0


# ---------------------------------------------------------------------------
def run_curve(args) -> int:
    """Empirical probability of deviation against the union bound."""
    import torch
    from host import M_P, LevEvaluator, describe, read_host
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=True)
    M = sum(int(W.shape[0]) for W in L.W)
    r_H = describe(sigma)
    B = args.batch
    cap = args.steps
    print(f"[curve] M = {M:,}; {B} noisy copies per level, {cap:,} steps each",
          flush=True)

    grid = [float(x) for x in args.grid.split(",")]
    rows = []
    for s in grid:
        # below 0.08 the bound is under 1e-11 per step, far below anything this
        # experiment can resolve, so a short run suffices to record that
        # nothing was seen
        cap = args.steps if s >= 0.08 else min(args.steps, 3000)
        gen = torch.Generator(device=args.device).manual_seed(7 + int(1000 * s))
        v = L._vec(0, M_P).unsqueeze(0).to(args.device).repeat(B + 1, 1)
        D = Devices(L, [r_H] * (B + 1), args.device)
        first = [None] * B
        t0 = time.perf_counter()
        live = cap
        for n in range(cap):
            y = v
            for W, b in zip(L.W, L.B):
                pre = y @ W.T + b
                noise = torch.randn(pre.shape, generator=gen,
                                    device=pre.device) * s
                noise[0] = 0.0                      # copy 0 is noise free
                y = ((pre + noise) >= -0.5).float()
            v = y
            diff = (v[1:] != v[0]).any(dim=1).cpu().tolist()
            for i, d in enumerate(diff):
                if d and first[i] is None:
                    first[i] = n + 1
            if all(f is not None for f in first):
                live = n + 1
                break
            D.step(v)
        dt = time.perf_counter() - t0
        deviated = [f for f in first if f is not None]
        exposure = sum(f if f is not None else live for f in first)
        rate = len(deviated) / exposure if exposure else 0.0
        p = math.erfc(1.0 / (2.0 * math.sqrt(2.0) * s))
        rows.append({"sigma": s, "copies": B, "steps": live,
                     "deviated": len(deviated),
                     "median_first_deviation":
                         sorted(deviated)[len(deviated) // 2] if deviated else None,
                     "step_exposure": exposure,
                     "empirical_rate_per_step": rate,
                     "bound_rate_per_step": M * p,
                     "looseness": (M * p / rate) if rate else None,
                     "first_deviation": first,
                     "seconds": dt})
        print(f"    s={s:.2f}: {len(deviated)}/{B} deviated within {live:,} "
              f"steps; empirical {rate:.3e} per step, bound {M * p:.3e}, "
              f"loose by {(M * p / rate) if rate else float('inf'):.1f}x "
              f"({dt / 60:.1f} min)", flush=True)
    json.dump({"M": M, "grid": grid, "rows": rows},
              open(runs_path("paper_noise_curve.json"), "w"), indent=1)
    return 0


# ---------------------------------------------------------------------------
def run_leak(args) -> int:
    """Static weight error including the leakage of the zero entries."""
    import torch
    from host import M_P, LevEvaluator, describe, read_host
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=True)
    M = sum(int(W.shape[0]) for W in L.W)
    nmax = max(int(W.shape[1]) for W in L.W)
    kmax = max(int((W != 0).sum(dim=1).max()) for W in L.W)
    r_H = describe(sigma)
    print(f"[leak] M = {M:,}, widest layer {nmax}, largest fan-in {kmax}; "
          f"Proposition 8.3 permits eps*k + gamma*(n-k) < 1/2", flush=True)

    def worst_bound(eps, gam):
        w = 0.0
        for W in L.W:
            k = (W != 0).sum(dim=1).float()
            n = float(W.shape[1])
            w = max(w, float((eps * k + gam * (n - k)).max()))
        return w

    def run(eps, gam, seed, steps):
        gen = torch.Generator(device=args.device).manual_seed(seed)
        Wp = []
        for W in L.W:
            u = torch.rand(W.shape, generator=gen, device=W.device) * 2 - 1
            scale = torch.where(W != 0,
                                torch.tensor(eps, device=W.device),
                                torch.tensor(gam, device=W.device))
            Wp.append(W + u * scale)
        v = L._vec(0, M_P).unsqueeze(0).to(args.device).repeat(2, 1)
        D = Devices(L, [r_H, r_H], args.device)
        first = None
        for n in range(steps):
            y = v
            for W, Wq, b in zip(L.W, Wp, L.B):
                pre = torch.stack([y[0] @ W.T, y[1] @ Wq.T]) + b
                y = (pre >= -0.5).float()
            v = y
            if bool((v[1] != v[0]).any()):
                first = n + 1
                break
            D.step(v)
        emitted = len(D.tapes[1].out)
        del Wp
        if args.device.startswith("cuda"):
            torch.cuda.empty_cache()
        return first, emitted

    gamma_star = 0.5 / (nmax - 1)
    rows = []
    grid = [(0.0, 0.0), (0.0, gamma_star / 2), (0.0, gamma_star * 0.9),
            (0.0, gamma_star * 1.5), (0.0, gamma_star * 4),
            (0.0, gamma_star * 16), (0.0, gamma_star * 64),
            (0.0, gamma_star * 256),
            (1 / 1024, gamma_star / 4), (1 / 512, 0.0), (1 / 256, 0.0),
            (1 / 64, 0.0), (1 / 16, 0.0), (1 / 4, 0.0)]
    for eps, gam in grid:
        t0 = time.perf_counter()
        first, emitted = run(eps, gam, 4242, args.steps)
        wb = worst_bound(eps, gam)
        certified = wb < 0.5
        rows.append({"epsilon": eps, "gamma": gam, "worst_bound": wb,
                     "certified_by_corollary": bool(certified),
                     "first_deviation": first,
                     "steps": first if first else args.steps,
                     "bytes": emitted, "seconds": time.perf_counter() - t0})
        print(f"    eps={eps:.3e} gamma={gam:.3e}: worst bound {wb:.4f}, "
              f"certified {'yes' if certified else 'no '}; "
              f"{'exact for all ' + format(args.steps, ',') + ' steps' if first is None else f'deviates at step {first:,}'}"
              f" ({time.perf_counter() - t0:.0f}s)", flush=True)
    json.dump({"M": M, "widest_layer": nmax, "largest_fanin": kmax,
               "gamma_star": gamma_star, "steps": args.steps, "rows": rows},
              open(runs_path("paper_noise_leak.json"), "w"), indent=1)
    return 0


# ---------------------------------------------------------------------------
def run_margin(args) -> int:
    """The margin of Theorem 8.2 along the trajectory of the instance."""
    import torch
    from host import M_STAR, LevEvaluator, Tape, read_host, tau_star
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=True)
    tau = tau_star(sigma, M_STAR)
    mem, pc, dev, states, n = list(M_STAR), 0, Tape(tau), [], 0
    want = args.states
    while pc != HALT_PC:
        if n % 2000 == 0 and len(states) < want:
            states.append((pc, list(mem)))
        A = mem[pc]
        B = mem[(pc + 1) & 0xFF]
        C = mem[(pc + 2) & 0xFF]
        r = (mem[B] - mem[A]) & 0xFF
        mem[B] = r
        pc = C if (r == 0 or r >= 0x80) else (pc + 3) & 0xFF
        dev.apply(mem)
        n += 1
    V = torch.stack([L._vec(p, m) for p, m in states]).to(args.device)
    worst, frac, y = float("inf"), 0.0, V
    for W, b in zip(L.W, L.B):
        pre = y @ W.T + b
        worst = min(worst, float((pre + 0.5).abs().min()))
        frac = max(frac, float((pre - pre.round()).abs().max()))
        y = (pre >= 0).float()
    print(f"[margin] {len(states)} states sampled along the {n:,} steps of the "
          f"run, all {L.info['layers']} layers: minimum distance of a "
          f"pre-activation from -1/2 is {worst:.6f}, maximum distance from an "
          f"integer is {frac:.3e}")
    json.dump({"states": len(states), "steps": n,
               "layers": int(L.info["layers"]), "min_margin": worst,
               "max_noninteger": frac},
              open(runs_path("paper_noise_margin.json"), "w"), indent=1)
    return 0 if abs(worst - 0.5) < 1e-9 and frac == 0.0 else 1


def upper_tail(x: float) -> float:
    """P(e > x) for a standard normal e."""
    return 0.5 * math.erfc(x / math.sqrt(2.0))


def poisson_cdf(k: int, lam: float) -> float:
    """P(N <= k) for N Poisson with mean lam."""
    if k < 0:
        return 0.0
    if lam <= 0.0:
        return 1.0
    logs = [i * math.log(lam) - lam - math.lgamma(i + 1) for i in range(k + 1)]
    m = max(logs)
    return min(1.0, math.exp(m) * sum(math.exp(t - m) for t in logs))


def poisson_interval(k: int, conf: float = 0.95) -> tuple:
    """The exact (Garwood) two-sided interval for the mean of a Poisson count k."""
    a = (1.0 - conf) / 2.0

    def bisect(f, lo, hi):
        for _ in range(200):
            mid = (lo + hi) / 2.0
            if f(mid):
                hi = mid
            else:
                lo = mid
        return (lo + hi) / 2.0

    top = 10.0 * (k + 10)
    lower = 0.0 if k == 0 else bisect(lambda lam: 1.0 - poisson_cdf(k - 1, lam) >= a, 0.0, top)
    upper = bisect(lambda lam: poisson_cdf(k, lam) <= a, 0.0, top)
    return lower, upper


def run_profile(args) -> int:
    """The pre-activations along the exact trajectory, and the one-sided bound."""
    import torch
    from host import M_P, LevEvaluator, describe, read_host
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=False)
    M = L.info["size"]
    r_H = describe(sigma)
    T = args.steps
    v = L._vec(0, M_P).unsqueeze(0).to(args.device)
    D = Devices(L, [r_H], args.device)
    lo, hi = -300, 300
    hist = torch.zeros(hi - lo + 1, dtype=torch.float64, device=args.device)
    step_hist = torch.zeros(T, hi - lo + 1, dtype=torch.float64, device=args.device)
    t0 = time.perf_counter()
    for n in range(T):
        x = v
        for groups, width in zip(L.plan, L.widths):
            y = torch.empty(1, width, device=args.device)
            for rows, idx, w, b in groups:
                pre = (x[:, idx] * w).sum(-1) + b
                z = pre.round().long().flatten()
                c = torch.bincount(z - lo, minlength=hi - lo + 1).double()
                hist += c
                step_hist[n] += c
                y[:, rows] = (pre >= 0).float()
            x = y
        v = x
        if D.step(v)[0]:
            T = n + 1
            break
    z = torch.arange(lo, hi + 1)
    per_step = (hist.cpu() / T)
    at = {int(k): float(per_step[k - lo]) for k in range(-3, 3)}
    curve = json.load(open(runs_path("paper_noise_curve.json")))
    rows = []
    steps_hist = step_hist[:T].cpu()
    for r in curve["rows"]:
        s = r["sigma"]
        q = torch.tensor([upper_tail(abs(int(k) + 0.5) / s) for k in z.tolist()],
                         dtype=torch.float64)
        b_t = steps_hist @ q                    # the bound at every step of the orbit
        cum = torch.cat([torch.zeros(1, dtype=torch.float64), b_t.cumsum(0)])
        # each copy is exposed from step 1 to its first deviation, or to the
        # end of the run; the bound on its expected number of deviations is
        # the sum of b_t over those steps
        exp_steps = [f if f is not None else r["steps"] for f in r["first_deviation"]]
        expected = float(sum(cum[min(e, T)] for e in exp_steps))
        exposure = sum(exp_steps)
        weighted = expected / exposure if exposure else 0.0
        # the rows of a layer read the same exact input and fail independently, so one minus the
        # product of their probabilities of holding is the probability that some row fails
        bp_t = -torch.expm1(steps_hist @ torch.log1p(-q))
        cum_p = torch.cat([torch.zeros(1, dtype=torch.float64), bp_t.cumsum(0)])
        expected_p = float(sum(cum_p[min(e, T)] for e in exp_steps))
        weighted_p = expected_p / exposure if exposure else 0.0
        lo, hi = poisson_interval(r["deviated"])
        p = math.erfc(1.0 / (2.0 * math.sqrt(2.0) * s))
        rows.append({"sigma": s, "measured": r["empirical_rate_per_step"],
                     "deviated": r["deviated"], "exposure": exposure,
                     "rate_interval_95": [lo / exposure, hi / exposure],
                     "bound_prop": M * p, "bound_profile_mean": float(b_t.mean()),
                     "bound_profile_exposure": weighted,
                     "expected_deviations_bound": expected,
                     "bound_product_exposure": weighted_p,
                     "expected_deviations_product": expected_p,
                     "ratio_profile": (weighted / r["empirical_rate_per_step"])
                     if r["empirical_rate_per_step"] else None,
                     "ratio_product": (weighted_p / r["empirical_rate_per_step"])
                     if r["empirical_rate_per_step"] else None})
        print(f"    s={s:.2f}: measured {r['empirical_rate_per_step']:.3e} "
              f"[{lo / exposure:.3e}, {hi / exposure:.3e}] ({r['deviated']} deviations), "
              f"profile bound {weighted:.3e}, product form {weighted_p:.3e}, M p {M * p:.3e}",
              flush=True)
    half = at[-1] + at[0]
    print(f"[profile] {T:,} steps of the exact orbit, M = {M:,} rows per step: "
          f"on average {at[-1]:,.0f} rows at z = -1 and {at[0]:,.0f} at z = 0, "
          f"{half / M:.4f} of all rows at distance 1/2 from the comparator "
          f"({time.perf_counter() - t0:.0f} s)")
    json.dump({"steps": T, "M": M, "rows_per_step_at": at,
               "fraction_at_half": half / M,
               "histogram_per_step": {int(k): float(c) for k, c in
                                      zip(z.tolist(), per_step.tolist()) if c},
               "rows": rows},
              open(runs_path("paper_noise_profile.json"), "w"), indent=1)
    return 0


def run_certify(args) -> int:
    """Active-input counts along an orbit."""
    import torch
    from host import M_P, M_STAR, LevEvaluator, describe, read_host, tau_star
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=False)
    M = L.info["size"]
    print(f"[certify] M = {M:,} rows in {L.info['layers']} layers", flush=True)
    leak = json.load(open(runs_path("paper_noise_leak.json"))) \
        if os.path.exists(runs_path("paper_noise_leak.json")) else None

    def walk(mem0, tape, cap, label):
        v = L._vec(0, mem0).unsqueeze(0).to(args.device)
        D = Devices(L, [tape], args.device)
        nl = len(L.plan)
        a_max = torch.zeros(nl, dtype=torch.long, device=args.device)
        k_max = torch.zeros(nl, dtype=torch.long, device=args.device)
        pop_max = torch.zeros(nl, dtype=torch.long, device=args.device)
        t0 = time.perf_counter()
        T = cap
        for n in range(cap):
            x = v
            for li, (groups, width) in enumerate(zip(L.plan, L.widths)):
                y = torch.empty(1, width, device=args.device)
                pop = x.sum().long()
                pop_max[li] = torch.maximum(pop_max[li], pop)
                for rows, idx, w, b in groups:
                    g = x[:, idx]
                    k1 = (g * (w != 0)).sum(-1).long()
                    k_max[li] = torch.maximum(k_max[li], k1.max())
                    a_max[li] = torch.maximum(a_max[li], (pop - k1).max())
                    y[:, rows] = (((g * w).sum(-1) + b) >= 0).float()
                x = y
            v = x
            if D.step(v)[0]:
                T = n + 1
                break
            if (n + 1) % 100000 == 0:
                print(f"    {label}: {n + 1:,} steps ({(n + 1) / (time.perf_counter() - t0):.0f}/s)",
                      flush=True)
        A_star, K_star = int(a_max.max()), int(k_max.max())
        rec = {"steps": T, "A_star": A_star, "K_star": K_star,
               "per_layer_active_zero_max": a_max.tolist(),
               "per_layer_active_programmed_max": k_max.tolist(),
               "per_layer_popcount_max": pop_max.tolist(),
               "gamma_certified": 0.5 / A_star, "epsilon_certified": 0.5 / K_star,
               "seconds": time.perf_counter() - t0}
        print(f"    {label}: {T:,} steps; at most {A_star} active zero-weight inputs and "
              f"{K_star} active programmed inputs at any row: gamma < {0.5 / A_star:.3e}, "
              f"eps < {0.5 / K_star:.3e} ({rec['seconds']:.0f} s)", flush=True)
        return rec

    out = {"M": M, "widest_layer": max(L.widths)}
    out["selfdesc_10000"] = walk(M_P, describe(sigma), args.steps, "self-describing orbit")
    if leak is not None:
        A_star, K_star = out["selfdesc_10000"]["A_star"], out["selfdesc_10000"]["K_star"]
        rows = []
        for r in leak["rows"]:
            cert = r["epsilon"] * K_star + r["gamma"] * A_star
            rows.append({"epsilon": r["epsilon"], "gamma": r["gamma"], "worst_bound": r["worst_bound"],
                         "orbit_bound": cert, "certified_by_orbit": cert < 0.5,
                         "first_deviation": r["first_deviation"]})
        out["leak_rows"] = rows
        n_cert = sum(r["certified_by_orbit"] for r in rows)
        n_exact = sum(r["first_deviation"] is None for r in rows)
        print(f"    of the {len(rows)} rows of the leak table, Proposition 8.3 certifies "
              f"{sum(r['worst_bound'] < 0.5 for r in rows)}, the orbit certificate {n_cert}, "
              f"and {n_exact} were exact")
    if not args.short:
        tau = tau_star(sigma, M_STAR)
        out["omega_star"] = walk(M_STAR, tau, 1 << 21, "self-reproducing run")
    json.dump(out, open(runs_path("paper_noise_certify.json"), "w"), indent=1)
    return 0


def run_masked(args) -> int:
    """The steps in which some row fails, against the product form b'(t), and the fraction of them
    that leave the state on the orbit; every noisy copy starts each step from the exact state."""
    import torch
    from host import M_P, LevEvaluator, describe, read_host
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=True)
    M = sum(int(W.shape[0]) for W in L.W)
    r_H = describe(sigma)
    B, T = args.batch, args.steps
    grid = [float(x) for x in args.grid.split(",")]
    print(f"[masked] M = {M:,}; {B} noisy copies restarted from the exact state at each of "
          f"{T:,} steps", flush=True)
    rows = []
    for s in grid:
        gen = torch.Generator(device=args.device).manual_seed(11 + int(1000 * s))
        v = L._vec(0, M_P).unsqueeze(0).to(args.device)
        D = Devices(L, [r_H], args.device)
        failed = deviated = 0
        expected = variance = 0.0
        t0 = time.perf_counter()
        steps = T
        for n in range(T):
            y = v.repeat(B + 1, 1)
            bad = torch.zeros(B, dtype=torch.bool, device=args.device)
            keep = torch.zeros((), dtype=torch.float64, device=args.device)
            for W, b in zip(L.W, L.B):
                pre = y @ W.T + b
                # copy 0 is noise free and reads the orbit's values, so its pre-activations are the
                # orbit's, and a row fails with probability Q(|z + 1/2| / s)
                q = 0.5 * torch.erfc((pre[0].double() + 0.5).abs() / (s * math.sqrt(2.0)))
                keep += torch.log1p(-q).sum()
                noise = torch.randn(pre.shape, generator=gen, device=pre.device) * s
                noise[0] = 0.0
                y = ((pre + noise) >= -0.5).float()
                bad |= (y[1:] != y[0]).any(dim=1)
            off = (y[1:] != y[0]).any(dim=1)
            failed += int(bad.sum())
            deviated += int(off.sum())
            bp = float(-torch.expm1(keep))          # b'(t): some row of the step fails
            expected += B * bp
            variance += B * bp * (1.0 - bp)
            v = y[0:1]
            if D.step(v)[0]:
                steps = n + 1
                break
        dt = time.perf_counter() - t0
        n_cs = B * steps
        rows.append({"sigma": s, "copies": B, "steps": steps, "copy_steps": n_cs,
                     "row_failure_steps": failed, "deviations": deviated,
                     "masked": failed - deviated,
                     "masked_fraction": (failed - deviated) / failed if failed else None,
                     "row_failure_rate": failed / n_cs, "deviation_rate": deviated / n_cs,
                     "product_form_mean": expected / n_cs,
                     "row_failures_expected": expected,
                     "row_failures_sd": math.sqrt(variance),
                     "seconds": dt})
        print(f"    s={s:.2f}: {failed:,} steps with a failing row against {expected:,.1f} "
              f"(sd {math.sqrt(variance):.1f}) from b'; {deviated:,} left the orbit, "
              f"{failed - deviated:,} masked"
              + (f" ({(failed - deviated) / failed:.2%})" if failed else "")
              + f" ({dt / 60:.1f} min)", flush=True)
    json.dump({"M": M, "grid": grid, "rows": rows},
              open(runs_path("paper_noise_masked.json"), "w"), indent=1)
    return 0


def run_intervals(args) -> int:
    """Exact Poisson intervals for the measured rates of the host and of the threefold host."""
    out = {"confidence": 0.95, "host": [], "amplified_r3": []}
    for r in json.load(open(runs_path("paper_noise_curve.json")))["rows"]:
        lo, hi = poisson_interval(r["deviated"])
        e = r["step_exposure"]
        out["host"].append({"sigma": r["sigma"], "deviated": r["deviated"], "exposure": e,
                            "rate": r["empirical_rate_per_step"], "rate_low": lo / e, "rate_high": hi / e})
    for r in json.load(open(runs_path("paper_redundant_restore_r3.json")))["rows"]:
        lo, hi = poisson_interval(r["failed"])
        e = r["exposure"]
        out["amplified_r3"].append({"sigma": r["sigma"], "failed": r["failed"], "exposure": e,
                                    "rate": r["failure_rate_per_step"], "rate_low": lo / e, "rate_high": hi / e})
    json.dump(out, open(runs_path("paper_noise_intervals.json"), "w"), indent=1)
    print(json.dumps(out, indent=1))
    return 0


def main_host() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["omega", "curve", "leak", "margin", "profile", "certify", "intervals",
                                     "masked"])
    ap.add_argument("--short", action="store_true", help="certify: the leak orbit only")
    ap.add_argument("--states", type=int, default=800)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--batch", type=int, default=256)
    ap.add_argument("--grid", default="0.02,0.04,0.06,0.08,0.09,0.10,0.11,0.12,0.15,0.20")
    ap.add_argument("--steps", type=int, default=4000)
    args = ap.parse_args()
    return {"omega": run_omega, "curve": run_curve, "leak": run_leak,
            "margin": run_margin, "profile": run_profile,
            "certify": run_certify, "intervals": run_intervals, "masked": run_masked}[args.mode](args)

SRC = os.path.join(REPO, "src")


def sha(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


# ---------------------------------------------------------------------------
# the construction
# ---------------------------------------------------------------------------

class Amplified:
    """N_r = A_r(Lev(N)) for the netlist N of a serialization, in flat form."""

    def __init__(self, sigma: bytes, r: int):
        assert r >= 1 and r % 2 == 1, "the redundancy must be odd"
        n, units, nxt = parse_sigma(sigma)
        ln, lunits, lnxt, info = levelize(n, units, nxt)
        assert ln == n
        self.r, self.n, self.sigma0 = r, n, sigma
        self.layers: List[int] = info["layer_sizes"]
        self.depth = info["depth"]
        self.lev_rows = len(lunits)
        self.lev_units = lunits
        self.lev_nxt = lnxt
        shift = (r - 1) // 2

        def pos(k: int, c: int) -> int:
            return c * n + k if k < n else r * n + r * (k - n) + c

        out: List[Tuple[int, List[Tuple[int, int]]]] = []
        for b, ins in lunits:
            ins_r = [(pos(k, c), w) for k, w in ins for c in range(r)]
            for c in range(r):
                out.append((r * b + shift, ins_r))
        self.units = out
        self.nxt = [r * n + r * (lnxt[i] - n) + c for c in range(r) for i in range(n)]
        self.state_bits = r * n
        self.edges = sum(len(ins) for _, ins in out)

    def sigma(self) -> bytes:
        bias, fanin, pred = [], [], []
        for b, ins in self.units:
            bias.append(b)
            fanin.append(len(ins))
            for k, w in ins:
                pred.append(w * (k + 1))
        lay = dict(STATE_LAYOUT)
        lay["replicas"] = self.r
        lay["copy_stride"] = self.n
        return encode(8, bias, fanin, pred, self.nxt, f"subleq8r{self.r}", lay, IO_CELLS)

    def write_flat(self, path: str) -> None:
        write_flat(path, self.state_bits, self.units, self.nxt)

    # -- the layered plan for the evaluator on the accelerator ----------------
    def plan(self, device: str, pad: bool):
        """Per layer, the rows, input positions, weights and biases of the amplified layer."""
        import torch
        r, n = self.r, self.n
        base = [n]
        for L in self.layers:
            base.append(base[-1] + L)
        plan = []
        j = 0
        for ell, L in enumerate(self.layers, start=1):
            rows_of = []
            for m in range(L):
                b, ins = self.lev_units[j]
                j += 1
                if ell == 1:
                    idx = [c * n + k for k, w in ins for c in range(r)]
                else:
                    idx = [r * (k - base[ell - 2]) + c for k, w in ins for c in range(r)]
                w = [wt for k, wt in ins for c in range(r)]
                for c in range(r):
                    rows_of.append((m * r + c, idx, w, r * b + (r - 1) // 2))
            groups = []
            if pad:
                K = max(len(x[1]) for x in rows_of)
                idx_t = torch.zeros(len(rows_of), K, dtype=torch.long)
                w_t = torch.zeros(len(rows_of), K)
                for q, (row, idx, w, b) in enumerate(rows_of):
                    idx_t[q, :len(idx)] = torch.tensor(idx)
                    w_t[q, :len(w)] = torch.tensor(w, dtype=torch.float32)
                rows_t = torch.tensor([x[0] for x in rows_of])
                b_t = torch.tensor([x[3] for x in rows_of], dtype=torch.float32)
                groups.append((rows_t.to(device), idx_t.to(device), w_t.to(device), b_t.to(device)))
            else:
                by_k: Dict[int, list] = {}
                for x in rows_of:
                    by_k.setdefault(len(x[1]), []).append(x)
                for K in sorted(by_k):
                    g = by_k[K]
                    rows_t = torch.tensor([x[0] for x in g])
                    idx_t = torch.tensor([x[1] for x in g], dtype=torch.long)
                    w_t = torch.tensor([x[2] for x in g], dtype=torch.float32)
                    b_t = torch.tensor([x[3] for x in g], dtype=torch.float32)
                    groups.append((rows_t.to(device), idx_t.to(device), w_t.to(device), b_t.to(device)))
            plan.append((r * L, groups))
        # the final layer holds copy c of state bit i at i r + c; the state
        # holds it at c n + i
        perm = torch.tensor([i * r + c for c in range(r) for i in range(n)]).to(device)
        return plan, perm


class Evaluator:
    """N_r on the accelerator, with read noise and the device on the copies."""

    def __init__(self, amp: Amplified, device: str = "cuda", pad: bool = True,
                 graph: bool = False):
        import torch
        self.torch = torch
        self.amp, self.device = amp, device
        self.r, self.n = amp.r, amp.n
        self.plan, self.perm = amp.plan(device, pad)
        self.M = sum(L for L, _ in self.plan)
        self.N = amp.state_bits
        cells = (IO_CELLS["c_wr"], IO_CELLS["c_eot"], IO_CELLS["c_rw"],
                 IO_CELLS["c_rd"], IO_CELLS["r_in"], IO_CELLS["r_out"])
        self.cells = cells
        # bit t of cell j of copy c
        self.cell_idx = torch.tensor([[c * self.n + 9 + 8 * j + t for j in cells for t in range(8)]
                                      for c in range(self.r)], device=device)
        self.pow2 = torch.tensor([1 << (7 - t) for t in range(8)], device=device)
        self.shift = torch.tensor([7 - t for t in range(8)])
        self.graph = None
        self.gen = None
        if graph:
            self.capture()

    # -- one application of the map, over a batch --------------------------
    def _forward(self, v, s: float = 0.0, gen=None):
        torch = self.torch
        x = v
        for width, groups in self.plan:
            y = torch.empty(x.shape[0], width, device=x.device)
            for rows, idx, w, b in groups:
                pre = (x[:, idx] * w).sum(-1) + b
                if s > 0:
                    pre = pre + torch.randn(pre.shape, generator=gen, device=pre.device) * s
                y[:, rows] = (pre >= -0.5).float()
            x = y
        return x[:, self.perm]

    def step(self, v):
        if self.graph is not None and v.shape[0] == 1 and self.gs == 0.0:
            return self._replay(v)
        return self._forward(v)

    def step_noisy(self, v, s: float, gen):
        if self.graph is not None and v.shape[0] == 1 and self.gs == s and gen is self.gen:
            return self._replay(v)
        return self._forward(v, s, gen)

    def capture(self, s: float = 0.0, seed: int = 0):
        """Capture one application, with noise of deviation s, as a CUDA graph."""
        torch = self.torch
        self.gs = s
        self.gin = torch.zeros(1, self.N, device=self.device)
        gen = None
        if s > 0:
            gen = torch.Generator(device=self.device).manual_seed(seed)
            self.gen = gen
        for _ in range(3):
            self._forward(self.gin, s, gen)
        torch.cuda.synchronize()
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self.gout = self._forward(self.gin, s, gen)
        torch.cuda.synchronize()
        if s == 0.0:
            probe = (torch.rand(1, self.N, device=self.device) < 0.5).float()
            self.gin.copy_(probe)
            self.graph.replay()
            torch.cuda.synchronize()
            assert bool((self.gout == self._forward(probe)).all()), \
                "the captured graph differs from the eager step"
        return self

    def _replay(self, v):
        if v.data_ptr() != self.gin.data_ptr():
            self.gin.copy_(v)
        self.graph.replay()
        self.gin.copy_(self.gout)
        return self.gin

    # -- states --------------------------------------------------------------
    def vec(self, pc: int, mem: List[int], copies: int = 1):
        torch = self.torch
        v0 = torch.zeros(self.n)
        for k in range(8):
            v0[k] = (pc >> (7 - k)) & 1
        for j in range(256):
            for k in range(8):
                v0[9 + j * 8 + k] = (mem[j] >> (7 - k)) & 1
        v = v0.repeat(self.r).unsqueeze(0).to(self.device)
        return v.repeat(copies, 1)

    def device_step(self, v, tapes: List[Tape], shadow: List[List[int]]):
        """The device on every copy: cells read by majority, every copy written."""
        torch = self.torch
        B = v.shape[0]
        bits = v[:, self.cell_idx]                       # (B, r, 48)
        maj = (bits.sum(1) * 2 > self.r).long()          # (B, 48)
        vals = (maj.reshape(B, len(self.cells), 8) * self.pow2).sum(-1).cpu()
        halts = (v[:, [c * self.n + 8 for c in range(self.r)]].sum(1) * 2 > self.r).cpu().tolist()
        new = torch.empty(B, len(self.cells), dtype=torch.long)
        for i in range(B):
            for k, j in enumerate(self.cells):
                shadow[i][j] = int(vals[i, k])
            tapes[i].apply(shadow[i])
            for k, j in enumerate(self.cells):
                new[i, k] = shadow[i][j]
        nb = ((new.unsqueeze(-1) >> self.shift) & 1).reshape(B, 1, -1).float().to(self.device)
        v[:, self.cell_idx] = nb.expand(B, self.r, -1)
        return halts


# ---------------------------------------------------------------------------
def build_amplified(r: int) -> Amplified:
    return Amplified(read_host(), r)


def compile_c(work: str, cc: str = "gcc") -> str:
    """The evaluator in C, the levels of a large netlist evaluated in parallel."""
    exe = os.path.join(work, "netlist_eval")
    subprocess.run([cc, "-O3", "-march=native", "-fopenmp", "-o", exe,
                    os.path.join(SRC, "netlist_eval.c")], check=True)
    os.environ.setdefault("OMP_WAIT_POLICY", "ACTIVE")
    os.environ.setdefault("OMP_PROC_BIND", "true")
    os.environ.setdefault("OMP_NUM_THREADS", str(max(1, min(16, (os.cpu_count() or 2) - 4))))
    return exe


def cmd_build(args) -> int:
    t0 = time.perf_counter()
    amp = build_amplified(args.r)
    s = amp.sigma()
    t_build = time.perf_counter() - t0
    rec = {"r": args.r, "state_bits": amp.state_bits, "units": len(amp.units),
           "edges": amp.edges, "depth": amp.depth, "lev_rows": amp.lev_rows,
           "layer_widths": [args.r * L for L in amp.layers],
           "sigma_bytes": len(s), "sigma_sha256": sha(s),
           "bias_min": min(b for b, _ in amp.units), "bias_max": max(b for b, _ in amp.units),
           "max_fanin": max(len(i) for _, i in amp.units), "build_seconds": t_build}
    print(f"[build] N_{args.r}: {amp.state_bits:,} state bits, {len(amp.units):,} units, "
          f"{amp.edges:,} edges, depth {amp.depth}; sigma {len(s):,} bytes, "
          f"sha {sha(s)[:16]} ({t_build:.1f} s)", flush=True)
    # the serialization denotes the netlist it was written from
    n2, u2, x2 = parse_sigma(s)
    rec["sigma_roundtrip"] = (n2 == amp.state_bits and u2 == amp.units and x2 == amp.nxt)
    print(f"    parses back to the same netlist: {rec['sigma_roundtrip']}")
    # one step against SUBLEQ on the evaluator in C, every copy checked
    work = tempfile.mkdtemp(prefix="redundant_")
    exe = compile_c(work, args.cc)
    flat = os.path.join(work, "net.bin")
    amp.write_flat(flat)
    p = subprocess.run([exe, "rand", flat, str(args.states), "1"], capture_output=True, text=True)
    rec["c_random"] = p.stdout.strip()
    rec["c_random_ok"] = p.returncode == 0 and " 0 disagree" in p.stdout
    print(f"    {p.stderr.strip()}")
    print(f"    {p.stdout.strip()}")
    # the margins: every pre-activation on a replicated state is r z + (r-1)/2
    import torch
    device = "cuda" if torch.cuda.is_available() else "cpu"
    E = Evaluator(amp, device=device, pad=True)
    gen = torch.Generator().manual_seed(11)
    worst = float("inf")
    for _ in range(4):
        base = (torch.rand(8, amp.n, generator=gen) < 0.5).float()
        x = base.repeat(1, args.r).to(device)
        for width, groups in E.plan:
            y = torch.empty(x.shape[0], width, device=device)
            for rows, idx, w, b in groups:
                pre = (x[:, idx] * w).sum(-1) + b
                worst = min(worst, float((pre + 0.5).abs().min()))
                y[:, rows] = (pre >= 0).float()
            x = y
    rec["min_margin_random_states"] = worst
    print(f"    minimum distance of a pre-activation from -1/2 on 32 replicated random "
          f"states: {worst}")
    # the evaluator on the accelerator agrees with the evaluator in C on a prefix
    tau = tau_star(read_host(), M_STAR)
    v = E.vec(0, M_STAR)
    tapes, shadow = [Tape(tau)], [[0] * 256]
    for _ in range(300):
        v = E.step(v)
        E.device_step(v, tapes, shadow)
    st = v[0].cpu().long().tolist()
    open(os.path.join(work, "mstar"), "wb").write(bytes(M_STAR))
    open(os.path.join(work, "tau"), "wb").write(tau)
    p = subprocess.run([exe, "run", flat, os.path.join(work, "mstar"), os.path.join(work, "tau"),
                        os.path.join(work, "out"), "net", "300"], capture_output=True, text=True)
    # replay the C state through the same 300 steps by the reference, then compare the
    # emitted prefix and the memory read by majority
    mem = list(M_STAR); pc = 0; dev = Tape(tau)
    for _ in range(300):
        a, b, c = mem[pc], mem[(pc + 1) & 255], mem[(pc + 2) & 255]
        rr = (mem[b] - mem[a]) & 255
        mem[b] = rr
        pc = c if (rr == 0 or rr >= 128) else (pc + 3) & 255
        dev.apply(mem)
    got_mem = [sum(st[c * amp.n + 9 + 8 * j + t] << (7 - t) for t in range(8))
               for j in range(256) for c in (0,)]
    all_copies_equal = all(st[c * amp.n + i] == st[i] for c in range(args.r) for i in range(amp.n))
    rec["accelerator_prefix_ok"] = (got_mem == mem and all_copies_equal
                                    and tapes[0].out == dev.out
                                    and open(os.path.join(work, "out"), "rb").read() == bytes(dev.out))
    print(f"    300 steps of Omega* on the accelerator and in C agree with the reference, "
          f"all copies equal: {rec['accelerator_prefix_ok']}")
    rec["ok"] = bool(rec["sigma_roundtrip"] and rec["c_random_ok"]
                     and abs(worst - args.r / 2) < 1e-9 and rec["accelerator_prefix_ok"])
    json.dump(rec, open(runs_path(f"paper_redundant_r{args.r}.json"), "w"), indent=1)
    print("REDUNDANT HOST:", "PASS" if rec["ok"] else "FAIL")
    return 0 if rec["ok"] else 1


def census(r: bytes) -> Dict[str, int]:
    """L, R, f_lit, f_rep of a recipe."""
    L = R = fl = fr = 0
    i = 0
    while True:
        t = r[i]
        if t == 0:
            break
        if 1 <= t <= 127:
            L += 1; fl += t; i += 1 + t
        else:
            R += 1; fr += 256 - t; i += 2
    return {"L": L, "R": R, "f_lit": fl, "f_rep": fr}


def closed_form(c: Dict[str, int], tape_len: int) -> Dict[str, int]:
    a = 5 * c["f_lit"] + 4 * c["f_rep"] + 7 * c["L"] + 9 * c["R"] + 7
    return {"A": a, "R": 1, "B": 6 * tape_len + 2, "total": a + 1 + 6 * tape_len + 2}


def cmd_selfrep(args) -> int:
    """Omega*_r on the evaluator in C in lockstep with the reference, to completion."""
    t0 = time.perf_counter()
    amp = build_amplified(args.r)
    s_r = amp.sigma()
    tau = tau_star(s_r, M_STAR)
    target = ser(s_r, bytes(M_STAR), tau)
    cen = census(tau)
    cf = closed_form(cen, len(tau))
    print(f"[selfrep] N_{args.r}: sigma {len(s_r):,} bytes, tau*_r {len(tau):,} bytes, "
          f"ser {len(target):,} bytes; closed form {cf['total']:,} steps", flush=True)
    work = tempfile.mkdtemp(prefix="redundant_")
    exe = compile_c(work, args.cc)
    flat = os.path.join(work, "net.bin")
    amp.write_flat(flat)
    open(os.path.join(work, "mstar"), "wb").write(bytes(M_STAR))
    open(os.path.join(work, "tau"), "wb").write(tau)
    t1 = time.perf_counter()
    p = subprocess.run([exe, "run", flat, os.path.join(work, "mstar"), os.path.join(work, "tau"),
                        os.path.join(work, "out"), args.cmode], capture_output=True, text=True)
    dt = time.perf_counter() - t1
    out = open(os.path.join(work, "out"), "rb").read()
    exact = out == target
    import re
    m = re.search(r"after (\d+) steps", p.stdout)
    steps = int(m.group(1)) if m else None
    sg, mb, tb = inst(out) if exact else (None, None, None)
    rec = {"r": args.r, "sigma_bytes": len(s_r), "sigma_sha256": sha(s_r),
           "tau_bytes": len(tau), "ser_bytes": len(target), "ser_sha256": sha(target),
           "census": cen, "closed_form": cf, "mode": args.cmode,
           "c_result": p.stdout.strip(), "steps": steps, "seconds": dt,
           "steps_per_second": steps / dt if steps and dt else None,
           "out_bytes": len(out), "out_sha256": sha(out), "equals_ser": exact,
           "closed_form_agrees": steps == cf["total"],
           "inst_recovers_instance": bool(exact and sg == s_r and mb == bytes(M_STAR) and tb == tau),
           "units": len(amp.units), "edges": amp.edges, "state_bits": amp.state_bits}
    rec["ok"] = bool(exact and rec["closed_form_agrees"] and rec["inst_recovers_instance"]
                     and p.returncode == 0)
    print(f"    {p.stderr.strip()}")
    print(f"    {p.stdout.strip()} ({dt / 60:.1f} min, {rec['steps_per_second'] or 0:,.0f} steps/s)")
    print(f"    Out = ser(Omega*_r): {exact}; closed form agrees: {rec['closed_form_agrees']}; "
          f"inst recovers the instance: {rec['inst_recovers_instance']}; sha {sha(out)[:16]}")
    rec["total_seconds"] = time.perf_counter() - t0
    json.dump(rec, open(runs_path(f"paper_redundant_selfrep_r{args.r}.json"), "w"), indent=1)
    print("REDUNDANT SELF-REPRODUCTION:", "PASS" if rec["ok"] else "FAIL")
    return 0 if rec["ok"] else 1


def tail_bound(s: float, r: int, M: int, T: int) -> float:
    """T M erfc(r / (2 sqrt 2 s)): the bound of the fault-tolerance theorem."""
    return T * M * math.erfc(r / (2.0 * math.sqrt(2.0) * s))


def cmd_omega(args) -> int:
    """Omega* run on N_r under read noise on every pre-activation, to completion."""
    import torch
    amp = build_amplified(args.r)
    sigma = read_host()
    tau = tau_star(sigma, M_STAR)
    target = ser(sigma, bytes(M_STAR), tau)
    E = Evaluator(amp, device=args.device, pad=True)
    M = E.M
    print(f"[omega] N_{args.r}: {amp.depth} layers, M = {M:,}, target {len(target):,} bytes",
          flush=True)
    grid = [float(x) for x in args.grid.split(",")]
    rows = []
    for s in grid:
        seed = 20260929 + int(1000 * s)
        try:
            E.capture(s, seed)
            gen = E.gen
        except Exception as e:                       # noqa: BLE001
            print(f"    (graph capture unavailable: {e}; running eagerly)", flush=True)
            E.graph = None
            gen = torch.Generator(device=args.device).manual_seed(seed)
        v = E.vec(0, M_STAR)
        tapes, shadow = [Tape(tau)], [[0] * 256]
        steps, wrong, halted = 0, None, False
        budget = args.budget
        t0 = time.perf_counter()
        while steps < budget:
            v = E.step_noisy(v, s, gen)
            steps += 1
            before = len(tapes[0].out)
            halted = E.device_step(v, tapes, shadow)[0]
            out = tapes[0].out
            if len(out) > before:
                k = len(out) - 1
                if k >= len(target) or out[k] != target[k]:
                    wrong = k
                    break
            if halted:
                break
            if steps % 200000 == 0:
                print(f"    s={s:.2f}: {steps:,} steps, {len(out):,} bytes "
                      f"({steps / (time.perf_counter() - t0):.0f}/s)", flush=True)
        dt = time.perf_counter() - t0
        emitted = len(tapes[0].out)
        complete = emitted == len(target) and wrong is None and halted
        outcome = ("emitted the whole instance" if complete else
                   f"a wrong byte after {wrong} correct" if wrong is not None
                   else "the halt bit was set" if halted else "the budget was exhausted")
        rows.append({"sigma": s, "steps": steps, "bytes": emitted, "complete": complete,
                     "first_wrong_byte": wrong, "halted": halted, "seconds": dt,
                     "bound": tail_bound(s, args.r, M, max(steps, 1))})
        print(f"    s={s:.2f}: {steps:,} steps, {emitted:,} bytes, {outcome} "
              f"({dt / 60:.1f} min); bound {rows[-1]['bound']:.3e}", flush=True)
        json.dump({"r": args.r, "M": M, "target_bytes": len(target), "rows": rows},
                  open(runs_path(f"paper_redundant_omega_r{args.r}.json"), "w"), indent=1)
    return 0


def upper_tail_amp(x: float) -> float:
    return 0.5 * math.erfc(x / math.sqrt(2.0))


def cmd_curve(args) -> int:
    """The rate at which noisy copies of N_r leave the orbit, against the profile and first-order bounds."""
    import torch
    amp = build_amplified(args.r)
    sigma = read_host()
    r_H = describe(sigma)
    E = Evaluator(amp, device=args.device, pad=False)
    M = E.M
    B = args.batch
    print(f"[curve] N_{args.r}: M = {M:,}; {B} noisy copies per level, {args.steps:,} steps each",
          flush=True)
    prof = None
    pp = runs_path("paper_noise_profile.json")
    if os.path.exists(pp):
        prof = json.load(open(pp))["histogram_per_step"]
    grid = [float(x) for x in args.grid.split(",")]
    rows = []
    for s in grid:
        cap = args.steps if s >= 0.24 else min(args.steps, 3000)
        gen = torch.Generator(device=args.device).manual_seed(7 + int(1000 * s))
        v = E.vec(0, M_P, copies=B + 1)
        tapes = [Tape(r_H) for _ in range(B + 1)]
        shadow = [[0] * 256 for _ in range(B + 1)]
        first = [None] * B
        live = cap
        t0 = time.perf_counter()
        for n in range(cap):
            x = v
            for width, groups in E.plan:
                y = torch.empty(x.shape[0], width, device=x.device)
                for rws, idx, w, b in groups:
                    pre = (x[:, idx] * w).sum(-1) + b
                    noise = torch.randn(pre.shape, generator=gen, device=pre.device) * s
                    noise[0] = 0.0
                    y[:, rws] = ((pre + noise) >= -0.5).float()
                x = y
            v = x[:, E.perm]
            diff = (v[1:] != v[0]).any(dim=1).cpu().tolist()
            for i, d in enumerate(diff):
                if d and first[i] is None:
                    first[i] = n + 1
            if all(f is not None for f in first):
                live = n + 1
                break
            E.device_step(v, tapes, shadow)
        dt = time.perf_counter() - t0
        deviated = [f for f in first if f is not None]
        exposure = sum(f if f is not None else live for f in first)
        rate = len(deviated) / exposure if exposure else 0.0
        bound_tail = M * math.erfc(args.r / (2.0 * math.sqrt(2.0) * s))
        bound_profile = None
        if prof is not None:
            bound_profile = args.r * sum(c * upper_tail_amp(args.r * abs(int(z) + 0.5) / s)
                                         for z, c in prof.items())
        rows.append({"sigma": s, "copies": B, "steps": live, "deviated": len(deviated),
                     "median_first_deviation":
                         sorted(deviated)[len(deviated) // 2] if deviated else None,
                     "step_exposure": exposure, "empirical_rate_per_step": rate,
                     "bound_tail_per_step": bound_tail,
                     "bound_profile_per_step": bound_profile,
                     "first_deviation": first, "seconds": dt})
        print(f"    s={s:.2f}: {len(deviated)}/{B} deviated within {live:,} steps; "
              f"empirical {rate:.3e} per step, profile bound "
              f"{bound_profile if bound_profile is None else format(bound_profile, '.3e')}, "
              f"M erfc {bound_tail:.3e} ({dt / 60:.1f} min)", flush=True)
        json.dump({"r": args.r, "M": M, "grid": grid, "rows": rows},
                  open(runs_path(f"paper_redundant_curve_r{args.r}.json"), "w"), indent=1)
    return 0


def cmd_restore(args) -> int:
    """One exact and B noisy copies of N_r stepped together, each read by majority."""
    import torch
    amp = build_amplified(args.r)
    sigma = read_host()
    r_H = describe(sigma)
    E = Evaluator(amp, device=args.device, pad=False)
    n, r = amp.n, args.r
    B = args.batch
    grid = [float(x) for x in args.grid.split(",")]
    rows = []
    print(f"[restore] N_{r}: {B} noisy copies, up to {args.steps:,} steps per level", flush=True)
    for s in grid:
        gen = torch.Generator(device=args.device).manual_seed(31 + int(1000 * s))
        v = E.vec(0, M_P, copies=B + 1)
        tapes = [Tape(r_H) for _ in range(B + 1)]
        shadow = [[0] * 256 for _ in range(B + 1)]
        first = [None] * B
        wrong_copy_steps = 0
        wrong_copies = 0
        worst = 0
        live = torch.ones(B, dtype=torch.bool)
        t0 = time.perf_counter()
        ran = args.steps
        for t in range(args.steps):
            x = v
            for width, groups in E.plan:
                y = torch.empty(x.shape[0], width, device=x.device)
                for rws, idx, w, b in groups:
                    pre = (x[:, idx] * w).sum(-1) + b
                    noise = torch.randn(pre.shape, generator=gen, device=pre.device) * s
                    noise[0] = 0.0
                    y[:, rws] = ((pre + noise) >= -0.5).float()
                x = y
            v = x[:, E.perm]
            ex = v[0].reshape(r, n)
            wrong = (v[1:].reshape(B, r, n) != ex.unsqueeze(0))
            k = wrong.sum(dim=(1, 2)).cpu()
            maj = (v[1:].reshape(B, r, n).sum(1) * 2 > r)
            bad = (maj != (ex[0] > 0.5).unsqueeze(0)).any(1).cpu()
            for i in range(B):
                if live[i]:
                    if int(k[i]):
                        wrong_copy_steps += 1
                        wrong_copies += int(k[i])
                        worst = max(worst, int(k[i]))
                    if bool(bad[i]):
                        first[i] = t + 1
                        live[i] = False
            if not bool(live.any()):
                ran = t + 1
                break
            E.device_step(v, tapes, shadow)
        dt = time.perf_counter() - t0
        failed = [f for f in first if f is not None]
        exposure = sum(f if f is not None else ran for f in first)
        rows.append({"sigma": s, "copies": B, "steps": ran, "failed": len(failed),
                     "median_first_failure": sorted(failed)[len(failed) // 2] if failed else None,
                     "exposure": exposure, "failure_rate_per_step": len(failed) / exposure,
                     "steps_with_wrong_copies": wrong_copy_steps, "wrong_copies": wrong_copies,
                     "max_wrong_copies_in_a_step": worst, "first_failure": first, "seconds": dt})
        print(f"    s={s:.2f}: {len(failed)}/{B} majority-read failures within {ran:,} steps, "
              f"exposure {exposure:,}, rate {len(failed) / exposure:.3e} per step; "
              f"{wrong_copies:,} wrong copies in {wrong_copy_steps:,} copy-steps before the "
              f"failures ({dt / 60:.1f} min)", flush=True)
        json.dump({"r": r, "copies": B, "cap": args.steps, "rows": rows},
                  open(runs_path(f"paper_redundant_restore_r{r}.json"), "w"), indent=1)
    return 0


# ---------------------------------------------------------------------------
# the restoration bound
# ---------------------------------------------------------------------------

def write_bound_input(amp: "Amplified", path: str, mem: List[int], tape: bytes) -> None:
    """Lev(N_host) layer by layer as CSR, with the device cells, memory and tape, for restore_bound.c."""
    import struct
    from host import C_WR, C_EOT, C_RW, C_RD, R_IN, R_OUT
    n = amp.n
    base = [n]
    for L in amp.layers:
        base.append(base[-1] + L)
    assert amp.layers[-1] == n and all(amp.lev_nxt[i] == base[-2] + i for i in range(n)), \
        "the last layer of Lev(N_host) is the state, in order"
    with open(path, "wb") as f:
        f.write(struct.pack("<ii", n, len(amp.layers)))
        j0 = 0
        for ell, L in enumerate(amp.layers, start=1):
            ip, ix, w, bias = [0], [], [], []
            for m in range(L):
                b, ins = amp.lev_units[j0 + m]
                for k, wt in sorted(ins):
                    ix.append(k if ell == 1 else k - base[ell - 2])
                    w.append(wt)
                ip.append(len(ix))
                bias.append(b)
            j0 += L
            cols = n if ell == 1 else amp.layers[ell - 2]
            f.write(struct.pack("<iii", L, cols, len(ix)))
            for arr in (ip, ix, w, bias):
                f.write(struct.pack(f"<{len(arr)}i", *arr))
        dev = sorted(9 + 8 * c + t for c in (C_WR, C_EOT, C_RW, C_RD, R_IN, R_OUT) for t in range(8))
        f.write(struct.pack("<i", len(dev)))
        f.write(struct.pack(f"<{len(dev)}i", *dev))
        f.write(bytes(mem))
        f.write(struct.pack("<i", len(tape)))
        f.write(bytes(tape))


def bound_series(exe: str, inp: str, s: float, total: int, chunks: int, workers: int,
                 work: str, warm: int = 20):
    """b(t) for t = 0 .. total-1, or to the halt, in chunks run in parallel."""
    import numpy as np
    from concurrent.futures import ThreadPoolExecutor
    edges = np.linspace(0, total, chunks + 1).astype(int)

    def one(k):
        a, b = int(edges[k]), int(edges[k + 1])
        out = os.path.join(work, f"b_{s:.3f}_{k}.bin")
        p = subprocess.run([exe, inp, repr(s), str(a), str(b - a), str(warm), "1e-20", out],
                           capture_output=True, text=True)
        assert p.returncode == 0, p.stderr
        vals = np.fromfile(out)
        os.remove(out)
        return vals, "halted=1" in p.stdout

    with ThreadPoolExecutor(workers) as ex:
        res = list(ex.map(one, range(chunks)))
    series, halted = [], False
    for vals, h in res:
        series.append(vals)
        if h:
            halted = True
            break
    return np.concatenate(series), halted


def cmd_bound(args) -> int:
    """The restoration bound of Proposition 8.14 along the exact orbit, by restore_bound.c."""
    import numpy as np
    amp = build_amplified(3)
    sigma = read_host()
    work = tempfile.mkdtemp()
    exe = os.path.join(work, "restore_bound")
    subprocess.run([args.cc, "-O3", "-march=native", "-o", exe, os.path.join(SRC, "restore_bound.c"), "-lm"],
                   check=True)
    desc_in, repr_in = os.path.join(work, "describing.bin"), os.path.join(work, "reproducing.bin")
    write_bound_input(amp, desc_in, M_P, describe(sigma))
    write_bound_input(amp, repr_in, M_STAR, tau_star(sigma, M_STAR))
    restore = {}
    rp = runs_path("paper_redundant_restore_r3.json")
    if os.path.exists(rp):
        for row in json.load(open(rp))["rows"]:
            restore[round(row["sigma"], 3)] = row
    workers = max(1, min(args.workers, os.cpu_count() or 1))
    out = {"r": 3, "eps_rel": 1e-20, "warm": 20, "steps": args.steps, "rows": [], "whole_run": []}
    for s in [float(x) for x in args.grid.split(",")]:
        t0 = time.perf_counter()
        b, _ = bound_series(exe, desc_in, s, args.steps, args.chunks, workers, work)
        cb = np.concatenate([[0.0], np.cumsum(b)])
        row = {"sigma": s, "mean_per_step": float(cb[-1] / len(b)), "copy_fails_within": float(cb[-1]),
               "max_per_step": float(b.max()),
               "first_order_per_step": 3 * sum(amp.layers) * math.erfc(3 / (2 * math.sqrt(2) * s))}
        m = restore.get(round(s, 3))
        if m:
            ran = [f if f is not None else m["steps"] for f in m["first_failure"]]
            row["weighted_per_step"] = float(sum(cb[e] for e in ran) / sum(ran))
            row["measured_per_step"] = m["failure_rate_per_step"]
            row["bound_over_measured"] = row["weighted_per_step"] / m["failure_rate_per_step"] \
                if m["failure_rate_per_step"] > 0 else None
        row["seconds"] = time.perf_counter() - t0
        out["rows"].append(row)
        print(f"    s={s:.2f}: {row.get('weighted_per_step', row['mean_per_step']):.3e} per step"
              + (f", measured {row['measured_per_step']:.3e}" if m else "")
              + f", first-order {row['first_order_per_step']:.3e} ({row['seconds']:.0f} s)", flush=True)
        json.dump(out, open(runs_path("paper_redundant_bound_r3.json"), "w"), indent=1)
    for s in [float(x) for x in args.whole.split(",") if x]:
        t0 = time.perf_counter()
        b, halted = bound_series(exe, repr_in, s, 1_664_939, args.whole_chunks, workers, work)
        row = {"sigma": s, "steps": len(b), "halted": halted, "sum": float(b.sum()),
               "first_order": len(b) * 3 * sum(amp.layers) * math.erfc(3 / (2 * math.sqrt(2) * s)),
               "seconds": time.perf_counter() - t0}
        out["whole_run"].append(row)
        print(f"    whole run, s={s:.2f}: {row['sum']:.3e} over {row['steps']:,} steps, first-order "
              f"{row['first_order']:.3e} ({row['seconds'] / 60:.1f} min)", flush=True)
        json.dump(out, open(runs_path("paper_redundant_bound_r3.json"), "w"), indent=1)
    return 0


def cmd_layers(args) -> int:
    """The profile bound of N_r split by layer."""
    import torch
    from host import M_P, LevEvaluator
    r = args.r
    sigma = read_host()
    L = LevEvaluator(sigma, device=args.device, dense=False)
    v = L._vec(0, M_P).unsqueeze(0).to(args.device)
    D = Devices(L, [describe(sigma)], args.device)
    lo, hi = -300, 300
    nl = len(L.plan)
    H = torch.zeros(nl, hi - lo + 1, dtype=torch.float64, device=args.device)
    T = args.steps
    for n in range(args.steps):
        x = v
        for li, (groups, width) in enumerate(zip(L.plan, L.widths)):
            y = torch.empty(1, width, device=args.device)
            for rows, idx, w, b in groups:
                pre = (x[:, idx] * w).sum(-1) + b
                H[li] += torch.bincount(pre.round().long().flatten() - lo,
                                        minlength=hi - lo + 1).double()
                y[:, rows] = (pre >= 0).float()
            x = y
        v = x
        if D.step(v)[0]:
            T = n + 1
            break
    H = (H / T).cpu()
    zs = list(range(lo, hi + 1))
    measured = {}
    cp = runs_path(f"paper_redundant_curve_r{r}.json")
    if os.path.exists(cp):
        for row in json.load(open(cp))["rows"]:
            measured[round(row["sigma"], 3)] = (row["empirical_rate_per_step"], row["deviated"])
    out = {"r": r, "steps": T,
           "per_layer_rows_at_half": [float(H[l][-1 - lo] + H[l][-lo]) for l in range(nl)],
           "last_layer_hist": {z: float(H[nl - 1][z - lo]) for z in zs if H[nl - 1][z - lo] > 0},
           "rows": []}
    for s in [float(x) for x in args.grid.split(",")]:
        q = torch.tensor([upper_tail_amp(r * abs(z + 0.5) / s) for z in zs], dtype=torch.float64)
        per_layer = [r * float(H[l] @ q) for l in range(nl)]
        m = measured.get(round(s, 3))
        out["rows"].append({"sigma": s, "measured": m[0] if m else None,
                            "deviated": m[1] if m else None, "full": sum(per_layer),
                            "last_layer": per_layer[-1], "per_layer": per_layer})
        print(f"    s={s:.2f}: last layer {per_layer[-1]:.3e}, all layers {sum(per_layer):.3e}"
              + (f", measured {m[0]:.3e} ({m[1]} deviations)" if m else ""), flush=True)
    json.dump(out, open(runs_path(f"paper_redundant_layers_r{r}.json"), "w"), indent=1)
    return 0


def main_amplified() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["build", "selfrep", "omega", "curve", "restore", "layers", "bound"])
    ap.add_argument("--r", type=int, default=3)
    ap.add_argument("--cc", default="gcc")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--states", type=int, default=1 << 16)
    ap.add_argument("--batch", type=int, default=256)
    ap.add_argument("--steps", type=int, default=None)
    ap.add_argument("--budget", type=int, default=1 << 21)
    ap.add_argument("--cmode", default="lock",
                    help="for selfrep: lock (compared with the reference every step) or net")
    ap.add_argument("--grid", default=None)
    ap.add_argument("--whole", default="0.20,0.22,0.25,0.27",
                    help="for bound: the noise levels of the sums over the self-reproducing run")
    ap.add_argument("--chunks", type=int, default=10, help="for bound: chunks of the 20,000 steps")
    ap.add_argument("--whole-chunks", type=int, default=36, help="for bound: chunks of the whole run")
    ap.add_argument("--workers", type=int, default=6, help="for bound: chunks run at once")
    args = ap.parse_args()
    if args.steps is None:
        args.steps = 20000 if args.mode == "bound" else 4000
    if args.grid is None:
        args.grid = {"omega": "0.10,0.15,0.20,0.25",
                     "curve": "0.06,0.12,0.18,0.24,0.27,0.30,0.33,0.36,0.45,0.60",
                     "restore": "0.27,0.30,0.33,0.36",
                     "layers": "0.06,0.12,0.18,0.24,0.27,0.30,0.33,0.36,0.45,0.60",
                     "bound": "0.10,0.15,0.20,0.22,0.25,0.27,0.30,0.33,0.36"}.get(args.mode, "")
    return {"build": cmd_build, "selfrep": cmd_selfrep, "omega": cmd_omega,
            "curve": cmd_curve, "restore": cmd_restore, "layers": cmd_layers,
            "bound": cmd_bound}[args.mode](args)

COMMANDS = {'host': main_host, 'amplified': main_amplified}


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        sys.exit('commands: ' + ', '.join(COMMANDS))
    cmd = sys.argv.pop(1)
    sys.argv[0] += ' ' + cmd
    sys.exit(COMMANDS[cmd]())