SteelXR/src/components/three/viewCubeBus.ts

import * as THREE from 'three';

/**
 * Shared bus between the main 3D viewer and the floating ViewCube widget.
 * - mainCameraRef / mainControlsRef: written by ModelViewer each frame
 * - viewAnim: smooth camera animation when a face is clicked
 * - calibration: orchestrates a 4-click (+ optional 2-click verify) procedure
 *   to align the active model's intrinsic axes to the world axes.
 */
export const mainCameraRef: { current: THREE.Camera | null } = { current: null };
export const mainControlsRef: { current: any | null } = { current: null };

// ── View animation ────────────────────────────────────────────────
export interface ViewAnim {
  active: boolean;
  startPos: THREE.Vector3;
  endPos: THREE.Vector3;
  startUp: THREE.Vector3;
  endUp: THREE.Vector3;
  t: number;
  duration: number;
}

export const viewAnim: ViewAnim = {
  active: false,
  startPos: new THREE.Vector3(),
  endPos: new THREE.Vector3(),
  startUp: new THREE.Vector3(0, 1, 0),
  endUp: new THREE.Vector3(0, 1, 0),
  t: 0,
  duration: 0.5,
};

export function requestView(dir: THREE.Vector3) {
  const cam = mainCameraRef.current;
  const controls = mainControlsRef.current;
  if (!cam || !controls) return;
  const target: THREE.Vector3 = controls.target ?? new THREE.Vector3();
  const distance = cam.position.distanceTo(target);
  const d = dir.clone().normalize();
  const endPos = target.clone().add(d.multiplyScalar(distance));
  const up = new THREE.Vector3(0, 1, 0);
  if (Math.abs(dir.y) > 0.99) {
    up.set(0, 0, dir.y > 0 ? -1 : 1);
  }
  viewAnim.startPos.copy(cam.position);
  viewAnim.endPos.copy(endPos);
  viewAnim.startUp.copy(cam.up);
  viewAnim.endUp.copy(up);
  viewAnim.t = 0;
  viewAnim.active = true;
}

// ── Calibration state machine ─────────────────────────────────────
export type CalibrationStep =
  | 'idle'
  | 'await-cube-1'
  | 'await-model-1'
  | 'await-cube-2'
  | 'await-model-2'
  | 'await-cube-3'   // optional verification
  | 'await-model-3'
  | 'done';

interface Pair { cube: THREE.Vector3; model: THREE.Vector3 }

interface CalState {
  step: CalibrationStep;
  modelId: string | null;
  pairs: Pair[];
  pendingCube: THREE.Vector3 | null;
  /** 0..1 progress for the button visual. */
  progress: number;
  /** Last verification angular error, in degrees. NaN until verify pair captured. */
  verifyErrorDeg: number;
  /** Listener notified on every state change so the UI re-renders. */
  listeners: Set<() => void>;
}

export const calibration: CalState = {
  step: 'idle',
  modelId: null,
  pairs: [],
  pendingCube: null,
  progress: 0,
  verifyErrorDeg: NaN,
  listeners: new Set(),
};

function notify() { calibration.listeners.forEach(fn => { try { fn(); } catch {} }); }

export function subscribeCalibration(fn: () => void): () => void {
  calibration.listeners.add(fn);
  return () => { calibration.listeners.delete(fn); };
}

export function startCalibration(modelId: string) {
  calibration.step = 'await-cube-1';
  calibration.modelId = modelId;
  calibration.pairs = [];
  calibration.pendingCube = null;
  calibration.progress = 0;
  calibration.verifyErrorDeg = NaN;
  notify();
}

export function cancelCalibration() {
  calibration.step = 'idle';
  calibration.modelId = null;
  calibration.pairs = [];
  calibration.pendingCube = null;
  calibration.progress = 0;
  calibration.verifyErrorDeg = NaN;
  notify();
}

/** Called when user clicks a cube face during calibration. */
export function pushCubeFace(dirWorld: THREE.Vector3) {
  if (calibration.step === 'await-cube-1') {
    calibration.pendingCube = dirWorld.clone().normalize();
    calibration.step = 'await-model-1';
  } else if (calibration.step === 'await-cube-2') {
    calibration.pendingCube = dirWorld.clone().normalize();
    calibration.step = 'await-model-2';
  } else if (calibration.step === 'await-cube-3') {
    calibration.pendingCube = dirWorld.clone().normalize();
    calibration.step = 'await-model-3';
  } else {
    return;
  }
  // Auto-rotate the main camera to look straight at the picked cube face
  // so the user sees the corresponding model face head-on (especially helpful
  // in orthographic mode where this guarantees a perpendicular click).
  try { requestView(dirWorld.clone()); } catch {}
  notify();
}

