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

/**
 * Popcount table for the 6-bit neighbour mask — same encoding used by
 * {@link v3_grid_solve_pressure}. Indexed by the mask byte (0-63), value is the
 * SOR / Laplacian denominator = number of fluid neighbours.
 *
 * @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;
})();

/**
 * MIC(0) modification weight. Bridson's recommendation is 0.97; the factor scales
 * how much "would-have-been-fill-in" we re-add onto the diagonal. Lower → closer
 * to plain IC(0). Higher → stronger preconditioning but riskier (can produce a
 * non-SPD matrix). 0.97 is the canonical sweet spot from the literature.
 *
 * @type {number}
 */
const MIC_TAU = 0.97;

/**
 * Diagonal safeguard threshold. If the computed Cholesky diagonal element falls
 * below `MIC_SIGMA * A_diag`, we discard the IC factorization for that cell and
 * fall back to `e = A_diag` (i.e. Jacobi for that cell). Prevents `1/sqrt(small)`
 * from blowing up.
 *
 * @type {number}
 */
const MIC_SIGMA = 0.25;

/**
 * Pressure Poisson solver: Modified Incomplete Cholesky preconditioned Conjugate
 * Gradient — MIC(0)-PCG.
 *
 * Solves `A * p = -div(v)` where A is the 7-point discrete Laplacian over fluid
 * cells (same sign convention as {@link v3_grid_solve_pressure} — that
 * function's SOR update is exactly the splitting `A * p = -div`). Per cell:
 *
 *     A[c, c]  = number of fluid neighbours (= popcount(neighbour_mask[c]))
 *     A[c, c'] = -1 if c' is a fluid neighbour of c, else 0
 *
 * PCG converges in roughly `O(√κ)` iterations where κ is the spectral condition
 * number of `M⁻¹A`. The MIC(0) preconditioner reduces κ from `O(N²)` (raw
 * Laplacian) to `O(N)`, so a single PCG iter does a lot more useful work than
 * a single SOR sweep.
 *
 * **Empirical finding (V8 / Node, single-thread, this implementation).** Each
 * PCG iter involves ~8 separate per-cell passes (As compute, axpy×2,
 * mean-subtract, fwd+bwd triangular solves, dot, s update) vs SOR's 2. In
 * native code each pass is bandwidth-bound and per-iter ratio approaches the
 * textbook 2× — yielding the canonical 1.5–2× wall-time win for PCG. In V8,
 * per-pass loop overhead dominates; the actual per-iter ratio is closer to
 * 4-6× and PCG ends up SLOWER than SOR at iso-quality on production grids:
 *
 *     res=64:  SOR  16 iter  /  75 ms,  max|div| 3.09e-1
 *              PCG   8 iter  / 181 ms,  max|div| 3.72e-3   (PCG 2.4× slower)
 *
 *     res=128: SOR  32 iter  / 862 ms,  max|div| 8.74e-1
 *              PCG   8 iter  /1452 ms,  max|div| 5.65e-3   (PCG 1.7× slower)
 *
 * (Reproduce via `bench: MIC-PCG vs SOR` in bench.spec.js.)
 *
 * Where PCG still wins: **convergence depth**. At the SAME iter count it
 * produces dramatically lower residual divergence (84× better max|div| at
 * 64³, 247× at 128³). If you need quality below the cell-centered operator
 * floor SOR plateaus at — e.g. quantitative pressure inspection, or
 * extreme-impulse scenes where the SOR floor is visible — PCG is the answer.
 * For visual-fluids steady state, **leave the default at SOR**.
 *
 * If this code is ever ported to WASM SIMD or WebGPU the perf comparison
 * flips: per-iter overhead becomes negligible in those environments and the
 * theoretical 1.5–2× PCG win re-emerges. Worth re-running the bench then.
 *
 * Algorithm (Bridson "Fluid Simulation for Computer Graphics" §5.6, Saad
 * §10.2):
 *
 *     1.  build_preconditioner: walk fluid cells in lex order, computing
 *         precon[c] = 1 / sqrt(e[c]) where e[c] is the modified Cholesky
 *         diagonal. Backward neighbours (-x, -y, -z) contribute to the
 *         subtraction.
 *
 *     2.  r₀ = -div - A * p₀     (initial residual; p₀ is the warm-start)
 *         z₀ = M⁻¹ * r₀          (apply preconditioner — two triangular solves)
 *         s₀ = z₀
 *         σ₀ = r₀ · z₀
 *
 *     3.  For k = 0, 1, …, max_iter-1:
 *             A·sₖ
 *             α = σₖ / (sₖ · A·sₖ)
 *             p ← p + α·sₖ
 *             r ← r - α·A·sₖ
 *             z ← M⁻¹ · r
 *             σₖ₊₁ = r · z
 *             β = σₖ₊₁ / σₖ
 *             s ← z + β·s
 *
 * Preconditioner apply is a forward sweep (L·y = r) followed by a backward
 * sweep (Lᵀ·z = y). Both sweeps are O(N) and inherently sequential — they're
 * the hot path that prevents MIC-PCG from parallelizing onto a GPU without
 * substantial reformulation. On single-thread JS that's not a concern.
 *
 * Cells with no fluid neighbours — solid cells (mask bit 7 set) and isolated
 * fluid (mask = 0), both identified by `(mask & 63) === 0` — are excluded from
 * the system: their pressure stays at whatever it was (typically 0), their
 * entries in r, z, s, As stay 0, and they contribute nothing to dot products.
 *
 * @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 `v3_grid_compute_divergence`.
 * @param {number} res_x
 * @param {number} res_y
 * @param {number} res_z
 * @param {number} max_iterations  Cap. Converged early-exits via implicit-zero
 *     dot product (sigma reaches 0 or numerical noise).
 * @param {Uint8Array} neighbour_mask  From `v3_grid_compute_solid_neighbour_mask`.
 * @param {Float32Array} scratch_r  Length ≥ rx*ry*rz. Residual.
 * @param {Float32Array} scratch_z  Length ≥ rx*ry*rz. Preconditioned residual.
 * @param {Float32Array} scratch_s  Length ≥ rx*ry*rz. Search direction.
 * @param {Float32Array} scratch_As  Length ≥ rx*ry*rz. A applied to s.
 * @param {Float32Array} scratch_precon  Length ≥ rx*ry*rz. Preconditioner
 *     diagonal (1/sqrt(e_c) per fluid cell, 0 for solid).
 */
