<?php
let div = document.createElement('div');
document.body.appendChild( div );
div.style['font-size'] = '20pt';
function log( s )
{
  console.log( s );
  let args = [...arguments].join(' ');
  div.innerHTML += args + '<br><br>';
}

log('WebGPU Compute Example');

if (!navigator.gpu) { log("WebGPU is not supported (or is it disabled? flags/settings)"); return; }

const adapter = await navigator.gpu.requestAdapter();
const device  = await adapter.requestDevice();


let canvas = document.createElement('canvas');
document.body.appendChild( canvas ); canvas.height = canvas.width = 500;

const context = canvas.getContext('webgpu');
const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); 
console.log('presentationFormat:', presentationFormat );


context.configure({ device: device, 
                    usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC | GPUTextureUsage.COPY_DST,
                    format: "rgba8unorm" /*presentationFormat*/  });

let canvasTexture = context.getCurrentTexture();

let imgWidth  = canvas.width;
let imgHeight = imgWidth;


const texture1 = device.createTexture({
  size: [imgWidth, imgHeight, 1],
  format: "rgba8unorm",
  usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.COPY_SRC | GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.STORAGE_BINDING
});

// ----------------------------------------------------------------

const bufferSize = imgWidth * imgHeight * 4 * 2; // for vec2<f32>

const velocityBuffer0 = device.createBuffer({ size: bufferSize, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });

const velocityBuffer1 = device.createBuffer({ size: bufferSize, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });

const densityBuffer0 = device.createBuffer({ size: bufferSize / 2,  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });

const densityBuffer1 = device.createBuffer({ size: bufferSize / 2,  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });

// ----------------------------------------------------------------

/*
// Circle density test pattern
function createDensityPattern() {
  // Create density data with test pattern
  const density = new Float32Array(imgWidth * imgWidth);
  const size = imgWidth;
  const densityData = new Float32Array(size * size);
  for (let y = 0; y < size; y++) {
    for (let x = 0; x < size; x++) {
      // Example test pattern (e.g., a circular density)
      const centerX = size / 2;
      const centerY = size / 2;
      const distance = Math.sqrt((x - centerX) ** 2 + (y - centerY) ** 2);
      const densityValue = distance < 50 ? 1.0 : 0.0; // Adjust threshold and shape as needed
      density[y * size + x] = densityValue * Math.random() * 100.0;
    }
  }
  return density;
}
*/

function createDensityPattern() {
  // Create temporary canvas to draw the text
  const canvas = document.createElement('canvas');
  canvas.width = imgWidth;
  canvas.height = imgWidth;
  const ctx = canvas.getContext('2d');

  // Set font and text properties
  ctx.font = 'bold 150px Arial';
  ctx.fillStyle = 'white';

  // Draw the text 'cat' onto the canvas
  ctx.fillText('xbdev', 30, 250);

  // Extract pixel data from the canvas
  const imageData = ctx.getImageData(0, 0, imgWidth, imgWidth);
  const pixels = imageData.data;

  // Convert pixel data to density array
  const density = new Float32Array(imgWidth * imgWidth);
  for (let i = 0; i < pixels.length; i += 4) {
    // Convert RGBA pixel values to grayscale
    const grayscale = (pixels[i] + pixels[i + 1] + pixels[i + 2]) / 3;
    // Normalize grayscale value to range [0, 1]
    const densityValue = grayscale / 255;
    density[i / 4] = densityValue;
  }

  return density;
}

const initialDensityData = createDensityPattern();

device.queue.writeBuffer(densityBuffer0, 0, initialDensityData);
device.queue.writeBuffer(densityBuffer1, 0, initialDensityData);

// ----------------------------------------------------------------

function createVortexPattern() {
  const data = new Float32Array(imgWidth * imgWidth * 2);
  const centerX = imgWidth / 2;
  const centerY = imgWidth / 2;
  const strength = 10.0;
  for (let y = 0; y < imgWidth; y++) {
    for (let x = 0; x < imgWidth; x++) {
      const dx = x - centerX;
      const dy = y - centerY;
      const distance = Math.sqrt(dx * dx + dy * dy);
      const angle = Math.atan2(dy, dx);
      const index = 2 * (y * imgWidth + x);
     // data[index + 0] = -strength * Math.sin(angle) * Math.exp(-distance / 100);
     // data[index + 1] =  strength * Math.cos(angle) * Math.exp(-distance / 100);
      data[index + 0 ] = -10.0 + Math.random() * 20.0;
      data[index + 1 ] = -10.0 + Math.random() * 20.0;
    }
  }
  return data;
}

const initialVelocityData = createVortexPattern();

device.queue.writeBuffer(velocityBuffer0, 0, initialVelocityData);
device.queue.writeBuffer(velocityBuffer1, 0, initialVelocityData);

// ----------------------------------------------------------------

