<script type="importmap">
  {
    "imports": {
      "three": "https://cdn.jsdelivr.net/npm/three@v0.172.0/build/three.webgpu.js",
      "three/webgpu": "https://cdn.jsdelivr.net/npm/three@v0.172.0/build/three.webgpu.js",
      "three/tsl": "https://cdn.jsdelivr.net/npm/three@v0.172.0/build/three.tsl.js",
      "three/addons/": "https://cdn.jsdelivr.net/npm/three@v0.171.0/examples/jsm/"
    }
  }
</script>

<div id="courses"><a href="https://niklever.com/courses" target="_blank">niklever.com/courses</a></div>

body {
  padding: 0;
  margin: 0;
}

#courses {
  font: bold 30px "Arial";
  position: fixed;
  left: 20px;
  top: 20px;
  color: #ffffff;
  text-decoration: none;
}

a:link {
  color: white;
  text-decoration: none;
}

a:hover {
  color: #dddd33;
  text-decoration: underline;
}

a:visited {
  color: white;
  text-decoration: none;
}

import * as THREE from "three";
import {
  uv,
  uvec2,
  float,
  uint,
  vec2,
  vec3,
  vec4,
  instanceIndex,
  Fn,
  abs,
  oneMinus,
  mix,
  texture,
  textureStore,
  time,
  floor,
  saturate,
  reciprocal,
  uniform,
  If,
  Loop,
  Break,
  positionWorld,
  cameraPosition,
  exp,
  fract,
  select,
  ceil
} from "three/tsl";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import {
  perlinFbm,
  worleyFbm,
  remap,
  pointInAABB
} from "https://assets.codepen.io/2666677/TSL-Common.js";
import { GUI } from "three/addons/libs/lil-gui.module.min.js";

let container;

let camera, scene, renderer, mesh, clock, controls;

init();

