// A column of water falls over in a box that rocks back and forth, run three times: a leader and
// two twins whose first particle starts one and two position units higher than the leader's.
// 1,024 particles move through a grid of 16 x 16 x 16 cells by MLS-MPM (the moving least squares
// material point method, Hu et al. 2018) as a weakly compressible fluid: each step scatters the
// particles' mass and momentum to the grid, moves the grid under gravity and the walls, and
// gathers the grid back into the particles. The rocking keeps the water sloshing, so the twins'
// one-unit difference grows until their water moves differently.
//
// Units: lengths in cells of 2^16 position units (Q), time in steps, velocities in Q per step.
// The hot loops multiply, add and divide by powers of two only, and every value stays below 2^53
// by the bounds noted below, so the sums are exact and the scatter gives the same grid, node by
// node, in any particle order.

const I64 N = 16;                        // cells per side
const I64 NP = 1_024;                    // particles per copy
const I64 NODES = 4_096;                 // N^3
const I64 Q = 65_536;                    // position units per cell
const I64 W = 2_048;                     // one weight, 2^11
const I64 W3 = 8_589_934_592;            // the 27 weights of a particle sum to 2^33
const I64 GRAV = 131;                    // 0.002 cells per step squared
const I64 TILT_X = 70;                   // the rocking: sideways gravity up to 70 and 50 Q per
const I64 TILT_Z = 50;                   // step squared, about 28 and 21 degrees of tilt
const I64 VMAX = 32_768;                 // half a cell per step: no grid velocity goes faster
const I64 XLO = 65_536;                  // particles stay between one cell and 14.5 cells, so
const I64 XHI = 950_271;                 // their 27 nodes are always inside the grid
const I64 PUSH_REACH = 65_536;           // the push acts within one cell of the leader's particle

// Particles, slot copy * NP + p: position, velocity, the affine velocity C (c0 to c8 row by row,
// scale Q per step) and the volume ratio J (scale Q). A module array holds at most 64 KiB
I64[3_072] px;
I64[3_072] py;
I64[3_072] pz;
I64[3_072] vx;
I64[3_072] vy;
I64[3_072] vz;
I64[3_072] c0;
I64[3_072] c1;
I64[3_072] c2;
I64[3_072] c3;
I64[3_072] c4;
I64[3_072] c5;
I64[3_072] c6;
I64[3_072] c7;
I64[3_072] c8;
I64[3_072] jv;
U64[3] digest;
I64 tether = 0;
I64 order = 0;
I64 steps = 0;

// The grid of one copy: mass (scale 2^33 per particle), then momentum (scale 2^25), then velocity
I64[4_096] gm;
I64[4_096] gx;
I64[4_096] gy;
I64[4_096] gz;

// One particle's quadratic B-spline weights (w[d * 3 + o], summing to W along each axis), the
// lowest node of its 3 x 3 x 3 block and its offset from that node
I64[9] w;
I64[3] base;
I64[3] fx;

// The weights of slot s: 0.5 (1.5 - f)^2, 0.75 - (f - 1)^2 and 0.5 (f - 0.5)^2 of the offset f in
// cells, the middle one taken as the rest so the three sum to W exactly
void weight_axis(I64 d, I64 x) {
    I64 b = (x - Q / 2) / Q;
    I64 f = x - b * Q;
    base[d] = b;
    fx[d] = f;
    I64 a = 3 * Q / 2 - f;
    I64 c = f - Q / 2;
    I64 w0 = (a * a + 2_097_152) / 4_194_304;
    I64 w2 = (c * c + 2_097_152) / 4_194_304;
    w[d * 3] = w0;
    w[d * 3 + 1] = W - w0 - w2;
    w[d * 3 + 2] = w2;
}

void weights(I64 s) {
    weight_axis(0, px[s]);
    weight_axis(1, py[s]);
    weight_axis(2, pz[s]);
}

// The particle scattered at turn i. Order 0 goes 0, 1, 2, ...; order 1 goes backward; order 2
// jumps by 389 from a start that moves every step. Each order visits every particle once
I64 visit(I64 i) {
    if (order == 1) {
        return NP - 1 - i;
    }
    if (order == 2) {
        return (i * 389 + steps * 97) % NP;
    }
    return i;
}