export function v3_grid_solve_pressure_pcg(
    pressure, divergence,
    res_x, res_y, res_z,
    max_iterations,
    neighbour_mask,
    scratch_r, scratch_z, scratch_s, scratch_As, scratch_precon
) {
    const cells = res_x * res_y * res_z;

    assert.notEqual(pressure, divergence, "pressure must not alias divergence");
    assert.greaterThanOrEqual(pressure.length, cells, "pressure covers grid");
    assert.greaterThanOrEqual(divergence.length, cells, "divergence covers grid");
    assert.greaterThanOrEqual(neighbour_mask.length, cells, "neighbour_mask covers grid");
    assert.greaterThanOrEqual(scratch_r.length, cells, "scratch_r covers grid");
    assert.greaterThanOrEqual(scratch_z.length, cells, "scratch_z covers grid");
    assert.greaterThanOrEqual(scratch_s.length, cells, "scratch_s covers grid");
    assert.greaterThanOrEqual(scratch_As.length, cells, "scratch_As covers grid");
    assert.greaterThanOrEqual(scratch_precon.length, cells, "scratch_precon covers grid");
    assert.isNonNegativeInteger(max_iterations, "max_iterations");

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

    const rx = res_x;
    const slice = res_x * res_y;

    const precon = scratch_precon;
    const r = scratch_r;
    const z = scratch_z;
    const s = scratch_s;
    const As = scratch_As;

    // Count fluid cells for the null-space projection (mean subtraction). With
    // all-Neumann boundary conditions and no Dirichlet anchor, the discrete
    // Laplacian has a 1-D null space (constant pressure). Plain (P)CG drifts
    // along that direction — iterates can explode within ~10 sweeps even
    // though the residual stays small in the orthogonal complement. Projecting
    // r back onto the mean-zero subspace after every update keeps the search
    // direction perpendicular to the null space and lets PCG converge as
    // theory predicts.
    //
    // (Standard fluids practice — see Bridson §5.5. The alternative is to pin
    // one cell to a Dirichlet value and exclude it from the system, but mean
    // subtraction has the same effect with simpler bookkeeping.)
    // A cell is a degree of freedom iff it has ≥1 fluid neighbour, i.e. any of
    // the low 6 bits set. Solid cells (bit 7) and isolated fluid (all-zero) are
    // both excluded by masking off bit 7 with & 63.
    let fluid_count = 0;
    for (let c = 0; c < cells; c++) {
        if ((neighbour_mask[c] & 63) !== 0) fluid_count++;
    }
    if (fluid_count === 0) {
        return;
    }

    // ─── 1. Build MIC(0) preconditioner ─────────────────────────────────────
    //
    // Walk fluid cells in lex order; for each, compute the modified Cholesky
    // diagonal e[c] from the backward (-x, -y, -z) fluid neighbours and store
    // precon[c] = 1 / sqrt(e[c]). Solid / isolated cells keep precon = 0.
    //
    // The "modified" piece adds tau * (off-direction terms) back onto the
    // diagonal — it compensates for fill-in that plain IC(0) drops, and makes
    // PCG converge in O(√N) instead of O(N²/³).
    precon.fill(0);
    for (let c = 0; c < cells; c++) {
        const mask = neighbour_mask[c];
        if ((mask & 63) === 0) {
            continue;
        }
        const a_diag = POPCOUNT_6[mask];
        let e = a_diag;

        if (mask & 1) {  // -x backward neighbour is fluid
            const c_prev = c - 1;
            const p_prev = precon[c_prev];
            const p_prev_sq = p_prev * p_prev;
            // (A[c, c_prev] * precon[c_prev])^2 = precon_prev^2 since A = -1.
            e -= p_prev_sq;
            // Modification: tau * A[c, c_prev] * (off-direction A entries from c_prev) * precon_prev^2.
            // A[c, c_prev] = -1. Off-directions from c_prev pointing away from c
            // are +y (bit 3 of c_prev's mask) and +z (bit 5).
            const m_prev = neighbour_mask[c_prev];
            const off = ((m_prev >> 3) & 1) + ((m_prev >> 5) & 1);
            // Contribution: tau * (-1) * (-(off)) * precon_prev^2 = tau * off * precon_prev^2
            e -= MIC_TAU * off * p_prev_sq;
        }
        if (mask & 4) {  // -y
            const c_prev = c - rx;
            const p_prev = precon[c_prev];
            const p_prev_sq = p_prev * p_prev;
            e -= p_prev_sq;
            const m_prev = neighbour_mask[c_prev];
            // Off-directions from c_prev: +x (bit 1), +z (bit 5).
            const off = ((m_prev >> 1) & 1) + ((m_prev >> 5) & 1);
            e -= MIC_TAU * off * p_prev_sq;
        }
        if (mask & 16) {  // -z
            const c_prev = c - slice;
            const p_prev = precon[c_prev];
            const p_prev_sq = p_prev * p_prev;
            e -= p_prev_sq;
            const m_prev = neighbour_mask[c_prev];
            // Off-directions from c_prev: +x (bit 1), +y (bit 3).
            const off = ((m_prev >> 1) & 1) + ((m_prev >> 3) & 1);
            e -= MIC_TAU * off * p_prev_sq;
        }

        // Safeguard: if e dropped below sigma*A_diag, the factorization is
        // becoming non-SPD-friendly; fall back to plain Jacobi for this cell.
        if (e < MIC_SIGMA * a_diag) {
            e = a_diag;
        }
        precon[c] = 1 / Math.sqrt(e);
    }

    // ─── 2. r = -div - A * pressure  (initial residual) ─────────────────────

    for (let c = 0; c < cells; c++) {
        const mask = neighbour_mask[c];
        if ((mask & 63) === 0) {
            r[c] = 0;
            continue;
        }
        let Ap = POPCOUNT_6[mask] * pressure[c];
        if (mask & 1)  Ap -= pressure[c - 1];
        if (mask & 2)  Ap -= pressure[c + 1];
        if (mask & 4)  Ap -= pressure[c - rx];
        if (mask & 8)  Ap -= pressure[c + rx];
        if (mask & 16) Ap -= pressure[c - slice];
        if (mask & 32) Ap -= pressure[c + slice];
        r[c] = -divergence[c] - Ap;
    }
    subtract_mean_fluid(r, neighbour_mask, cells, fluid_count);

    // ─── 3. z = M⁻¹ * r, s = z, sigma = r · z ───────────────────────────────

    apply_preconditioner(r, z, precon, neighbour_mask, cells, rx, slice);

    let sigma = 0;
    for (let c = 0; c < cells; c++) {
        s[c] = z[c];
        sigma += r[c] * z[c];
    }

    if (sigma === 0) {
        // Initial residual was already zero (or pathological). Either way,
        // nothing to do.
        return;
    }

    // ─── 4. PCG outer loop ──────────────────────────────────────────────────

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

        // As = A * s, accumulating s · As in the same pass.
        let s_dot_As = 0;
        for (let c = 0; c < cells; c++) {
            const mask = neighbour_mask[c];
            if ((mask & 63) === 0) {
                As[c] = 0;
                continue;
            }
            let val = POPCOUNT_6[mask] * s[c];
            if (mask & 1)  val -= s[c - 1];
            if (mask & 2)  val -= s[c + 1];
            if (mask & 4)  val -= s[c - rx];
            if (mask & 8)  val -= s[c + rx];
            if (mask & 16) val -= s[c - slice];
            if (mask & 32) val -= s[c + slice];
            As[c] = val;
            s_dot_As += s[c] * val;
        }

        if (s_dot_As === 0) {
            // Degenerate search direction; stop.
            break;
        }

        const alpha = sigma / s_dot_As;

        // pressure += alpha * s
        // r        -= alpha * As
        for (let c = 0; c < cells; c++) {
            pressure[c] += alpha * s[c];
            r[c]        -= alpha * As[c];
        }
        // Project r back onto the mean-zero (range(A)) subspace. Round-off
        // accumulates a constant component each iter; without this, alpha
        // explodes within ~8 iters once s aligns with the null space.
        subtract_mean_fluid(r, neighbour_mask, cells, fluid_count);

        // z = M⁻¹ * r
        apply_preconditioner(r, z, precon, neighbour_mask, cells, rx, slice);

        // sigma_new = r · z
        let sigma_new = 0;
        for (let c = 0; c < cells; c++) {
            sigma_new += r[c] * z[c];
        }

        if (sigma_new === 0) {
            break;
        }

        const beta = sigma_new / sigma;

        // s = z + beta * s
        for (let c = 0; c < cells; c++) {
            s[c] = z[c] + beta * s[c];
        }

        sigma = sigma_new;
    }
}

