GPU Linked Particles in React Three Fiber with WebGPU

'use client'

import { useGSAP } from '@gsap/react'
import {
  type CreatorState,
  useFrame,
  useLocalNodes,
  useThree,
  useUniforms,
} from '@react-three/fiber/webgpu'
import gsap from 'gsap'
import { type FC, useCallback, useEffect, useMemo, useRef } from 'react'
import {
  Break,
  Fn,
  If,
  Loop,
  PI2,
  array,
  bool,
  color,
  cos,
  deltaTime,
  float,
  hash,
  instanceIndex,
  instancedArray,
  mix,
  mx_noise_float,
  mx_noise_vec3,
  positionLocal,
  sin,
  smoothstep,
  step,
  storage,
  time,
  uv,
  vec2,
  vec3,
  vec4,
} from 'three/tsl'
import {
  BufferGeometry,
  DoubleSide,
  type Node,
  StorageBufferAttribute,
  type UniformNode,
} from 'three/webgpu'

import useIsPortraitSize from '../layout/useIsPortraitSize'

gsap.registerPlugin(useGSAP)

const PARTICLE_COUNT = 1024
const INITIAL_VELOCITY = 80
const LINKED_PARTICLES_UNIFORM_SCOPE = 'linkedParticles'
const RING_RADIUS = 3

const PARTICLE_PALETTE = [
  '#ffe071',
  '#ff9183',
  '#ff4b93',
  '#5f00aa',
  '#9a00a3',
  '#fff636',
  '#f6cc00',
  '#ffdf3a',
  '#f8d200',
  '#ffe81f',
  '#fff132',
  '#fffc5f',
  '#edc740',
  '#d5c8a3',
  '#ff9b91',
  '#ffa798',
  '#ff978b',
  '#ffbaa7',
  '#ff957a',
  '#ffa59d',
  '#ffc3a6',
  '#ffceb2',
  '#bf897e',
  '#ff9be2',
  '#ff85d4',
  '#ffb1ff',
  '#ff8ddc',
  '#ff6ecd',
  '#ff6fc7',
  '#ffb1f2',
  '#ff96e4',
  '#ca92a3',
  '#ffbfff',
  '#faa7ff',
  '#ffb6ff',
  '#ffc8ff',
  '#ffbcff',
  '#ffccff',
  '#e2c0ff',
  '#ffc9ff',
  '#ccb0d9',
  '#ffc3ff',
  '#ffb0ff',
  '#ffafff',
  '#ffb5ff',
  '#ff90ff',
  '#ffd5ff',
  '#ffa9ff',
  '#d6abda',
] as const

const PARTICLE_PALETTE_SIZE = PARTICLE_PALETTE.length
const particlePaletteNodes = array(PARTICLE_PALETTE.map((hex) => color(hex)))

const getLinksIndices = (particleCount: number): number[] => {
  const linksIndices = []
  // Each link is a quad (2 triangles), using 4 vertices.
  for (let i = 0; i < particleCount; i++) {
    // Calculate the base index for the *vertices* of the i-th link strip.
    // Since each link uses 4 vertices, the base vertex index is i * 4.
    const baseVertexIndex = i * 4
    // Define the two triangles for the quad using these vertex indices.
    // Triangle 1: vertices base, base + 1, base + 2
    // Triangle 2: vertices base, base + 2, base + 3
    linksIndices.push(baseVertexIndex, baseVertexIndex + 1, baseVertexIndex + 2)
    linksIndices.push(baseVertexIndex, baseVertexIndex + 2, baseVertexIndex + 3)
  }
  return linksIndices
}
type LinkedParticlesUniforms = {
  uOpacity: UniformNode<'float', number>
  uRingRadius: UniformNode<'float', number>
  uVelocityMultiplier: UniformNode<'float', number>
}

type Props = {
  particleCount?: number
}

