Ambient Occlusion

Video Lecture

Section	Video Links
Ambient Occlusion

Description

We will implement Screen Space Ambient Occlusion (AO).

This will make areas darker based on the surrounding objects.

We can use AO independently of shadows.

Instead of casting a shadow ray, AO checks how much nearby geometry is blocking light.

It will sample small distances around the object, and if any are blocked, then the position is darkened.

Areas near the ground, near other objects, or inside corners will appear darker due to AO.

Surfaces will feel more realistic as occluded regions will receive less light.

Working Example

<>

Start Script

./src/main.ts

import './style.css'
import * as THREE from 'three/webgpu'
import {
  positionLocal,
  Fn,
  If,
  Loop,
  Break,
  float,
  vec2,
  vec3,
  min,
  max,
  length,
  normalize,
  dot,
  time,
  sin,
  cos,
  abs,
  negate,
  uniform,
  pow,
  reflect,
  clamp,
  mod,
  exp,
  mix,
  smoothstep,
} from 'three/tsl'
import { OrbitControls } from 'three/addons/controls/OrbitControls.js'
import { GUI } from 'three/addons/libs/lil-gui.module.min.js'

const scene = new THREE.Scene()

const camera = new THREE.PerspectiveCamera(
  53,
  window.innerWidth / window.innerHeight,
  0.1,
  1000
)
camera.position.set(5, 1.5, 4)

const renderer = new THREE.WebGPURenderer()
renderer.setSize(window.innerWidth, window.innerHeight)
document.body.appendChild(renderer.domElement)
renderer.setAnimationLoop(animate)
renderer.setPixelRatio(0.25) // start low for slower cards

window.addEventListener('resize', function () {
  camera.aspect = window.innerWidth / window.innerHeight
  camera.updateProjectionMatrix()
  renderer.setSize(window.innerWidth, window.innerHeight)
})

const controls = new OrbitControls(camera, renderer.domElement)
controls.enableDamping = true
controls.target.y = 1
controls.autoRotateSpeed = -0.5
//controls.autoRotate = true
controls.maxPolarAngle = Math.PI / 1.9

const options = {
  zoom: 0,
  maxSteps: 256,
  surfaceDistance: 0.0001,
  cameraNear: 0.0001,
  cameraFar: 128.0,
  shadowSoftness: 10,
  shadowIntensity: 0.1,
  shadowMaxSteps: 256,
  shadowFar: 128,
  maxReflections: 1,
  reflectivity: 0.5,
  // samples: 5,
  // spread: 0.5,
  objectHeight: 1,
}

const camPos = uniform(new THREE.Vector3())
const camTarget = uniform(new THREE.Vector3())
const zoom = uniform(options.zoom)
const maxSteps = uniform(options.maxSteps)
const surfaceDistance = uniform(options.surfaceDistance)
const cameraNear = uniform(options.cameraNear)
const cameraFar = uniform(options.cameraFar)
const shadowSoftness = uniform(options.shadowSoftness)
const shadowIntensity = uniform(options.shadowIntensity)
const shadowMaxSteps = uniform(options.shadowMaxSteps)
const shadowFar = uniform(options.shadowFar)
const maxReflections = uniform(options.maxReflections)
const reflectivity = uniform(options.reflectivity)
// const samples = uniform(options.samples)
// const spread = uniform(options.spread)
const objectHeight = uniform(options.objectHeight)

// @ts-ignore
const Floor = Fn(([position]) => {
  const floor = position.y

  const radius = 10.0
  const circle = length(position.xz).sub(radius)

  return max(floor, circle)
})

// @ts-ignore
const Sphere = Fn(([position, radius]) => {
  return length(position).sub(radius)
})

// @ts-ignore
const Box = Fn(([position, dimensions]) => {
  const distance = abs(position).sub(dimensions)
  return length(max(distance, 0.0)).add(
    min(max(distance.x, max(distance.y, distance.z)), 0.0)
  )
})

