"""Code written from the paper's text alone: H_A, the levelization, noise and weight error."""
import hashlib
import json
import math
import os
import struct
import subprocess
import sys
import time

import numpy as np
import torch


RUNS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'paper', 'runs')


class Net:
    def __init__(self, n_state):
        self.n = n_state
        self.bias = []
        self.ins = []
        self.block = []
        self.cur = None

    def u(self, b, ins):
        assert all(w in (-1, 1) for _, w in ins)
        assert len({s for s, _ in ins}) == len(ins)
        self.bias.append(b)
        self.ins.append(ins)
        self.block.append(self.cur)
        return self.n + len(self.bias) - 1

    def AND(self, *lits):
        return self.u(-sum(1 for _, pos in lits if pos), [(s, 1 if pos else -1) for s, pos in lits])

    def OR(self, *lits):
        return self.u(sum(1 for _, pos in lits if not pos) - 1, [(s, 1 if pos else -1) for s, pos in lits])


def build(A, fold=False):
    S = 1 << A
    P_IDX = list(range(A))
    H = A
    MU = lambda j, t: A + 1 + j * A + t
    n = A + 1 + A * S
    N = Net(n)

    def sel(bits, j, extra=()):
        # the decode line [bits == j] as one conjunction of literals, Lemma 5.6
        lits = [(bits[i], bool((j >> i) & 1)) for i in range(A)] + list(extra)
        return N.AND(*lits)

    N.cur = 'P'
    P = [sel(P_IDX, j) for j in range(S)]
    N.cur = 'HP'
    HP = [N.AND((H, True), (P_IDX[i], True)) for i in range(A)]
    N.cur = 'FA'
    if fold:
        FA = [[[sel(P_IDX, j, [(MU((j + o) % S, t), True)]) for j in range(S)] for t in range(A)] for o in range(3)]
    else:
        FA = [[[N.AND((MU((j + o) % S, t), True), (P[j], True)) for j in range(S)] for t in range(A)] for o in range(3)]
    N.cur = 'Q'
    Q = [N.OR(*[(P[j], True) for j in range(S) if ((j + 3) % S >> i) & 1]) for i in range(A)]
    N.cur = 'F'
    F = [[N.OR(*[(FA[o][t][j], True) for j in range(S)]) for t in range(A)] for o in range(3)]
    if not fold:
        N.cur = 'PAB'
        PA = [sel(F[0], j) for j in range(S)]
        PB = [sel(F[1], j) for j in range(S)]
    N.cur = 'AC'
    AC = N.AND(*[(F[2][t], True) for t in range(A)])
    N.cur = 'W'
    if fold:
        W = [sel(F[1], j, [(H, False)]) for j in range(S)]
    else:
        W = [N.AND((PB[j], True), (H, False)) for j in range(S)]
    N.cur = 'RXY'
    if fold:
        RX = [[sel(F[0], j, [(MU(j, t), True)]) for j in range(S)] for t in range(A)]
        RY = [[sel(F[1], j, [(MU(j, t), True)]) for j in range(S)] for t in range(A)]
    else:
        RX = [[N.AND((MU(j, t), True), (PA[j], True)) for j in range(S)] for t in range(A)]
        RY = [[N.AND((MU(j, t), True), (PB[j], True)) for j in range(S)] for t in range(A)]
    N.cur = 'KP'
    KP = [[N.AND((MU(j, t), True), (W[j], False)) for t in range(A)] for j in range(S)]
    N.cur = 'XY'
    X = [N.OR(*[(RX[t][j], True) for j in range(S)]) for t in range(A)]
    Y = [N.OR(*[(RY[t][j], True) for j in range(S)]) for t in range(A)]
    # the carry-lookahead subtractor of Lemma 5.3, r = y - x
    N.cur = 'sub'
    pp = [N.u(0, [(Y[i], 1), (X[i], -1)]) for i in range(A - 1)]
    Tstar, Tj, c = {}, {}, {}
    for k in range(1, A):
        Tstar[k] = N.u(-k, [(pp[i], 1) for i in range(k)])
        for j in range(k):
            Tj[k, j] = N.u(-(k - j), [(Y[j], 1), (X[j], -1)] + [(pp[i], 1) for i in range(j + 1, k)])
    for k in range(1, A):
        c[k] = N.u(-1, [(Tstar[k], 1)] + [(Tj[k, j], 1) for j in range(k)])
    hm = {}
    for k in range(1, A):
        for m in (1, 2, 3):
            hm[m, k] = N.u(1 - m, [(Y[k], 1), (X[k], -1), (c[k], 1)])
    o1 = N.OR((X[0], True), (Y[0], True))
    o2 = N.u(1, [(X[0], -1), (Y[0], -1)])
    r = [N.AND((o1, True), (o2, True))]
    for k in range(1, A):
        r.append(N.u(-1, [(hm[1, k], 1), (hm[2, k], -1), (hm[3, k], 1)]))
    N.cur = 'eq'
    eqs = []
    for t in range(A):
        a1 = N.AND((X[t], True), (Y[t], True))
        a0 = N.AND((X[t], False), (Y[t], False))
        eqs.append(N.OR((a1, True), (a0, True)))
    e = N.AND(*[(q, True) for q in eqs])
    N.cur = 'LEQ'
    leq = N.u(-1, [(hm[1, A - 1], 1), (hm[2, A - 1], -1), (hm[3, A - 1], 1), (e, 1)])
    N.cur = 'NP'
    TA = [N.AND((H, False), (leq, True), (F[2][i], True)) for i in range(A)]
    TB = [N.AND((H, False), (leq, False), (Q[i], True)) for i in range(A)]
    NP = [N.OR((HP[i], True), (TA[i], True), (TB[i], True)) for i in range(A)]
    N.cur = 'NH'
    HA = N.AND((H, False), (leq, True), (AC, True))
    HB = N.AND((H, False), (leq, False), (P[S - 4], True))
    NH = N.OR((H, True), (HA, True), (HB, True))
    N.cur = 'NM'
    NM = {}
    for j in range(S):
        for t in range(A):
            wr = N.AND((r[t], True), (W[j], True))
            NM[j, t] = N.OR((KP[j][t], True), (wr, True))
    nxt = [None] * n
    for i in range(A):
        nxt[P_IDX[i]] = NP[i]
    nxt[H] = NH
    for j in range(S):
        for t in range(A):
            nxt[MU(j, t)] = NM[j, t]
    return N, nxt