const timerUniformBuffer = device.createBuffer({
  size: 4, 
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
});

const timestep  = new Float32Array( [0.0] );

device.queue.writeBuffer(timerUniformBuffer,   0, timestep             );

// ----------------------------------------------------------


// Bind group layout and bind group
const bindGroupLayout = device.createBindGroupLayout({
  entries: [  
    {binding: 0, visibility: GPUShaderStage.COMPUTE,  buffer: { type: "uniform" }   },
    {binding: 1, visibility: GPUShaderStage.COMPUTE,  buffer: { type: "storage"  }  },
    {binding: 2, visibility: GPUShaderStage.COMPUTE,  buffer: { type: "storage"  }  },
    {binding: 3, visibility: GPUShaderStage.COMPUTE,  buffer: { type: "storage"  }  },
    {binding: 4, visibility: GPUShaderStage.COMPUTE,  buffer: { type: "storage"  }  },
    {binding: 5, visibility: GPUShaderStage.COMPUTE,  storageTexture: {format:"rgba8unorm", access:"write-only", viewDimension:"2d"}   }
           ]
});

const bindGroup0 = device.createBindGroup({ layout: bindGroupLayout,
                                            entries: [  {  binding: 0,  resource: { buffer: timerUniformBuffer} },
                                                       {  binding: 1,  resource: { buffer: velocityBuffer0 } },
                                                       {  binding: 2,  resource: { buffer: velocityBuffer1 } },
                                                       {  binding: 3,  resource: { buffer: densityBuffer0 } },
                                                       {  binding: 4,  resource: { buffer: densityBuffer1 } },
                                                       {  binding: 5,  resource: texture1.createView()  },
                                                    ] });

const bindGroup1 = device.createBindGroup({ layout: bindGroupLayout,
                                            entries: [  {  binding: 0,  resource: { buffer: timerUniformBuffer} },
                                                       {  binding: 1,  resource: { buffer: velocityBuffer1 } },
                                                       {  binding: 2,  resource: { buffer: velocityBuffer0 } },
                                                       {  binding: 3,  resource: { buffer: densityBuffer1 } },
                                                       {  binding: 4,  resource: { buffer: densityBuffer0 } },
                                                       {  binding: 5,  resource: texture1.createView()  },
                                                    ] });


// Compute shader code
const computeShader = ` 
@binding(0) @group(0) var<uniform> mytimer : f32;
@binding(1) @group(0) var<storage, read_write> velocity0: array<vec2<f32>>;
@binding(2) @group(0) var<storage, read_write> velocity1: array<vec2<f32>>;
@binding(3) @group(0) var<storage, read_write> density0: array<f32>;
@binding(4) @group(0) var<storage, read_write> density1: array<f32>;
@binding(5) @group(0) var myTexture1: texture_storage_2d<rgba8unorm, write>;

// Scene parameters
const imgWidth:  u32 = ${canvas.width};
const imgHeight: u32 = ${canvas.width};
const gridSize:  i32 = ${canvas.width};

fn random(uv: vec2<f32>) -> f32 {
    return fract(sin(dot(uv, vec2<f32>(12.9898, 78.233))) * 43758.5453);
}

const dt: f32 = 0.1; // Time step
const diffusion: f32 = 0.1; // Diffusion coefficient
const vdiff : f32 = 0.1;
const curlCoefficient = 0.4;
const size: u32 = ${canvas.width}; // Simulation grid size

fn index(x: u32, y: u32) -> u32 {
    return y * size + x;
}

fn diffuseDensity( factor:f32, x:u32, y:u32 )
{
    let center = density0[index(x, y)];
    let left   = density0[index(max(x, 1u) - 1u, y)];
    let right  = density0[index(min(x + 1u, size - 1u), y)];
    let up     = density0[index(x, max(y, 1u) - 1u)];
    let down   = density0[index(x, min(y + 1u, size - 1u))];
    density1[index(x,y)] = ( center + factor * (left + right + up + down - 4.0 * center) );
}

fn diffuseVelocity( factor:f32, x:u32, y:u32 )
{
    let center = velocity0[index(x, y)];
    let left   = velocity0[index(max(x, 1u) - 1u, y)];
    let right  = velocity0[index(min(x + 1u, size - 1u), y)];
    let up     = velocity0[index(x, max(y, 1u) - 1u)];
    let down   = velocity0[index(x, min(y + 1u, size - 1u))];
    velocity1[index(x,y)] = ( center + factor * (left + right + up + down - 4.0 * center) );
}

@compute @workgroup_size(8, 8)
fn diffuse_step(@builtin(global_invocation_id) id: vec3<u32>) {
    let x = id.x;
    let y = id.y;
    if (x >= size || y >= size) {
        return;
    }

    diffuseDensity( diffusion, x, y );
    
    diffuseVelocity( vdiff, x, y );
}

@compute @workgroup_size(8, 8)
fn advect_step(@builtin(global_invocation_id) id: vec3<u32>) {
    let x = id.x;
    let y = id.y;
    if (x >= size || y >= size) {
        return;
    }

    // Simple advection
    let uv = velocity0[index(x, y)];
    let pos = vec2<f32>(f32(x), f32(y)) - dt * uv;
    let x0 = clamp(floor(pos.x), 0.0, f32(size - 1));
    let x1 = clamp(ceil(pos.x), 0.0, f32(size - 1));
    let y0 = clamp(floor(pos.y), 0.0, f32(size - 1));
    let y1 = clamp(ceil(pos.y), 0.0, f32(size - 1));
    let s1 = pos.x - x0;
    let s0 = 1.0 - s1;
    let t1 = pos.y - y0;
    let t0 = 1.0 - t1;
    let advectedDensity = s0 * (t0 * density0[index(u32(x0), u32(y0))] + t1 * density0[index(u32(x0), u32(y1))]) +
                          s1 * (t0 * density0[index(u32(x1), u32(y0))] + t1 * density0[index(u32(x1), u32(y1))]);
    density1[index(x, y)] = advectedDensity;
}


@compute @workgroup_size(8, 8)
fn curl_step(@builtin(global_invocation_id) id: vec3<u32>) {
    let x = id.x;
    let y = id.y;
    if (x >= size || y >= size) {
        return;
    }

    
    let du_dy = ( velocity0[index(x    ,y + 1)].x  - velocity0[index(x    ,y - 1)].x ) / (2.0);
    let dv_dx = ( velocity0[index(x + 1,y    )].y  - velocity0[index(x - 1,y    )].y ) / (2.0);
    let curl = dv_dx - du_dy;
    
    velocity1[index(x,y)].x += -curl * dt * curlCoefficient;
    velocity1[index(x,y)].y +=  curl * dt * curlCoefficient;

}

@compute @workgroup_size(8, 8)
fn display(@builtin(global_invocation_id) globalId: vec3<u32>) {
    let x = globalId.x;
    let y = globalId.y;
    if (x >= imgWidth || y >= imgHeight) {
        return;
    }

    let d = density0[index(x, y)];
    var finalColor = vec4<f32>(d, d, d, 1.0);
    textureStore(myTexture1, vec2<i32>(globalId.xy), finalColor);
}


`;
  

