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

/**
 * Implicit (Gauss-Seidel-like) Jacobi diffusion on a 3D cell-centered scalar grid.
 *
 * Solves the implicit diffusion equation `(I - a*L) x = x0` for `x`, where L is the
 * discrete Laplacian. Solving implicitly is unconditionally stable for any `a >= 0`,
 * which is why Stam's "Stable Fluids" uses it instead of an explicit step.
 *
 * Each iteration writes into a scratch buffer and swaps; after `iterations` rounds the
 * final result is copied into `output`. Increasing `iterations` increases accuracy at
 * linear cost; 4-20 is the usual range. Boundary cells use Neumann (reflective)
 * conditions — missing neighbours are treated as equal to the cell itself, which
 * reduces the effective denominator.
 *
 * Solid cells (marked non-zero in `solid`) are skipped: their value in `output` is
 * left at the corresponding `x0` value.
 *
 * @param {Float32Array} output      Destination (size = res_x * res_y * res_z). Mutated.
 * @param {Float32Array} x0          Source values to diffuse. NOT mutated.
 * @param {Float32Array} scratch     Same size as output; used for ping-pong. Mutated.
 * @param {number} res_x
 * @param {number} res_y
 * @param {number} res_z
 * @param {number} a                 Implicit diffusion coefficient (= dt * diffusion_rate).
 * @param {number} iterations        Jacobi iteration count.
 * @param {Uint8Array} solid         Required cell flags (zero-filled for a wall-free
 *                                   domain); non-zero cells are not updated.
 */
export function v3_grid_apply_diffusion(output, x0, scratch, res_x, res_y, res_z, a, iterations, solid) {
    assert.notEqual(output, x0, "output must differ from x0 (x0 is the source and must not be mutated)");
    assert.notEqual(output, scratch, "output must differ from scratch");
    assert.notEqual(x0, scratch, "x0 must differ from scratch");
    assert.greaterThanOrEqual(output.length, res_x * res_y * res_z, "output covers grid");
    assert.greaterThanOrEqual(x0.length, res_x * res_y * res_z, "x0 covers grid");
    assert.greaterThanOrEqual(scratch.length, res_x * res_y * res_z, "scratch covers grid");
    assert.greaterThanOrEqual(a, 0, "a");
    assert.isNonNegativeInteger(iterations, "iterations");

    // Seed first iterate with x0.
    output.set(x0);

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

    const slice_size = res_x * res_y;
    const last_x = res_x - 1;
    const last_y = res_y - 1;
    const last_z = res_z - 1;

    let prev = output;
    let next = scratch;

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

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

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

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

                    if (solid[c] !== 0) {
                        next[c] = x0[c];
                        continue;
                    }

                    let sum = 0;
                    let count = 0;

                    if (x > 0)       { sum += prev[c - 1];          count++; }
                    if (x < last_x)  { sum += prev[c + 1];          count++; }
                    if (has_ym)      { sum += prev[c - res_x];      count++; }
                    if (has_yp)      { sum += prev[c + res_x];      count++; }
                    if (has_zm)      { sum += prev[c - slice_size]; count++; }
                    if (has_zp)      { sum += prev[c + slice_size]; count++; }

                    // Implicit Jacobi update with Neumann boundary correction:
                    // missing neighbours contribute as if equal to the cell, which is
                    // equivalent to dropping them from both the numerator and the
                    // denominator factor on `a`.
                    next[c] = (x0[c] + a * sum) / (1 + a * count);
                }
            }
        }

        const tmp = prev;
        prev = next;
        next = tmp;
    }

    if (prev !== output) {
        output.set(prev);
    }
}