// @ts-ignore
const IntersectedSphereBox = Fn(([position, radius]) => {
  const circle = length(position).sub(radius.add(0.33))
  const box = Box(position, vec3(radius))
  return max(circle, box)
})

// @ts-ignore
const SubtractedSphereBox = Fn(([position, radius]) => {
  const circle = length(position).sub(radius.add(0.33))
  const box = Box(position, vec3(radius))
  return max(negate(circle), box)
})

// @ts-ignore
const sdfScene = Fn(([position]) => {
  const floor = Floor(position)

  const sphere = Sphere(position.sub(vec3(1.5, objectHeight, 1.5)), 1)

  const box = Box(position.sub(vec3(1.5, objectHeight, -1.5)), vec3(1))

  const intersectedSphereBox = IntersectedSphereBox(
    position.sub(vec3(-1.5, objectHeight, -1.5)),
    1
  )
  const subtractedSphereBox = SubtractedSphereBox(
    position.sub(vec3(-1.5, objectHeight, 1.5)),
    1
  )

  const distance = vec2(floor, 0).toVar()
  If(distance.x.greaterThan(sphere), () => {
    distance.assign(vec2(sphere, 1))
  })
  If(distance.x.greaterThan(box), () => {
    distance.assign(vec2(box, 2))
  })
  If(distance.x.greaterThan(intersectedSphereBox), () => {
    distance.assign(vec2(intersectedSphereBox, 3))
  })
  If(distance.x.greaterThan(subtractedSphereBox), () => {
    distance.assign(vec2(subtractedSphereBox, 4))
  })

  return distance
})

// @ts-ignore
const getNormal = Fn(([position, distance]) => {
  const offset = vec2(0.0001, 0)

  return normalize(
    distance.sub(
      vec3(
        sdfScene(position.sub(offset.xyy)).x,
        sdfScene(position.sub(offset.yxy)).x,
        sdfScene(position.sub(offset.yyx)).x
      )
    )
  )
})

// @ts-ignore
const shadowMarcher = Fn(([rayOrigin, rayDirection]) => {
  const shadow = float(1).toVar()
  const accumulatedDistance = float(cameraNear).toVar()

  // @ts-ignore
  Loop({ start: 0, end: shadowMaxSteps }, () => {
    const distance = sdfScene(
      rayOrigin.add(rayDirection.mul(accumulatedDistance))
    )

    If(abs(distance.x).lessThan(surfaceDistance), () => {
      shadow.assign(0)
      Break()
    }).Else(() => {
      shadow.assign(
        min(shadow, shadowSoftness.mul(distance.x).div(accumulatedDistance))
      ) //softens shadow
      accumulatedDistance.addAssign(clamp(distance.x, surfaceDistance, 0.025)) // works better for soft shadows
      //accumulatedDistance.addAssign(distance) // works better for hard shadows

      If(accumulatedDistance.greaterThan(shadowFar), () => {
        Break()
      })
    })
  })

  return shadow
})

// @ts-ignore
const atmosphericScattering = Fn(([position, direction]) => {
  const topColour = vec3(0.1, 0.2, 0.5) // Deep blue (upper sky)
  const midColour = vec3(1.0, 0.4, 0.2) // Orange-red (mid sky)
  const bottomColour = vec3(0.9, 0.6, 0.3) // Yellow near horizon

  const t = clamp(position.y.mul(0.5).add(0.5), 0.0, 1.0) // Horizon-based gradient
  const skyColour = mix(mix(bottomColour, midColour, t), vec3(0), t.mul(1.25))

  const radius = 0.01
  const distance = length(position.sub(normalize(direction)))
  const sun = exp(distance.negate().mul(distance.div(radius)))
  skyColour.addAssign(vec3(1.0, 0.8, 0.5).mul(sun))

  // Mie scattering (orange glow near the sun)
  const mie = exp(distance.mul(3.0).pow(2.0).negate()).mul(0.5)
  const mieColor = vec3(midColour).mul(mie)
  skyColour.addAssign(mieColor.mul(1.5))

  // Rayleigh scattering (blue tint in the upper sky)
  const rayleigh = exp(position.y.mul(2.5)).mul(0.3)
  const rayleighColor = rayleigh.mul(topColour)
  skyColour.addAssign(rayleighColor)

  // Night Transition
  const nightFactor = smoothstep(-0.1, 0.1, direction.y)
  skyColour.mulAssign(nightFactor)

  return skyColour
})