def stats(N, nxt):
    depth = [0] * (N.n + len(N.bias))
    for k, ins in enumerate(N.ins):
        depth[N.n + k] = 1 + max((depth[s] for s, _ in ins), default=0)
    blocks = {}
    for k, b in enumerate(N.block):
        d = depth[N.n + k]
        lo, hi = blocks.get(b, (d, d))
        blocks[b] = (min(lo, d), max(hi, d))
    return dict(units=len(N.bias), edges=sum(len(i) for i in N.ins),
                depth=max(depth[s] for s in nxt), bias_range=[min(N.bias), max(N.bias)],
                block_layers=blocks)


def evaluate(N, nxt, x):
    """x: (n, batch) uint8; returns the next state (n, batch)."""
    sig = np.zeros((N.n + len(N.bias), x.shape[1]), dtype=np.int32)
    sig[:N.n] = x
    for k, (b, ins) in enumerate(zip(N.bias, N.ins)):
        idx = np.fromiter((s for s, _ in ins), dtype=np.int64, count=len(ins))
        w = np.fromiter((w for _, w in ins), dtype=np.int32, count=len(ins))
        if len(ins) == 0:
            pre = np.full(x.shape[1], b, dtype=np.int32)
        else:
            pre = w @ sig[idx] + b
        sig[N.n + k] = pre >= 0
    return sig[nxt].astype(np.uint8)


