opendance-web/web-server/src/preprocess/mmd-pose.ts

import { Mat3, Transform, Vec3, applyTransform, matVec, quatToMat3, vAdd, vLen, vNorm, vScale, vSub } from './math3d';
import { LEG_IK, LEG_LENGTHS, MMD_SKELETON, MmdBoneDef } from './mmd-skeleton';
import { VmdMotion, VmdBoneTrack, sampleBoneTrack, trackHasPositionMotion } from './vmd-parser';

/**
 * VMD 数值的解释方式。
 *
 * MMD 使用左手坐标系，而本项目的重定向通路（与 BVH 一致）是右手系。
 * 两种解释只有在真实数据上跑一遍才能判断哪种是"物理正确"的：
 * `zMirrored` 相当于把动画在 Z 轴上做一次镜像（旋转四元数取
 * (x,y,z,w) → (-x,-y,z,w)，位移 z 取反）。
 * 见 mmd-vmd.spec.ts 中基于真实 VMD 的"膝盖朝前"对照实验。
 */
export type MmdRotationMode = 'asIs' | 'zMirrored';

export interface MmdPoseOptions {
  rotationMode?: MmdRotationMode;
  /** 采样帧率，默认 VMD 固定的 30 */
  fps?: number;
  /** 最大采样帧数，用于防御异常文件（默认 30000 帧 ≈ 16 分钟） */
  maxFrames?: number;
  /** 是否启用腿部 IK 解算（默认自动判断；诊断与测试可强制关闭） */
  legIk?: boolean;
}

/** 单帧解算结果：boneName（标准名）→ 世界变换 */
export type Pose = Map<string, Transform>;

const IDENTITY_TRANSFORM: Transform = { r: [1, 0, 0, 0, 1, 0, 0, 0, 1], t: [0, 0, 0] };

/** S·R·S，S = diag(1,1,-1)：把左手系旋转搬到右手系 */
function mirrorRotationZ(r: Mat3): Mat3 {
  return [
    r[0], r[1], -r[2],
    r[3], r[4], -r[5],
    -r[6], -r[7], r[8],
  ];
}

function convertRotation(
  quat: readonly [number, number, number, number],
  mode: MmdRotationMode,
): Mat3 {
  const matrix = quatToMat3(quat);
  return mode === 'zMirrored' ? mirrorRotationZ(matrix) : matrix;
}

function convertPosition(position: Vec3, mode: MmdRotationMode): Vec3 {
  return mode === 'zMirrored' ? [position[0], position[1], -position[2]] : [...position] as Vec3;
}

/** 骨架的拓扑顺序（父骨骼一定排在子骨骼之前） */
const SKELETON_ORDER: readonly MmdBoneDef[] = (() => {
  const remaining = [...MMD_SKELETON];
  const ordered: MmdBoneDef[] = [];
  const done = new Set<string>();
  while (remaining.length > 0) {
    const index = remaining.findIndex((bone) => bone.parent === null || done.has(bone.parent));
    if (index < 0) throw new Error('MMD 标准骨架存在环，无法拓扑排序');
    const [bone] = remaining.splice(index, 1);
    ordered.push(bone);
    done.add(bone.name);
  }
  return ordered;
})();

/**
 * 解算单帧：标准骨架 FK + 腿部 IK。
 *
 * MMD 的位移通道是**世界轴**上的偏移（不被父骨骼旋转带走），
 * 因此平移单独叠加：`t = parentT + parentR * restOffset + animPos`。
 */
export function solveFrame(
  motion: VmdMotion,
  frame: number,
  mode: MmdRotationMode,
  useLegIk: { 左: boolean; 右: boolean },
): Pose {
  const pose: Pose = new Map();

  for (const bone of SKELETON_ORDER) {
    const parent = bone.parent ? pose.get(bone.parent) ?? IDENTITY_TRANSFORM : IDENTITY_TRANSFORM;
    const track: VmdBoneTrack | undefined = motion.tracks.get(bone.name);
    const sampled = track ? sampleBoneTrack(track, frame) : null;

    const localR = sampled
      ? convertRotation(sampled.rotation, mode)
      : IDENTITY_TRANSFORM.r;
    const animPos = sampled ? convertPosition(sampled.position, mode) : ([0, 0, 0] as Vec3);

    pose.set(bone.name, {
      r: applyTransform(parent, bone.offset, localR).r,
      t: vAdd(vAdd(parent.t, matVec(parent.r, bone.offset)), animPos),
    });
  }

  for (const side of ['左', '右'] as const) {
    if (!useLegIk[side]) continue;
    solveLegIk(pose, side);
  }

  return pose;
}

