js/fluid.js

/**
 * 🌿 Ivy's GPU Art Studio
 * Tab 2: Fluid Simulation
 *
 * GPU-accelerated fluid dynamics using compute shaders
 * Based on Jos Stam's "Stable Fluids" algorithm
 * Enhanced with styles, palettes, and effects!
 */

class FluidRenderer {
    constructor() {
        this.device = null;
        this.context = null;
        this.format = null;

        // Simulation parameters
        this.params = {
            style: 0, // 0=classic, 1=ivy, 2=ink, 3=smoke, 4=plasma, 5=watercolor
            palette: 0, // 0=ivy, 1=rainbow, 2=fire, 3=ocean, 4=neon, 5=sunset, 6=cosmic, 7=mono
            viscosity: 0.1,
            diffusion: 0.0001,
            force: 100,
            curl: 30,
            pressure: 0.8,
            bloom: true,
            vortex: false
        };

        // Simulation state
        this.gridSize = 256;
        this.velocityBuffers = [];
        this.densityBuffers = [];
        this.currentBuffer = 0;

        this.input = null;
        this.animationLoop = null;
        this.isActive = false;
        this.time = 0;

        // Previous mouse position for velocity
        this.prevMouseX = 0.5;
        this.prevMouseY = 0.5;
    }

    async init(device, context, format, canvas) {
        this.device = device;
        this.context = context;
        this.format = format;
        this.canvas = canvas;

        // Create simulation buffers
        this.createBuffers();

        // Create pipelines
        await this.createPipelines();

        // Setup input
        this.input = new WebGPUUtils.InputHandler(canvas);

        // Animation loop
        this.animationLoop = new WebGPUUtils.AnimationLoop((dt, totalTime) => {
            this.time = totalTime;
            this.simulate(dt);
            this.render();
        });
    }

    createBuffers() {
        const size = this.gridSize * this.gridSize;

        // Double buffering for velocity (vec2) and density (f32)
        for (let i = 0; i < 2; i++) {
            this.velocityBuffers.push(
                this.device.createBuffer({
                    label: `Velocity Buffer ${i}`,
                    size: size * 8, // vec2f = 8 bytes
                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
                })
            );

            this.densityBuffers.push(
                this.device.createBuffer({
                    label: `Density Buffer ${i}`,
                    size: size * 16, // vec4f for RGBA = 16 bytes
                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
                })
            );
        }

        // Uniform buffer
        this.uniformBuffer = this.device.createBuffer({
            label: "Fluid Uniforms",
            size: 64,
            usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
        });

        // Initialize with zeros
        const zeroVelocity = new Float32Array(size * 2);
        const zeroDensity = new Float32Array(size * 4);

        for (let i = 0; i < 2; i++) {
            this.device.queue.writeBuffer(this.velocityBuffers[i], 0, zeroVelocity);
            this.device.queue.writeBuffer(this.densityBuffers[i], 0, zeroDensity);
        }
    }

