import { assert } from "../../../../core/assert.js";

/**
 * Popcount lookup for the 6-bit neighbour mask. Index = mask value (0–63 used in
 * practice; the table is sized 64 so the read is a single bounded indexed access).
 * Values are the SOR denominator (number of fluid neighbours) for each mask.
 *
 * @type {Uint8Array}
 */
const POPCOUNT_6 = (function () {
    const t = new Uint8Array(64);
    for (let i = 0; i < 64; i++) {
        let c = 0;
        for (let b = 0; b < 6; b++) if (i & (1 << b)) c++;
        t[i] = c;
    }
    return t;
})();

/**
 * Solve the pressure Poisson equation `∇²p = div(v)` by red-black Successive
 * Over-Relaxation (SOR), gated by a pre-baked per-cell neighbour mask.
 *
 * Per cell, the update is:
 *
 *     p_jacobi(c) = ( sum_of_fluid_neighbour_pressures - div(c) ) / count
 *     p(c) = (1 - ω) * p(c) + ω * p_jacobi(c)
 *
 * "Which neighbours count as fluid" is determined entirely by `neighbour_mask`,
 * a Uint8Array allocated and populated by {@link FluidField} via
 * {@link FluidField.recomputeSolidNeighbourMask}. The encoding is:
 *
 *     bit 0 (= 1)   : -x neighbour is fluid
 *     bit 1 (= 2)   : +x neighbour is fluid
 *     bit 2 (= 4)   : -y neighbour is fluid
 *     bit 3 (= 8)   : +y neighbour is fluid
 *     bit 4 (= 16)  : -z neighbour is fluid
 *     bit 5 (= 32)  : +z neighbour is fluid
 *     bit 7 (= 128) : this cell is itself solid (consumed by
 *                     {@link v3_grid_subtract_pressure_gradient}, ignored here)
 *
 * The low 6 bits are the fluid-neighbour set. A solid cell encodes as `128` and
 * an isolated fluid cell as `0`, so the inner loop's `(mask & 63) === 0 ? skip`
 * covers both self-solid AND fluid-with-no-fluid-neighbours in one comparison.
 * This eliminates six boundary checks plus six solid-neighbour checks per cell
 * per iteration, replacing them with a single mask read and six bit-tests on a
 * register-resident value. Boundary handling (cells at x=0, x=last, etc.) is
 * implicit — the recompute step writes 0 into the boundary-facing bit, so the
 * SOR loop never reads out-of-range memory.
 *
 * Cells are visited in two passes per iteration: first "red" (`(x+y+z) % 2 === 0`)
 * then "black". Adjacent cells have opposite colour, so each pass reads only
 * cells of the other colour and writes are independent — same in-place red-black
 * scheme as the un-masked variant.
 *
 * @param {Float32Array|Float16Array} pressure  Mutated in place. Pre-fill with 0
 *     or with the previous step's solution to warm-start.
 * @param {Float32Array|Float16Array} divergence  From
 *     {@link v3_grid_compute_divergence}. Same dtype as pressure.
 * @param {number} res_x
 * @param {number} res_y
 * @param {number} res_z
 * @param {number} iterations  Number of red-black sweeps.
 * @param {number} omega  SOR relaxation factor in (0, 2). Near-optimal values
 *     come from {@link optimal_sor_omega}.
 * @param {Uint8Array} neighbour_mask  Length ≥ res_x*res_y*res_z. Each byte
 *     encodes the fluid-neighbour bitmask (see above). MUST be recomputed
 *     whenever the underlying solid mask changes.
 */
export function v3_grid_solve_pressure(pressure, divergence, res_x, res_y, res_z, iterations, omega, neighbour_mask) {
    const cell_count = res_x * res_y * res_z;

    assert.notEqual(pressure, divergence, "pressure must not alias divergence");
    assert.greaterThanOrEqual(pressure.length, cell_count, "pressure covers grid");
    assert.greaterThanOrEqual(divergence.length, cell_count, "divergence covers grid");
    assert.greaterThanOrEqual(neighbour_mask.length, cell_count, "neighbour_mask covers grid");
    assert.isNonNegativeInteger(iterations, "iterations");
    assert.greaterThan(omega, 0, "omega");
    assert.lessThan(omega, 2, "omega");

    if (iterations === 0) {
        return;
    }

    const slice_size = res_x * res_y;
    const one_minus_omega = 1 - omega;

    for (let iter = 0; iter < iterations; iter++) {

        // Two passes per iteration: red (parity 0), then black (parity 1). Inside each
        // pass the cell colour determines which x-positions we visit per scanline.
        for (let pass = 0; pass < 2; pass++) {

            for (let z = 0; z < res_z; z++) {
                const z_off = z * slice_size;

                for (let y = 0; y < res_y; y++) {
                    const y_off = y * res_x;

                    // First x for this scanline of the requested colour.
                    const x_start = ((z + y) & 1) ^ pass;

                    for (let x = x_start; x < res_x; x += 2) {
                        const c = z_off + y_off + x;

                        const mask = neighbour_mask[c];
                        if ((mask & 63) === 0) {
                            // No fluid neighbours: solid cell (bit 7 set) OR isolated
                            // fluid. Either way, no degrees of freedom — skip. Cells
                            // that reach past here have bit 7 clear, so `mask` indexes
                            // POPCOUNT_6 (sized 64) in range.
                            continue;
                        }

                        let sum = 0;
                        if (mask & 1)  sum += pressure[c - 1];
                        if (mask & 2)  sum += pressure[c + 1];
                        if (mask & 4)  sum += pressure[c - res_x];
                        if (mask & 8)  sum += pressure[c + res_x];
                        if (mask & 16) sum += pressure[c - slice_size];
                        if (mask & 32) sum += pressure[c + slice_size];

                        const count = POPCOUNT_6[mask];
                        const p_target = (sum - divergence[c]) / count;
                        pressure[c] = one_minus_omega * pressure[c] + omega * p_target;
                    }
                }
            }
        }
    }
}
