import { Color, DirectionalLight, HemisphereLight, InstancedMesh, MeshLambertMaterial, Object3D, OrthographicCamera, PCFSoftShadowMap, PerspectiveCamera, Quaternion, Scene, SphereGeometry, Vector3, WebGLRenderer } from "https://esm.sh/three@0.169.0"


const N = 6000
const MIN_R = 0.005
const MAX_R = 0.5


//////// SETUP

const scene = new Scene()
scene.background = new Color(0.01, 0.01, 0.01)

let asp = innerWidth / innerHeight
const camera = new OrthographicCamera(-asp, asp, 1, -1, -1, 1)

const renderer = new WebGLRenderer({ canvas })
renderer.setSize(innerWidth, innerHeight)
renderer.setPixelRatio(Math.min(2, devicePixelRatio))
renderer.shadowMap.enabled = true
renderer.shadowMap.type = PCFSoftShadowMap


//////// LIGHT

const ambient = new HemisphereLight(
  new Color(1, 0.6, 0.2),
  new Color(0.2, 0.6, 1),
  2,
)
scene.add(ambient)

const light = new DirectionalLight(
  new Color(1, 0.9, 0.8),
  3,
)
light.position.set(-1, 1, 1)
scene.add(light)

light.castShadow = true
light.shadow.mapSize.width  = 1024
light.shadow.mapSize.height = 1024
light.shadow.camera.near   =  0
light.shadow.camera.far    =  4
light.shadow.camera.left   = -1
light.shadow.camera.right  =  1
light.shadow.camera.top    = -1
light.shadow.camera.bottom =  1
light.shadow.bias = -0.001
light.shadow.intensity = 0.62


//////// INTSTANCED SPHERES

const mat = new MeshLambertMaterial()

const new_spheres = (ws: number, hs: number) => {
  const geom = new SphereGeometry(1, ws, hs)
  const spheres = new InstancedMesh(geom, mat, N)
  spheres.castShadow = true
  spheres.receiveShadow = true
  spheres.count = 0
  return spheres
}

const spheres = [
  // small spheres
  { num: 0, mesh: new_spheres(9, 6) },
  // medium spheres
  { num: 0, mesh: new_spheres(21, 14) },
  // large spheres
  { num: 0, mesh: new_spheres(72, 48) },
]

const obj = new Object3D
obj.scale.multiplyScalar(0.5)
for (const i of spheres) obj.add(i.mesh)
scene.add(obj)
const qt = obj.quaternion


const dummy = new Object3D()

type iTransfer = {
  n: number
  tests_n: number
  data: Float32Array
}

const set_instances = (e: MessageEvent<iTransfer>) => {
  const { n, tests_n, data } = e.data

  for (let i = 0; i < data.length; i += 4) {
    const x = data[i + 0]
    const y = data[i + 1]
    const z = data[i + 2]
    const r = data[i + 3]

    dummy.position.set(x, y, z)
    dummy.scale.setScalar(r)
    dummy.updateMatrix()

    const size_id = r > 0.03 ? r > 0.1 ? 2 : 1 : 0
    const inst = spheres[size_id]
    const { mesh } = inst

    mesh.setMatrixAt(inst.num, dummy.matrix)
    mesh.setColorAt(inst.num, new Color(
      random(0.2, 1.0),
      random(0.2, 0.9),
      random(0.2, 0.8),
    ))

    mesh.instanceMatrix.addUpdateRange(inst.num * 16, 16)
    mesh.instanceColor!.addUpdateRange(inst.num *  3,  3)
    mesh.instanceMatrix.needsUpdate = true
    mesh.instanceColor!.needsUpdate = true

    mesh.count = ++inst.num
  }

  log.textContent = `${str(tests_n)} tests\n${str(n)} spheres`
}


//////// ANIMATION LOOP

const HPI   = Math.PI / 2
const DUR   = 3000
const DELAY = 1000


const get_next_qt = (() => {
  const new_qt = (
    v: Vector3,
    a: number,
  ) => (
    new Quaternion().setFromAxisAngle(v, a)
  )

  const vx = new Vector3(1, 0, 0)
  const qtxs = [
    new_qt(vx, -HPI),
    new_qt(vx, 0),
    new_qt(vx, HPI),
  ]

  const vy = new Vector3(0, 1, 0)
  const qtys = [
    new_qt(vy, -HPI),
    new_qt(vy, 0),
    new_qt(vy, HPI),
  ]

  // 0b0101 gives zero rotation, qtxs[1] * qtys[1]
  let prev_qt_id = 5

  const get_rnd_qt = () => {
    const id = (random(0, 3) << 2) + (random(0, 3) & 3)
    // don't return zero or same rotation
    if (id === prev_qt_id || id === 5) {
      return get_rnd_qt()
    }
    prev_qt_id = id
    return [ qtxs[id >> 2], qtys[id & 3] ]
  }

  return (prev_qt: Quaternion) => {
    const [ xqt, yqt ] = get_rnd_qt()
    return prev_qt.clone()
      .premultiply(xqt)
      .premultiply(yqt)
  }
})()


