import { assert } from "../../../core/assert.js";
import { optimal_sor_omega } from "./solver/optimal_sor_omega.js";
import { v3_grid_apply_advection_forward } from "./solver/v3_grid_apply_advection_forward.js";
import { v3_grid_apply_diffusion } from "./solver/v3_grid_apply_diffusion.js";
import { v3_grid_apply_scalar_advection } from "./solver/v3_grid_apply_scalar_advection.js";
import { v3_grid_compute_divergence } from "./solver/v3_grid_compute_divergence.js";
import { v3_grid_solve_pressure } from "./solver/v3_grid_solve_pressure.js";
import { v3_grid_solve_pressure_pcg } from "./solver/v3_grid_solve_pressure_pcg.js";
import { v3_grid_subtract_pressure_gradient } from "./solver/v3_grid_subtract_pressure_gradient.js";

/**
 * Pressure solver selection. Both solve the same A·p = -div system over the
 * same neighbour-mask topology; the difference is convergence rate and per-iter
 * cost.
 *
 *   - `"sor"` — red-black Successive Over-Relaxation, our long-standing
 *     default. One sweep is cheap (~3-5 µs/cell at 128³), convergence is
 *     O(N) iterations for fixed accuracy.
 *
 *   - `"mic-pcg"` — Modified Incomplete Cholesky preconditioned Conjugate
 *     Gradient. ~2× the per-iter cost of SOR, but converges in O(√N) iters
 *     (typically 8-12 vs SOR's 30+ at 128³), so a net 1.5-2× speedup at
 *     quality parity for production grids. Allocates 5 additional N-sized
 *     scratch buffers (r, z, s, As, precon).
 *
 * @enum {string}
 */
export const PressureSolver = {
    SOR: "sor",
    MICPCG: "mic-pcg",
};

/**
 * Cell-centered 3D Stable Fluids solver (Stam, 1999/2003).
 *
 * Holds tuning knobs plus transient working memory (per-step velocity snapshots,
 * divergence, ping-pong scratch). None of that working memory carries information
 * across step boundaries — its value is overwritten before being read on the next
 * step. Replacing the simulator instance between two steps of the same field
 * produces an identical result (the only cross-step state, `pressure`, lives on
 * the field). Scratch is cached by cell-count so per-step allocation stays at zero.
 *
 * Per step:
 *   1. Apply external forces (effectors mutate field velocity).
 *   2. Diffuse velocity (viscosity; skipped when `velocity_diffusion_rate == 0`).
 *   3. Project (pressure solve + gradient subtract) to enforce ∇·v ≈ 0.
 *   4. Self-advect velocity (semi-Lagrangian back-trace).
 *   5. Project again (advection introduces small divergence).
 *   6. For each scalar attribute on the field: optionally diffuse, then advect.
 *
 * Pressure is warm-started from the previous step's solution (`field.pressure`),
 * which converges much faster than starting from zero once the flow reaches a
 * quasi-steady state.
 *
 * Compute model: synchronous, main-thread. Cost scales linearly with cell count and
 * with iteration counts. The pressure SOR sweep count is auto-scaled to grid size
 * (see {@link pressure_iteration_scale}) so the per-cell budget stays roughly flat
 * across resolutions instead of over-spending at 32³ and under-spending at 256³.
 */
/**
 * Minimum iteration count when `pressure_iteration_scale > 0`. Picked so very
 * small grids (8³, 16³) still get enough sweeps to converge — at those sizes
 * `scale * max_res` underflows the practical "useful work" zone. For grids
 * larger than `floor / scale` the scale formula dominates and this constant
 * stops mattering. Caller may inspect / override if a per-deployment minimum
 * makes sense.
 *
 * @type {number}
 */
export const PRESSURE_ITERATIONS_FLOOR = 8;

export class FluidSimulator {

    /**
     * Implicit-diffusion coefficient (= viscosity * dt-like scaling). 0 disables.
     * @type {number}
     */
    velocity_diffusion_rate = 0;

    /**
     * Jacobi iterations for the velocity diffusion solve. Only consulted when
     * `velocity_diffusion_rate > 0`.
     * @type {number}
     */
    velocity_diffusion_iterations = 4;

    /**
     * Implicit-diffusion coefficient applied uniformly to all scalar attributes. 0
     * disables. Semi-Lagrangian advection already introduces visible numerical
     * diffusion, so leaving this at 0 is usually fine.
     * @type {number}
     */
    scalar_diffusion_rate = 0;

    /**
     * @type {number}
     */
    scalar_diffusion_iterations = 4;

