/* The restoration bound of the amplified host N_3 = Lev(N_host)^(3), computed along the exact
 * orbit of an instance (src/redundant.py bound).
 *
 *   restore_bound IN.bin s start len warm eps_rel PERSTEP.bin|- [skip_rel]
 *
 * IN.bin holds Lev(N_host) layer by layer (CSR over the layer before, weights +-1, biases), the
 * state bits the device rewrites, the memory image and the tape. The program follows the
 * orbit on the SUBLEQ reference to step start - warm, then on Lev(N_host) with its device,
 * and from the pre-activation z of every row computes, per step t of start .. start + len - 1,
 * the bound b(t) as the sum of
 *   t2   sum_w min(1, 3 q2(w)) sum_{u in T(w)} P2(u)
 *   t3   sum_w min(1, 3 q3(w)) sum_{u in T(w)} P3(u)
 *   tind sum_w sum_{u1<u2 in T(w)} P1(u1) P1(u2)
 *   tK   sum_w sum_{u1<u2 in T(w)} K(u1,u2)                 (sparse part)
 *   tR   sum_w (1/2) ((A R)_w)^2                            (majorant R R^T)
 *   tab  sum_w (A a)_w (A b)_w                              (majorant a b^T + b a^T)
 *   tB   sum over state rows v of P2(v)
 * with X = A P1, P1 = 3 q0 + 3 (q1-q0) X, P2 = 3 q0^2 + 3 (q1^2-q0^2) X,
 * P3 = q0^3 + (q1^3-q0^3) X, q_j = Q((D-j)/s), D = 3|z+1/2|, A the incidence of the edges
 * that push toward the comparator, and K = 9 diag(q1-q0) A (diag P1 + K) A^T diag(q1-q0) off
 * the diagonal. Pairs reached through a row with q1-q0 below DC, sources with P1 below
 * skip_rel and entries below eps_rel, both relative to 3 Q(3/(2s)), are carried by the
 * majorants, which propagate exactly, so the computed b(t) bounds the exact recursion. A run
 * that starts after step 0 begins from probability one on every state bit the device does not
 * rewrite: every quantity of the recursion bounds a probability. */
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define ZOFF 400
#define ZN (2 * ZOFF + 1)
#define DC 1e-14
static double SKIP = 1e-20;

typedef struct {
    int rows, cols, nnz;
    int *ip, *ix, *w, *bias;  /* CSR */
    int *cp, *ce, *erow;      /* children of each source: edges, sorted by child */
} Layer;

static int n, L, ndev, tlen, maxw;
static Layer *Ly;
static int *devbits;
static unsigned char mem0[256], *tape;

static void rd(void *p, size_t sz, size_t k, FILE *f) {
    if (fread(p, sz, k, f) != k) { fprintf(stderr, "short read\n"); exit(2); }
}

static void load(const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) { perror(path); exit(2); }
    rd(&n, 4, 1, f); rd(&L, 4, 1, f);
    Ly = calloc(L, sizeof(Layer));
    maxw = n;
    for (int l = 0; l < L; l++) {
        Layer *y = &Ly[l];
        rd(&y->rows, 4, 1, f); rd(&y->cols, 4, 1, f); rd(&y->nnz, 4, 1, f);
        y->ip = malloc(4 * (y->rows + 1)); y->ix = malloc(4 * y->nnz);
        y->w = malloc(4 * y->nnz); y->bias = malloc(4 * y->rows);
        rd(y->ip, 4, y->rows + 1, f); rd(y->ix, 4, y->nnz, f);
        rd(y->w, 4, y->nnz, f); rd(y->bias, 4, y->rows, f);
        if (y->rows > maxw) maxw = y->rows;
        y->erow = malloc(4 * y->nnz);
        for (int v = 0; v < y->rows; v++)
            for (int e = y->ip[v]; e < y->ip[v + 1]; e++) y->erow[e] = v;
        y->cp = calloc(y->cols + 1, 4); y->ce = malloc(4 * y->nnz);
        for (int e = 0; e < y->nnz; e++) y->cp[y->ix[e] + 1]++;
        for (int u = 0; u < y->cols; u++) y->cp[u + 1] += y->cp[u];
        int *fill = malloc(4 * y->cols);
        memcpy(fill, y->cp, 4 * y->cols);
        for (int e = 0; e < y->nnz; e++) y->ce[fill[y->ix[e]]++] = e;  /* rows ascending */
        free(fill);
    }
    rd(&ndev, 4, 1, f); devbits = malloc(4 * ndev); rd(devbits, 4, ndev, f);
    rd(mem0, 1, 256, f); rd(&tlen, 4, 1, f); tape = malloc(tlen); rd(tape, 1, tlen, f);
    fclose(f);
}