// // @ts-ignore
// const computeAO = Fn(([position, normal]) => {
//   const occlusion = float(0.0).toVar()

//   Loop({ start: 0, end: samples, condition: '<=' }, ({ i }) => {
//     const spacer = float(i).div(samples)
//     const samplePos = position.add(normal.mul(spacer.mul(spread)))
//     const distance = sdfScene(samplePos)
//     occlusion.addAssign(smoothstep(0.0, 0.1, distance))
//   })
//   return clamp(occlusion.div(samples), 0, 1.0)
//   //return clamp(occlusion.div(samples).oneMinus(), 0, 1.0)
// })

// @ts-ignore
const main = Fn(() => {
  const p = positionLocal

  const rayOrigin = camPos.toVar()
  const lookAt = camTarget

  const forward = normalize(lookAt.sub(rayOrigin))
  const rayDirection = normalize(p.add(forward.mul(zoom))).toVar()

  //const t = time.div(5)
  const lightPosition = vec3(-25, 7.5, -25)
  //const lightPosition = vec3(-25, sin(t).mul(10).add(7), -25)
  //const lightPosition = vec3(sin(t).mul(25), 2, cos(t).mul(25))
  //const lightPosition = vec3(sin(t).mul(25), sin(t).mul(10).add(8), cos(t).mul(25))

  const skyColour = atmosphericScattering(p, normalize(lightPosition))
  const finalColour = skyColour.toVar()
  const reflectivityFactor = float(1.0).toVar()

  Loop({ start: 0, end: maxReflections, condition: '<=' }, () => {
    const accumulatedDistance = float(cameraNear).toVar()
    const distance = vec2(0).toVar()
    const position = vec3(0).toVar()

    Loop({ start: 0, end: maxSteps, condition: '<' }, () => {
      position.assign(rayOrigin.add(rayDirection.mul(accumulatedDistance)))
      distance.assign(sdfScene(position))

      If(
        abs(distance.x)
          .lessThan(surfaceDistance)
          .or(accumulatedDistance.greaterThan(cameraFar)),
        () => {
          Break()
        }
      )

      accumulatedDistance.addAssign(distance.x)
    })

    If(accumulatedDistance.lessThan(cameraFar), () => {
      const lightDirection = normalize(lightPosition.sub(position))

      const normal = getNormal(position, distance.x)

      const diffuse = clamp(dot(normal, lightDirection), 0.75, 1).toVar()

      const diffuseColour = atmosphericScattering(
        reflect(rayDirection, normal),
        normalize(lightPosition)
      ).toVar()

      const shadow = shadowMarcher(
        position.add(normal.mul(surfaceDistance)),
        vec3(lightDirection)
      )
      shadow.assign(max(shadowIntensity, shadow))
      diffuse.mulAssign(shadow)

      // const ao = computeAO(position, normal)
      // diffuse.mulAssign(ao)

      finalColour.assign(
        mix(finalColour, diffuseColour.mul(diffuse), reflectivityFactor)
      )

      rayOrigin.assign(position.add(normal.mul(surfaceDistance)))
      rayDirection.assign(reflect(rayDirection, normal))
      reflectivityFactor.mulAssign(reflectivity)
    })
  })

  return finalColour
})

scene.backgroundNode = main()

const gui = new GUI()