/**
 * 两骨骼解析 IK：把小腿末端拉到 IK 目标点，膝盖朝 +Z（角色正面）弯。
 * 目标超出可达范围时按比例收敛到伸直/极限，不产生 NaN。
 */
function solveLegIk(pose: Pose, side: '左' | '右'): void {
  const names = LEG_IK[side];
  const hip = pose.get(`${side}足`);
  const ik = pose.get(names.ik);
  if (!hip || !ik) return;

  const upper = LEG_LENGTHS.thigh;
  const lower = LEG_LENGTHS.shin;
  const hipPos = hip.t;
  const target = ik.t;

  let delta = vSub(target, hipPos);
  let distance = vLen(delta);
  const minReach = Math.abs(upper - lower) + 1e-4;
  const maxReach = upper + lower - 1e-4;

  if (distance < 1e-6) {
    // 退化：目标与髋重合，直接落地
    delta = [0, -1, 0];
    distance = 1;
  }
  const axis = vNorm(delta);
  let ankle = target;
  if (distance > maxReach) {
    ankle = vAdd(hipPos, vScale(axis, maxReach));
    distance = maxReach;
  } else if (distance < minReach) {
    ankle = vAdd(hipPos, vScale(axis, minReach));
    distance = minReach;
  }

  // 髋到膝在髋-踝连线上的投影长度
  const along = (upper * upper - lower * lower + distance * distance) / (2 * distance);
  const heightSq = Math.max(0, upper * upper - along * along);
  const height = Math.sqrt(heightSq);

  // 膝盖参考方向：优先 +Z（正面），与连线正交化；退化时退回 +Y
  const forward: Vec3 = [0, 0, 1];
  let perp = vSub(forward, vScale(axis, forward[0] * axis[0] + forward[1] * axis[1] + forward[2] * axis[2]));
  if (vLen(perp) < 1e-4) {
    const up: Vec3 = [0, 1, 0];
    perp = vSub(up, vScale(axis, up[0] * axis[0] + up[1] * axis[1] + up[2] * axis[2]));
  }
  perp = vNorm(perp);

  const kneePos = vAdd(vAdd(hipPos, vScale(axis, along)), vScale(perp, height));

  const thighDir = vNorm(vSub(kneePos, hipPos));
  const shinDir = vNorm(vSub(ankle, kneePos));
  const knee = pose.get(names.knee);
  const ankleBone = pose.get(names.ankle);

  if (knee) {
    knee.t = kneePos;
    if (vLen(thighDir) > 1e-6) {
      // 让膝的旋转与"指向踝"的方向一致，供下游（如趾尖）使用
      knee.r = alignRotation(knee.r, thighDir);
    }
  }
  if (ankleBone) {
    ankleBone.t = ankle;
    // MMD 把 IK 骨骼的旋转直接赋给足首
    ankleBone.r = ik.r;
    if (vLen(shinDir) > 1e-6) ankleBone.r = alignRotation(ik.r, shinDir);
  }
  const toe = pose.get(names.toe);
  if (toe && ankleBone) {
    const restOffset: Vec3 = [0, -0.2, LEG_LENGTHS.toe > 0 ? 1.3 : 1.3];
    toe.t = vAdd(ankle, matVec(ankleBone.r, restOffset));
    toe.r = ankleBone.r;
  }
}

/**
 * 让骨骼的 Y 轴（MMD 骨骼默认朝向子骨骼的方向）对齐给定方向，
 * 保留绕该方向的扭转分量。用于 IK 解算后给出合理的局部朝向。
 */