def subleq_step(A, x):
    """Definition 5.2 on the state layout above, columns of x being states."""
    S = 1 << A
    batch = x.shape[1]
    p = np.zeros(batch, dtype=np.int64)
    for i in range(A):
        p |= x[i].astype(np.int64) << i
    h = x[A].astype(bool)
    mem = np.zeros((S, batch), dtype=np.int64)
    for t in range(A):
        mem |= x[A + 1 + np.arange(S) * A + t].astype(np.int64) << t
    cols = np.arange(batch)
    a = mem[p, cols]
    b = mem[(p + 1) % S, cols]
    c = mem[(p + 2) % S, cols]
    newb = (mem[b, cols] - mem[a, cols]) % S
    signed = np.where(newb >= S // 2, newb - S, newb)
    np_ = np.where(signed <= 0, c, (p + 3) % S)
    mem2 = mem.copy()
    mem2[b, cols] = newb
    mem2 = np.where(h[None, :], mem, mem2)
    p2 = np.where(h, p, np_)
    h2 = h | (p2 == S - 1)
    y = np.zeros_like(x)
    for i in range(A):
        y[i] = (p2 >> i) & 1
    y[A] = h2
    for t in range(A):
        y[A + 1 + np.arange(S) * A + t] = (mem2 >> t) & 1
    return y


def check(A, fold, states, batch, seed):
    N, nxt = build(A, fold)
    st = stats(N, nxt)
    rng = np.random.default_rng(seed)
    n = N.n
    bad = 0
    halted_bad = 0
    done = 0
    while done < states:
        bsz = min(batch, states - done)
        x = rng.integers(0, 2, size=(n, bsz), dtype=np.uint8)
        x[A] = 0
        x[A, : bsz // 8] = 1                  # an eighth of the states halted
        y = evaluate(N, nxt, x)
        ref = subleq_step(A, x)
        bad += int(np.any(y != ref, axis=0).sum())
        halted = x[A] == 1
        halted_bad += int(np.any(y[:, halted] != x[:, halted], axis=0).sum())
        done += bsz
    formula = (8 * A + 4) * (1 << A) + (A * A + 39 * A + 4) // 2
    return dict(A=A, fold=fold, states=states, disagree=bad, halted_not_fixed=halted_bad,
                formula_units=None if fold else formula, **{k: v for k, v in st.items() if k != 'block_layers'},
                block_layers={k: list(v) for k, v in st['block_layers'].items()})

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

SIGMA = os.path.join(REPO, 'variants', 'neural_subleq8io_netlist.safetensors')

WORK = os.path.join(REPO, 'src', '.work')

OUT = os.path.join(WORK, 'host_lev.bin')


def parse(path):
    b = open(path, 'rb').read()
    n = struct.unpack('<Q', b[:8])[0]
    hdr = json.loads(b[8:8 + n])
    base = 8 + n
    md = hdr.pop('__metadata__')
    T = {}
    for k, v in hdr.items():
        dt = {'I8': np.int8, 'I16': np.int16, 'I32': np.int32}[v['dtype']]
        lo, hi = v['data_offsets']
        T[k] = np.frombuffer(b[base + lo:base + hi], dtype=dt).astype(np.int64)
    return b, md, T


def levelize(nstate, bias, fanin, pred, nxt):
    U = len(bias)
    starts = np.concatenate([[0], np.cumsum(fanin)])
    ins = [[(abs(int(e)) - 1, 1 if e > 0 else -1) for e in pred[starts[u]:starts[u + 1]]] for u in range(U)]
    depth = [0] * (nstate + U)
    for u in range(U):
        depth[nstate + u] = 1 + max((depth[s] for s, _ in ins[u]), default=0)
    d = max(depth[s] for s in nxt)
    # units no next value depends on
    need = [False] * (nstate + U)
    stack = [s for s in nxt]
    while stack:
        s = stack.pop()
        if need[s]:
            continue
        need[s] = True
        if s >= nstate:
            stack.extend(t for t, _ in ins[s - nstate])
    # lambda(u): the largest layer that consumes u; next values consumed at d + 1
    lam = [-1] * (nstate + U)
    for u in range(U):
        if need[nstate + u]:
            for s, _ in ins[u]:
                lam[s] = max(lam[s], depth[nstate + u])
    for s in nxt:
        lam[s] = d + 1
    layers = []                                    # layer l: list of (key, bias, [(src_key, w)])
    pos = [{('s', i): i for i in range(nstate)}]   # key -> row index in layer
    for l in range(1, d + 1):
        rows = []
        if l < d:
            for s in range(nstate + U):
                if not need[s] and s >= nstate:
                    continue
                if s >= nstate and depth[s] == l:
                    u = s - nstate
                    rows.append((('s', s), int(bias[u]), [(('s', t), w) for t, w in ins[u]]))
                elif depth[s] < l < lam[s]:
                    rows.append((('s', s), -1, [(('s', s), 1)]))
        else:
            for i, s in enumerate(nxt):
                if s >= nstate and depth[s] == d:
                    u = s - nstate
                    rows.append((('o', i), int(bias[u]), [(('s', t), w) for t, w in ins[u]]))
                else:
                    rows.append((('o', i), -1, [(('s', s), 1)]))
        layers.append(rows)
        pos.append({k: r for r, (k, _, _) in enumerate(rows)})
    # resolve each input to its row in the layer before
    flat = []
    for l, rows in enumerate(layers, start=1):
        prev = pos[l - 1]
        flat.append([(b, [(prev[k], w) for k, w in sl]) for _, b, sl in rows])
    return d, flat


def frame(u):
    return struct.pack('<Q', len(u)) + bytes(u)


def nu(x):
    return (-x) % 256


def describe(f):
    out, i = [], 0
    while i < len(f):
        ell = 1
        while i + ell < len(f) and f[i + ell] == f[i] and ell < 127:
            ell += 1
        if ell >= 4:
            out += [256 - ell, nu(f[i])]
            i += ell
            continue
        k = i + 1
        while k < len(f) and k < i + 127:
            if k + 3 < len(f) and f[k] == f[k + 1] == f[k + 2] == f[k + 3]:
                break
            k += 1
        out += [k - i] + [nu(c) for c in f[i:k]]
        i = k
    return bytes(out + [0])


Z, ONE, T1, T2, NEG1, EOK = 0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5

C_WR, C_EOT, C_RW, C_RD, R_IN, R_OUT = 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE

TABLE9 = [
    (T2, T2, 0x03), (T1, T1, 0x06), (NEG1, C_RD, 0x09), (R_IN, T1, 0x0C), (T1, T2, 0x0F),
    (Z, T2, 0x15), (Z, Z, 0x1B), (Z, T1, 0x3C), (Z, Z, 0x2D), (NEG1, C_RD, 0x1E),
    (R_IN, R_OUT, 0x21), (NEG1, C_WR, 0x24), (ONE, T2, 0x2A), (Z, Z, 0x1B), (Z, Z, 0x00),
    (NEG1, C_RD, 0x30), (R_IN, R_OUT, 0x33), (NEG1, C_WR, 0x36), (ONE, T1, 0x2A), (Z, Z, 0x30),
    (Z, Z, 0xFF)]

TABLE10 = [(NEG1, C_RW, 0x3F), (NEG1, C_RD, 0x42), (C_EOT, EOK, 0xFF), (R_IN, Z, 0x48),
           (Z, R_OUT, 0x4B), (NEG1, C_WR, 0x4E), (Z, Z, 0x3F)]


def image(prog):
    mem = [0] * 256
    for k, ins in enumerate(prog):
        mem[3 * k:3 * k + 3] = list(ins)
    mem[ONE], mem[NEG1], mem[EOK] = 1, 255, 1
    return bytes(mem)


def main_levelize():
    raw, md, T = parse(SIGMA)
    nstate = int(T['manifest.state_bits'][0])
    d, layers = levelize(nstate, T['net.bias'], T['net.fanin'], T['net.pred'], [int(v) for v in T['net.next']])
    widths = [nstate] + [len(r) for r in layers]
    M = sum(widths[1:])
    nnz = sum(len(sl) for rows in layers for _, sl in rows)
    ident = sum(1 for rows in layers for b, sl in rows if b == -1 and len(sl) == 1 and sl[0][1] == 1)
    reg = json.loads(md['signal_registry'])
    io = json.loads(md['io'])
    m_p = image(TABLE9)
    m_star = image(TABLE9[:20] + TABLE10)
    r_h = describe(raw)
    tau = describe(frame(raw) + frame(m_star))
    ser = frame(raw) + frame(m_star) + tau
    info = dict(depth=d, state_bits=nstate, M=M, nnz=nnz, max_width=max(widths[1:]), identity_like=ident,
                dense_entries=sum(widths[i] * widths[i + 1] for i in range(d)),
                r_H=len(r_h), tau=len(tau), ser=len(ser), ser_sha=hashlib.sha256(ser).hexdigest(),
                sigma_sha=hashlib.sha256(raw).hexdigest(), registry=reg, io=io)
    print(json.dumps(info))
    os.makedirs(WORK, exist_ok=True)
    json.dump(info, open(os.path.join(REPO, 'paper', 'runs', 'independent_levelize.json'), 'w'), indent=1)
    with open(OUT, 'wb') as f:
        w = lambda *v: f.write(struct.pack('<%di' % len(v), *v))
        w(nstate, d, M, nnz)
        w(*widths)
        for rows in layers:
            for b, sl in rows:
                w(b, len(sl))
        for rows in layers:
            for b, sl in rows:
                for src, wt in sl:
                    w(src, wt)
        w(reg['pc'][0], reg['halt'][0], reg['mem'][0])
        w(io['c_wr'], io['c_eot'], io['c_rw'], io['c_rd'], io['r_in'], io['r_out'], io['halt_pc'])
        for blob in (m_p, r_h, m_star, tau, ser):
            w(len(blob))
            f.write(blob)

NSIM, BIN, HIST = os.path.join(WORK, 'nsim'), os.path.join(WORK, 'host_lev.bin'), os.path.join(WORK, 'hist_sd.bin')

M = 41967

Q = lambda u: 0.5 * math.erfc(u / math.sqrt(2.0))


def nsim(*args):
    return subprocess.run([NSIM, BIN] + [str(a) for a in args], capture_output=True, text=True, check=True).stdout


def profile_series(s, r=1, steps=4000):
    """b(t) along the exact orbit: the sum over rows of r Q(r (k + 1/2) / s)."""
    h = np.fromfile(HIST, dtype=np.int64).reshape(-1, 64)[:steps]
    q = np.array([r * Q(r * (k + 0.5) / s) for k in range(64)])
    return h @ q


def curve(s, copies, T, seed, r=1):
    first = json.loads(nsim('curve', 0, s, copies, T, seed, r))
    dev = [f for f in first if f > 0]
    runs = [f if f > 0 else T for f in first]
    exposure = sum(runs)
    cum = np.concatenate([[0.0], np.cumsum(profile_series(s, r))])
    expected = float(sum(cum[e] for e in runs))
    return dict(s=s, r=r, deviated=len(dev), copies=copies, T=T,
                median=float(np.median(dev)) if dev else None, exposure=exposure,
                rate=len(dev) / exposure, profile_bound=expected / exposure, bound_on_count=expected,
                rMp=r * M * math.erfc(r / (2 * math.sqrt(2) * s)))


def main_noise():
    os.makedirs(WORK, exist_ok=True)
    subprocess.run(['gcc', '-O3', '-march=native', '-fopenmp', '-o', NSIM, os.path.join(REPO, 'src', 'nsim.c'), '-lm'], check=True)
    out = {}
    t0 = time.time()
    out['exact_self_describing'] = json.loads(nsim('exact', 0, 10000, 4000, HIST))
    out['exact_self_reproducing'] = json.loads(nsim('exact', 1, 2000000, 0, '/dev/null'))
    h = np.fromfile(HIST, dtype=np.int64).reshape(-1, 64)[:4000].mean(axis=0)
    out['profile'] = dict(at_one_half=float(h[0]), at_three_halves=float(h[1]), at_five_halves=float(h[2]),
                          beyond=float(h[3:].sum()), fraction_at_one_half=float(h[0] / M))
    out['curve'] = [curve(s, 256, 3000 if s <= 0.06 else 4000, int(s * 1000) + 17)
                    for s in (0.02, 0.04, 0.06, 0.08, 0.09, 0.10, 0.11, 0.12, 0.15, 0.20)]
    out['curve_s008_seeds'] = [curve(0.08, 256, 4000, seed) for seed in range(101, 109)]
    out['deviation_N3'] = [curve(s, 256, 4000, int(s * 1000) + 29, r=3) for s in (0.24, 0.27, 0.30, 0.33)]
    out['failure_N3'] = []
    for s in (0.27, 0.30, 0.33, 0.36):
        first = json.loads(nsim('fail', s, 64, 20000, int(s * 1000) + 41, 3))
        runs = [f if f > 0 else 20000 for f in first]
        failed = sum(1 for f in first if f > 0)
        out['failure_N3'].append(dict(s=s, failed=failed, copies=64, exposure=sum(runs), rate=failed / sum(runs)))
    procs = [subprocess.Popen([NSIM, BIN, 'run', s, r, seed], stdout=subprocess.PIPE, text=True)
             for s, r, seed in (('0.05', '1', '906'), ('0.10', '1', '911'), ('0.15', '1', '916'), ('0.20', '1', '921'),
                                ('0.10', '3', '913'), ('0.15', '3', '918'), ('0.20', '3', '923'), ('0.25', '3', '928'))]
    out['complete_runs'] = [json.loads(p.communicate()[0]) for p in procs]
    out['seconds'] = round(time.time() - t0, 1)
    json.dump(out, open(os.path.join(REPO, 'paper', 'runs', 'independent_noise.json'), 'w'), indent=1)
    print(json.dumps({k: v for k, v in out.items() if k in ('exact_self_reproducing', 'profile', 'seconds')}))


def load():
    b = open(BIN, 'rb').read()
    pos = 0

    def ints(k):
        nonlocal pos
        v = struct.unpack_from('<%di' % k, b, pos)
        pos += 4 * k
        return v
    n, L, M, nnz = ints(4)
    widths = ints(L + 1)
    rows = [ints(2) for _ in range(M)]
    edges = [ints(2) for _ in range(nnz)]
    PCB, HALT, MEMB = ints(3)
    io = ints(7)
    blobs = []
    for _ in range(5):
        (ln,) = ints(1)
        blobs.append(b[pos:pos + ln])
        pos += ln
    layers, r0, e0 = [], 0, 0
    for l in range(1, L + 1):
        W = np.zeros((widths[l], widths[l - 1]), dtype=np.float32)
        bias = np.zeros(widths[l], dtype=np.float32)
        for i in range(widths[l]):
            bb, f = rows[r0 + i]
            bias[i] = bb
            for k in range(f):
                src, w = edges[e0 + k]
                W[i, src] = w
            e0 += f
        r0 += widths[l]
        layers.append((W, bias))
    return n, widths, layers, (PCB, HALT, MEMB), io, blobs


class Device:
    def __init__(self, io, MEMB, tape):
        self.C_WR, self.C_EOT, self.C_RW, self.C_RD, self.R_IN, self.R_OUT, _ = io
        self.MEMB, self.tape, self.head = MEMB, tape, 0

    def get(self, x, j):
        v = 0
        for k in range(8):
            v = (v << 1) | int(x[self.MEMB + 8 * j + k])
        return v

    def put(self, x, j, v):
        for k in range(8):
            x[self.MEMB + 8 * j + k] = (v >> (7 - k)) & 1

    def __call__(self, x):
        if self.get(x, self.C_WR) == 1:
            self.put(x, self.R_OUT, 0)
        if self.get(x, self.C_RD) == 1:
            if self.head < len(self.tape):
                self.put(x, self.R_IN, self.tape[self.head]); self.put(x, self.C_EOT, 0); self.head += 1
            else:
                self.put(x, self.C_EOT, 1)
        if self.get(x, self.C_RW) == 1:
            self.head = 0
        for c in (self.C_RD, self.C_RW, self.C_WR):
            self.put(x, c, 0)


def main_weight_error():
    n, widths, layers, (PCB, HALT, MEMB), io, blobs = load()
    dev = torch.device('cuda')
    Wt = [torch.from_numpy(W).to(dev) for W, _ in layers]
    bt = [torch.from_numpy(b).to(dev) for _, b in layers]
    steps = 10000
    settings = [(0, 0.25 / 8209), (0, 0.45 / 8209), (0, 0.75 / 8209), (0, 2 / 8209), (0, 8 / 8209),
                (0, 32 / 8209), (0, 128 / 8209), (0.25 / 256, 0.125 / 8209), (0.5 / 256, 0), (1 / 256, 0),
                (4 / 256, 0), (16 / 256, 0), (64 / 256, 0)]
    out = []
    for si, (eps, gam) in enumerate(settings):
        t0 = time.time()
        g = torch.Generator(device=dev).manual_seed(4242 + si)
        Wp = []
        for W in Wt:
            u = torch.rand(W.shape, generator=g, device=dev) * 2 - 1
            Wp.append(W + torch.where(W != 0, u * eps, u * gam))
        x = np.zeros(n, dtype=np.uint8)
        dx = Device(io, MEMB, blobs[1]); dp = Device(io, MEMB, blobs[1])
        for j, v in enumerate(blobs[0]):
            dx.put(x, j, v)
        xp = x.copy()
        first = 0
        for t in range(steps):
            a = torch.from_numpy(x.astype(np.float32)).to(dev)
            ap = torch.from_numpy(xp.astype(np.float32)).to(dev)
            for W, Wq, b in zip(Wt, Wp, bt):
                a = ((W @ a + b) >= -0.5).float()
                ap = ((Wq @ ap + b) >= -0.5).float()
            x = a.to(torch.uint8).cpu().numpy(); xp = ap.to(torch.uint8).cpu().numpy()
            dx(x); dp(xp)
            if not np.array_equal(x, xp):
                first = t + 1
                break
        res = dict(eps=eps, gamma=gam, outcome=('exact for %d steps' % steps) if not first else 'deviates at step %d' % first,
                   seconds=round(time.time() - t0, 1))
        print(json.dumps(res), flush=True)
        out.append(res)
        del Wp
        torch.cuda.empty_cache()
    json.dump(out, open(os.path.join(REPO, 'paper', 'runs', 'independent_weight_error.json'), 'w'), indent=1)


def cmd_host_family():
    out = []
    widths = [int(a) for a in sys.argv[1].split(',')] if len(sys.argv) > 1 else list(range(3, 13))
    for A in widths:
        for fold in (False, True):
            states = 1 << 16 if A <= 10 else (1 << 15 if A == 11 else 1 << 14)
            batch = 4096 if A <= 9 else (1024 if A <= 11 else 512)
            t0 = time.time()
            res = check(A, fold, states, batch, seed=1000 + A)
            res['seconds'] = round(time.time() - t0, 1)
            print(json.dumps({k: v for k, v in res.items() if k != 'block_layers'}), flush=True)
            out.append(res)
    json.dump(out, open(os.path.join(RUNS, 'independent_host_family.json'), 'w'), indent=1)


COMMANDS = {'host-family': cmd_host_family, 'levelize': main_levelize, 'noise': main_noise, 'weight-error': main_weight_error}


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]())