const raymarchFolder = gui.addFolder('Raymarching')
raymarchFolder
  .add(options, 'zoom', 0, 10, 0.001)
  .name('Zoom')
  .onChange((v) => {
    zoom.value = v
  })
raymarchFolder
  .add(options, 'maxSteps', 1, 512, 1)
  .name('Raymarch Max Steps')
  .onChange((v) => {
    maxSteps.value = v
  })
raymarchFolder
  .add(options, 'surfaceDistance', 0, 0.001, 0.0001)
  .name('Surface Distance')
  .onChange((v) => {
    surfaceDistance.value = v
  })
raymarchFolder
  .add(options, 'cameraNear', 0.0001, 10, 0.1)
  .name('Camera Near')
  .onChange((v) => {
    cameraNear.value = v
  })
raymarchFolder
  .add(options, 'cameraFar', 1, 512, 0.1)
  .name('Camera Far')
  .onChange((v) => {
    cameraFar.value = v
  })
raymarchFolder.close()

// const lightingFolder = gui.addFolder('Lighting')
// lightingFolder
//     .add(options, 'shininess', 0, 100, 0.1)
//     .name('Shininess')
//     .onChange((v) => {
//         shininess.value = v
//     })

// lightingFolder
//     .add(options, 'specularStrength', 0, 1, 0.01)
//     .name('SpecularStrength')
//     .onChange((v) => {
//         specularStrength.value = v
//     })
// lightingFolder.close()

const shadowFolder = gui.addFolder('Shadows')
shadowFolder
  .add(options, 'shadowSoftness', 0, 100, 0.01)
  .name('Shadow Softness')
  .onChange((v) => {
    shadowSoftness.value = v
  })
shadowFolder
  .add(options, 'shadowIntensity', 0, 1, 0.01)
  .name('Shadow Intensity')
  .onChange((v) => {
    shadowIntensity.value = v
  })
shadowFolder
  .add(options, 'shadowMaxSteps', 1, 512, 1)
  .name('Shadow Max Steps')
  .onChange((v) => {
    shadowMaxSteps.value = v
  })
shadowFolder
  .add(options, 'shadowFar', 1, 512, 0.1)
  .name('Shadow Far')
  .onChange((v) => {
    shadowFar.value = v
  })
shadowFolder.close()

const reflectionFolder = gui.addFolder('Reflections')
reflectionFolder
  .add(options, 'maxReflections', 0, 3, 1)
  .name('Max Reflections')
  .onChange((v) => {
    maxReflections.value = v
  })
reflectionFolder
  .add(options, 'reflectivity', 0, 1, 0.1)
  .name('Reflectivity')
  .onChange((v) => {
    reflectivity.value = v
  })
reflectionFolder.close()

// const aoFolder = gui.addFolder('Ambient Occlusion')
// aoFolder
//   .add(options, 'samples', 1, 10, 1)
//   .name('Samples')
//   .onChange((v) => {
//     samples.value = v
//   })
// aoFolder
//   .add(options, 'spread', 0.1, 4, 0.01)
//   .name('spread')
//   .onChange((v) => {
//     spread.value = v
//   })
// aoFolder
//   .add(options, 'objectHeight', 0, 2, 0.01)
//   .name('Object Height')
//   .onChange((v) => {
//     objectHeight.value = v
//   })

//Adaptive DPR
const clock = new THREE.Clock()
let targetFPS = 45 // see notes
let frameCount = 0
let elapsed = 0

function adjustDPR(renderer: THREE.WebGPURenderer) {
  elapsed += clock.getDelta()
  frameCount++

  if (elapsed >= 0.094) {
    let fps = frameCount * elapsed * 100
    frameCount = 0
    elapsed -= 0.1

    if (fps < targetFPS * 0.95) {
      renderer.setPixelRatio(
        Math.min(1, Math.max(0.1, renderer.getPixelRatio() * 0.9))
      )
    } else if (fps > targetFPS) {
      renderer.setPixelRatio(Math.min(1, renderer.getPixelRatio() * 1.1))
    }

    //console.log(`DPR: ${renderer.getPixelRatio()} (FPS: ${fps})`)
  }
}