function alignRotation(base: Mat3, direction: Vec3): Mat3 {
  const dir = vNorm(direction);
  if (vLen(dir) < 1e-6) return base;
  // 取一个与 dir 不平行的参考轴构造正交基（列优先：X, Y=dir, Z）
  const reference: Vec3 = Math.abs(dir[1]) > 0.9 ? [1, 0, 0] : [0, 1, 0];
  const xAxis = vNorm([
    reference[1] * dir[2] - reference[2] * dir[1],
    reference[2] * dir[0] - reference[0] * dir[2],
    reference[0] * dir[1] - reference[1] * dir[0],
  ]);
  if (vLen(xAxis) < 1e-6) return base;
  const zAxis: Vec3 = [
    xAxis[1] * dir[2] - xAxis[2] * dir[1],
    xAxis[2] * dir[0] - xAxis[0] * dir[2],
    xAxis[0] * dir[1] - xAxis[1] * dir[0],
  ];
  return [
    xAxis[0], dir[0], zAxis[0],
    xAxis[1], dir[1], zAxis[1],
    xAxis[2], dir[2], zAxis[2],
  ];
}

/** 关节映射表使用的骨骼名（与 joint-map.ts 的候选名一致） */
const HEAD_KEY = 'head';
const JOINT_KEYS: ReadonlyArray<{ key: string; bone: string; offset?: Vec3 }> = [
  { key: HEAD_KEY, bone: '頭', offset: [0, 0.6, 0] },
  { key: 'leftarm', bone: '左腕' },
  { key: 'rightarm', bone: '右腕' },
  { key: 'leftforearm', bone: '左ひじ' },
  { key: 'rightforearm', bone: '右ひじ' },
  { key: 'lefthand', bone: '左手首' },
  { key: 'righthand', bone: '右手首' },
  { key: 'leftupleg', bone: '左足' },
  { key: 'rightupleg', bone: '右足' },
  { key: 'leftleg', bone: '左ひざ' },
  { key: 'rightleg', bone: '右ひざ' },
  { key: 'leftfoot', bone: '左足首' },
  { key: 'rightfoot', bone: '右足首' },
  { key: 'lefttoebase', bone: '左つま先' },
  { key: 'righttoebase', bone: '右つま先' },
];

/** 把一帧的姿态转成重定向模块使用的"骨骼名 → 世界坐标"表 */
function poseToWorld(pose: Pose): Map<string, Vec3> {
  const world = new Map<string, Vec3>();
  for (const { key, bone, offset } of JOINT_KEYS) {
    const transform = pose.get(bone);
    if (!transform) continue;
    world.set(key, offset ? vAdd(transform.t, offset) : transform.t);
  }
  return world;
}

/**
 * VMD → 逐帧世界坐标（键名与 `joint-map.ts` 的候选名一致）。
 * 输出可直接交给 `worldFramesToKeypointFrames` 复用 BVH 通路的重定向。
 */
export function vmdToWorldFrames(
  motion: VmdMotion,
  options: MmdPoseOptions = {},
): { frames: Map<string, Vec3>[]; fps: number } {
  const mode = options.rotationMode ?? 'zMirrored';
  const fps = options.fps && options.fps > 0 ? options.fps : 30;
  const maxFrames = options.maxFrames ?? 30000;
  const frameCount = Math.min(motion.maxFrame + 1, maxFrames);

  // 腿部是否由 IK 驱动：IK 目标有位移 → 用 IK；否则用 FK 旋转（部分动作数据是烘焙好的 FK）
  const autoIk = {
    左: trackHasPositionMotion(motion.tracks.get(LEG_IK.左.ik)),
    右: trackHasPositionMotion(motion.tracks.get(LEG_IK.右.ik)),
  };
  const useLegIk =
    options.legIk === false ? { 左: false, 右: false } : autoIk;

  const frames: Map<string, Vec3>[] = [];
  for (let frame = 0; frame < frameCount; frame += 1) {
    frames.push(poseToWorld(solveFrame(motion, frame, mode, useLegIk)));
  }
  return { frames, fps };
}

/** 诊断用：返回本帧被 IK 驱动的腿（测试与日志使用） */
export function legIkDriven(motion: VmdMotion): { 左: boolean; 右: boolean } {
  return {
    左: trackHasPositionMotion(motion.tracks.get(LEG_IK.左.ik)),
    右: trackHasPositionMotion(motion.tracks.get(LEG_IK.右.ik)),
  };
}