// DeepFracture live web runtime: Three.js rendering + Rapier physics +
// prebaked VQ-VAE codebook fragments.
//
// Pipeline mirror of TEBP-DeepFracture 04.Run-time. The decoder output only
// depends on the quantized codebook entry, so every entry was decoded and
// watershed-segmented OFFLINE (bake.py); at runtime the Siren encoder and
// the nearest-neighbour codebook lookup run live in JS and the matching
// pre-segmented fragment set enters the simulation.

import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
import RAPIER from '@dimforge/rapier3d-compat';

const ASSETS_ROOT = './assets';
const SHAPES = ['squirrel', 'bunny', 'pot', 'base', 'lion'];
const BALL_RADIUS = 0.16;
const BALL_MASS = 1.0;
const IMPULSE_MAX = 10000;
const IMPULSE_SCALE = 250;      // imp_raw ~= |v| * 1/dt, dt = 1/250 (TEBP timestep)
const GRAVITY = -6.0;

// display orientation: rotate the model +90° about Y; the encoder still
// receives coordinates in the original (un-rotated) training frame.
const MODEL_ROT = new THREE.Matrix4().makeRotationX(-Math.PI / 2);
const MODEL_ROT_INV = MODEL_ROT.clone().invert();

const $ = (id) => document.getElementById(id);
const status = (html) => { $('status').innerHTML = html; };
const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
const fmt3 = (v) => `[${v.x.toFixed(2)}, ${v.y.toFixed(2)}, ${v.z.toFixed(2)}]`;

// ---------------------------------------------------------------- setup
await RAPIER.init();
const world = new RAPIER.World({ x: 0, y: GRAVITY, z: 0 });
const eventQueue = new RAPIER.EventQueue(true);

const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setSize(innerWidth, innerHeight);
renderer.setPixelRatio(Math.min(devicePixelRatio, 2));
renderer.shadowMap.enabled = true;
renderer.shadowMap.type = THREE.PCFSoftShadowMap;
$('app').appendChild(renderer.domElement);

const scene = new THREE.Scene();
scene.background = new THREE.Color(0x10141b);
scene.fog = new THREE.Fog(0x10141b, 12, 30);

const camera = new THREE.PerspectiveCamera(45, innerWidth / innerHeight, 0.05, 100);
camera.position.set(2.6, 1.6, 3.4);

const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 0.1, 0);
controls.enableDamping = true;
controls.maxDistance = 12;
controls.minDistance = 1.2;

scene.add(new THREE.HemisphereLight(0xbcd2ff, 0x3a3126, 1.35));
const fill = new THREE.DirectionalLight(0x9db8ff, 0.7);
fill.position.set(-5, 3, -4);
scene.add(fill);
const sun = new THREE.DirectionalLight(0xfff2df, 3.0);
sun.position.set(4, 7, 3);
sun.castShadow = true;
sun.shadow.mapSize.set(2048, 2048);
sun.shadow.camera.left = sun.shadow.camera.bottom = -4;
sun.shadow.camera.right = sun.shadow.camera.top = 4;
scene.add(sun);

// ground
const GROUND_Y = -1.05;
const ground = new THREE.Mesh(
  new THREE.CylinderGeometry(9, 9, 0.12, 64),
  new THREE.MeshStandardMaterial({ color: 0x232a35, roughness: 0.9 }));
ground.position.y = GROUND_Y - 0.06;
ground.receiveShadow = true;
scene.add(ground);
world.createCollider(
  RAPIER.ColliderDesc.cuboid(9, 0.06, 9).setTranslation(0, GROUND_Y - 0.06, 0)
    .setFriction(0.8).setRestitution(0.15));

// ---------------------------------------------------------------- assets
const loader = new GLTFLoader();
const b64f32 = (b64) => new Float32Array(
  Uint8Array.from(atob(b64), (c) => c.charCodeAt(0)).buffer);

const shapeCache = new Map();   // shape -> {meta, encW, encB, cookbook, mesh}
const glbCache = new Map();     // `${shape}/${code}` -> gltf

function firstMesh(root) {
  let found = null;
  root.traverse((c) => { if (!found && c.isMesh) found = c; });
  return found;
}

