path: root/public/js/components/vrm-avatar.js

/* global AFRAME, THREE, CurrentWorld */
// A remote person's body: a VRM 0.x avatar that hangs from their tracked head,
// turns its body to follow where they look, reaches its arms to their tracked
// hands (two-bone IK), opens its mouth while they talk, and blinks. When they
// sit, its hips go on their seat, its back leans toward their head and its
// feet rest on the ground; when they jump, its knees tuck up.
//
// Lives on the remote copy of each person's rig. The head and hands are the
// networked child entities NAF creates inside that rig.
(function () {
  const AVATAR_DIR = 'assets/avatars/';
  const AVATAR_SLUG = /^[a-z0-9-]{1,40}$/;

  // Avatars are scaled so their head bone sits this high when standing.
  const HEAD_BONE_HEIGHT = 1.52;
  // From the eyes (tracked camera) to the head bone's pivot, in camera space.
  const EYES_TO_HEAD_PIVOT = new THREE.Vector3(0, -0.1, 0.07);
  // From the tracked controller grip back to the wrist, in hand space.
  const GRIP_TO_WRIST = new THREE.Vector3(0, -0.01, 0.07);
  // Body turns to follow the head, never lagging more than this (radians).
  const MAX_BODY_TWIST = 0.6;
  // Sitting: the hip joints are this far above the seat, and the back leans
  // at most this much (radians) to bring the head to where theirs is.
  const HIP_ABOVE_SEAT = 0.08;
  const MAX_LEAN = 0.6;

  const UP = new THREE.Vector3(0, 1, 0);
  const v1 = new THREE.Vector3();
  const v2 = new THREE.Vector3();
  const v3 = new THREE.Vector3();
  const cross = new THREE.Vector3();
  const shoulder = new THREE.Vector3();
  const elbow = new THREE.Vector3();
  const wrist = new THREE.Vector3();
  const target = new THREE.Vector3();
  const pole = new THREE.Vector3();
  const pivot = new THREE.Vector3();
  const forward = new THREE.Vector3();
  const headPos = new THREE.Vector3();
  const headQuat = new THREE.Quaternion();
  const bodyQuat = new THREE.Quaternion();
  const q1 = new THREE.Quaternion();
  const q2 = new THREE.Quaternion();
  const leanQuat = new THREE.Quaternion();
  const noRotation = new THREE.Quaternion();

  const loader = new THREE.GLTFLoader();

  function wrapAngle(a) {
    return Math.atan2(Math.sin(a), Math.cos(a));
  }

  function setBoneWorldQuaternion(bone, worldQuat) {
    bone.parent.getWorldQuaternion(q1).invert();
    bone.quaternion.copy(q1).multiply(worldQuat);
    bone.updateMatrixWorld(true);
  }

  function rotateBoneWorld(bone, deltaQuat) {
    bone.getWorldQuaternion(q2).premultiply(deltaQuat);
    setBoneWorldQuaternion(bone, q2);
  }

  // Bend the elbow (or knee) so shoulder->wrist spans the target distance, aim
  // the limb at the target, then roll it so the elbow points toward `poleDir`.
  function solveLimb(upper, lower, end, targetPos, poleDir) {
    upper.getWorldPosition(shoulder);
    lower.getWorldPosition(elbow);
    end.getWorldPosition(wrist);
    const a = shoulder.distanceTo(elbow);
    const b = elbow.distanceTo(wrist);
    const reach = Math.min(Math.max(shoulder.distanceTo(targetPos), Math.abs(a - b) + 1e-3), (a + b) * 0.999);

    // 1. Elbow bend (law of cosines), bending away from the pole.
    const wantInterior = Math.acos(THREE.MathUtils.clamp((a * a + b * b - reach * reach) / (2 * a * b), -1, 1));
    const haveInterior = v1.subVectors(shoulder, elbow).angleTo(v2.subVectors(wrist, elbow));
    const upperDir = v1.subVectors(elbow, shoulder).normalize();
    const bendDir = v2.copy(poleDir).negate();
    bendDir.addScaledVector(upperDir, -bendDir.dot(upperDir)).normalize();
    const axis = v3.crossVectors(upperDir, bendDir).normalize();
    rotateBoneWorld(lower, q1.setFromAxisAngle(axis, haveInterior - wantInterior));

    // 2. Aim the whole limb at the target.
    end.getWorldPosition(wrist);
    v1.subVectors(wrist, shoulder).normalize();
    v2.subVectors(targetPos, shoulder).normalize();
    rotateBoneWorld(upper, q1.setFromUnitVectors(v1, v2));

    // 3. Roll around the shoulder->target axis so the elbow faces the pole.
    const n = v3.subVectors(targetPos, shoulder).normalize();
    lower.getWorldPosition(elbow);
    const toElbow = v1.subVectors(elbow, shoulder);
    toElbow.addScaledVector(n, -toElbow.dot(n));
    const toPole = v2.copy(poleDir).addScaledVector(n, -poleDir.dot(n));
    if (toElbow.lengthSq() > 1e-8 && toPole.lengthSq() > 1e-8) {
      const roll = Math.atan2(n.dot(cross.crossVectors(toElbow, toPole)), toElbow.dot(toPole));
      rotateBoneWorld(upper, q1.setFromAxisAngle(n, roll));
    }
  }

  // No tracked hand (desktop, or controllers set down): let the arm hang.
  function relaxArm(arm, bodyQuaternion) {
    arm.upper.getWorldPosition(shoulder);
    arm.hand.getWorldPosition(wrist);
    v1.subVectors(wrist, shoulder).normalize();
    v2.set(arm.side * 0.22, -1, -0.04).normalize().applyQuaternion(bodyQuaternion);
    rotateBoneWorld(arm.upper, q1.setFromUnitVectors(v1, v2));
  }

  // VRM 0.x stores morph-target "blend shape" presets (a, blink, ...) by glTF mesh index.
  function blendShape(gltf, meshesByIndex, preset) {
    const groups = (gltf.parser.json.extensions.VRM.blendShapeMaster || {}).blendShapeGroups || [];
    const group = groups.find((g) => g.presetName === preset);
    if (!group) return [];
    const binds = [];
    for (const bind of group.binds || []) {
      for (const mesh of meshesByIndex.get(bind.mesh) || []) {
        if (mesh.morphTargetInfluences) binds.push({ mesh, index: bind.index, weight: bind.weight / 100 });
      }
    }
    return binds;
  }

  function setBlendShape(binds, value) {
    for (const b of binds) b.mesh.morphTargetInfluences[b.index] = value * b.weight;
  }

  // Lambert lighting lets the campfire light avatars; a little emissive keeps
  // them readable in the dark.
  function toLambert(material) {
    const lambert = new THREE.MeshLambertMaterial({
      map: material.map || null,
      color: material.color ? material.color.clone() : new THREE.Color(0xffffff),
      emissive: new THREE.Color(0xffffff),
      emissiveMap: material.map || null,
      emissiveIntensity: 0.3,
      transparent: material.transparent,
      alphaTest: material.alphaTest,
      side: material.side,
    });
    material.dispose();
    return lambert;
  }

  // Rest-pose positions of a VRM's humanoid bones, scaled like the drawn
  // avatar, read from the glTF JSON alone.
  function restBones(json) {
    const nodes = json.nodes || [];
    const parents = new Map();
    nodes.forEach((n, i) => (n.children || []).forEach((c) => parents.set(c, i)));
    const world = new Map();
    const worldMatrix = (i) => {
      if (!world.has(i)) {
        const n = nodes[i];
        const m = n.matrix ? new THREE.Matrix4().fromArray(n.matrix) : new THREE.Matrix4().compose(
          new THREE.Vector3().fromArray(n.translation || [0, 0, 0]),
          new THREE.Quaternion().fromArray(n.rotation || [0, 0, 0, 1]),
          new THREE.Vector3().fromArray(n.scale || [1, 1, 1]),
        );
        world.set(i, parents.has(i) ? worldMatrix(parents.get(i)).clone().multiply(m) : m);
      }
      return world.get(i);
    };
    const pos = (name) => {
      const hb = json.extensions.VRM.humanoid.humanBones.find((b) => b.bone === name);
      return hb ? new THREE.Vector3().setFromMatrixPosition(worldMatrix(hb.node)) : null;
    };
    const head = pos('head');
    const left = pos('leftUpperLeg');
    const right = pos('rightUpperLeg');
    if (!head || !left || !right) return null;
    const scale = HEAD_BONE_HEIGHT / Math.max(0.2, head.y);
    return {
      head: head.multiplyScalar(scale),
      hips: left.add(right).multiplyScalar(scale / 2), // between the hip joints
    };
  }

  // Where your eyes go when you sit as this avatar, so that its hips land on
  // the seat: how high above the seat, and how far in front of it.
  const seatedEyesCache = new Map();
  function seatedEyes(slug) {
    if (!AVATAR_SLUG.test(slug || '')) return Promise.resolve(null);
    if (!seatedEyesCache.has(slug)) {
      seatedEyesCache.set(slug, fetch(AVATAR_DIR + slug + '.vrm')
        .then((res) => {
          if (!res.ok) throw new Error(`HTTP ${res.status}`);
          return res.arrayBuffer();
        })
        .then((buf) => {
          // GLB: 12-byte header, then the JSON chunk's length, type, and bytes.
          const length = new DataView(buf).getUint32(12, true);
          const bones = restBones(JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 20, length))));
          if (!bones) return null;
          return {
            height: THREE.MathUtils.clamp(HIP_ABOVE_SEAT + bones.head.y - bones.hips.y - EYES_TO_HEAD_PIVOT.y, 0.55, 1.25),
            forward: bones.hips.z - bones.head.z + EYES_TO_HEAD_PIVOT.z, // the avatar faces -Z
          };
        })
        .catch((err) => {
          console.warn(`Couldn't measure avatar "${slug}"`, err);
          return null;
        }));
    }
    return seatedEyesCache.get(slug);
  }

  window.VrmBody = { seatedEyes };

  AFRAME.registerComponent('vrm-avatar', {
    init() {
      this.root = null;
      this.slug = '';
      this.loadToken = 0;
      this.bodyYaw = null;
      this.nextBlink = 1 + Math.random() * 3;
      this.blinkTime = -1;
      this.seatIndex = -1;
      this.tuck = 0;

      this.onInfo = (evt) => {
        this.seatIndex = evt.detail.seat;
        this.setAvatar(evt.detail.avatar);
      };
      this.el.addEventListener('player-info-changed', this.onInfo);
      const info = this.el.getAttribute('player-info');
      if (info) {
        this.seatIndex = info.seat;
        if (info.avatar) this.setAvatar(info.avatar);
      }
    },

    remove() {
      this.el.removeEventListener('player-info-changed', this.onInfo);
      this.loadToken++;
      this.disposeAvatar();
    },

    setAvatar(slug) {
      if (!AVATAR_SLUG.test(slug || '') || slug === this.slug) return;
      this.slug = slug;
      const token = ++this.loadToken;
      loader.load(AVATAR_DIR + slug + '.vrm', (gltf) => {
        if (token !== this.loadToken || !this.el.isConnected) return;
        this.disposeAvatar();
        this.buildAvatar(gltf);
      }, undefined, (err) => console.warn(`Couldn't load avatar "${slug}"`, err));
    },

    buildAvatar(gltf) {
      const vrm = gltf.parser.json.extensions && gltf.parser.json.extensions.VRM;
      if (!vrm) {
        console.warn(`Avatar "${this.slug}" is not a VRM 0.x file`);
        return;
      }

      const nodes = new Map();
      const meshesByIndex = new Map();
      gltf.scene.traverse((obj) => {
        const assoc = gltf.parser.associations.get(obj);
        if (assoc && assoc.nodes !== undefined) nodes.set(assoc.nodes, obj);
        if (obj.isMesh) {
          obj.frustumCulled = false; // skinned bounds don't follow the bones
          obj.material = Array.isArray(obj.material) ? obj.material.map(toLambert) : toLambert(obj.material);
          if (assoc && assoc.meshes !== undefined) {
            if (!meshesByIndex.has(assoc.meshes)) meshesByIndex.set(assoc.meshes, []);
            meshesByIndex.get(assoc.meshes).push(obj);
          }
        }
      });
      const bone = (name) => {
        const hb = vrm.humanoid.humanBones.find((b) => b.bone === name);
        return hb ? nodes.get(hb.node) : null;
      };

      const root = new THREE.Group();
      root.add(gltf.scene);
      this.el.sceneEl.object3D.add(root);
      root.updateMatrixWorld(true);

      this.head = bone('head');
      if (!this.head) {
        console.warn(`Avatar "${this.slug}" has no head bone`);
        this.el.sceneEl.object3D.remove(root);
        return;
      }

      // Rest pose, measured with the root at the origin (VRM 0.x faces -Z, like the camera).
      this.headRestPos = this.head.getWorldPosition(new THREE.Vector3());
      this.headRestQuat = this.head.getWorldQuaternion(new THREE.Quaternion());
      this.scale = HEAD_BONE_HEIGHT / Math.max(0.2, this.headRestPos.y);

      this.arms = [['left', -1], ['right', 1]].map(([side, sign]) => {
        const upper = bone(`${side}UpperArm`);
        const lower = bone(`${side}LowerArm`);
        const hand = bone(`${side}Hand`);
        if (!upper || !lower || !hand) return null;
        return {
          name: side, side: sign, upper, lower, hand,
          upperRest: upper.quaternion.clone(),
          lowerRest: lower.quaternion.clone(),
          handRestScale: hand.scale.clone(),
        };
      }).filter(Boolean);

      this.legs = ['left', 'right'].map((side) => {
        const upper = bone(`${side}UpperLeg`);
        const lower = bone(`${side}LowerLeg`);
        const foot = bone(`${side}Foot`);
        if (!upper || !lower || !foot) return null;
        const hip = upper.getWorldPosition(new THREE.Vector3());
        const knee = lower.getWorldPosition(new THREE.Vector3());
        const ankle = foot.getWorldPosition(new THREE.Vector3());
        return {
          upper, lower, foot,
          upperRest: upper.quaternion.clone(),
          lowerRest: lower.quaternion.clone(),
          thigh: hip.distanceTo(knee),
          length: hip.distanceTo(knee) + knee.distanceTo(ankle),
          ankleHeight: ankle.y,
        };
      }).filter(Boolean);
      this.spine = bone('spine');
      this.spineRest = this.spine ? this.spine.quaternion.clone() : null;
      const bones = restBones(gltf.parser.json);
      this.hipsRestPos = bones ? bones.hips.divideScalar(this.scale) : null;

      this.mouth = blendShape(gltf, meshesByIndex, 'a');
      this.blink = blendShape(gltf, meshesByIndex, 'blink');
      this.root = root;
      this.root.visible = false;
    },

    disposeAvatar() {
      if (!this.root) return;
      this.el.sceneEl.object3D.remove(this.root);
      this.root.traverse((obj) => {
        if (!obj.isMesh) return;
        obj.geometry.dispose();
        for (const m of [].concat(obj.material)) {
          if (m.map) m.map.dispose();
          m.dispose();
        }
      });
      this.root = null;
    },

    // The seat they're on, if any.
    currentSeat() {
      if (this.seatIndex < 0 || !this.hipsRestPos) return null;
      const world = CurrentWorld(this.el.sceneEl);
      return (world && world.seats[this.seatIndex]) || null;
    },

    findChild(predicate) {
      for (const child of this.el.children) {
        if (predicate(child)) return child;
      }
      return null;
    },

    findHand(side) {
      return this.findChild((c) => {
        const hc = c.components && c.components['networked-hand-controls'];
        return hc && hc.data.hand === side;
      });
    },

    tick(time, delta) {
      if (!this.root) return;
      const headEl = this.findChild((c) => c.classList.contains('remote-head'));
      if (!headEl) {
        this.root.visible = false;
        return;
      }
      this.root.visible = true;

      const dt = Math.min(delta, 100) / 1000;
      const seat = this.currentSeat();

      // Body: face the fire when seated, otherwise turn toward where they're
      // looking, lazily.
      headEl.object3D.getWorldPosition(headPos);
      headEl.object3D.getWorldQuaternion(headQuat);
      v1.set(0, 0, -1).applyQuaternion(headQuat);
      const lookYaw = Math.atan2(-v1.x, -v1.z);
      if (seat) {
        this.bodyYaw = seat.yaw;
      } else {
        if (this.bodyYaw === null) this.bodyYaw = lookYaw;
        let twist = wrapAngle(lookYaw - this.bodyYaw);
        this.bodyYaw += twist * Math.min(1, dt * 2.5);
        twist = wrapAngle(lookYaw - this.bodyYaw);
        if (Math.abs(twist) > MAX_BODY_TWIST) this.bodyYaw += twist - Math.sign(twist) * MAX_BODY_TWIST;
      }
      bodyQuat.setFromAxisAngle(UP, this.bodyYaw);

      // Hang the avatar from the head pivot, or sit its hips on the seat.
      pivot.copy(EYES_TO_HEAD_PIVOT).applyQuaternion(headQuat).add(headPos);
      if (seat) {
        v3.copy(this.hipsRestPos).multiplyScalar(this.scale).applyQuaternion(bodyQuat);
        this.root.position.set(seat.x, seat.y + HIP_ABOVE_SEAT, seat.z).sub(v3);
      } else {
        v3.copy(this.headRestPos).multiplyScalar(this.scale).applyQuaternion(bodyQuat);
        this.root.position.subVectors(pivot, v3);
      }
      this.root.quaternion.copy(bodyQuat);
      this.root.scale.setScalar(this.scale);
      for (const limb of this.arms.concat(this.legs)) {
        limb.upper.quaternion.copy(limb.upperRest);
        limb.lower.quaternion.copy(limb.lowerRest);
      }
      if (this.spine) this.spine.quaternion.copy(this.spineRest);
      this.root.updateMatrixWorld(true);

      // Seated: lean the back so the head goes where their head is.
      if (seat && this.spine) {
        this.spine.getWorldPosition(v1);
        this.head.getWorldPosition(v2).sub(v1).normalize();
        v3.subVectors(pivot, v1).normalize();
        leanQuat.setFromUnitVectors(v2, v3);
        const lean = 2 * Math.acos(Math.min(1, Math.abs(leanQuat.w)));
        if (lean > MAX_LEAN) leanQuat.slerp(noRotation, 1 - MAX_LEAN / lean);
        rotateBoneWorld(this.spine, leanQuat);
      }

      // Head: same rotation as their real head.
      setBoneWorldQuaternion(this.head, q2.copy(headQuat).multiply(this.headRestQuat));

      // Arms: reach for tracked hands, otherwise hang.
      for (const arm of this.arms) {
        const handEl = this.findHand(arm.name);
        if (handEl && handEl.object3D.visible) {
          handEl.object3D.getWorldQuaternion(q2);
          handEl.object3D.getWorldPosition(target);
          target.add(v1.copy(GRIP_TO_WRIST).applyQuaternion(q2));
          pole.set(arm.side * 0.5, -1, 0.6).normalize().applyQuaternion(bodyQuat);
          solveLimb(arm.upper, arm.lower, arm.hand, target, pole);
          arm.hand.scale.setScalar(0.001); // their networked hand model shows instead
        } else {
          relaxArm(arm, bodyQuat);
          arm.hand.scale.copy(arm.handRestScale);
        }
        arm.hand.updateMatrixWorld(true);
      }

      // Legs: feet on the ground in front of the seat, knees tucked up while
      // they're in the air (their rig only rises off the ground in a jump),
      // otherwise straight down.
      const world = CurrentWorld(this.el.sceneEl);
      const rigPos = this.el.object3D.position;
      const floor = world && world.floorAt ? world.floorAt(seat ? seat.x : rigPos.x, seat ? seat.z : rigPos.z) : 0;
      const airborne = !seat && rigPos.y > floor + 0.02;
      this.tuck += ((airborne ? 1 : 0) - this.tuck) * Math.min(1, dt * 12);
      if (seat || this.tuck > 0.01) {
        forward.set(0, 0, -1).applyQuaternion(bodyQuat);
        pole.set(0, 0.3, -1).normalize().applyQuaternion(bodyQuat);
        for (const leg of this.legs) {
          leg.foot.getWorldPosition(target);
          if (seat) {
            target.addScaledVector(forward, leg.thigh * this.scale);
            target.y = floor + leg.ankleHeight * this.scale;
          } else {
            target.addScaledVector(forward, 0.2 * leg.length * this.scale * this.tuck);
            target.y += 0.35 * leg.length * this.scale * this.tuck;
          }
          solveLimb(leg.upper, leg.lower, leg.foot, target, pole);
        }
      }

      // Face: talk and blink.
      const voice = headEl.components['voice-level'];
      setBlendShape(this.mouth, voice ? voice.level : 0);
      const t = time / 1000;
      if (this.blinkTime < 0 && t > this.nextBlink) this.blinkTime = t;
      if (this.blinkTime >= 0) {
        const p = (t - this.blinkTime) / 0.16;
        setBlendShape(this.blink, p < 1 ? Math.sin(p * Math.PI) : 0);
        if (p >= 1) {
          this.blinkTime = -1;
          this.nextBlink = t + 2 + Math.random() * 4;
        }
      }
    },
  });
})();