// Scatter: every particle adds its mass and momentum to its 27 nodes. With A = stress * I + C, the
// momentum at a node o (0 to 2 along each axis) is v + A (o Q - f) / Q = m0 + A o, m0 taken once
// per particle. A node adds weight * momentum / 2^8: the weight is below 2^32 and the momentum
// below 2^18, so each product is below 2^50, and a node's sums stay below 2^46
void scatter(I64 c) {
    for n in 0..NODES {
        gm[n] = 0;
        gx[n] = 0;
        gy[n] = 0;
        gz[n] = 0;
    }
    for i in 0..NP {
        I64 s = c * NP + visit(i);
        weights(s);
        // pressure of the weakly compressible fluid: 4 E (1 - J) with E = 0.15
        I64 st = 3 * (Q - jv[s]) / 5;
        I64 a0 = c0[s] + st;
        I64 a1 = c1[s];
        I64 a2 = c2[s];
        I64 a3 = c3[s];
        I64 a4 = c4[s] + st;
        I64 a5 = c5[s];
        I64 a6 = c6[s];
        I64 a7 = c7[s];
        I64 a8 = c8[s] + st;
        I64 mx = vx[s] - (a0 * fx[0] + a1 * fx[1] + a2 * fx[2] + Q / 2) / Q;
        I64 my = vy[s] - (a3 * fx[0] + a4 * fx[1] + a5 * fx[2] + Q / 2) / Q;
        I64 mz = vz[s] - (a6 * fx[0] + a7 * fx[1] + a8 * fx[2] + Q / 2) / Q;
        for i0 in 0..3 {
            for i1 in 0..3 {
                I64 wxy = w[i0] * w[3 + i1];
                I64 n0 = ((base[0] + i0) * N + base[1] + i1) * N + base[2];
                I64 rx = mx + a0 * i0 + a1 * i1;
                I64 ry = my + a3 * i0 + a4 * i1;
                I64 rz = mz + a6 * i0 + a7 * i1;
                I64 wt0 = wxy * w[6];
                gm[n0] = gm[n0] + wt0;
                gx[n0] = gx[n0] + wt0 * rx / 256;
                gy[n0] = gy[n0] + wt0 * ry / 256;
                gz[n0] = gz[n0] + wt0 * rz / 256;
                I64 wt1 = wxy * w[7];
                gm[n0 + 1] = gm[n0 + 1] + wt1;
                gx[n0 + 1] = gx[n0 + 1] + wt1 * (rx + a2) / 256;
                gy[n0 + 1] = gy[n0 + 1] + wt1 * (ry + a5) / 256;
                gz[n0 + 1] = gz[n0 + 1] + wt1 * (rz + a8) / 256;
                I64 wt2 = wxy * w[8];
                gm[n0 + 2] = gm[n0 + 2] + wt2;
                gx[n0 + 2] = gx[n0 + 2] + wt2 * (rx + 2 * a2) / 256;
                gy[n0 + 2] = gy[n0 + 2] + wt2 * (ry + 2 * a5) / 256;
                gz[n0 + 2] = gz[n0 + 2] + wt2 * (rz + 2 * a8) / 256;
            }
        }
    }
}

// A grid velocity: no faster than VMAX, and no motion into a wall within three nodes of it
I64 wall(I64 u, I64 at) {
    I64 v = clamp(u, -VMAX, VMAX);
    if (at < 3 && v < 0) {
        return 0;
    }
    if (at > N - 4 && v > 0) {
        return 0;
    }
    return v;
}

// The rocking: a triangle wave from -amp to amp and back over `period` steps
I64 rock(I64 t, I64 period, I64 amp) {
    I64 ph = t % period;
    I64 half = period / 2;
    if (ph < half) {
        return -amp + 2 * amp * ph / half;
    }
    return amp - 2 * amp * (ph - half) / half;
}

// Momentum to velocity, plus gravity and the rocking, on every node that holds at least 2^-25 of
// a particle. Rounding: (2 a + b) / (2 b) is a / b to the nearest, halves up
void move_grid() {
    I64 tx = rock(steps, 480, TILT_X);
    I64 tz = rock(steps + 100, 330, TILT_Z);
    for n in 0..NODES {
        I64 mass = gm[n] / 256;
        if (mass > 0) {
            gx[n] = wall((2 * gx[n] + mass) / (2 * mass) + tx, n / 256);
            gy[n] = wall((2 * gy[n] + mass) / (2 * mass) - GRAV, (n / 16) % 16);
            gz[n] = wall((2 * gz[n] + mass) / (2 * mass) + tz, n % 16);
        } else {
            gx[n] = 0;
            gy[n] = 0;
            gz[n] = 0;
        }
    }
}

