<script async src="https://unpkg.com/es-module-shims@1.6.3/dist/es-module-shims.js"></script>

<script type="importmap">

  {

    "imports": {

      "three": "https://unpkg.com/three@0.158/build/three.module.js", 

      "three/addons/": "https://unpkg.com/three@0.157/examples/jsm/",

      "three/nodes": "https://unpkg.com/three@0.157/examples/jsm/nodes/Nodes.js"

    }

  }

</script>

body{

  overflow: hidden;

  margin: 0;

}

/*The example here has to be run locally because CodePen apparently has difficulties with texture_storage_2D. I can't really say why that is because I don't know. It works locally for me.
*/

import * as THREE from "three";
import {texture, MeshBasicNodeMaterial, uniform, textureStore, instanceIndex, wgslFn, /*StorageTexture*/ } from 'three/nodes';
import {OrbitControls} from "three/addons/controls/OrbitControls.js";
import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';


let computeTest, size, stepUniform, material, testTexture;


let scene = new THREE.Scene();
scene.background = new THREE.Color(0x00001f);
let renderer = new WebGPURenderer();
renderer.setPixelRatio(window.devicePixelRatio);
renderer.setSize(window.innerWidth,window.innerHeight);
document.body.appendChild(renderer.domElement);
let camera = new THREE.PerspectiveCamera(50.0, window.innerWidth/window.innerHeight, 0.1, 10000);
camera.position.set(3, 3, 3);
let controls = new OrbitControls(camera, renderer.domElement);

window.addEventListener("resize", onWindowResize, false);

init();
render();


function init() {
  
  const computeShader = wgslFn(`
    fn mainWGSL(
      writeTex: texture_storage_2d<rgba8unorm, write>, 
      index: u32,
      size: f32,
      step: f32,
    ) -> void {
    
      var posX = f32(index) % size;
      var posY = f32(index) / size;
      var idx = vec2i(i32(posX), i32(posY));
         
      //loop controlColors
      if(step == 0){color = vec4<f32>(0,0.2,0.1,0)};
      if(step == 1){color = vec4<f32>(0,0,1,0)};
      if(step == 2){color = vec4<f32>(0,1,0,0)};
      if(step == 3){color = vec4<f32>(0,1,1,0)};
      if(step == 4){color = vec4<f32>(1,0,0,0)};
      if(step == 5){color = vec4<f32>(1,0,1,0)};
      if(step == 6){color = vec4<f32>(1,1,0,0)};
      if(step == 7){color = vec4<f32>(1,1,1,0)};
     
      textureStore(writeTex, idx, color);
    } 
  `);
  
  size = 256;
  
  //testTexture = new StorageTexture(size, size); //not work in codePen
  testTexture = new THREE.Texture();
  testTexture.image = {width: size, height: size};
 
  stepUniform = uniform(0); //initiaize with 0
  
  const computeTest = computeShader({
    size: size,
    step: stepUniform,
    index: instanceIndex,
    writeTex: textureStore(testTexture),
  }).compute(size ** 2);
  
  
	material = new MeshBasicNodeMaterial();
	material.colorNode = texture(testTexture);

	const geometry = new THREE.BoxGeometry();
	let cube = new THREE.Mesh(geometry, material);
	scene.add(cube);
}

function updateCompute(){

  const iterations = 6; //In this case I expect a purple cube because 5 is the last iteration step. But I get the default gray.
  
  /*The analog mechanism in webgl2 delivers correct iteration values in the shader and always the color value of the last iteration step. In webgl2 I save textures with renderTargets but the mechanism itself is analog. If I replace the "i" on the uniform node with a number I see the corresponding color. If I start the for loop with a number higher than 0, I see the color that corresponds to the number with which the for loop starts. The current iteration steps do not arrive in the shader.*/
  
  //iterations = 1 -> dark green
  //iterations = 2 -> blue
  //iterations = 3 -> green
  //iterations = 4 -> cyan
  //iterations = 5 -> red
  //iterations = 6 -> purple
  //iterations = 7 -> yellow
  //iterations = 8 -> white
  
  for(let i = 0; i < iterations; i++){
   
    stepUniform.value = i;
    renderer.compute(computeTest);
  }
  
}  

function render() {
  requestAnimationFrame(render);
  
  updateCompute();
  
  renderer.render(scene, camera);
}

function onWindowResize() {
  camera.aspect = window.innerWidth / window.innerHeight;
  camera.updateProjectionMatrix();
  renderer.setSize(window.innerWidth, window.innerHeight);
}