    async createPipelines() {
        // Compute shader for simulation
        const computeShader = this.device.createShaderModule({
            label: "Fluid Compute Shader",
            code: this.getComputeShaderCode()
        });

        // Render shader for display
        const renderShader = this.device.createShaderModule({
            label: "Fluid Render Shader",
            code: this.getRenderShaderCode()
        });

        // Bind group layouts
        this.computeBindGroupLayout = this.device.createBindGroupLayout({
            entries: [
                { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
                { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
                { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } }
            ]
        });

        this.renderBindGroupLayout = this.device.createBindGroupLayout({
            entries: [
                { binding: 0, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "uniform" } },
                { binding: 1, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "read-only-storage" } },
                { binding: 2, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "read-only-storage" } }
            ]
        });

        // Compute pipeline
        this.computePipeline = this.device.createComputePipeline({
            label: "Fluid Compute Pipeline",
            layout: this.device.createPipelineLayout({
                bindGroupLayouts: [this.computeBindGroupLayout]
            }),
            compute: {
                module: computeShader,
                entryPoint: "main"
            }
        });

        // Render pipeline
        this.renderPipeline = this.device.createRenderPipeline({
            label: "Fluid Render Pipeline",
            layout: this.device.createPipelineLayout({
                bindGroupLayouts: [this.renderBindGroupLayout]
            }),
            vertex: {
                module: renderShader,
                entryPoint: "vertexMain"
            },
            fragment: {
                module: renderShader,
                entryPoint: "fragmentMain",
                targets: [{ format: this.format }]
            },
            primitive: {
                topology: "triangle-list"
            }
        });

        // Create bind groups
        this.updateBindGroups();
    }

    updateBindGroups() {
        const curr = this.currentBuffer;
        const next = 1 - curr;

        this.computeBindGroup = this.device.createBindGroup({
            layout: this.computeBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[curr] } },
                { binding: 2, resource: { buffer: this.velocityBuffers[next] } },
                { binding: 3, resource: { buffer: this.densityBuffers[curr] } },
                { binding: 4, resource: { buffer: this.densityBuffers[next] } }
            ]
        });

        this.renderBindGroup = this.device.createBindGroup({
            layout: this.renderBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[next] } },
                { binding: 2, resource: { buffer: this.densityBuffers[next] } }
            ]
        });
    }

    start() {
        this.isActive = true;
        console.log("🌊 FluidRenderer started!");
        this.animationLoop.start();
    }

    stop() {
        this.isActive = false;
        console.log("🌊 FluidRenderer stopped");
        this.animationLoop.stop();
    }

    reset() {
        const size = this.gridSize * this.gridSize;
        const zeroVelocity = new Float32Array(size * 2);
        const zeroDensity = new Float32Array(size * 4);

        for (let i = 0; i < 2; i++) {
            this.device.queue.writeBuffer(this.velocityBuffers[i], 0, zeroVelocity);
            this.device.queue.writeBuffer(this.densityBuffers[i], 0, zeroDensity);
        }
    }

    setViscosity(value) {
        this.params.viscosity = value;
    }

    setDiffusion(value) {
        this.params.diffusion = value;
    }

    setForce(value) {
        this.params.force = value;
    }

    setColorMode(mode) {
        const modes = { ink: 0, fire: 1, rainbow: 2, smoke: 3, ivy: 4 };
        this.params.colorMode = modes[mode] || 0;
    }

    setStyle(style) {
        const styles = { classic: 0, ivy: 1, ink: 2, smoke: 3, plasma: 4, watercolor: 5 };
        this.params.style = styles[style] ?? 0;
    }

    setPalette(palette) {
        const palettes = { ivy: 0, rainbow: 1, fire: 2, ocean: 3, neon: 4, sunset: 5, cosmic: 6, monochrome: 7 };
        this.params.palette = palettes[palette] ?? 0;
    }

    setCurl(value) {
        this.params.curl = value;
    }

    setPressure(value) {
        this.params.pressure = value;
    }

    setBloom(enabled) {
        this.params.bloom = enabled;
    }

    setVortex(enabled) {
        this.params.vortex = enabled;
    }

    simulate(dt) {
        if (!this.isActive) return;

        // Auto-spawn some fluid for visual feedback even without mouse
        const autoSpawn = !this.input.isPressed;
        let mouseX = this.input.mouseX;
        let mouseY = this.input.mouseY;
        let isPressed = this.input.isPressed;

        // Auto animation when not interacting
        if (autoSpawn && this.time > 0) {
            // Create swirling patterns automatically
            const t = this.time * 0.5;
            mouseX = 0.5 + 0.3 * Math.sin(t);
            mouseY = 0.5 + 0.3 * Math.cos(t * 0.7);
            isPressed = true; // Simulate mouse press for auto-spawn
        }

        // Calculate mouse velocity
        const dx = (mouseX - this.prevMouseX) * this.params.force;
        const dy = (mouseY - this.prevMouseY) * this.params.force;
        this.prevMouseX = mouseX;
        this.prevMouseY = mouseY;

        // Update uniforms - expanded for new params
        const uniforms = new Float32Array([
            this.gridSize, // 0: grid size
            dt, // 1: delta time
            this.params.viscosity, // 2: viscosity
            this.params.diffusion, // 3: diffusion
            mouseX, // 4: mouse X
            mouseY, // 5: mouse Y
            dx, // 6: velocity X
            dy, // 7: velocity Y
            isPressed ? 1.0 : 0.0, // 8: is mouse pressed
            this.params.style, // 9: style
            this.params.palette, // 10: palette
            this.params.curl, // 11: curl/vorticity
            this.params.pressure, // 12: pressure
            this.params.bloom ? 1.0 : 0.0, // 13: bloom
            this.params.vortex ? 1.0 : 0.0, // 14: vortex
            this.time // 15: time
        ]);

        this.device.queue.writeBuffer(this.uniformBuffer, 0, uniforms);

        // Update bind groups with current buffer state
        this.updateBindGroups();

        // Run compute shader
        const commandEncoder = this.device.createCommandEncoder();
        const computePass = commandEncoder.beginComputePass();

        computePass.setPipeline(this.computePipeline);
        computePass.setBindGroup(0, this.computeBindGroup);
        computePass.dispatchWorkgroups(Math.ceil(this.gridSize / 8), Math.ceil(this.gridSize / 8));

        computePass.end();
        this.device.queue.submit([commandEncoder.finish()]);

        // Swap buffers
        this.currentBuffer = 1 - this.currentBuffer;
    }

    render() {
        if (!this.isActive) return;

        WebGPUUtils.resizeCanvasToDisplaySize(this.canvas, window.devicePixelRatio);

        // IMPORTANT: Create render bind group to read the LATEST buffer (after compute)
        const curr = this.currentBuffer;
        const renderBindGroup = this.device.createBindGroup({
            layout: this.renderBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[curr] } },
                { binding: 2, resource: { buffer: this.densityBuffers[curr] } }
            ]
        });

        const commandEncoder = this.device.createCommandEncoder();
        const renderPass = commandEncoder.beginRenderPass({
            colorAttachments: [
                {
                    view: this.context.getCurrentTexture().createView(),
                    clearValue: { r: 0, g: 0, b: 0, a: 1 },
                    loadOp: "clear",
                    storeOp: "store"
                }
            ]
        });

        renderPass.setPipeline(this.renderPipeline);
        renderPass.setBindGroup(0, renderBindGroup);
        renderPass.draw(3);
        renderPass.end();

        this.device.queue.submit([commandEncoder.finish()]);
    }

    getComputeShaderCode() {
        return /* wgsl */ `
            struct Uniforms {
                gridSize: f32,
                dt: f32,
                viscosity: f32,
                diffusion: f32,
                mouseX: f32,
                mouseY: f32,
                velX: f32,
                velY: f32,
                mousePressed: f32,
                style: f32,
                palette: f32,
                curl: f32,
                pressure: f32,
                doBloom: f32,
                doVortex: f32,
                time: f32,
            }

            @group(0) @binding(0) var<uniform> u: Uniforms;
            @group(0) @binding(1) var<storage, read> velIn: array<vec2f>;
            @group(0) @binding(2) var<storage, read_write> velOut: array<vec2f>;
            @group(0) @binding(3) var<storage, read> densIn: array<vec4f>;
            @group(0) @binding(4) var<storage, read_write> densOut: array<vec4f>;

            fn idx(x: i32, y: i32) -> u32 {
                let size = i32(u.gridSize);
                let cx = clamp(x, 0, size - 1);
                let cy = clamp(y, 0, size - 1);
                return u32(cy * size + cx);
            }

            fn getPaletteColor(t: f32, paletteId: i32) -> vec3f {
                let tt = fract(t);

                if (paletteId == 0) { // Ivy Green
                    return vec3f(0.13 * tt + 0.05, 0.77 * tt + 0.2, 0.37 * tt + 0.1);
                } else if (paletteId == 1) { // Rainbow
                    return vec3f(
                        0.5 + 0.5 * sin(tt * 6.28 + 0.0),
                        0.5 + 0.5 * sin(tt * 6.28 + 2.094),
                        0.5 + 0.5 * sin(tt * 6.28 + 4.188)
                    );
                } else if (paletteId == 2) { // Fire
                    return vec3f(tt, tt * 0.4, tt * 0.1);
                } else if (paletteId == 3) { // Ocean
                    return vec3f(0.1 * tt, 0.4 * tt + 0.1, 0.9 * tt + 0.1);
                } else if (paletteId == 4) { // Neon
                    return vec3f(
                        0.5 + 0.5 * sin(tt * 12.0),
                        0.5 + 0.5 * sin(tt * 12.0 + 2.0),
                        0.5 + 0.5 * sin(tt * 12.0 + 4.0)
                    );
                } else if (paletteId == 5) { // Sunset
                    return vec3f(0.9 * tt + 0.1, 0.4 * tt, 0.3 * tt + 0.1);
                } else if (paletteId == 6) { // Cosmic
                    return vec3f(0.3 * tt + 0.1, 0.1 * tt + 0.05, 0.8 * tt + 0.2);
                } else { // Monochrome
                    return vec3f(tt * 0.9 + 0.1);
                }
            }

            @compute @workgroup_size(8, 8)
            fn main(@builtin(global_invocation_id) gid: vec3u) {
                let size = i32(u.gridSize);
                let x = i32(gid.x);
                let y = i32(gid.y);

                if (x >= size || y >= size) {
                    return;
                }

                let i = idx(x, y);
                let paletteId = i32(u.palette);

                // Read previous state
                var newVel = velIn[i];
                var newDens = densIn[i];

                // Get neighbors for diffusion
                let vL = velIn[idx(x - 1, y)];
                let vR = velIn[idx(x + 1, y)];
                let vU = velIn[idx(x, y + 1)];
                let vD = velIn[idx(x, y - 1)];

                let dL = densIn[idx(x - 1, y)];
                let dR = densIn[idx(x + 1, y)];
                let dU = densIn[idx(x, y + 1)];
                let dD = densIn[idx(x, y - 1)];

                // Apply diffusion (controlled by diffusion parameter)
                let diffAmount = u.diffusion * 1000.0;
                newVel = mix(newVel, (vL + vR + vU + vD) * 0.25, diffAmount);
                newDens = mix(newDens, (dL + dR + dU + dD) * 0.25, diffAmount);

                // Apply viscosity (dampens velocity)
                newVel *= (1.0 - u.viscosity * 0.1);

                // Vorticity / curl effect
                if (u.doVortex > 0.5) {
                    let curlAmount = u.curl * 0.0005;
                    let vortex = (vR.y - vL.y) - (vU.x - vD.x);
                    newVel += vec2f(-vortex, vortex) * curlAmount;
                }

                // Add forces from mouse
                let fx = f32(x) / f32(size);
                let fy = f32(y) / f32(size);
                let dist = distance(vec2f(fx, fy), vec2f(u.mouseX, u.mouseY));
                let radius = 0.02 + (u.pressure * 0.1); // Pressure affects brush size

                if (dist < radius && u.mousePressed > 0.5) {
                    let strength = 1.0 - dist / radius;
                    // Force affects velocity strength
                    let forceMultiplier = u.velX * u.velX + u.velY * u.velY;
                    newVel += vec2f(u.velX, u.velY) * strength * u.dt * 2.0;

                    // Add density/color using palette
                    let colorHue = strength + u.time * 0.1;
                    let color = getPaletteColor(colorHue, paletteId);
                    newDens += vec4f(color * strength * 3.0, strength * 3.0);
                }

                // Apply pressure (affects how much velocity is preserved)
                newVel *= u.pressure;

                // Decay
                newVel *= 0.995;
                newDens *= 0.992;

                // Boundary conditions
                if (x <= 1 || x >= size - 2 || y <= 1 || y >= size - 2) {
                    newVel *= 0.5;
                }

                velOut[i] = newVel;
                densOut[i] = newDens;
            }
        `;
    }

    getRenderShaderCode() {
        return /* wgsl */ `
            struct Uniforms {
                gridSize: f32,
                dt: f32,
                viscosity: f32,
                diffusion: f32,
                mouseX: f32,
                mouseY: f32,
                velX: f32,
                velY: f32,
                mousePressed: f32,
                style: f32,
                palette: f32,
                curl: f32,
                pressure: f32,
                doBloom: f32,
                doVortex: f32,
                time: f32,
            }

            @group(0) @binding(0) var<uniform> u: Uniforms;
            @group(0) @binding(1) var<storage, read> velocity: array<vec2f>;
            @group(0) @binding(2) var<storage, read> density: array<vec4f>;

            struct VertexOutput {
                @builtin(position) position: vec4f,
                @location(0) uv: vec2f,
            }

            @vertex
            fn vertexMain(@builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
                var pos = array<vec2f, 3>(
                    vec2f(-1.0, -1.0),
                    vec2f(3.0, -1.0),
                    vec2f(-1.0, 3.0)
                );

                var output: VertexOutput;
                output.position = vec4f(pos[vertexIndex], 0.0, 1.0);
                output.uv = pos[vertexIndex] * 0.5 + 0.5;
                return output;
            }

            fn getPaletteColor(t: f32, paletteId: i32) -> vec3f {
                let tt = fract(t);

                if (paletteId == 0) { // Ivy Green
                    return vec3f(0.1 + 0.2 * tt, 0.5 + 0.5 * tt, 0.2 + 0.3 * tt);
                } else if (paletteId == 1) { // Rainbow
                    return vec3f(
                        0.5 + 0.5 * cos(6.28318 * (tt + 0.0)),
                        0.5 + 0.5 * cos(6.28318 * (tt + 0.33)),
                        0.5 + 0.5 * cos(6.28318 * (tt + 0.67))
                    );
                } else if (paletteId == 2) { // Fire
                    return vec3f(min(1.0, tt * 2.5), tt * tt, tt * tt * tt * 0.3);
                } else if (paletteId == 3) { // Ocean
                    return vec3f(0.0 + 0.2 * tt, 0.3 + 0.4 * tt, 0.6 + 0.4 * tt);
                } else if (paletteId == 4) { // Neon
                    return vec3f(
                        0.5 + 0.5 * sin(tt * 12.56),
                        0.5 + 0.5 * sin(tt * 12.56 + 2.094),
                        0.5 + 0.5 * sin(tt * 12.56 + 4.188)
                    );
                } else if (paletteId == 5) { // Sunset
                    return vec3f(0.9 - 0.2 * tt, 0.3 + 0.4 * tt, 0.3 + 0.5 * tt);
                } else if (paletteId == 6) { // Cosmic
                    return vec3f(
                        0.2 + 0.5 * sin(tt * 6.28),
                        0.1 + 0.3 * sin(tt * 6.28 + 2.0),
                        0.5 + 0.5 * sin(tt * 6.28 + 4.0)
                    );
                } else { // Monochrome
                    return vec3f(tt, tt, tt);
                }
            }

            @fragment
            fn fragmentMain(input: VertexOutput) -> @location(0) vec4f {
                let size = i32(u.gridSize);
                let x = i32(input.uv.x * f32(size));
                let y = i32(input.uv.y * f32(size));
                let i = u32(clamp(y, 0, size - 1) * size + clamp(x, 0, size - 1));

                let d = density[i];
                let v = velocity[i];

                let style = i32(u.style);
                let paletteId = i32(u.palette);
                let speed = length(v);
                let dens = length(d.rgb);

                var color = vec3f(0.0);

                // Show mouse position as a dot for visual feedback
                let mouseDist = distance(input.uv, vec2f(u.mouseX, u.mouseY));
                let mouseGlow = smoothstep(0.08, 0.0, mouseDist) * 0.5;

                // Style-based rendering
                if (style == 0) { // Classic - use density color directly
                    color = d.rgb;
                } else if (style == 1) { // Ivy Flow - organic green tones
                    let hue = dens * 0.3 + speed * 0.1;
                    color = getPaletteColor(hue, paletteId);
                    color *= dens * 1.5;
                } else if (style == 2) { // Ink Drop - high contrast
                    color = getPaletteColor(dens + speed * 0.2, paletteId);
                    color = pow(color * dens, vec3f(0.8));
                } else if (style == 3) { // Smoke - soft gradient
                    let smoke = smoothstep(0.0, 1.0, dens);
                    color = mix(vec3f(0.02), getPaletteColor(speed * 0.5, paletteId), smoke);
                } else if (style == 4) { // Plasma - vibrant swirls
                    let plasma = sin(dens * 10.0 + u.time) * 0.5 + 0.5;
                    color = getPaletteColor(plasma + speed * 0.3, paletteId);
                    color *= dens * 2.0;
                } else { // Watercolor - soft bleeding edges
                    let wc = smoothstep(0.0, 0.5, dens);
                    color = getPaletteColor(dens * 0.5 + u.time * 0.05, paletteId) * wc;
                    color = mix(color, vec3f(1.0), (1.0 - wc) * 0.1);
                }

                // Velocity-based highlights
                color += getPaletteColor(0.8, paletteId) * speed * 0.15;

                // Add mouse indicator
                color += getPaletteColor(u.time * 0.2, paletteId) * mouseGlow;

                // Vortex visualization
                if (u.doVortex > 0.5) {
                    // Approximate curl from velocity
                    let curlVis = abs(v.x - v.y) * 0.5;
                    color += vec3f(curlVis * 0.3, curlVis * 0.1, curlVis * 0.4);
                }

                // Bloom effect
                if (u.doBloom > 0.5) {
                    let bloom = max(0.0, dens - 0.5) * 2.0;
                    color += color * bloom * 0.5;
                    color = color / (1.0 + color * 0.3); // Tone mapping
                }

                return vec4f(color, 1.0);
            }
        `;
    }
}

// Export
window.FluidRenderer = FluidRenderer;