async function loadShape(shape) {
  if (shapeCache.has(shape)) return shapeCache.get(shape);
  const res = await fetch(`${ASSETS_ROOT}/${shape}/meta.json`);
  if (!res.ok) throw new Error(`assets for "${shape}" not published yet`);
  const meta = await res.json();
  const gltf = await loader.loadAsync(`${ASSETS_ROOT}/${shape}/target.glb`);
  const mesh = firstMesh(gltf.scene);
  mesh.geometry = mesh.geometry.clone();
  mesh.geometry.applyMatrix4(MODEL_ROT);
  if (!mesh.geometry.attributes.normal) mesh.geometry.computeVertexNormals();
  mesh.material = new THREE.MeshStandardMaterial({ color: 0x76c893, roughness: 0.65 });
  mesh.castShadow = mesh.receiveShadow = true;
  const entry = {
    meta,
    encW: b64f32(meta.encoderW),
    encB: b64f32(meta.encoderB),
    cookbook: b64f32(meta.cookbook),
    mesh,
  };
  shapeCache.set(shape, entry);
  return entry;
}

// ---------------------------------------------------------------- state
let SHAPE = 'squirrel';
let shapeData = null;
let target = null;          // { mesh, body, collider }
let fractured = false;
let paused = false;         // world freezes while the pipeline panel plays
const balls = [];           // { mesh, body, collider }
const fragments = [];       // { mesh, body }

function spawnTarget() {
  const mesh = shapeData.mesh.clone();
  mesh.position.set(0, 0, 0);
  scene.add(mesh);
  const geo = mesh.geometry;
  const verts = new Float32Array(geo.attributes.position.array);
  const idx = geo.index
    ? new Uint32Array(geo.index.array)
    : new Uint32Array([...Array(verts.length / 3).keys()]);
  const body = world.createRigidBody(RAPIER.RigidBodyDesc.fixed());
  const collider = world.createCollider(
    RAPIER.ColliderDesc.trimesh(verts, idx).setActiveEvents(
      RAPIER.ActiveEvents.COLLISION_EVENTS), body);
  target = { mesh, body, collider };
  fractured = false;
}

function clearDynamic() {
  for (const list of [balls, fragments]) {
    for (const o of list) { scene.remove(o.mesh); world.removeRigidBody(o.body); }
    list.length = 0;
  }
}

function reset() {
  clearDynamic();
  if (target) { scene.remove(target.mesh); world.removeRigidBody(target.body); target = null; }
  paused = false;
  $('pipeline').style.display = 'none';
  spawnTarget();
  status(`Ready — <b>click the ${SHAPE}</b> to fire.`);
}
$('reset').onclick = reset;

async function switchShape(shape) {
  try {
    status(`⏳ Loading <b>${shape}</b>…`);
    shapeData = await loadShape(shape);
    SHAPE = shape;
    reset();
  } catch (e) {
    $('shape').value = SHAPE;
    status(`⚠️ ${e.message}`);
  }
}
$('shape').onchange = (e) => switchShape(e.target.value);

