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,
        energy: f32,
      };
      @group(0) @binding(0) var<storage, read_write> particles: array<Particle>;

      @compute @workgroup_size(${workgroup_size}) fn kick( 
        @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 j = 0u; j<${Nparticles}-1; j++){
          let pj = select(j+1,j,j<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(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;
      }

      @compute @workgroup_size(${workgroup_size}) fn energy( 
        @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
        let v = length(particles[pi].vel);
        particles[pi].energy = ${mass} * v*v;
        for (var j = 0u; j<${Nparticles}-1; j++){
          let pj = select(j+1,j,j<pi);
          let rel_pos = particles[pj].pos-particles[pi].pos;
          let softening = 0.01;
          let d = length(rel_pos)+ 0.01;
          particles[pi].energy += ${mass} * ${mass} / d ;
        }
      }
    `,
  });

  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 energyComputePipeline = device.createComputePipeline({
    label: 'energy compute pipeline',
    layout: computePipelineLayout,
    compute: {
      entryPoint: "energy",
      module: computeModule,
    },
  });

  let rng_state = 20;
  const initialConditions = new Float32Array(8*Nparticles);
  for (var i = 0; i<Nparticles*8; i++){
    rng_state = (1103515245 * rng_state + 12345) % 0x80000000; 
    initialConditions[i] = rng_state / (0x80000000 - 1) - 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.setPipeline(energyComputePipeline);
  // computePass.dispatchWorkgroups(Nworkgroups);
  
  computePass.end();
  const computeCommandBuffer = computeEncoder.finish();
  device.queue.submit([computeCommandBuffer]);


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


  // Copy to CPU
  console.log("transfer started");
  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()

}

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

main();