    /**
     * Which pressure solver to use. `PressureSolver.SOR` (default) is the
     * baseline red-black Successive Over-Relaxation. `PressureSolver.MICPCG`
     * switches to Modified Incomplete Cholesky preconditioned Conjugate
     * Gradient, which typically converges in 1/3 the iterations at 2× per-iter
     * cost (net ~1.5-2× faster wall time at quality parity on production
     * grids). MIC-PCG allocates an additional 5 N-sized scratch buffers, so
     * peak memory is higher.
     *
     * Switching at runtime is supported but flushes the warm-started pressure
     * meaning for one step (both solvers warm-start from `field.pressure`,
     * but the convergence trajectory is different — first frame after a
     * switch will look "fresh", not warm).
     *
     * @type {string}
     */
    pressure_solver = PressureSolver.SOR;

    /**
     * Run the pressure projection *before* self-advection in addition to after.
     *
     * The Stam (1999/2003) reference solver projects twice per step:
     *
     *     [forces] → [diffuse] → **project₁** → [advect] → **project₂** → [scalars]
     *
     * `project₁` ensures the velocity field that's *fed into* semi-Lagrangian
     * advection is divergence-free, so back-traced trajectories follow a physical
     * flow. `project₂` cleans up the small divergence advection re-introduces.
     *
     * Skipping `project₁` (this flag = false) cuts the pressure-solve cost roughly
     * in half — at default settings, pressure is ~80% of step() time, and the two
     * projections share that ~evenly. The cost is that effectors' splatted
     * impulses (which inject divergence) and any post-diffusion residual get
     * carried into advection un-projected, then have to be cleaned up by the
     * single post-advection projection.
     *
     * In practice:
     *   - If diffusion is OFF (rate=0) and effectors are gentle (gentle wind, no
     *     big impulses), the visual quality cost is small. The post-advection
     *     project handles everything in roughly the same number of iterations.
     *   - If diffusion is ON, or you have strong impulses (explosions, large
     *     wakes), the second-only solver is asked to fix everything in one pass
     *     and may visibly under-converge — flicker, "pumping", or divergent
     *     vortices.
     *
     * Default `true` (= Stam reference behavior) for correctness; flip to `false`
     * if profiling shows pressure as a bottleneck and quality A/B looks acceptable.
     *
     * @type {boolean}
     */
    project_before_advection = true;

    /**
     * Adaptive scaling factor for the pressure Poisson SOR sweep count. The actual
     * iteration count per `project()` call is
     *
     *     iters = max(PRESSURE_ITERATIONS_FLOOR, ceil(pressure_iteration_scale * max(rx, ry, rz)))
     *
     * unless `pressure_iteration_scale <= 0`, in which case projection is skipped
     * entirely (useful for tests that need to inspect un-projected velocity or
     * tear-down scenarios).
     *
     * Why scaled by resolution? Red-black SOR with the optimal omega converges
     * geometrically with a rate that depends on grid size — the iterations needed
     * to reach a given residual scale roughly as `O(N)` where N is the max grid
     * extent. A fixed iteration count therefore over-spends at small grids and
     * under-spends at large ones; the default value of 30 (from the pre-adaptive
     * era) was right for ~128³ but wasteful at 32³ and marginal at 256³.
     *
     * Default scale `0.25` plus floor `8` chosen empirically against the bench
     * sweep in `bench.spec.js` (`pressure_iterations convergence + step() cost`)
     * — it's the smallest setting that keeps a stress-test divergent velocity
     * field stable under continuous stepping across the 32³ → 128³ range. For a
     * gentler real-world flow you can drop the scale further; for cold-start
     * spawn frames or scenes with heavy impulses, bump it up.
     *
     * @type {number}
     */
    pressure_iteration_scale = 0.25;

    // ────────────────────────────────────────────────────────────────────────────
    // Transient scratch — overwritten before read on every step. Cached by
    // cell-count so we don't allocate per step. Their value never crosses a step
    // boundary; if it did it would belong on the field (see `field.pressure`).
    // ────────────────────────────────────────────────────────────────────────────

    #scratch_cells = 0;

    /** @type {Float32Array|null} */ #prev_x = null;
    /** @type {Float32Array|null} */ #prev_y = null;
    /** @type {Float32Array|null} */ #prev_z = null;
    /** @type {Float32Array|null} */ #divergence = null;
    /** @type {Float32Array|null} */ #diffusion_scratch = null;
    /** @type {Float32Array|null} */ #scalar_scratch = null;

    // PCG-only scratch (allocated lazily on first MICPCG-enabled project call).
    #pcg_scratch_cells = 0;
    /** @type {Float32Array|null} */ #pcg_r = null;
    /** @type {Float32Array|null} */ #pcg_z = null;
    /** @type {Float32Array|null} */ #pcg_s = null;
    /** @type {Float32Array|null} */ #pcg_As = null;
    /** @type {Float32Array|null} */ #pcg_precon = null;