function animate() {
  adjustDPR(renderer)

  controls.update()

  camPos.value.copy(camera.position)
  camTarget.value.copy(controls.target)

  renderer.render(scene, camera)
}

Final Script

./src/main.ts

import './style.css'
import * as THREE from 'three/webgpu'
import {
  positionLocal,
  Fn,
  If,
  Loop,
  Break,
  float,
  vec2,
  vec3,
  min,
  max,
  length,
  normalize,
  dot,
  time,
  sin,
  cos,
  abs,
  negate,
  uniform,
  pow,
  reflect,
  clamp,
  mod,
  exp,
  mix,
  smoothstep,
} from 'three/tsl'
import { OrbitControls } from 'three/addons/controls/OrbitControls.js'
import { GUI } from 'three/addons/libs/lil-gui.module.min.js'

const scene = new THREE.Scene()

const camera = new THREE.PerspectiveCamera(
  53,
  window.innerWidth / window.innerHeight,
  0.1,
  1000
)
camera.position.set(5, 1.5, 4)

const renderer = new THREE.WebGPURenderer()
renderer.setSize(window.innerWidth, window.innerHeight)
document.body.appendChild(renderer.domElement)
renderer.setAnimationLoop(animate)
renderer.setPixelRatio(0.25) // start low for slower cards

window.addEventListener('resize', function () {
  camera.aspect = window.innerWidth / window.innerHeight
  camera.updateProjectionMatrix()
  renderer.setSize(window.innerWidth, window.innerHeight)
})

const controls = new OrbitControls(camera, renderer.domElement)
controls.enableDamping = true
controls.target.y = 1
controls.autoRotateSpeed = -0.5
//controls.autoRotate = true
controls.maxPolarAngle = Math.PI / 1.9

const options = {
  zoom: 0,
  maxSteps: 256,
  surfaceDistance: 0.0001,
  cameraNear: 0.0001,
  cameraFar: 128.0,
  shadowSoftness: 10,
  shadowIntensity: 0.1,
  shadowMaxSteps: 256,
  shadowFar: 128,
  maxReflections: 1,
  reflectivity: 0.5,
  samples: 5,
  spread: 0.5,
  objectHeight: 1,
}

const camPos = uniform(new THREE.Vector3())
const camTarget = uniform(new THREE.Vector3())
const zoom = uniform(options.zoom)
const maxSteps = uniform(options.maxSteps)
const surfaceDistance = uniform(options.surfaceDistance)
const cameraNear = uniform(options.cameraNear)
const cameraFar = uniform(options.cameraFar)
const shadowSoftness = uniform(options.shadowSoftness)
const shadowIntensity = uniform(options.shadowIntensity)
const shadowMaxSteps = uniform(options.shadowMaxSteps)
const shadowFar = uniform(options.shadowFar)
const maxReflections = uniform(options.maxReflections)
const reflectivity = uniform(options.reflectivity)
const samples = uniform(options.samples)
const spread = uniform(options.spread)
const objectHeight = uniform(options.objectHeight)

// @ts-ignore
const Floor = Fn(([position]) => {
  const floor = position.y

  const radius = 10.0
  const circle = length(position.xz).sub(radius)

  return max(floor, circle)
})

// @ts-ignore
const Sphere = Fn(([position, radius]) => {
  return length(position).sub(radius)
})

// @ts-ignore
const Box = Fn(([position, dimensions]) => {
  const distance = abs(position).sub(dimensions)
  return length(max(distance, 0.0)).add(
    min(max(distance.x, max(distance.y, distance.z)), 0.0)
  )
})

// @ts-ignore
const IntersectedSphereBox = Fn(([position, radius]) => {
  const circle = length(position).sub(radius.add(0.33))
  const box = Box(position, vec3(radius))
  return max(circle, box)
})

