import TimingHelper from 'https://webgpufundamentals.org/webgpu/resources/js/timing-helper.js';

async function main() {
  const adapter = await navigator.gpu?.requestAdapter();
  const canTimestamp = adapter?.features.has('timestamp-query');
  if (!canTimestamp) {
    fail('need a browser that supports timestamp-query');
    return;
  }
  const device = await adapter?.requestDevice({
    requiredFeatures: [
      ...(canTimestamp ? ['timestamp-query'] : []),
    ],
  });
  const timingHelper = new TimingHelper(device);

  if (!device) {
    fail('need a browser that supports WebGPU');
    return;
  }
  
  let Nparticles = 2**16;
  let dt = 0.04;
  let mass = 1.0/Nparticles;

  let workgroup_size = 64;
  let Nworkgroups = Nparticles/workgroup_size;

  const computeModule = device.createShaderModule({
    label: 'nbody compute module',
    code: `
    	struct Particle {
        pos: vec3f,
        vel: vec3f,
      };
      @group(0) @binding(0) var<storage, read_write> particles: array<Particle>;

      @compute @workgroup_size(${workgroup_size}) fn kick( 
        @builtin(global_invocation_id) id: vec3u,
        @builtin(local_invocation_index) local_invocation_index: u32,
        @builtin(num_workgroups) num_workgroups: vec3<u32>,
        @builtin(workgroup_id) workgroup_id : vec3<u32>,
      ) {
        let pi = workgroup_id.x * ${workgroup_size} + local_invocation_index; // global_invocation_index
        for (var pj = 0u; pj<${Nparticles}; pj++){
          if (pj!=pi){
            let rel_pos = particles[pj].pos-particles[pi].pos;
            let softening = 0.01;
            let d = length(rel_pos)+ 0.01;
            particles[pi].vel += ${mass} * ${dt} * rel_pos/(d*d*d);
          }
        }
      }

      @compute @workgroup_size(${workgroup_size}) fn drift( 
        @builtin(global_invocation_id) id: vec3u,
        @builtin(local_invocation_index) local_invocation_index: u32,
        @builtin(num_workgroups) num_workgroups: vec3<u32>,
        @builtin(workgroup_id) workgroup_id : vec3<u32>,
      ) {
        let pi = workgroup_id.x * ${workgroup_size} + local_invocation_index; // global_invocation_index
        particles[pi].pos += 0.5*${dt} * particles[pi].vel;
      }
    `,
  });

  const computeBindGroupLayout = device.createBindGroupLayout({
    entries: [
      {
        binding: 0,
        visibility: GPUShaderStage.COMPUTE,
        buffer: {
          type: 'storage',
          minBindingSize: 0,
        },
      },
    ],
  });

  const computePipelineLayout = device.createPipelineLayout({
    bindGroupLayouts: [ computeBindGroupLayout ],
  });

  const kickComputePipeline = device.createComputePipeline({
    label: 'kick compute pipeline',
    layout: computePipelineLayout,
    compute: {
      entryPoint: "kick",
      module: computeModule,
    },
  });

  const driftComputePipeline = device.createComputePipeline({
    label: 'drift compute pipeline',
    layout: computePipelineLayout,
    compute: {
      entryPoint: "drift",
      module: computeModule,
    },
  });

  const initialConditions = new Float32Array(8*Nparticles);
  for (var i = 0; i<Nparticles*8; i++){
    initialConditions[i] = Math.random()-0.5;
  }

  const workBuffer = device.createBuffer({
    label: 'work buffer',
    size: initialConditions.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(workBuffer, 0, initialConditions);

  const computeBindGroup = device.createBindGroup({
    label: 'bindGroup for compute',
    layout: computeBindGroupLayout,
    entries: [ { binding: 0, resource: { buffer: workBuffer } } ] 
  });


  const computeEncoder = device.createCommandEncoder();    
  const computePass = timingHelper.beginComputePass(computeEncoder);
  computePass.setBindGroup(0, computeBindGroup);
  computePass.setPipeline(driftComputePipeline);
  computePass.dispatchWorkgroups(Nworkgroups);
  computePass.setPipeline(kickComputePipeline);
  computePass.dispatchWorkgroups(Nworkgroups);
  computePass.setPipeline(driftComputePipeline);
  computePass.dispatchWorkgroups(Nworkgroups);
  computePass.end();
  const computeCommandBuffer = computeEncoder.finish();
  device.queue.submit([computeCommandBuffer]);


  timingHelper.getResult().then(duration => {
    console.log(`duration: ${duration/1e9}s`);
  });

  workBuffer.destroy()
}

function fail(msg) {
  // eslint-disable-next-line no-alert
  alert(msg);
}

main();


// // Copy to CPU
//   const resultBuffer = device.createBuffer({
//     label: 'result buffer',
//     size: initialConditions.byteLength,
//     usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
//   });
// 
//   const encoder = device.createCommandEncoder();
//   encoder.copyBufferToBuffer(workBuffer, 0, resultBuffer, 0, resultBuffer.size);
//   const commandBuffer = encoder.finish();
//   device.queue.submit([commandBuffer]);
// 
//   await resultBuffer.mapAsync(GPUMapMode.READ);
//   const result = new Float32Array(resultBuffer.getMappedRange().slice());
//   console.log('energy: ', result[7]);
//   resultBuffer.unmap();
// 
//   workBuffer.destroy()
//   resultBuffer.destroy()