// ---------------------------------------------------------------- network (JS mirror of Cook)
let seedState = 0;
function seededNormal(seed) {           // Box-Muller with mulberry32
  let a = seed >>> 0;
  const rand = () => {
    a |= 0; a = (a + 0x6D2B79F5) | 0;
    let t = Math.imul(a ^ (a >>> 15), 1 | a);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
  return () => {
    const u = Math.max(rand(), 1e-9), v = rand();
    return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
  };
}

function selectCode(pos, dir, impNorm, seed) {
  const { meta, encW, encB, cookbook } = shapeData;
  const x = [pos.x, pos.y, pos.z, dir.x, dir.y, dir.z, impNorm];
  const F = meta.featDim;
  const feat = new Float32Array(F + meta.latentDim);
  for (let i = 0; i < F; i++) {
    let s = encB[i];
    for (let j = 0; j < 7; j++) s += encW[i * 7 + j] * x[j];
    feat[i] = Math.sin(meta.w0 * s);
  }
  // xavier-normal latent (std = sqrt(2 / (1 + latentDim)))
  const nrm = seededNormal(seed + (seedState++));
  const std = Math.sqrt(2 / (1 + meta.latentDim));
  for (let i = 0; i < meta.latentDim; i++) feat[F + i] = nrm() * std;

  const D = F + meta.latentDim;
  let best = -1, bestDist = Infinity;
  for (const c of meta.codes) {          // restrict to baked codes
    let d = 0;
    for (let i = 0; i < D; i++) {
      const t = feat[i] - cookbook[c * D + i];
      d += t * t;
    }
    if (d < bestDist) { bestDist = d; best = c; }
  }
  return best;
}

// ---------------------------------------------------------------- fracture
function pipelineStep(html) {
  const li = document.createElement('li');
  li.innerHTML = html;
  li.style.margin = '3px 0';
  $('pipeline-steps').appendChild(li);
  return li;
}

async function fracture(hitPoint, hitDir, impRaw) {
  fractured = true;
  paused = true;                       // freeze time at the moment of impact
  $('pipeline-steps').innerHTML = '';
  $('pipeline-bar-wrap').style.display = 'none';
  $('pipeline').style.display = 'block';

  // back to the network's training frame (un-rotate the display transform)
  const posModel = hitPoint.clone().applyMatrix4(MODEL_ROT_INV);
  const dirModel = hitDir.clone().transformDirection(MODEL_ROT_INV);

  const impNorm = Math.min(impRaw, IMPULSE_MAX) / IMPULSE_MAX * 2 - 1;
  pipelineStep(`💥 <b style="color:#e8eaed">Impact captured</b> (PyBullet-style contact)<br>
    <small>pos ${fmt3(posModel)} · dir ${fmt3(dirModel)}<br>
    impulse ${impRaw.toFixed(0)} → normalized <b>${impNorm.toFixed(2)}</b></small>`);
  await sleep(700);

  pipelineStep(`🧠 <b style="color:#e8eaed">Siren encoder</b> — live in JS
    <small>7-D condition → ${shapeData.meta.featDim}-D feature = sin(Wx+b)</small>`);
  const seed = Number($('seed').value) || 42;
  const code = selectCode(posModel, dirModel, impNorm, seed);
  await sleep(700);

  pipelineStep(`📖 <b style="color:#e8eaed">VQ codebook lookup</b> — live in JS
    <small>nearest of ${shapeData.meta.nCodes} entries → <b style="color:#e0574f">#${code}</b></small>`);
  await sleep(700);

  const key = `${SHAPE}/${code}`;
  const cached = glbCache.has(key);
  pipelineStep(`🌊 <b style="color:#e8eaed">Decoder + watershed flooding</b>
    <small>128³ GS-SDF → h-minima + watershed → fragment meshes —
    precomputed offline for every codebook entry (identical output, zero wait)</small>`);
  await sleep(700);

  pipelineStep(`📦 <b style="color:#e8eaed">Fragment set #${code}</b>
    <small>${cached ? 'cached' : 'downloading'}…</small>`);
  let gltf = glbCache.get(key);
  if (!gltf) {
    $('pipeline-bar-wrap').style.display = 'block';
    gltf = await new Promise((resolve, reject) => {
      loader.load(`${ASSETS_ROOT}/${SHAPE}/codes/${String(code).padStart(3, '0')}.glb`,
        resolve,
        (ev) => { if (ev.total) $('pipeline-bar').style.width = `${(ev.loaded / ev.total) * 100}%`; },
        reject);
    });
    $('pipeline-bar').style.width = '100%';
    glbCache.set(key, gltf);
  }
  await sleep(400);

  scene.remove(target.mesh);
  world.removeRigidBody(target.body);
  target = null;

  let count = 0;
  gltf.scene.traverse((child) => {
    if (!child.isMesh) return;
    count += 1;
    const mesh = child.clone();
    mesh.geometry = child.geometry.clone();
    mesh.geometry.applyMatrix4(MODEL_ROT);
    if (!mesh.geometry.attributes.normal) mesh.geometry.computeVertexNormals();
    mesh.material = new THREE.MeshStandardMaterial({
      vertexColors: !!mesh.geometry.attributes.color, roughness: 0.6,
      color: mesh.geometry.attributes.color ? 0xffffff : 0xd08770 });
    mesh.castShadow = mesh.receiveShadow = true;
    scene.add(mesh);

    // shrink the hull slightly toward the centroid so touching fragments
    // don't start deeply interpenetrated (which would pop them apart)
    const src = mesh.geometry.attributes.position.array;
    const n = src.length / 3;
    let cx = 0, cy = 0, cz = 0;
    for (let i = 0; i < src.length; i += 3) { cx += src[i]; cy += src[i + 1]; cz += src[i + 2]; }
    cx /= n; cy /= n; cz /= n;
    const pts = new Float32Array(src.length);
    const SHRINK = 0.96;
    for (let i = 0; i < src.length; i += 3) {
      pts[i] = cx + (src[i] - cx) * SHRINK;
      pts[i + 1] = cy + (src[i + 1] - cy) * SHRINK;
      pts[i + 2] = cz + (src[i + 2] - cz) * SHRINK;
    }
    const body = world.createRigidBody(
      RAPIER.RigidBodyDesc.dynamic().setTranslation(0, 0, 0));
    const desc = RAPIER.ColliderDesc.convexHull(pts);
    if (desc) {
      desc.setDensity(1.0).setFriction(0.7).setRestitution(0.2);
      world.createCollider(desc, body);
    }
    fragments.push({ mesh, body });
  });

  pipelineStep(`🪨 <b style="color:#e8eaed">${count} fragments</b> re-enter the simulation`);
  paused = false;
  status(`💥 <b>${count} fragments</b> · codebook <b>#${code}</b> — ↺ Reset and hit a <b>different spot</b> for a different pattern!`);
  setTimeout(() => { $('pipeline').style.display = 'none'; }, 8000);
}

// ---------------------------------------------------------------- shooting
const raycaster = new THREE.Raycaster();
let dragging = false, downPos = null;

renderer.domElement.addEventListener('pointerdown', (e) => {
  downPos = [e.clientX, e.clientY]; dragging = false;
});
renderer.domElement.addEventListener('pointermove', (e) => {
  if (downPos && Math.hypot(e.clientX - downPos[0], e.clientY - downPos[1]) > 6) dragging = true;
});
renderer.domElement.addEventListener('pointerup', (e) => {
  if (dragging || !downPos) { downPos = null; return; }
  downPos = null;
  shoot(e);
});

function shoot(e) {
  if (paused) return;
  const ndc = new THREE.Vector2(
    (e.clientX / innerWidth) * 2 - 1, -(e.clientY / innerHeight) * 2 + 1);
  raycaster.setFromCamera(ndc, camera);
  const dir = raycaster.ray.direction.clone().normalize();
  const origin = camera.position.clone().add(dir.clone().multiplyScalar(0.4));
  const speed = Number($('speed').value);

  const mesh = new THREE.Mesh(
    new THREE.SphereGeometry(BALL_RADIUS, 24, 16),
    new THREE.MeshStandardMaterial({ color: 0xe0574f, roughness: 0.4 }));
  mesh.castShadow = true;
  mesh.position.copy(origin);
  scene.add(mesh);

  const body = world.createRigidBody(
    RAPIER.RigidBodyDesc.dynamic()
      .setTranslation(origin.x, origin.y, origin.z)
      .setLinvel(dir.x * speed, dir.y * speed, dir.z * speed)
      .setCcdEnabled(true));
  const col = world.createCollider(
    RAPIER.ColliderDesc.ball(BALL_RADIUS).setDensity(
      BALL_MASS / ((4 / 3) * Math.PI * BALL_RADIUS ** 3))
      .setActiveEvents(RAPIER.ActiveEvents.COLLISION_EVENTS)
      .setFriction(0.6).setRestitution(0.3), body);
  balls.push({ mesh, body, collider: col });
  if (!fractured) status('🔴 projectile away…');
}

// ---------------------------------------------------------------- impact detection
function handleCollisions() {
  eventQueue.drainCollisionEvents((h1, h2, started) => {
    if (!started || fractured || !target) return;
    const th = target.collider.handle;
    if (h1 !== th && h2 !== th) return;
    const otherHandle = h1 === th ? h2 : h1;
    const ball = balls.find((b) => b.collider.handle === otherHandle);
    if (!ball) return;

    const v = ball.body.linvel();
    const speed = Math.hypot(v.x, v.y, v.z);
    const impRaw = speed * IMPULSE_SCALE * BALL_MASS;
    if (impRaw < 2000) return;   // below fracture threshold (TEBP)

    // contact point: cast the ball's motion ray onto the target mesh
    const p = ball.body.translation();
    const vd = new THREE.Vector3(v.x, v.y, v.z).normalize();
    raycaster.set(new THREE.Vector3(p.x, p.y, p.z).sub(vd.clone().multiplyScalar(BALL_RADIUS * 3)), vd);
    const hits = raycaster.intersectObject(target.mesh, false);
    const hit = hits[0] ? hits[0].point
      : new THREE.Vector3(p.x, p.y, p.z).sub(vd.clone().multiplyScalar(BALL_RADIUS));

    fracture(hit, vd, impRaw);
  });
}

// ---------------------------------------------------------------- loop
const clock = new THREE.Clock();
function tick() {
  requestAnimationFrame(tick);
  const dt = Math.min(clock.getDelta(), 0.05);
  if (!paused) {
    world.timestep = dt;
    world.step(eventQueue);
    handleCollisions();
  }

  for (const list of [balls, fragments]) {
    for (const o of list) {
      const t = o.body.translation(), r = o.body.rotation();
      o.mesh.position.set(t.x, t.y, t.z);
      o.mesh.quaternion.set(r.x, r.y, r.z, r.w);
    }
  }
  // recycle far-away balls
  for (let i = balls.length - 1; i >= 0; i--) {
    if (balls[i].mesh.position.length() > 40) {
      scene.remove(balls[i].mesh); world.removeRigidBody(balls[i].body);
      balls.splice(i, 1);
    }
  }
  controls.update();
  renderer.render(scene, camera);
}

// ---------------------------------------------------------------- boot
shapeData = await loadShape(SHAPE);
spawnTarget();
$('loading').style.display = 'none';
status(`Ready — <b>click the ${SHAPE}</b> to fire.`);
tick();

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