    /**
     * Grow the cached scratch buffers if the requested size exceeds what's allocated.
     * Idempotent and monotonic — never shrinks. Called at the top of {@link step} and
     * {@link project} so user code never has to think about it.
     * @param {number} cell_count
     */
    #ensure_scratch(cell_count) {
        if (this.#scratch_cells >= cell_count) {
            return;
        }
        this.#scratch_cells = cell_count;
        this.#prev_x = new Float32Array(cell_count);
        this.#prev_y = new Float32Array(cell_count);
        this.#prev_z = new Float32Array(cell_count);
        this.#divergence = new Float32Array(cell_count);
        this.#diffusion_scratch = new Float32Array(cell_count);
        this.#scalar_scratch = new Float32Array(cell_count);
    }

    /**
     * Lazily allocate the PCG-specific scratch (5 N-sized Float32 buffers).
     * Idempotent and monotonic.
     * @param {number} cell_count
     */
    #ensure_pcg_scratch(cell_count) {
        if (this.#pcg_scratch_cells >= cell_count) {
            return;
        }
        this.#pcg_scratch_cells = cell_count;
        this.#pcg_r = new Float32Array(cell_count);
        this.#pcg_z = new Float32Array(cell_count);
        this.#pcg_s = new Float32Array(cell_count);
        this.#pcg_As = new Float32Array(cell_count);
        this.#pcg_precon = new Float32Array(cell_count);
    }

    /**
     * Re-allocate the divergence scratch buffer if its element type doesn't match
     * the field's pressure type. Lets `half_precision_pressure` on the field carry
     * through into the SOR inner-loop bandwidth (the solver reads divergence and
     * writes pressure of the same dtype).
     *
     * @param {FluidField} field
     */
    #ensure_divergence_matches(field) {
        const desired_ctor = field.pressure.constructor;
        const cell_count = field.cellCount();
        if (this.#divergence === null
            || this.#divergence.constructor !== desired_ctor
            || this.#divergence.length < cell_count) {
            this.#divergence = new desired_ctor(cell_count);
        }
    }

    /**
     * Resolve the per-call SOR iteration count from {@link pressure_iteration_scale}
     * and the field's resolution. Public so callers can budget frame time or query
     * the actual cost in advance.
     *
     *   - `scale <= 0` → 0 (projection disabled).
     *   - `scale > 0`  → `max(PRESSURE_ITERATIONS_FLOOR, ceil(scale * max(rx, ry, rz)))`.
     *
     * @param {number} res_x
     * @param {number} res_y
     * @param {number} res_z
     * @return {number} non-negative integer iteration count.
     */
    resolvePressureIterations(res_x, res_y, res_z) {
        const scale = this.pressure_iteration_scale;
        if (scale <= 0) {
            return 0;
        }
        const max_res = res_x > res_y ? (res_x > res_z ? res_x : res_z) : (res_y > res_z ? res_y : res_z);
        const scaled = Math.ceil(scale * max_res);
        return scaled > PRESSURE_ITERATIONS_FLOOR ? scaled : PRESSURE_ITERATIONS_FLOOR;
    }

    /**
     * Project the field's velocity to be (approximately) divergence-free. Exposed
     * because tests and external code may need to project on demand (e.g., after
     * splatting an out-of-band impulse). `field.pressure` is read as the warm-start
     * initial guess and written with the new solution.
     * @param {FluidField} field
     */
    project(field) {
        this.#ensure_scratch(field.cellCount());
        this.#ensure_divergence_matches(field);

        const res = field.getResolution();
        const solid = field.solid;
        const iterations = this.resolvePressureIterations(res[0], res[1], res[2]);

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

        // Pre-bake the per-cell neighbourhood mask consumed by the pressure SOR
        // loop. One O(N) byte-pass amortized across `iterations` SOR sweeps;
        // saves six boundary checks plus six solid-neighbour checks per cell per
        // iteration. See FluidField.solid_neighbour_mask.
        field.recomputeSolidNeighbourMask();

        v3_grid_compute_divergence(
            this.#divergence,
            field.velocity_x, field.velocity_y, field.velocity_z,
            res[0], res[1], res[2],
            solid
        );

        if (this.pressure_solver === PressureSolver.MICPCG) {
            this.#ensure_pcg_scratch(field.cellCount());
            v3_grid_solve_pressure_pcg(
                field.pressure, this.#divergence,
                res[0], res[1], res[2],
                iterations,
                field.solid_neighbour_mask,
                this.#pcg_r, this.#pcg_z, this.#pcg_s, this.#pcg_As, this.#pcg_precon
            );
        } else {
            v3_grid_solve_pressure(
                field.pressure, this.#divergence,
                res[0], res[1], res[2],
                iterations,
                optimal_sor_omega(res[0], res[1], res[2]),
                field.solid_neighbour_mask
            );
        }

        v3_grid_subtract_pressure_gradient(
            field.velocity_x, field.velocity_y, field.velocity_z,
            field.pressure,
            res[0], res[1], res[2],
            field.solid_neighbour_mask
        );
    }

    /**
     * Diffuse a single velocity component in-place, using `source_copy_buf` as the
     * "before" snapshot and `#diffusion_scratch` as the ping-pong scratch.
     * @param {Float32Array} component
     * @param {Float32Array} source_copy_buf
     * @param {number[]} res
     * @param {Uint8Array} solid
     */
    #diffuse_velocity_component(component, source_copy_buf, res, solid) {
        source_copy_buf.set(component);
        v3_grid_apply_diffusion(
            component, source_copy_buf, this.#diffusion_scratch,
            res[0], res[1], res[2],
            this.velocity_diffusion_rate,
            this.velocity_diffusion_iterations,
            solid
        );
    }

    /**
     * Advance the simulation by one timestep.
     *
     * @param {FluidField} field
     * @param {number} time_delta            seconds (or grid-cells-per-velocity-unit;
     *                                       caller defines the convention)
     * @param {AbstractFluidEffector[]} effectors    Pass `[]` for no external forces.
     * @param {Float32Array|number[]} world_to_grid  Column-major 4×4 affine sending
     *                                               world-space points into this
     *                                               field's grid space. Passed through
     *                                               to each effector's `apply`. Pass
     *                                               `M4_IDENTITY` when effectors are
     *                                               already configured in grid space.
     */
    step(field, time_delta, effectors, world_to_grid) {
        assert.defined(field, "field");
        assert.notNull(field, "field");
        assert.isNumber(time_delta, "time_delta");
        assert.greaterThanOrEqual(time_delta, 0, "time_delta");
        assert.isArray(effectors, "effectors");
        assert.defined(world_to_grid, "world_to_grid");
        assert.notNull(world_to_grid, "world_to_grid");

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

        const cell_count = field.cellCount();
        assert.greaterThan(cell_count, 0, "field has not been built");

        this.#ensure_scratch(cell_count);

        const res = field.getResolution();

        // 1. External forces.
        for (let i = 0; i < effectors.length; i++) {
            effectors[i].apply(field, time_delta, world_to_grid);
        }

        const solid = field.solid;

        // 2. Velocity diffusion (viscosity). Cheap to skip when disabled.
        if (this.velocity_diffusion_rate > 0 && this.velocity_diffusion_iterations > 0) {
            this.#diffuse_velocity_component(field.velocity_x, this.#prev_x, res, solid);
            this.#diffuse_velocity_component(field.velocity_y, this.#prev_y, res, solid);
            this.#diffuse_velocity_component(field.velocity_z, this.#prev_z, res, solid);
        }

        // 3. Project to enforce incompressibility before advection. Skipping this
        //    is a perf knob (~halves the pressure cost) at the price of advection
        //    seeing a slightly-divergent input. See `project_before_advection`.
        if (this.project_before_advection) {
            this.project(field);
        }

        // 4. Self-advect velocity (semi-Lagrangian back-trace through the projected field).
        this.#prev_x.set(field.velocity_x);
        this.#prev_y.set(field.velocity_y);
        this.#prev_z.set(field.velocity_z);

        v3_grid_apply_advection_forward(
            [field.velocity_x, field.velocity_y, field.velocity_z],
            [this.#prev_x, this.#prev_y, this.#prev_z],
            res[0], res[1], res[2],
            time_delta,
            solid
        );

        // 5. Re-project after advection.
        this.project(field);

        // 6. Transport each passive scalar by the now-finalized velocity field.
        const scalars = field.getScalarAttributes();
        for (let i = 0; i < scalars.length; i++) {
            const attr = scalars[i];

            if (this.scalar_diffusion_rate > 0 && this.scalar_diffusion_iterations > 0) {
                this.#scalar_scratch.set(attr.data);
                v3_grid_apply_diffusion(
                    attr.data, this.#scalar_scratch, this.#diffusion_scratch,
                    res[0], res[1], res[2],
                    this.scalar_diffusion_rate,
                    this.scalar_diffusion_iterations,
                    solid
                );
            }

            this.#scalar_scratch.set(attr.data);
            v3_grid_apply_scalar_advection(
                attr.data, this.#scalar_scratch,
                field.velocity_x, field.velocity_y, field.velocity_z,
                res[0], res[1], res[2],
                time_delta,
                solid
            );
        }
    }
}