/** Snap a unit vector to the nearest principal axis (±X/±Y/±Z) when within
 *  `maxDeg` degrees of it. Returns the snapped vector or the input unchanged. */
function snapToPrincipalAxis(v: THREE.Vector3, maxDeg = 12): THREE.Vector3 {
  const cosT = Math.cos(THREE.MathUtils.degToRad(maxDeg));
  const axes = [
    new THREE.Vector3(1, 0, 0), new THREE.Vector3(-1, 0, 0),
    new THREE.Vector3(0, 1, 0), new THREE.Vector3(0, -1, 0),
    new THREE.Vector3(0, 0, 1), new THREE.Vector3(0, 0,-1),
  ];
  let best = v;
  let bestDot = cosT;
  for (const a of axes) {
    const d = v.dot(a);
    if (d > bestDot) { bestDot = d; best = a.clone(); }
  }
  return best;
}

/**
 * Called when user clicks a model face during calibration.
 * @param normalWorld world-space face normal of the clicked face
 * @param calGroupWorldQuat current world quaternion of the calibration group
 *   (used to map both cube and model directions into the same local frame).
 */
export function pushModelFaceNormal(normalWorld: THREE.Vector3, calGroupWorldQuat: THREE.Quaternion) {
  if (!calibration.pendingCube) return;
  const cubeWorld = calibration.pendingCube;
  const inv = calGroupWorldQuat.clone().invert();
  const cubeLocal = cubeWorld.clone().applyQuaternion(inv).normalize();
  const modelLocalRaw = normalWorld.clone().applyQuaternion(inv).normalize();
  // Snap near-axial picks (chamfered or slightly off faces) to the closest
  // principal axis. This dramatically improves robustness when picking faces
  // from an orthographic head-on view.
  const modelLocal = snapToPrincipalAxis(modelLocalRaw, 12);
  calibration.pairs.push({ cube: cubeLocal, model: modelLocal });
  calibration.pendingCube = null;

  if (calibration.step === 'await-model-1') {
    calibration.step = 'await-cube-2';
    calibration.progress = 0.33;
  } else if (calibration.step === 'await-model-2') {
    calibration.progress = 0.75;
    calibration.step = 'await-cube-3';
  } else if (calibration.step === 'await-model-3') {
    calibration.progress = 1;
    calibration.verifyErrorDeg = computeVerifyError(calibration.pairs);
    calibration.step = 'done';
  }
  notify();
}

/** Computes a quaternion R such that R*pairs[i].model ≈ pairs[i].cube using
 *  Gram-Schmidt to ortho-normalize both bases. Returns null if degenerate. */
export function computeCalibrationQuaternion(pairs: Pair[]): THREE.Quaternion | null {
  if (pairs.length < 2) return null;
  const n1 = pairs[0].model.clone().normalize();
  const n2raw = pairs[1].model.clone();
  let n2 = n2raw.sub(n1.clone().multiplyScalar(n2raw.dot(n1)));
  if (n2.lengthSq() < 1e-6) return null;
  n2.normalize();
  const n3 = new THREE.Vector3().crossVectors(n1, n2);

  const c1 = pairs[0].cube.clone().normalize();
  const c2raw = pairs[1].cube.clone();
  let c2 = c2raw.sub(c1.clone().multiplyScalar(c2raw.dot(c1)));
  if (c2.lengthSq() < 1e-6) return null;
  c2.normalize();
  const c3 = new THREE.Vector3().crossVectors(c1, c2);

  // R = C * Nᵀ where columns of N are (n1,n2,n3) and C are (c1,c2,c3).
  // In three.js Matrix4: makeBasis sets columns.
  const N = new THREE.Matrix4().makeBasis(n1, n2, n3);
  const C = new THREE.Matrix4().makeBasis(c1, c2, c3);
  const Nt = N.clone().transpose();
  const R = new THREE.Matrix4().multiplyMatrices(C, Nt);
  const q = new THREE.Quaternion().setFromRotationMatrix(R);
  return q;
}

/** Predicted angular error of the verify pair, in degrees. */
export function computeVerifyError(pairs: Pair[]): number {
  if (pairs.length < 3) return NaN;
  const q = computeCalibrationQuaternion(pairs.slice(0, 2));
  if (!q) return NaN;
  const predicted = pairs[2].model.clone().applyQuaternion(q).normalize();
  const target = pairs[2].cube.clone().normalize();
  const cos = THREE.MathUtils.clamp(predicted.dot(target), -1, 1);
  return THREE.MathUtils.radToDeg(Math.acos(cos));
}