const LinkedParticles: FC<Props> = ({ particleCount = PARTICLE_COUNT }) => {
  const renderer = useThree((s) => s.renderer)
  const isPortrait = useIsPortraitSize()
  const ringRadiusMultiplier = isPortrait ? 0.75 : 1.0
  const hasInitializedRef = useRef(false)

  const { linksGeometry, linksVerticesSBA } = useMemo(() => {
    const geometry = new BufferGeometry()
    const vertices = new StorageBufferAttribute(particleCount * 4, 3)

    geometry.setIndex(getLinksIndices(particleCount))
    geometry.setAttribute('position', vertices)

    return {
      linksGeometry: geometry,
      linksVerticesSBA: vertices,
    }
  }, [particleCount])

  const particleBuffers = useMemo(() => {
    const seeds = new Float32Array(particleCount)
    for (let i = 0; i < particleCount; i++) {
      seeds[i] = Math.random()
    }

    return {
      particleVelocityBuffer: instancedArray(particleCount, 'vec3'),
      particleColourBuffer: instancedArray(particleCount, 'vec4'),
      particleSeedBuffer: instancedArray(seeds, 'float'),
      particlePositionBuffer: instancedArray(particleCount, 'vec4'),
    }
  }, [particleCount])

  useUniforms(
    {
      uOpacity: 0,
      uRingRadius: RING_RADIUS,
      uVelocityMultiplier: INITIAL_VELOCITY,
    },
    LINKED_PARTICLES_UNIFORM_SCOPE,
  )

  const createSimulationNodes = useCallback(
    ({ uniforms }: CreatorState) => {
      const { uOpacity, uRingRadius, uVelocityMultiplier } =
        uniforms.scope<LinkedParticlesUniforms>(LINKED_PARTICLES_UNIFORM_SCOPE)
      const {
        particleColourBuffer,
        particlePositionBuffer,
        particleSeedBuffer,
        particleVelocityBuffer,
      } = particleBuffers

      const linksPositions = storage(linksVerticesSBA, 'vec3', linksVerticesSBA.count).setPBO(
        true,
      )

      const computeParticlePositions = Fn(() => {
        const position = particlePositionBuffer.element(instanceIndex).xyz
        const life = particlePositionBuffer.element(instanceIndex).w
        const seed = particleSeedBuffer.element(instanceIndex)

        // Assign a random lifetime
        life.assign(mix(0.0, 1.0, hash(seed)))
        // Spread around the ring...

        const angle = hash(instanceIndex.add(2)).mul(PI2).toVar()
        const ringNoise = mx_noise_float(seed.add(time)).mul(0.2).toVar()
        const newPos = vec3(
          sin(angle.add(ringNoise)).mul(uRingRadius),
          cos(angle).mul(uRingRadius).add(ringNoise),
          0.0,
        )
        position.assign(newPos)
      })().compute(particleCount)

      const computeColors = Fn(() => {
        const c = particleColourBuffer.element(instanceIndex)
        const colorIndex = hash(instanceIndex.add(2)).mul(PARTICLE_PALETTE_SIZE).floor()
        const paletteColor = particlePaletteNodes.element(colorIndex) as unknown as Node<'vec3'>
        c.assign(vec4(paletteColor, 1))
      })().compute(particleCount)

      const positionNode = particlePositionBuffer.toAttribute().xyz

      const colorNode = Fn(() => {
        const colour = particleColourBuffer.element(instanceIndex).rgb
        const centeredUv = uv().distance(vec2(0.5))
        const circle = step(0.5, centeredUv).oneMinus()
        return vec4(colour, circle)
      })()

      const opacityNode = Fn(() => {
        const distNorm = positionLocal.xy.length().div(uRingRadius).toVar()
        const distNormSq = distNorm.mul(distNorm)
        const fadeToEdges = smoothstep(6.0, 2.5, distNormSq)
        return fadeToEdges.mul(uOpacity)
      })()

      const linkOpacityNode = opacityNode

      const thresholdSq = float(0.64)
      const halfWidth = float(0.005)

      const updateParticles = Fn(() => {
        // Update the particle position using noise.
        const position = particlePositionBuffer.element(instanceIndex).xyz
        const life = particlePositionBuffer.element(instanceIndex).w
        const vel = particleVelocityBuffer.element(instanceIndex).xy
        const dt = deltaTime.mul(0.01).mul(uVelocityMultiplier).toVar()

        If(life.greaterThan(0.0), () => {
          // Deduct life
          life.assign(life.sub(dt))
          // Generate some noise
          const noise = mx_noise_vec3(position.add(time.mul(0.3)))
            .mul(0.02)
            .toVar()
          // Update velocity with noise
          vel.addAssign(vec2(noise.x, noise.y))

          // Apply attraction force to velocity with a bit of turbulence for visual interest
          const turbulence = vec2(
            mx_noise_float(position.x.add(time)).mul(0.1),
            mx_noise_float(position.y.add(time)).mul(0.1),
          )

          vel.addAssign(turbulence.mul(0.32))

          // Update position based on velocity and time.
          position.addAssign(vel.mul(dt))

          // LINKS
          const foundLink = bool(false).toVar()
          const toPosition = vec3(0.0).toVar()

          Loop(particleCount, ({ i }) => {
            const other = particlePositionBuffer.element(i)
            // // skip self
            If(float(i).equal(float(instanceIndex)), () => {
              Break()
            })

            If(foundLink, () => {
              // If we already found a link, skip the rest
              Break()
            })

            const d2 = position.sub(other.xyz).lengthSq()

            If(d2.lessThan(thresholdSq), () => {
              toPosition.assign(other.xyz)
              foundLink.assign(bool(true))
            })
          })

          // --- Calculate Link Vertex Positions ---
          const baseVertexIndex = instanceIndex.mul(4)

          // --- Link 1 (Forced to next particle) ---
          // No If(foundLink) needed as we force it to true
          const linkStart = position
          const linkEnd = toPosition // Use the forced 'toPosition'

          If(foundLink, () => {
            // direction in XY only, then its 2‑D perpendicular
            const dir2D = vec3(linkEnd.x.sub(linkStart.x), linkEnd.y.sub(linkStart.y), 0.0)
              .normalize()
              .toVar()
            const perp2D = vec3(dir2D.y.negate(), dir2D.x, 0.0).mul(halfWidth).toVar()

            linksPositions.element(baseVertexIndex).assign(linkStart.add(perp2D))
            linksPositions.element(baseVertexIndex.add(1)).assign(linkStart.sub(perp2D))
            linksPositions.element(baseVertexIndex.add(2)).assign(linkEnd.add(perp2D))
            linksPositions.element(baseVertexIndex.add(3)).assign(linkEnd.sub(perp2D))
          }).Else(() => {
            // Collapse
            linksPositions.element(baseVertexIndex).assign(linkStart)
            linksPositions.element(baseVertexIndex.add(1)).assign(linkStart)
            linksPositions.element(baseVertexIndex.add(2)).assign(linkStart)
            linksPositions.element(baseVertexIndex.add(3)).assign(linkStart)
          })
        })
      })().compute(particleCount)

      const spawnParticles = Fn(() => {
        const position = particlePositionBuffer.element(instanceIndex).xyz
        const life = particlePositionBuffer.element(instanceIndex).w

        const vel = particleVelocityBuffer.element(instanceIndex).xy
        const seed = particleSeedBuffer.element(instanceIndex)

        // Spawn new particle if the current one is dead and the system is alive
        If(life.lessThanEqual(0.0).and(uOpacity.greaterThan(0.0)), () => {
          // Assign a random lifetime
          life.assign(mix(0.12, 0.8, hash(float(instanceIndex).add(time))))
          // TODO: extract ring position logic into a separate function
          const angle = hash(instanceIndex.add(1)).mul(PI2).toVar()
          const ringNoise = mx_noise_float(seed.add(time)).mul(0.2).toVar()
          const velocityDirection = vec2(sin(angle), cos(angle)).toVar()
          const newPos = vec3(
            sin(angle.add(ringNoise)).mul(uRingRadius),
            velocityDirection.y.mul(uRingRadius).add(ringNoise),
            0.0,
          )
          position.assign(newPos)

          // Generate an outward velocity based on the angle
          vel.assign(velocityDirection.mul(seed.mul(10.0)))
        })
      })().compute(particleCount)

      return {
        colorNode,
        computeColors,
        computeParticlePositions,
        linkOpacityNode,
        opacityNode,
        positionNode,
        setUOpacity: (value: number) => {
          uOpacity.value = value
        },
        setURingRadius: (value: number) => {
          uRingRadius.value = value
        },
        setUVelocityMultiplier: (value: number) => {
          uVelocityMultiplier.value = value
        },
        spawnParticles,
        updateParticles,
      }
    },
    [linksVerticesSBA, particleBuffers, particleCount],
  )

  const simulation = useLocalNodes(createSimulationNodes)

  useEffect(() => {
    let isCancelled = false
    hasInitializedRef.current = false

    void renderer
      .computeAsync([simulation.computeParticlePositions, simulation.computeColors])
      .then(() => {
        if (!isCancelled) hasInitializedRef.current = true
      })

    return () => {
      isCancelled = true
    }
  }, [renderer, ringRadiusMultiplier, simulation])

  useGSAP(
    () => {
      const ringRadius = RING_RADIUS * ringRadiusMultiplier
      const animation = {
        opacity: 0,
        velocityMultiplier: INITIAL_VELOCITY,
      }

      simulation.setUOpacity(animation.opacity)
      simulation.setURingRadius(ringRadius)
      simulation.setUVelocityMultiplier(animation.velocityMultiplier)

      gsap
        .timeline()
        .to(
          animation,
          {
            opacity: 1,
            duration: 2,
            ease: 'power2.out',
            onUpdate: () => simulation.setUOpacity(animation.opacity),
          },
          0,
        )
        .to(
          animation,
          {
            velocityMultiplier: 30,
            duration: 5,
            ease: 'power2.out',
            onUpdate: () => simulation.setUVelocityMultiplier(animation.velocityMultiplier),
          },
          0,
        )
    },
    {
      dependencies: [ringRadiusMultiplier, simulation],
      revertOnUpdate: true,
    },
  )

  useFrame(
    () => {
      if (!hasInitializedRef.current) return
      renderer.compute([simulation.spawnParticles, simulation.updateParticles])
    },
    { fps: 60 },
  )

  return (
    <group>
      {/* Mesh displaying the links */}
      <mesh geometry={linksGeometry} frustumCulled={false}>
        <meshBasicNodeMaterial
          attach="material"
          color="#9a00a3"
          side={DoubleSide}
          depthTest={false}
          opacityNode={simulation.linkOpacityNode}
          transparent={true}
        />
      </mesh>
      {/* InstancedMesh displaying the particles */}
      <instancedMesh args={[undefined, undefined, particleCount]} frustumCulled={false}>
        <planeGeometry args={[0.06, 0.06]} />
        <spriteNodeMaterial
          key={simulation.colorNode.uuid}
          positionNode={simulation.positionNode}
          colorNode={simulation.colorNode}
          depthTest={false}
          opacityNode={simulation.opacityNode}
          transparent={true}
        />
      </instancedMesh>
    </group>
  )
}

export default LinkedParticles

// Future update...(leave for now)
// Attract particles that are near the pointer to the pointer - within a given threshold
// const pointerDist = uPointer.sub(position.xy)
// const pointerDir = pointerDist.normalize()
// const pointerDistSq = pointerDist.lengthSq()
// Calculate attraction strength with increased effect
// Using a steeper curve for more dramatic effect
// const attractionStrength = smoothstep(0.0, pointerDistThresholdSq, pointerDistSq).oneMinus() // Increased from 20.0 to 80.0 for stronger effect