// @ts-ignore
const SubtractedSphereBox = Fn(([position, radius]) => {
  const circle = length(position).sub(radius.add(0.33))
  const box = Box(position, vec3(radius))
  return max(negate(circle), box)
})

// @ts-ignore
const sdfScene = Fn(([position]) => {
  const floor = Floor(position)

  const sphere = Sphere(position.sub(vec3(1.5, objectHeight, 1.5)), 1)

  const box = Box(position.sub(vec3(1.5, objectHeight, -1.5)), vec3(1))

  const intersectedSphereBox = IntersectedSphereBox(
    position.sub(vec3(-1.5, objectHeight, -1.5)),
    1
  )
  const subtractedSphereBox = SubtractedSphereBox(
    position.sub(vec3(-1.5, objectHeight, 1.5)),
    1
  )

  const distance = vec2(floor, 0).toVar()
  If(distance.x.greaterThan(sphere), () => {
    distance.assign(vec2(sphere, 1))
  })
  If(distance.x.greaterThan(box), () => {
    distance.assign(vec2(box, 2))
  })
  If(distance.x.greaterThan(intersectedSphereBox), () => {
    distance.assign(vec2(intersectedSphereBox, 3))
  })
  If(distance.x.greaterThan(subtractedSphereBox), () => {
    distance.assign(vec2(subtractedSphereBox, 4))
  })

  return distance
})

// @ts-ignore
const getNormal = Fn(([position, distance]) => {
  const offset = vec2(0.0001, 0)

  return normalize(
    distance.sub(
      vec3(
        sdfScene(position.sub(offset.xyy)).x,
        sdfScene(position.sub(offset.yxy)).x,
        sdfScene(position.sub(offset.yyx)).x
      )
    )
  )
})

// @ts-ignore
const shadowMarcher = Fn(([rayOrigin, rayDirection]) => {
  const shadow = float(1).toVar()
  const accumulatedDistance = float(cameraNear).toVar()

  // @ts-ignore
  Loop({ start: 0, end: shadowMaxSteps }, () => {
    const distance = sdfScene(
      rayOrigin.add(rayDirection.mul(accumulatedDistance))
    )

    If(abs(distance.x).lessThan(surfaceDistance), () => {
      shadow.assign(0)
      Break()
    }).Else(() => {
      shadow.assign(
        min(shadow, shadowSoftness.mul(distance.x).div(accumulatedDistance))
      ) //softens shadow
      accumulatedDistance.addAssign(clamp(distance.x, surfaceDistance, 0.025)) // works better for soft shadows
      //accumulatedDistance.addAssign(distance) // works better for hard shadows

      If(accumulatedDistance.greaterThan(shadowFar), () => {
        Break()
      })
    })
  })

  return shadow
})

// @ts-ignore
const atmosphericScattering = Fn(([position, direction]) => {
  const topColour = vec3(0.1, 0.2, 0.5) // Deep blue (upper sky)
  const midColour = vec3(1.0, 0.4, 0.2) // Orange-red (mid sky)
  const bottomColour = vec3(0.9, 0.6, 0.3) // Yellow near horizon

  const t = clamp(position.y.mul(0.5).add(0.5), 0.0, 1.0) // Horizon-based gradient
  const skyColour = mix(mix(bottomColour, midColour, t), vec3(0), t.mul(1.25))

  const radius = 0.01
  const distance = length(position.sub(normalize(direction)))
  const sun = exp(distance.negate().mul(distance.div(radius)))
  skyColour.addAssign(vec3(1.0, 0.8, 0.5).mul(sun))

  // Mie scattering (orange glow near the sun)
  const mie = exp(distance.mul(3.0).pow(2.0).negate()).mul(0.5)
  const mieColor = vec3(midColour).mul(mie)
  skyColour.addAssign(mieColor.mul(1.5))

  // Rayleigh scattering (blue tint in the upper sky)
  const rayleigh = exp(position.y.mul(2.5)).mul(0.3)
  const rayleighColor = rayleigh.mul(topColour)
  skyColour.addAssign(rayleighColor)

  // Night Transition
  const nightFactor = smoothstep(-0.1, 0.1, direction.y)
  skyColour.mulAssign(nightFactor)

  return skyColour
})

