<?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');

async function loadTexture(fileName = "https://webgpulab.xbdev.net/var/images/test512.png", 
                           width=512, height=512) {
  console.log('loading image:', fileName);
  // Load image 
  const img = document.createElement("img");
  img.src = fileName;

  await img.decode();

  const originalWidth = img.width;
  const originalHeight = img.height;

  const imageCanvas = document.createElement('canvas');
  imageCanvas.width = width;
  imageCanvas.height = height;
  const imageCanvasContext = imageCanvas.getContext('2d');

  // Draw the image onto the canvas, resizing it to the specified width and height
  imageCanvasContext.drawImage(img, 0, 0, width, height);

  const imageData = imageCanvasContext.getImageData(0, 0, width, height);
  const textureData = imageData.data;
  console.log('textureData.byteLength:', textureData.byteLength);

  const basicTexture = device.createTexture({
    size: [width, height, 1],
    format: "rgba8unorm",
    usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.TEXTURE_BINDING
  });

  await device.queue.writeTexture(
    { texture: basicTexture },
    textureData,
    { bytesPerRow: width * 4 },
    [width, height, 1]
  );

  return { w: width, h: height, t: basicTexture };
}



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();

const imgWidth = 512;
const imgHeight = imgWidth;

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

const texture0         = await loadTexture( 'https://webgpulab.xbdev.net/var/images/test512.png', imgWidth );
//const texture1         = await loadTexture( 'https://webgpulab.xbdev.net/var/images/avatar.png', imgWidth);

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

// Basic canvas which will be used to display the output from the compute shader

let canvasa = document.createElement('canvas');
document.body.appendChild( canvasa ); canvasa.height = canvasa.width = imgWidth;
const context = canvasa.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();

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

// Output texture - output from the compute shader written to this texture
// Copy this texutre to the 'canvas' - needs to be the same size as the output
// canvas size
const texture1 = device.createTexture({
  size: [imgWidth, imgHeight, 1],
  format: "rgba8unorm",
  usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.COPY_SRC | GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.STORAGE_BINDING
});

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


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

const timestep  = new Float32Array( [0.0] );

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


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

const GCOMPUTE = GPUShaderStage.COMPUTE;

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

const bindGroup = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [  {   binding: 0,  resource: texture0.t.createView()         },
                  {   binding: 1,  resource: { buffer: timerUniformBuffer  } },
                {   binding: 2,  resource: texture1.createView()           }
    ]
});


// Compute shader code
const computeShader = ` 
@group(0) @binding(0) var myTexture0:  texture_2d<f32>; // input texture image
@group(0) @binding(1) var<uniform>     mytimer   : f32; // timer increments each frame
@group(0) @binding(2) var myTexture1:  texture_storage_2d<rgba8unorm, write>; // output image


@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) globalId      : vec3<u32>,
        @builtin(local_invocation_id)  localId       : vec3<u32>,
        @builtin(workgroup_id)         workgroupId   : vec3<u32>,
        @builtin(num_workgroups)       workgroupSize : vec3<u32>
        ) 
{
    var imgWidth  = f32( ${imgWidth}  ); 
    var imgHeight = f32( ${imgHeight} ); 

    var coords = vec2<f32>( f32(globalId.x), f32(globalId.y) );
    
    var uv = coords / vec2<f32>(imgWidth, imgHeight); // normalize coordinates to 0.0 - 1.0 range

    // Calculate displacement for water drop effect
    var dropCenter = vec2<f32>(0.5, 0.5); // center of the effect, can be modified
    var dropFrequency = 20.0; // frequency of the ripples
    var dropAmplitude = 0.01; // amplitude of the ripples
    var distance = distance(uv, dropCenter);

    var displacement = sin(distance * dropFrequency - mytimer * 5.0) * dropAmplitude;

    // Apply displacement
    var displacedUv = uv + normalize(uv - dropCenter) * displacement;

    // Clamp UV coordinates to the edges of the texture
    displacedUv = clamp(displacedUv, vec2<f32>(0.0), vec2<f32>(1.0));

    // Sample color from the texture with displaced coordinates
    var texCol0 = textureLoad(myTexture0, vec2<i32>(displacedUv * vec2<f32>(imgWidth, imgHeight)), 0);
    
    // Store the result in the output texture
    textureStore(myTexture1, vec2<i32>(globalId.xy), texCol0);
    
/*
    var coords = vec2<f32>( f32(globalId.x), f32(globalId.y) ) * 3.0;
    
    var uv = vec2<f32>( f32(globalId.x), f32(globalId.y) ); // uvs * 2.0 - 1.0; 
    uv = uv / ${imgWidth}; // 0.0 - 1.0
   
    // Sample color from the texture
    var texCol0 = textureLoad( myTexture0, vec2<i32>( i32(uv.x*${imgWidth}), i32(uv.y*${imgHeight}) ), 0 );
    
    // .....
    
    
    textureStore(myTexture1, vec2<i32>( i32(globalId.x) , i32(globalId.y)  ), texCol0 );
*/
}
`;
  

// Pipeline setup
const computePipeline = device.createComputePipeline({
    layout :   device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]}),
    compute: { module    : device.createShaderModule({code:computeShader}),
               entryPoint: "main" }
});


async function frame()
{
  
  // Commands submission
  const commandEncoder = device.createCommandEncoder();
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup);
  passEncoder.dispatchWorkgroups( imgWidth/8, imgWidth/8 );
  await passEncoder.end();


  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();