// 3rd order smoothstep
const ease = (t: number) => (
    35 * t ** 4
  - 84 * t ** 5
  + 70 * t ** 6
  - 20 * t ** 7
)


const rnd_rot = () => {
  const prev_qt = qt.clone()
  const next_qt = get_next_qt(prev_qt)

  const start = performance.now()

  const tick = () => {
    const now = performance.now()
    const t = ease(Math.min(1, (now - start) / DUR))
    qt.slerpQuaternions(prev_qt, next_qt, t)
    if (t < 1) requestAnimationFrame(tick)
    else setTimeout(rnd_rot, DELAY)
  }

  requestAnimationFrame(tick)
}

setTimeout(rnd_rot, DELAY)


//////// RENDER

renderer.setAnimationLoop(() => {
  renderer.render(scene, camera)
})


//////// PACKING WORKER

const worker_src = `
const N = ${N}
const MIN_R = ${MIN_R}
const MAX_R = ${MAX_R}
${worker_self.toString().replace(/^.+\n|\n.+$/g, "")}`

const worker = new Worker(URL.createObjectURL(new Blob(
  [ worker_src ],
  { type: "text/javascript" },
)))

worker.onmessage = set_instances

worker.postMessage(1) // init packing


function worker_self() {
  const $ = self as any as Worker

  let n = 0
  const packed = new Float32Array(N * 4)

  let chunk_n = 0
  const N_PER_CHUNK = 20

  let tests_n = 0
  const MAX_TESTS = N * 1e6

  const transfer: iTransfer = {
    n,
    tests_n,
    data: null!,
  }

  onmessage = () => {
    while (n < N) {
      let x = 0, y = 0, z = 0, r = 0

      while (r === 0) {
        x = rnd_pos()
        y = rnd_pos()
        z = rnd_pos()
        r = get_radius(x, y, z, n)

        if (r === -1) {
          return send_data()
        }
      }

      const i = n++ * 4
      packed[i + 0] = x
      packed[i + 1] = y
      packed[i + 2] = z
      packed[i + 3] = r

      if (++chunk_n === N_PER_CHUNK || n === N) {
        chunk_n = 0
        send_data()
      }
    }
  }

  const send_data = () => {
    transfer.n = n
    transfer.tests_n = tests_n
    transfer.data = packed.slice((n - N_PER_CHUNK) * 4, n * 4)
    $.postMessage(transfer, [ transfer.data.buffer ])
  }

  const rnd_pos = () => -1 + Math.random() * 2

  function dist(
    x0: number,
    y0: number,
    x1: number,
    y1: number,
  ) {
    return Math.hypot(x1 - x0, y1 - y0)
  }

  function get_radius(
    x: number,
    y: number,
    z: number,
    n: number,
  ) {
    let r = Math.min(
      1 - Math.abs(x),
      1 - Math.abs(y),
      1 - Math.abs(z),
      MAX_R,
    )

    if (r < MIN_R) {
      return 0
    }

    if (n === 0) {
      return r
    }

    for (let i = 0; i < n * 4; i += 4) {
      if (tests_n === MAX_TESTS) {
        return -1
      }

      const xo = packed[i + 0]
      const yo = packed[i + 1]
      const zo = packed[i + 2]
      const ro = packed[i + 3]

      const ryz = dist(y, z, yo, zo) - ro
      const rxz = dist(x, z, xo, zo) - ro
      const rxy = dist(x, y, xo, yo) - ro

      if (
        ryz < MIN_R ||
        rxz < MIN_R ||
        rxy < MIN_R
      ) {
        r = 0
        break
      }

      r = Math.min(ryz, rxz, rxy, r)

      tests_n++
    }

    return r
  }
}


//////// UTILS

function str(num: number) {
  return num.toLocaleString()
}

function random(from = 0, to = 1) {
  return from + Math.random() * (to - from)
}

onresize = () => {
  asp = innerWidth / innerHeight
  camera.left  = -asp
  camera.right =  asp
  camera.updateProjectionMatrix()

  renderer.setSize(innerWidth, innerHeight)
}