// Gather: each particle's velocity is the weighted sum S of its nodes' velocities over 2^33. Its
// C is 4 / 2^33 times the weighted sum of velocity times offset (o Q - f) / Q, which is
// T / 2^31 - (S / 2^31) f / 2^16 with T the weighted sum of velocity times o. Weights sum to 2^33 and
// velocities stay below 2^15, so S stays below 2^48 and T below 2^49
void gather(I64 c) {
    for p in 0..NP {
        I64 s = c * NP + p;
        weights(s);
        I64 sx = 0;
        I64 sy = 0;
        I64 sz = 0;
        I64[9] t;
        for i0 in 0..3 {
            for i1 in 0..3 {
                I64 wxy = w[i0] * w[3 + i1];
                I64 n0 = ((base[0] + i0) * N + base[1] + i1) * N + base[2];
                I64 wt0 = wxy * w[6];
                I64 wt1 = wxy * w[7];
                I64 wt2 = wxy * w[8];
                I64 ux1 = wt1 * gx[n0 + 1];
                I64 ux2 = wt2 * gx[n0 + 2];
                I64 ux = wt0 * gx[n0] + ux1 + ux2;
                I64 uy1 = wt1 * gy[n0 + 1];
                I64 uy2 = wt2 * gy[n0 + 2];
                I64 uy = wt0 * gy[n0] + uy1 + uy2;
                I64 uz1 = wt1 * gz[n0 + 1];
                I64 uz2 = wt2 * gz[n0 + 2];
                I64 uz = wt0 * gz[n0] + uz1 + uz2;
                sx = sx + ux;
                sy = sy + uy;
                sz = sz + uz;
                t[0] = t[0] + ux * i0;
                t[1] = t[1] + ux * i1;
                t[2] = t[2] + ux1 + 2 * ux2;
                t[3] = t[3] + uy * i0;
                t[4] = t[4] + uy * i1;
                t[5] = t[5] + uy1 + 2 * uy2;
                t[6] = t[6] + uz * i0;
                t[7] = t[7] + uz * i1;
                t[8] = t[8] + uz1 + 2 * uz2;
            }
        }
        vx[s] = (sx + W3 / 2) / W3;
        vy[s] = (sy + W3 / 2) / W3;
        vz[s] = (sz + W3 / 2) / W3;
        I64 hx = sx / 2_147_483_648;
        I64 hy = sy / 2_147_483_648;
        I64 hz = sz / 2_147_483_648;
        c0[s] = t[0] / 2_147_483_648 - hx * fx[0] / 65_536;
        c1[s] = t[1] / 2_147_483_648 - hx * fx[1] / 65_536;
        c2[s] = t[2] / 2_147_483_648 - hx * fx[2] / 65_536;
        c3[s] = t[3] / 2_147_483_648 - hy * fx[0] / 65_536;
        c4[s] = t[4] / 2_147_483_648 - hy * fx[1] / 65_536;
        c5[s] = t[5] / 2_147_483_648 - hy * fx[2] / 65_536;
        c6[s] = t[6] / 2_147_483_648 - hz * fx[0] / 65_536;
        c7[s] = t[7] / 2_147_483_648 - hz * fx[1] / 65_536;
        c8[s] = t[8] / 2_147_483_648 - hz * fx[2] / 65_536;
        I64 trace = c0[s] + c4[s] + c8[s];
        jv[s] = jv[s] + (jv[s] * trace + Q / 2) / Q;
    }
}