// @ts-ignore
const computeAO = Fn(([position, normal]) => {
  const occlusion = float(0.0).toVar()

  Loop({ start: 0, end: samples, condition: '<=' }, ({ i }) => {
    const spacer = float(i).div(samples)
    const samplePos = position.add(normal.mul(spacer.mul(spread)))
    const distance = sdfScene(samplePos)
    occlusion.addAssign(smoothstep(0.0, 0.1, distance))
  })
  return clamp(occlusion.div(samples), 0, 1.0)
  //return clamp(occlusion.div(samples).oneMinus(), 0, 1.0)
})

// @ts-ignore
const main = Fn(() => {
  const p = positionLocal

  const rayOrigin = camPos.toVar()
  const lookAt = camTarget

  const forward = normalize(lookAt.sub(rayOrigin))
  const rayDirection = normalize(p.add(forward.mul(zoom))).toVar()

  //const t = time.div(5)
  const lightPosition = vec3(-25, 7.5, -25)
  //const lightPosition = vec3(-25, sin(t).mul(10).add(7), -25)
  //const lightPosition = vec3(sin(t).mul(25), 2, cos(t).mul(25))
  //const lightPosition = vec3(sin(t).mul(25), sin(t).mul(10).add(8), cos(t).mul(25))

  const skyColour = atmosphericScattering(p, normalize(lightPosition))
  const finalColour = skyColour.toVar()
  const reflectivityFactor = float(1.0).toVar()

  Loop({ start: 0, end: maxReflections, condition: '<=' }, () => {
    const accumulatedDistance = float(cameraNear).toVar()
    const distance = vec2(0).toVar()
    const position = vec3(0).toVar()

    Loop({ start: 0, end: maxSteps, condition: '<' }, () => {
      position.assign(rayOrigin.add(rayDirection.mul(accumulatedDistance)))
      distance.assign(sdfScene(position))

      If(
        abs(distance.x)
          .lessThan(surfaceDistance)
          .or(accumulatedDistance.greaterThan(cameraFar)),
        () => {
          Break()
        }
      )

      accumulatedDistance.addAssign(distance.x)
    })

    If(accumulatedDistance.lessThan(cameraFar), () => {
      const lightDirection = normalize(lightPosition.sub(position))

      const normal = getNormal(position, distance.x)

      const diffuse = clamp(dot(normal, lightDirection), 0.75, 1).toVar()

      const diffuseColour = atmosphericScattering(
        reflect(rayDirection, normal),
        normalize(lightPosition)
      ).toVar()

      // const shadow = shadowMarcher(
      //   position.add(normal.mul(surfaceDistance)),
      //   vec3(lightDirection)
      // )
      // shadow.assign(max(shadowIntensity, shadow))
      // diffuse.mulAssign(shadow)

      const ao = computeAO(position, normal)
      diffuse.mulAssign(ao)

      finalColour.assign(
        mix(finalColour, diffuseColour.mul(diffuse), reflectivityFactor)
      )

      rayOrigin.assign(position.add(normal.mul(surfaceDistance)))
      rayDirection.assign(reflect(rayDirection, normal))
      reflectivityFactor.mulAssign(reflectivity)
    })
  })

  return finalColour
})

scene.backgroundNode = main()

const gui = new GUI()

const raymarchFolder = gui.addFolder('Raymarching')
raymarchFolder
  .add(options, 'zoom', 0, 10, 0.001)
  .name('Zoom')
  .onChange((v) => {
    zoom.value = v
  })
raymarchFolder
  .add(options, 'maxSteps', 1, 512, 1)
  .name('Raymarch Max Steps')
  .onChange((v) => {
    maxSteps.value = v
  })