/**
 * Apply the MIC(0) preconditioner: solve M · z = r where M = L · Lᵀ.
 *
 * Two sequential sweeps:
 *   - forward (L · y = r): walk fluid cells in lex order, using backward
 *     fluid neighbours' already-computed y values.
 *   - backward (Lᵀ · z = y): walk in reverse lex order, using forward
 *     fluid neighbours' already-computed z values.
 *
 * Both write into `z` in place; the forward pass stores y in z, the backward
 * pass overwrites with the final z. Solid cells stay at 0 throughout.
 *
 * For our normalized Laplacian (A = -1 off-diagonal, popcount diagonal), the
 * L factor has L[c, c-1] = -precon[c-1] etc., L[c, c] = 1/precon[c], so each
 * triangular solve degenerates to ~4 floating-point ops per cell.
 *
 * @param {Float32Array} r
 * @param {Float32Array} z
 * @param {Float32Array} precon
 * @param {Uint8Array} mask_array
 * @param {number} cells
 * @param {number} rx
 * @param {number} slice
 */
/**
 * Subtract the mean (over fluid cells) from a buffer in place. Used to project
 * a vector onto range(A) = mean-zero subspace; required for CG to converge on
 * a system with a non-trivial null space (pure Neumann pressure Poisson).
 *
 * @param {Float32Array} v
 * @param {Uint8Array} mask_array
 * @param {number} cells
 * @param {number} fluid_count
 */