/* ---- the exact orbit -------------------------------------------------- */
static int getword(const char *x, int j) {
    int v = 0;
    for (int k = 0; k < 8; k++) v = (v << 1) | x[9 + 8 * j + k];
    return v;
}
static void putword(char *x, int j, int v) {
    for (int k = 0; k < 8; k++) x[9 + 8 * j + k] = (v >> (7 - k)) & 1;
}

/* ---- a hash map from pairs to values ------------------------------------ */
typedef struct { uint64_t key; double val; } Slot;
static Slot *H; static int *Hused; static int Hn, Hcap;
static void hclear(void) { for (int i = 0; i < Hn; i++) H[Hused[i]].key = 0; Hn = 0; }
static int hslot(uint64_t key) {
    uint64_t h = key * 0x9E3779B97F4A7C15ull;
    int i = (int)(h >> 40) & (Hcap - 1);
    while (H[i].key && H[i].key != key) i = (i + 1) & (Hcap - 1);
    return i;
}
static void hgrow(void) {  /* double the table, keeping its entries */
    Slot *old = H; int *oldu = Hused, on = Hn;
    Hcap *= 2; H = calloc(Hcap, sizeof(Slot)); Hused = malloc(4 * Hcap); Hn = 0;
    for (int k = 0; k < on; k++) {
        Slot *o = &old[oldu[k]];
        int i = hslot(o->key);
        H[i] = *o; Hused[Hn++] = i;
    }
    free(old); free(oldu);
}
static void hadd(int a, int b, double v) {
    if (a > b) { int t = a; a = b; b = t; }
    uint64_t key = ((uint64_t)(a + 1) << 32) | (uint64_t)(b + 1);
    int i = hslot(key);
    if (!H[i].key) {
        if (Hn >= Hcap / 2) { hgrow(); i = hslot(key); }
        H[i].key = key; H[i].val = 0; Hused[Hn++] = i;
    }
    H[i].val += v;
}