// Consensus tether on the twins' particles. Pulling (k > 0) is a critically damped spring toward
// the leader's particle with stiffness k^4 / 10^10 per step squared, which at k = 100 settles in
// about 60 steps. Pushing (k < 0) is a weak spring away from the leader's particle that lets go one
// cell out
void pull_twins() {
    I64 k = tether;
    if (k == 0) {
        return;
    }
    for c in 1..3 {
        for p in 0..NP {
            I64 t = c * NP + p;
            I64 dx = px[p] - px[t];
            I64 dy = py[p] - py[t];
            I64 dz = pz[p] - pz[t];
            if (k > 0) {
                I64 k4 = k * k * k * k;
                I64 k2 = 2 * k * k;
                vx[t] = vx[t] + muldiv(I64, dx, k4, 10_000_000_000, half_even) + muldiv(I64, vx[p] - vx[t], k2, 100_000, half_even);
                vy[t] = vy[t] + muldiv(I64, dy, k4, 10_000_000_000, half_even) + muldiv(I64, vy[p] - vy[t], k2, 100_000, half_even);
                vz[t] = vz[t] + muldiv(I64, dz, k4, 10_000_000_000, half_even) + muldiv(I64, vz[p] - vz[t], k2, 100_000, half_even);
            } else if (max(abs(dx), max(abs(dy), abs(dz))) < PUSH_REACH) {
                vx[t] = vx[t] - muldiv(I64, dx, k * k, 10_000_000, half_even);
                vy[t] = vy[t] - muldiv(I64, dy, k * k, 10_000_000, half_even);
                vz[t] = vz[t] - muldiv(I64, dz, k * k, 10_000_000, half_even);
            }
        }
    }
}

// One step of all three copies, then the tether, then every particle moves
void step() {
    for c in 0..3 {
        scatter(c);
        move_grid();
        gather(c);
    }
    pull_twins();
    for c in 0..3 {
        I64 check = 0;
        for p in 0..NP {
            I64 s = c * NP + p;
            px[s] = clamp(px[s] + vx[s], XLO, XHI);
            py[s] = clamp(py[s] + vy[s], XLO, XHI);
            pz[s] = clamp(pz[s] + vz[s], XLO, XHI);
            check = check + px[s] * (2 * p + 1) + py[s] * (2 * p + 2_049) + pz[s] * (2 * p + 4_097);
        }
        digest[c] = (digest[c] *% 1_099_511_628_211) ^ (check as% U64);
    }
    steps = steps + 1;
}

// Loads the dam: a block 4 cells wide, 8 tall and 4 deep in a corner of the box, two particles per
// cell along each axis, at rest. Twin 1's particle 0 starts 1 Q higher, twin 2's 2 Q higher.
// Preset 1 starts the block in the middle of the box instead
export I64 start(I64 preset) {
    I64 ox = 2 * Q;
    if (preset == 1) {
        ox = 6 * Q;
    }
    for c in 0..3 {
        for i in 0..8 {
            for j in 0..16 {
                for k in 0..8 {
                    I64 s = c * NP + (i * 16 + j) * 8 + k;
                    px[s] = ox + (2 * i + 1) * Q / 4;
                    py[s] = 2 * Q + (2 * j + 1) * Q / 4;
                    pz[s] = ox + (2 * k + 1) * Q / 4;
                    vx[s] = 0;
                    vy[s] = 0;
                    vz[s] = 0;
                    jv[s] = Q;
                    c0[s] = 0;
                    c1[s] = 0;
                    c2[s] = 0;
                    c3[s] = 0;
                    c4[s] = 0;
                    c5[s] = 0;
                    c6[s] = 0;
                    c7[s] = 0;
                    c8[s] = 0;
                }
            }
        }
        digest[c] = 14_695_981_039_346_656_037;
    }
    py[NP] = py[NP] + 1;
    py[2 * NP] = py[2 * NP] + 2;
    steps = 0;
    return preset;
}

// The page drives the program through these exports. A state slot is field * 3072 + copy * 1024 +
// particle, fields px py pz vx vy vz j, then c0 to c8; slot 49152 is the step count, which sets
// the rocking and the scramble
export I64 load(I64 slot, I64 value) {
    if (slot == 49_152) {
        steps = value;
        return value;
    }
    I64 s = slot % 3_072;
    I64 f = slot / 3_072;
    if (f == 0) { px[s] = value; }
    if (f == 1) { py[s] = value; }
    if (f == 2) { pz[s] = value; }
    if (f == 3) { vx[s] = value; }
    if (f == 4) { vy[s] = value; }
    if (f == 5) { vz[s] = value; }
    if (f == 6) { jv[s] = value; }
    if (f == 7) { c0[s] = value; }
    if (f == 8) { c1[s] = value; }
    if (f == 9) { c2[s] = value; }
    if (f == 10) { c3[s] = value; }
    if (f == 11) { c4[s] = value; }
    if (f == 12) { c5[s] = value; }
    if (f == 13) { c6[s] = value; }
    if (f == 14) { c7[s] = value; }
    if (f == 15) { c8[s] = value; }
    return value;
}

