<?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

// Function to generate random float based on globalId
fn randnum(globalId: vec2<u32>, seed: f32) -> f32 {
    return fract(sin(dot(vec2<f32>(f32(globalId.x*0), f32(globalId.y)), vec2<f32>(12.9898, 78.233) )) * 43758.5453);
}

@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

    // Glitch Noise
    var glitchIntensity = 0.1; // Adjust this value to control the intensity of the glitch noise
    var glitchThreshold = 0.80; // Adjust this value to control the threshold for applying glitch noise

    var seed = f32(globalId.x) + f32(globalId.y)*123*345859; // Use globalId as seed
    var glitchColor = vec4<f32>(
        randnum(globalId.xy, seed) * glitchIntensity,
        randnum(globalId.yx, seed) * glitchIntensity,
        randnum(globalId.xy + vec2<u32>(10u, 20u), seed) * glitchIntensity, // Additional offsets for variety
        0.0
    );
    
    
    var glitchedColor = mix(textureLoad(myTexture0, vec2<i32>(uv * vec2<f32>(imgWidth, imgHeight)), 0),
    
                            textureLoad(myTexture0, vec2<i32>((uv + glitchColor.xy) * vec2<f32>(imgWidth, imgHeight)), 0).xyzw,
                            step(glitchThreshold, randnum(globalId.xy,0.0)  ) );
                            
    
    // Store the result in the output texture
    textureStore(myTexture1, vec2<i32>(globalId.xy), glitchedColor);
}
`;
  

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


timestep[0] = timestep[0] + 0.1;
device.queue.writeBuffer(timerUniformBuffer,   0, timestep             );



// 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]);