let shaderModule = device.createShaderModule({code:computeShader});
// Pipeline setup

var pipelines = {};
function addPipeline( pipelineName )
{
    let p = device.createComputePipeline({
              layout :   device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]}),
              compute: { module: shaderModule,
                         entryPoint: pipelineName }
    });
    pipelines[ pipelineName ] = p;
  
}

addPipeline( 'diffuse_step' );
addPipeline( 'advect_step' );
addPipeline( 'curl_step' );
addPipeline( 'display' );

async function frame()
{
  // Commands submission
  const commandEncoder = device.createCommandEncoder();

  /*
  {
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup0);
  // workgroup size defined on the wgsl shader
  passEncoder.dispatchWorkgroups( imgWidth/8, imgWidth/8 );
  await passEncoder.end();
  }
  */
  async function computePass( theCommandEncoder, thePipeline, theBindGroup )
  {
    const passEncoder = theCommandEncoder.beginComputePass();
    passEncoder.setPipeline(thePipeline);
    passEncoder.setBindGroup(0, theBindGroup);
    // workgroup size defined on the wgsl shader
    passEncoder.dispatchWorkgroups( imgWidth/8, imgWidth/8 );
    await passEncoder.end();
  }
  
  await computePass( commandEncoder, pipelines['diffuse_step'], bindGroup0 );
  await computePass( commandEncoder, pipelines['diffuse_step'], bindGroup1 );
  
  await computePass( commandEncoder, pipelines['advect_step'], bindGroup0 );
  await computePass( commandEncoder, pipelines['advect_step'], bindGroup1 );
  
  await computePass( commandEncoder, pipelines['curl_step'], bindGroup0 );
  await computePass( commandEncoder, pipelines['curl_step'], bindGroup1 );
  


  await computePass( commandEncoder, pipelines['display'], bindGroup0 );
  
  
  canvasTexture = context.getCurrentTexture();
  
  await
  commandEncoder.copyTextureToTexture( { texture:texture1 },
                                       { texture: canvasTexture },
                                       { width:imgWidth, height:imgHeight, depthOrArrayLayers:1} );
                                        
  // Submit GPU commands.
  const gpuCommands = commandEncoder.finish();
  await device.queue.submit([gpuCommands]);
  
  timestep[0] = timestep[0] + 0.01;
  device.queue.writeBuffer(timerUniformBuffer,   0, timestep             );
  
  requestAnimationFrame(frame);
}


frame();