export I64 read(I64 slot) {
    if (slot == 49_152) {
        return steps;
    }
    I64 s = slot % 3_072;
    I64 f = slot / 3_072;
    if (f == 0) { return px[s]; }
    if (f == 1) { return py[s]; }
    if (f == 2) { return pz[s]; }
    if (f == 3) { return vx[s]; }
    if (f == 4) { return vy[s]; }
    if (f == 5) { return vz[s]; }
    if (f == 6) { return jv[s]; }
    if (f == 7) { return c0[s]; }
    if (f == 8) { return c1[s]; }
    if (f == 9) { return c2[s]; }
    if (f == 10) { return c3[s]; }
    if (f == 11) { return c4[s]; }
    if (f == 12) { return c5[s]; }
    if (f == 13) { return c6[s]; }
    if (f == 14) { return c7[s]; }
    return c8[s];
}

// A particle's position for drawing, in 1/1024 of a cell: x * 2^32 + y * 2^16 + z
export I64 spot(I64 s) {
    return (px[s] / 64) * 4_294_967_296 + (py[s] / 64) * 65_536 + pz[s] / 64;
}

export U64 seed(I64 c, U64 value) {
    digest[c] = value;
    return value;
}

export U64 digest_of(I64 c) {
    return digest[c];
}

export I64 pull(I64 k) {
    tether = k;
    return k;
}

export I64 scramble(I64 o) {
    order = o;
    return o;
}

export I64 advance(I64 n) {
    for i in 0..n {
        step();
    }
    return n;
}

// The largest distance in Q from a particle of copy c to the same particle of the leader
export I64 spread(I64 c) {
    I64 most = 0;
    for p in 0..NP {
        I64 dx = px[c * NP + p] - px[p];
        I64 dy = py[c * NP + p] - py[p];
        I64 dz = pz[c * NP + p] - pz[p];
        most = max(most, isqrt(dx * dx + dy * dy + dz * dz));
    }
    return most;
}

// The grid of copy 0 scattered in one order: its mass and momentum summed with node weights
U64 grid_digest() {
    U64 h = 14_695_981_039_346_656_037;
    for n in 0..NODES {
        h = (h *% 1_099_511_628_211) ^ (gm[n] as% U64) ^ ((gx[n] as% U64) <<% 17) ^ ((gy[n] as% U64) <<% 31) ^ ((gz[n] as% U64) <<% 45);
    }
    return h;
}

U64 grid_in_order(I64 o) {
    order = o;
    scatter(0);
    order = 0;
    return grid_digest();
}

// The total mass on the grid of copy 0, in particles
I64 grid_mass() {
    I64 total = 0;
    for n in 0..NODES {
        total = total + gm[n];
    }
    return total / W3;
}

void main() {
    _ = start(0);
    _ = advance(200);
    "after 200 steps, dam break:\n";
    "leader digest={digest[0]}\n";
    "twin 1 digest={digest[1]}\n";
    "twin 2 digest={digest[2]}\n";
    "spread twin 1={spread(1)} Q, twin 2={spread(2)} Q\n";
    "grid forward={grid_in_order(0)}\n";
    "grid backward={grid_in_order(1)}\n";
    "grid jumping={grid_in_order(2)}\n";
    "grid mass={grid_mass()} particles\n";
}

test "the scatter gives the same grid in any particle order" {
    _ = start(0);
    _ = advance(40);
    U64 a = grid_in_order(0);
    assert(grid_in_order(1) == a && grid_in_order(2) == a);
}

test "the grid holds every particle's mass exactly" {
    _ = start(0);
    _ = advance(10);
    _ = grid_in_order(2);
    I64 total = 0;
    for n in 0..NODES {
        total = total + gm[n];
    }
    assert(total == NP * W3);
}

test "a scrambled run gives the same digests" {
    _ = start(0);
    _ = scramble(2);
    _ = advance(30);
    U64 d = digest[0];
    _ = start(0);
    _ = scramble(0);
    _ = advance(30);
    assert(digest[0] == d);
}