function init() {
  renderer = new THREE.WebGPURenderer({ antialias: true });
  renderer.setPixelRatio(window.devicePixelRatio);
  renderer.setSize(window.innerWidth, window.innerHeight);
  renderer.setAnimationLoop(render);
  document.body.appendChild(renderer.domElement);

  scene = new THREE.Scene();
  scene.background = new THREE.Color(0x0561a0);

  camera = new THREE.PerspectiveCamera(
    40,
    window.innerWidth / window.innerHeight,
    0.1,
    100
  );
  camera.position.set(0, 0, 30);

  controls = new OrbitControls(camera, renderer.domElement);
  controls.autoRotate = true;
  
  clock = new THREE.Clock();

  tsl();

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

async function createNoiseTexture(size = 256, cellsX = 16, cellsY = 8) {
  const storageTexture = new THREE.StorageTexture(size * cellsX, size * cellsY);
  storageTexture.minFilter = THREE.NearestFilter;
  storageTexture.maxFilter = THREE.NearestFilter;
  storageTexture.generateMipMaps = false;
  
  const computeTexture = Fn(({ storageTexture }) => {
    const posX = instanceIndex.modInt(size * cellsX).toVar();
    const posY = instanceIndex.div(size * cellsX).toVar();
    const indexUV = uvec2(posX, posY);

    const slices = cellsX * cellsY;
    const row = uint(posY.div(size)).toVar();
    const col = uint(posX.div(size)).toVar();
    const slice = row.mul(cellsX).add(col).toVar(); //uint(cellsX).sub(col) ).toVar();

    const pt = vec3(posX.sub(col.mul(size)), posY.sub(row.mul(size)), 0).div(
      size
    );
    pt.z = float(slice).div(slices);

    const freq = float(4);

    const pfbm = mix(1, perlinFbm(pt, 4, 7), 0.5).toVar();
    pfbm.assign(abs(pfbm.mul(2).sub(1))); // billowy perlin noise

    const g = worleyFbm(pt, freq);
    const b = worleyFbm(pt, freq.mul(2));
    const a = worleyFbm(pt, freq.mul(4));
    const r = remap(pfbm, 0, 1, g, 1); // perlin-worley

    textureStore(storageTexture, indexUV, vec4(r, g, b, a)).toWriteOnly();
  })({ storageTexture }).compute(size * size * cellsY * cellsX);

  await renderer.computeAsync(computeTexture);

  return storageTexture;
}

async function tsl() {
  const geometry = new THREE.BoxGeometry(40, 20, 40);
  const material = new THREE.MeshBasicNodeMaterial({ transparent: true });
  mesh = new THREE.Mesh(geometry, material);
  const bbox = new THREE.Box3();
  mesh.geometry.computeBoundingBox();
  bbox.copy(mesh.geometry.boundingBox).applyMatrix4(mesh.matrixWorld);
  const bbmin = uniform(vec3(bbox.min));
  const bbmax = uniform(vec3(bbox.max));

  const cellsX = 16;
  const cellsY = 8;
  const slices = cellsX * cellsY;

  const storageTexture = await createNoiseTexture(256, cellsX, cellsY);

  const stepSize = uniform(float(1));
  const stepCount = uniform(uint(30));
  const intensity = uniform(float(2.8));
  const pwctrl = uniform( float(0.85) );
  const wfbmctrl = uniform( float(1) );
  const useExp = uniform( uint(0) );

  const getDensityTSL = Fn(({ storageTexture, pt, next }) => {
    const slices = cellsX * cellsY;
    const slice = select(next, ceil(pt.z.mul(slices)), floor(pt.z.mul(slices)));
    If(slice.greaterThanEqual(slices), () => {
      slice.subAssign(slices);
    });

    const col = slice.modInt(cellsX).toVar();
    const row = uint(slice.div(cellsX)).toVar();
    const origin = vec2(float(col).div(cellsX), float(row).div(cellsY));

    const uv1 = vec2(pt.xy).div(vec2(cellsX, cellsY)).toVar();
    uv1.addAssign(origin);

    const texel = texture(storageTexture, uv1, 0).toVar();
    const perlinWorley = texel.x.toVar();
    const worley = texel.yzw.toVar();
    
    // worley fbms with different frequencies
    const wfbm = worley.x
      .mul(0.625)
      .add(worley.y.mul(0.125))
      .add(worley.z.mul(0.25))
      .toVar();

    // cloud shape modeled after the GPU Pro 7 chapter
    const cloud = remap(perlinWorley, wfbm.sub(wfbmctrl), 1, 0, 1).toVar();
    cloud.assign(saturate(remap(cloud, pwctrl, 1, 0, 1))); // fake cloud coverage

    return cloud;
  });

  const samplePositionToUV = Fn(({ pos, bbmin, bbmax }) => {
    const uv = pos.sub(bbmin).div(bbmax.sub(bbmin)).toVar();
    uv.x.subAssign(time.mul(0.02));
    uv.y.subAssign(time.mul(0.02));
    uv.x = fract(uv.x);
    uv.y = fract(uv.y);
    return uv;
  });

  const raymarchTSL = /*@__PURE__*/ Fn(({ storageTexture }) => {
    const rayDirection = positionWorld.sub(cameraPosition).normalize().toVar();
    const samplePosition = positionWorld.toVar();
    const stepVec = rayDirection.mul(stepSize);
    const slices = cellsX * cellsY;

    const density = float().toVar();
    const count = uint(0).toVar();
    const pt = vec3().toVar();
    const s1 = float().toVar();
    const s2 = float().toVar();

    Loop(
      { start: uint(0), end: stepCount, type: "uint", condition: "<" },
      ({ i }) => {
        If(pointInAABB(samplePosition, bbmin, bbmax), () => {
          pt.assign(
            samplePositionToUV({ pos: samplePosition, bbmin, bbmax }) );

          s1.assign(getDensityTSL({ storageTexture, pt, next: false }));
          s2.assign(getDensityTSL({ storageTexture, pt, next: true }));

          density.addAssign(mix(s1, s2, fract(pt.z.mul(slices))));
          count.addAssign(1);
        });

        samplePosition.addAssign(stepVec);
      }
    );

   density.divAssign(count);
    
    If( ( useExp.equal( 1 ) ), () => {
      density.assign(oneMinus(exp(density.negate())));
    });

    return vec4(1, 1, 1, saturate(density.mul(intensity)));
  });

  material.fragmentNode = raymarchTSL({ storageTexture });

  scene.add(mesh);

  const options = {
    useExp: false
  }
  const gui = new GUI();
  gui.add(stepSize, "value", 0, 5).name("stepSize");
  gui.add(stepCount, "value", 1, 128).name("stepCount");
  gui.add(intensity, "value", 0, 3).name("intensity");
  gui.add(pwctrl, "value", 0, 1).name("pwctrl");
  gui.add(wfbmctrl, "value", 0, 1).name("wfbmctrl");
  gui.add(options, "useExp").onChange( value => {
    useExp.value = ( value ) ? 1 : 0;
  });
}

function onWindowResize() {
  camera.aspect = window.innerWidth / window.innerHeight;
  camera.updateProjectionMatrix();

  renderer.setSize(window.innerWidth, window.innerHeight);
}

//

function render() {
  controls.update( clock.getDelta() );
  renderer.render(scene, camera);
}