int main(int argc, char **argv) {
    if (argc < 8) {
        fprintf(stderr, "usage: restore_bound IN.bin s start len warm eps_rel PERSTEP.bin|- [skip_rel]\n");
        return 2;
    }
    load(argv[1]);
    double s = atof(argv[2]);
    long start = atol(argv[3]), len = atol(argv[4]), warm = atol(argv[5]);
    double eps = atof(argv[6]);
    FILE *fo = strcmp(argv[7], "-") ? fopen(argv[7], "wb") : NULL;
    if (argc > 8) SKIP = atof(argv[8]);
    long first = start > warm ? start - warm : 0;

    /* tables over pre-activations */
    static double q0[ZN], q1[ZN], q2[ZN], q3[ZN], dl[ZN];
    for (int i = 0; i < ZN; i++) {
        double D = 3.0 * fabs((i - ZOFF) + 0.5), c = 1.0 / (s * sqrt(2.0));
        q0[i] = 0.5 * erfc(D * c); q1[i] = 0.5 * erfc((D - 1) * c);
        q2[i] = 0.5 * erfc((D - 2) * c); q3[i] = 0.5 * erfc((D - 3) * c);
        dl[i] = q1[i] - q0[i];
    }
    eps *= 3 * q0[ZOFF];    /* cutoffs relative to the start probability at margin 3/2 */
    SKIP *= 3 * q0[ZOFF];

    /* orbit buffers */
    char **y = malloc(sizeof(char *) * (L + 1));
    int **z = malloc(sizeof(int *) * L);
    char **tow = malloc(sizeof(char *) * L);
    y[0] = calloc(n, 1);
    for (int l = 0; l < L; l++) {
        y[l + 1] = calloc(Ly[l].rows, 1); z[l] = malloc(4 * Ly[l].rows); tow[l] = malloc(Ly[l].nnz);
    }
    /* reach the first step on the SUBLEQ reference with its device */
    int M[256], pc = 0;
    long h = 0, out = 0;
    for (int j = 0; j < 256; j++) M[j] = mem0[j];
    for (long k = 0; k < first; k++) {
        int a = M[pc], b = M[(pc + 1) & 255], c = M[(pc + 2) & 255];
        int v = (M[b] - M[a]) & 255;
        M[b] = v;
        pc = (v == 0 || v >= 128) ? c : (pc + 3) & 255;
        if (M[0xF9] == 1) { M[0xFE] = 0; out++; }
        if (M[0xFC] == 1) { if (h < tlen) { M[0xFD] = tape[h]; M[0xFA] = 0; h++; } else M[0xFA] = 1; }
        if (M[0xFB] == 1) h = 0;
        M[0xFC] = M[0xFB] = M[0xF9] = 0;
        if (pc == 255) { printf("halted before step %ld\n", first); return 0; }
    }
    char *x = calloc(n, 1);
    for (int k = 0; k < 8; k++) x[k] = (pc >> (7 - k)) & 1;
    for (int j = 0; j < 256; j++) putword(x, j, M[j]);
    char *isdev = calloc(n, 1);
    for (int i = 0; i < ndev; i++) isdev[devbits[i]] = 1;

    /* bound buffers: previous and next layer */
    double *P1p = calloc(maxw, 8), *P2p = calloc(maxw, 8), *P3p = calloc(maxw, 8);
    double *Rp = calloc(maxw, 8), *ap = calloc(maxw, 8), *bp = calloc(maxw, 8);
    double *P1n = calloc(maxw, 8), *P2n = calloc(maxw, 8), *P3n = calloc(maxw, 8);
    double *Rn = calloc(maxw, 8), *an = calloc(maxw, 8), *bn = calloc(maxw, 8);
    double *s1 = calloc(maxw, 8), *sk = calloc(maxw, 8), *kmax = calloc(maxw, 8), *gmax = calloc(maxw, 8);
    int kcap = 1 << 20, nk = 0, nk2 = 0;
    int *ku1 = malloc(4 * kcap), *ku2 = malloc(4 * kcap), *nu1 = malloc(4 * kcap), *nu2 = malloc(4 * kcap);
    double *kv = malloc(8 * kcap), *nv = malloc(8 * kcap);
    Hcap = 1 << 15; H = calloc(Hcap, sizeof(Slot)); Hused = malloc(4 * Hcap);
    int *big = malloc(4 * maxw), *big2 = malloc(4 * maxw);

    /* a restarted chunk begins from probability one on every state bit the device does not write */
    if (first > 0)
        for (int i = 0; i < n; i++) P1p[i] = isdev[i] ? 0.0 : 1.0;
    double tot[7] = {0}, total = 0, bmax = 0;
    long steps, counted = 0;
    int halted = 0, kpeak = 0;
    double Rpeak = 0, abpeak = 0;
    for (steps = first; steps < start + len; steps++) {
        int count = steps >= start;
        /* exact step on the orbit, recording pre-activations and toward edges */
        memcpy(y[0], x, n);
        for (int l = 0; l < L; l++) {
            Layer *ly = &Ly[l];
            for (int v = 0; v < ly->rows; v++) {
                int acc = ly->bias[v];
                for (int e = ly->ip[v]; e < ly->ip[v + 1]; e++) acc += ly->w[e] * y[l][ly->ix[e]];
                z[l][v] = acc; y[l + 1][v] = acc >= 0;
                for (int e = ly->ip[v]; e < ly->ip[v + 1]; e++) {
                    int d = ly->w[e] * (1 - 2 * y[l][ly->ix[e]]);
                    tow[l][e] = acc >= 0 ? d < 0 : d > 0;
                }
            }
        }
        /* the bound for this step */
        double t[7] = {0};  /* t2 t3 tind tK tR tab tB */
        memset(P2p, 0, 8 * n); memset(P3p, 0, 8 * n);
        for (int l = 0; l < L; l++) {
            Layer *ly = &Ly[l];
            int R = ly->rows, C = ly->cols;
            /* row maxima of K and the source majorant s1 = sqrt(P1) + sqrt(max K) */
            for (int i = 0; i < nk; i++) {
                if (kv[i] > kmax[ku1[i]]) kmax[ku1[i]] = kv[i];
                if (kv[i] > kmax[ku2[i]]) kmax[ku2[i]] = kv[i];
            }
            for (int u = 0; u < C; u++) {
                s1[u] = sqrt(P1p[u]) + sqrt(kmax[u]);
                sk[u] = P1p[u] < SKIP ? sqrt(P1p[u]) : 0.0;
            }
            for (int i = 0; i < nk; i++) { kmax[ku1[i]] = 0; kmax[ku2[i]] = 0; }
            /* reader terms and the vector recursions */
            for (int v = 0; v < R; v++) {
                int zi = z[l][v] + ZOFF;
                double X = 0, Y2 = 0, Y3 = 0, S2 = 0, AR = 0, Aa = 0, Ab = 0, As1 = 0, Ask = 0;
                for (int e = ly->ip[v]; e < ly->ip[v + 1]; e++) {
                    if (!tow[l][e]) continue;
                    int u = ly->ix[e];
                    X += P1p[u]; Y2 += P2p[u]; Y3 += P3p[u]; S2 += P1p[u] * P1p[u];
                    AR += Rp[u]; Aa += ap[u]; Ab += bp[u]; As1 += s1[u]; Ask += sk[u];
                }
                t[0] += fmin(1.0, 3 * q2[zi]) * Y2;
                t[1] += fmin(1.0, 3 * q3[zi]) * Y3;
                t[2] += 0.5 * (X * X - S2);
                t[4] += 0.5 * AR * AR;
                t[5] += Aa * Ab;
                double d = dl[zi], a0 = q0[zi], a1 = q1[zi];
                P1n[v] = 3 * a0 + 3 * d * X;
                P2n[v] = 3 * a0 * a0 + 3 * (a1 * a1 - a0 * a0) * X;
                P3n[v] = a0 * a0 * a0 + (a1 * a1 * a1 - a0 * a0 * a0) * X;
                double hh = 3 * d * As1;
                Rn[v] = 3 * d * (AR + Ask);
                an[v] = 3 * d * Aa + (d < DC ? hh : 0.0);
                bn[v] = 3 * d * Ab + hh;
            }
            /* tK: each stored pair times its number of common toward readers */
            for (int i = 0; i < nk; i++) {
                int a = ku1[i], b = ku2[i], ea = ly->cp[a], eb = ly->cp[b];
                int c = 0;
                while (ea < ly->cp[a + 1] && eb < ly->cp[b + 1]) {
                    int ra = ly->erow[ly->ce[ea]], rb = ly->erow[ly->ce[eb]];
                    if (ra < rb) ea++;
                    else if (rb < ra) eb++;
                    else { c += tow[l][ly->ce[ea]] && tow[l][ly->ce[eb]]; ea++; eb++; }
                }
                t[3] += kv[i] * c;
            }
            /* the sparse part of K at this layer, from big toward children */
            hclear();
            for (int u = 0; u < C; u++) {
                if (P1p[u] < SKIP) continue;
                int nb = 0;
                for (int q = ly->cp[u]; q < ly->cp[u + 1]; q++) {
                    int e = ly->ce[q], v = ly->erow[e];
                    if (tow[l][e] && dl[z[l][v] + ZOFF] >= DC) big[nb++] = v;
                }
                for (int i = 0; i < nb; i++)
                    for (int j = i + 1; j < nb; j++)
                        hadd(big[i], big[j], 9 * dl[z[l][big[i]] + ZOFF] * dl[z[l][big[j]] + ZOFF] * P1p[u]);
            }
            for (int k = 0; k < nk; k++) {
                int a = ku1[k], b = ku2[k], na = 0, nb = 0;
                for (int q = ly->cp[a]; q < ly->cp[a + 1]; q++) {
                    int e = ly->ce[q], v = ly->erow[e];
                    if (tow[l][e] && dl[z[l][v] + ZOFF] >= DC) big[na++] = v;
                }
                for (int q = ly->cp[b]; q < ly->cp[b + 1]; q++) {
                    int e = ly->ce[q], v = ly->erow[e];
                    if (tow[l][e] && dl[z[l][v] + ZOFF] >= DC) big2[nb++] = v;
                }
                for (int i = 0; i < na; i++)
                    for (int j = 0; j < nb; j++)
                        if (big[i] != big2[j])
                            hadd(big[i], big2[j], 9 * dl[z[l][big[i]] + ZOFF] * dl[z[l][big2[j]] + ZOFF] * kv[k]);
            }
            nk2 = 0;
            for (int i = 0; i < Hn; i++) {
                Slot *sl = &H[Hused[i]];
                int a = (int)(sl->key >> 32) - 1, b = (int)(sl->key & 0xffffffffu) - 1;
                if (sl->val >= eps) {
                    if (nk2 >= kcap) { fprintf(stderr, "K full\n"); return 3; }
                    nu1[nk2] = a; nu2[nk2] = b; nv[nk2++] = sl->val;
                } else {
                    if (sl->val > gmax[a]) gmax[a] = sl->val;
                    if (sl->val > gmax[b]) gmax[b] = sl->val;
                }
            }
            for (int v = 0; v < R; v++) if (gmax[v] > 0) { Rn[v] += sqrt(gmax[v]); gmax[v] = 0; }
            /* the next layer becomes the previous one */
            { double *tq;
              tq = P1p; P1p = P1n; P1n = tq; tq = P2p; P2p = P2n; P2n = tq; tq = P3p; P3p = P3n; P3n = tq;
              tq = Rp; Rp = Rn; Rn = tq; tq = ap; ap = an; an = tq; tq = bp; bp = bn; bn = tq; }
            { int *ti; ti = ku1; ku1 = nu1; nu1 = ti; ti = ku2; ku2 = nu2; nu2 = ti; }
            { double *tv = kv; kv = nv; nv = tv; }
            nk = nk2;
            if (nk > kpeak) kpeak = nk;
        }
        for (int v = 0; v < n; v++) t[6] += P2p[v];
        for (int v = 0; v < n; v++) {
            if (Rp[v] > Rpeak) Rpeak = Rp[v];
            if (ap[v] * bp[v] > abpeak) abpeak = ap[v] * bp[v];
        }
        /* the device restores its cells in every copy */
        for (int i = 0; i < ndev; i++) { int d = devbits[i]; P1p[d] = 0; Rp[d] = 0; ap[d] = 0; bp[d] = 0; }
        nk2 = 0;
        for (int i = 0; i < nk; i++)
            if (!isdev[ku1[i]] && !isdev[ku2[i]]) { ku1[nk2] = ku1[i]; ku2[nk2] = ku2[i]; kv[nk2++] = kv[i]; }
        nk = nk2;
        double b = 0;
        for (int k = 0; k < 7; k++) b += t[k];
        if (count) {
            for (int k = 0; k < 7; k++) tot[k] += t[k];
            total += b; counted++;
            if (b > bmax) bmax = b;
            if (fo) fwrite(&b, 8, 1, fo);
        }
        /* advance the orbit: next state and the device */
        memcpy(x, y[L], n);
        if (getword(x, 0xF9) == 1) { out++; putword(x, 0xFE, 0); }
        if (getword(x, 0xFC) == 1) {
            if (h < tlen) { putword(x, 0xFD, tape[h]); putword(x, 0xFA, 0); h++; }
            else putword(x, 0xFA, 1);
        }
        if (getword(x, 0xFB) == 1) h = 0;
        putword(x, 0xFC, 0); putword(x, 0xFB, 0); putword(x, 0xF9, 0);
        if (x[8]) { halted = 1; break; }
    }
    if (fo) fclose(fo);
    printf("s=%.3f start=%ld steps=%ld halted=%d bytes=%ld total=%.17e mean=%.6e max=%.6e "
           "t2=%.3e t3=%.3e tind=%.3e tK=%.3e tR=%.3e tab=%.3e tB=%.3e Kpeak=%d Rpeak=%.2e abpeak=%.2e\n",
           s, start, counted, halted, out, total, counted ? total / counted : 0.0, bmax,
           tot[0], tot[1], tot[2], tot[3], tot[4], tot[5], tot[6], kpeak, Rpeak, abpeak);
    return 0;
}