function subtract_mean_fluid(v, mask_array, cells, fluid_count) {
    let sum = 0;
    for (let c = 0; c < cells; c++) {
        if ((mask_array[c] & 63) !== 0) sum += v[c];
    }
    const mean = sum / fluid_count;
    for (let c = 0; c < cells; c++) {
        if ((mask_array[c] & 63) !== 0) v[c] -= mean;
    }
}

function apply_preconditioner(r, z, precon, mask_array, cells, rx, slice) {
    // Forward solve L · y = r. Store y in z (we won't read z[c'] for c' > c).
    for (let c = 0; c < cells; c++) {
        const mask = mask_array[c];
        if ((mask & 63) === 0) {
            z[c] = 0;
            continue;
        }
        let t = r[c];
        // L[c, c-1] = -precon[c-1] when -x neighbour fluid → subtracting L·y is
        // adding precon[c-1] * y[c-1].
        if (mask & 1)  t += precon[c - 1]     * z[c - 1];
        if (mask & 4)  t += precon[c - rx]    * z[c - rx];
        if (mask & 16) t += precon[c - slice] * z[c - slice];
        // L[c, c] = 1 / precon[c] → dividing by it is multiplying by precon[c].
        z[c] = t * precon[c];
    }

    // Backward solve Lᵀ · z = y, in place. When we read z[c+1], z[c+rx],
    // z[c+slice], those were written in this same backward pass (higher
    // indices, processed first in reverse).
    for (let c = cells - 1; c >= 0; c--) {
        const mask = mask_array[c];
        if ((mask & 63) === 0) {
            continue;
        }
        let t = z[c]; // = y[c] from forward pass
        // L[c+1, c] = -precon[c] when c+1's -x neighbour (= c) is fluid, which
        // is equivalent to "c's +x neighbour (= c+1) is fluid", i.e. mask & 2.
        if (mask & 2)  t += precon[c] * z[c + 1];
        if (mask & 8)  t += precon[c] * z[c + rx];
        if (mask & 32) t += precon[c] * z[c + slice];
        z[c] = t * precon[c];
    }
}