raymarchFolder
  .add(options, 'surfaceDistance', 0, 0.001, 0.0001)
  .name('Surface Distance')
  .onChange((v) => {
    surfaceDistance.value = v
  })
raymarchFolder
  .add(options, 'cameraNear', 0.0001, 10, 0.1)
  .name('Camera Near')
  .onChange((v) => {
    cameraNear.value = v
  })
raymarchFolder
  .add(options, 'cameraFar', 1, 512, 0.1)
  .name('Camera Far')
  .onChange((v) => {
    cameraFar.value = v
  })
raymarchFolder.close()

// const lightingFolder = gui.addFolder('Lighting')
// lightingFolder
//     .add(options, 'shininess', 0, 100, 0.1)
//     .name('Shininess')
//     .onChange((v) => {
//         shininess.value = v
//     })

// lightingFolder
//     .add(options, 'specularStrength', 0, 1, 0.01)
//     .name('SpecularStrength')
//     .onChange((v) => {
//         specularStrength.value = v
//     })
// lightingFolder.close()

const shadowFolder = gui.addFolder('Shadows')
shadowFolder
  .add(options, 'shadowSoftness', 0, 100, 0.01)
  .name('Shadow Softness')
  .onChange((v) => {
    shadowSoftness.value = v
  })
shadowFolder
  .add(options, 'shadowIntensity', 0, 1, 0.01)
  .name('Shadow Intensity')
  .onChange((v) => {
    shadowIntensity.value = v
  })
shadowFolder
  .add(options, 'shadowMaxSteps', 1, 512, 1)
  .name('Shadow Max Steps')
  .onChange((v) => {
    shadowMaxSteps.value = v
  })
shadowFolder
  .add(options, 'shadowFar', 1, 512, 0.1)
  .name('Shadow Far')
  .onChange((v) => {
    shadowFar.value = v
  })
shadowFolder.close()

const reflectionFolder = gui.addFolder('Reflections')
reflectionFolder
  .add(options, 'maxReflections', 0, 3, 1)
  .name('Max Reflections')
  .onChange((v) => {
    maxReflections.value = v
  })
reflectionFolder
  .add(options, 'reflectivity', 0, 1, 0.1)
  .name('Reflectivity')
  .onChange((v) => {
    reflectivity.value = v
  })
reflectionFolder.close()

const aoFolder = gui.addFolder('Ambient Occlusion')
aoFolder
  .add(options, 'samples', 1, 10, 1)
  .name('Samples')
  .onChange((v) => {
    samples.value = v
  })
aoFolder
  .add(options, 'spread', 0.1, 4, 0.01)
  .name('spread')
  .onChange((v) => {
    spread.value = v
  })
aoFolder
  .add(options, 'objectHeight', 0, 2, 0.01)
  .name('Object Height')
  .onChange((v) => {
    objectHeight.value = v
  })

//Adaptive DPR
const clock = new THREE.Clock()
let targetFPS = 45 // see notes
let frameCount = 0
let elapsed = 0

function adjustDPR(renderer: THREE.WebGPURenderer) {
  elapsed += clock.getDelta()
  frameCount++

  if (elapsed >= 0.094) {
    let fps = frameCount * elapsed * 100
    frameCount = 0
    elapsed -= 0.1

    if (fps < targetFPS * 0.95) {
      renderer.setPixelRatio(
        Math.min(1, Math.max(0.1, renderer.getPixelRatio() * 0.9))
      )
    } else if (fps > targetFPS) {
      renderer.setPixelRatio(Math.min(1, renderer.getPixelRatio() * 1.1))
    }

    //console.log(`DPR: ${renderer.getPixelRatio()} (FPS: ${fps})`)
  }
}

function animate() {
  adjustDPR(renderer)

  controls.update()

  camPos.value.copy(camera.position)
  camTarget.value.copy(controls.target)

  renderer.render(scene, camera)
}