engine/src/session.js

// session.js — layer 2: the one object a UI talks to.
//
//   const s = createSession();
//   s.on('stage',   e => ...)   // { stage, done, total, detail }
//   s.on('metrics', e => ...)   // { iter, splats, itersPerSec, psnrTrain, psnrHold }
//   s.on('event',   e => ...)   // { kind: 'refine' | 'train-complete', ... }
//   s.on('log',     m => ...)   // prose, for consoles
//
//   await s.load(files);        // File[] / Blob[] -> Frames
//   await s.solve();            // SfM: cameras + sparse points
//   await s.seed();             // Gaussians + WebGPU trainer
//   s.start(); s.pause();       // the training loop (policy lives HERE)
//
//   s.view.attach(canvas);      // render target
//   s.view.lookThrough(i);      // camera = training frame i
//   s.view.setCamera(pose);     // free camera
//   await s.exportPlyBlob();    // standard 3DGS .ply (opacity comp baked)
//
// All numeric policy that used to live in the demo page is here: iteration
// batching, auto-stop, refine cadence, blur exclusion, holdout choice, PSNR
// conversion, and the hidden-tab watchdog.

import { decodeFrames } from './io/frames.js';
import { isEquirect, sliceEquirect, faceSizeFor, probeImageSize, FACE_ROTS } from './io/pano.js';
import { runSfM } from './sfm/sfm.js';
import { initGaussians } from './gs/init.js';
import { GSTrainer } from './gs/trainer.js';
import { createGpu } from './gpu/context.js';
import { gaussiansToPly, bakeOpacityCompensation } from './io/ply.js';

/** world-space position of a camera pose {R, t} */
export function camPosition({ R, t }) {
  return [
    -(R[0] * t[0] + R[3] * t[1] + R[6] * t[2]),
    -(R[1] * t[0] + R[4] * t[1] + R[7] * t[2]),
    -(R[2] * t[0] + R[5] * t[1] + R[8] * t[2]),
  ];
}

/** Undistort training images in place (bilinear resample) so the pinhole
 *  trainer matches BA's distortion-corrected geometry. Returns true when a
 *  resample actually happened (|k| above the noise floor). Out-of-frame
 *  samples get a -1 sentinel the trainer excludes from the loss. */
export function undistortFrames(frames, recon) {
  const { k1, k2 } = recon;
  // below ~0.01 the estimated distortion is noise (synthetic GT k1=0 comes
  // back as ~0.005) and resampling would only soften the targets
  if (Math.abs(k1) < 0.01 && Math.abs(k2) < 0.01) return false;
  for (const im of frames) {
    const f = recon.fFeat * (im.tw / im.fw);
    const cx = im.tw / 2, cy = im.th / 2;
    const src = im.rgb;
    const dst = new Float32Array(src.length);
    for (let y = 0; y < im.th; y++) {
      for (let x = 0; x < im.tw; x++) {
        const xp = (x + 0.5 - cx) / f, yp = (y + 0.5 - cy) / f;
        const r2 = xp * xp + yp * yp;
        const D = 1 + k1 * r2 + k2 * r2 * r2;
        const rx = f * xp * D + cx - 0.5;
        const ry = f * yp * D + cy - 0.5;
        const o = (y * im.tw + x) * 3;
        if (rx < 0 || ry < 0 || rx > im.tw - 1.001 || ry > im.th - 1.001) {
          dst[o] = -1; dst[o + 1] = -1; dst[o + 2] = -1;
          continue;
        }
        const x0 = rx | 0, y0 = ry | 0;
        const fx = rx - x0, fy = ry - y0;
        const i00 = (y0 * im.tw + x0) * 3, i01 = i00 + 3;
        const i10 = i00 + im.tw * 3, i11 = i10 + 3;
        for (let c = 0; c < 3; c++) {
          dst[o + c] =
            src[i00 + c] * (1 - fx) * (1 - fy) + src[i01 + c] * fx * (1 - fy) +
            src[i10 + c] * (1 - fx) * fy + src[i11 + c] * fx * fy;
        }
      }
    }
    im.rgb = dst;
  }
  return true;
}

/**
 * @typedef {object} SessionOptions
 * @property {GPUDevice} [device]        host-owned WebGPU device to share
 * @property {number} [maxIters=60000]   auto-stop; trainer schedules scale to it
 * @property {number} [itersPerFrame=15] training batch per animation frame
 * @property {number} [initTarget=60000] initial Gaussian count (grows during training)
 * @property {number|'auto'|null} [holdout='auto']  frame excluded from training
 *   and scored as the novel-view metric; 'auto' picks a sharp mid-sequence frame
 * @property {number} [evalHoldEvery=4000]  holdout PSNR cadence (iterations)
 * @property {number} [evalSplit=0]  the standard benchmark protocol: every Nth
 *   frame of the input order is excluded from training and scored together via
 *   evalTestPsnr() (papers use N=8); 0 = off
 * @property {object} [sfm]              SfmOptions passed to solve()
 * @property {object} [trainer]          extra GSTrainer options (shDeg, ...)
 * @property {object} [frames]           FrameOptions passed to load()
 */

class Emitter {
  constructor() { this.map = new Map(); }
  on(type, fn) {
    if (!this.map.has(type)) this.map.set(type, new Set());
    this.map.get(type).add(fn);
    return () => this.map.get(type).delete(fn);
  }
  emit(type, e) {
    const s = this.map.get(type);
    if (s) for (const fn of s) fn(e);
  }
}

export class Session {
  /** @param {SessionOptions} [opts] */
  constructor(opts = {}) {
    this.opts = opts;
    this.frames = [];
    this.recon = null;      // solve() result: { cams, points, k1, k2, ... }
    this.model = null;      // seed() result: { data, n, center, radius }
    this.trainer = null;
    this.gpu = null;
    this.holdout = -1;
    this.testCams = [];
    this.training = false;
    this.lossHistory = [];  // [iter, psnrTrain]
    this._fences = [];      // in-flight batch fences (see _frameLoop)
    // Per-frame loop timings, recorded by default — the cost is a handful of
    // performance.now() calls and one small array row per submitted batch.
    // opts.perf: false disables.
    this.perf = opts.perf === false ? null : { frames: [], marks: [] };
    this._em = new Emitter();
    this._debug = null;     // solver internals (feats/tracks) for UI beats
    this.view = new SessionView(this);
  }

  on(type, fn) { return this._em.on(type, fn); }
  _log(m) { this._em.emit('log', m); }
  _stage(e) { this._em.emit('stage', e); }

  /** Decode photographs into Frames. files: File[]/Blob[]/{source,name}[]
   *  360 equirectangular panos (2:1 aspect) are detected here and sliced
   *  into six-face cubemap rigs — the solver then treats each pano as ONE
   *  camera pose and the trainer sees ordinary pinhole faces. */
  async load(files) {
    this._stage({ stage: 'decode', done: 0, total: files.length });
    const { prepared, rigs, focalNative, faceSize } = await this._slicePanos(files);
    this.frames = await decodeFrames(prepared, {
      ...this.opts.frames, log: (m) => this._log(m),
    });
    this._stage({ stage: 'decode', done: this.frames.length, total: files.length });
    if (this.frames.length < 2) throw new Error('need at least 2 decodable images');
    if (rigs) {
      // the sliced focal is exact but lives in face-native pixels; the
      // solver works at feature scale
      this.rigInfo = rigs;
      this.rigFocalPx = focalNative * this.frames[0].fw / faceSize;
      this._log(`camera rigs active: ${new Set(rigs.filter(Boolean).map((r) => r.id)).size} panos, ` +
        `focal ${this.rigFocalPx.toFixed(1)}px (known from the slice)`);
    }
    return this.frames;
  }

  /** Detect equirect panos in the input and slice them into rig faces.
   *  Detection reads image HEADERS only (probeImageSize) — a normal photo
   *  set pays nothing here; panos are decoded once, sliced, released. */
  async _slicePanos(files) {
    const meta = [];
    let panos = 0, sized = 0, maxW = 0;
    for (const f of files) {
      const src = f && f.source !== undefined ? f.source : f;
      const name = (f && f.name) || (src && src.name) || `image_${meta.length}`;
      const dims = await probeImageSize(src);
      const pano = !!dims && isEquirect(dims.w, dims.h);
      if (dims) sized++;
      if (pano) { panos++; maxW = Math.max(maxW, dims.w); }
      meta.push({ src, name, pano });
    }
    if (!panos) return { prepared: files, rigs: null };
    if (panos < sized) {
      throw new Error('mixed input: 360 panoramas and regular photos in one set are not ' +
        'supported yet — drop either panos or photos, not both');
    }
    // one shared face size + fov for the whole set (the solver's focal is shared)
    const size = faceSizeFor(maxW);
    this._log(`360 input: slicing ${panos} equirectangular panos into ${size}px cubemap rigs`);
    const prepared = [], rigs = [];
    let f = 0, rigId = 0;
    for (const b of meta) {
      const bmp = await createImageBitmap(b.src);
      const sl = sliceEquirect(bmp, size);
      bmp.close?.();
      f = sl.f;
      const base = b.name.replace(/\.[^.]+$/, '');
      sl.faces.forEach((cv, k) => {
        prepared.push({ source: cv, name: `${base}_f${k}` });
        rigs.push({ id: rigId, R: FACE_ROTS[k] });
      });
      rigId++;
      this._stage({ stage: 'decode', done: rigId, total: meta.length });
      await new Promise((r) => setTimeout(r, 0));
    }
    return { prepared, rigs, focalNative: f, faceSize: size };
  }

  /** Use frames prepared elsewhere (e.g. processSource on fetched bitmaps). */
  useFrames(frames) { this.frames = frames; return frames; }

  /** Structure from motion: camera poses + sparse points from the frames. */
  async solve(extra = {}) {
    const opts = { ...this.opts.sfm, ...extra };
    if (this.rigInfo && !opts.rigs) {
      opts.rigs = this.rigInfo;
      opts.focalPx = opts.focalPx ?? this.rigFocalPx;
    }
    // the GPU matcher shares the session device (created here rather than at
    // seed) — it carries the raised buffer limits big feature sets need
    if (!this.gpu) {
      this.gpu = await createGpu({ device: this.opts.device });
      this.gpu.onLost = (info) => this._deviceLost(info);
    }
    this.recon = await runSfM(
      this.frames,
      (m) => this._log(m),
      (imgIdx, x, y) => this.frames[imgIdx].sampleColor(x, y),
      {
        gpu: this.gpu,
        ...opts,
        onEvent: (e) => { this._stage(e); if (opts.onEvent) opts.onEvent(e); },
        debug: (d) => { this._debug = d; if (opts.debug) opts.debug(d); },
      });
    if (undistortFrames(this.frames, this.recon)) {
      this._log(`undistorted training images (k1 ${this.recon.k1.toFixed(4)}, k2 ${this.recon.k2.toFixed(4)})`);
    }
    if (this.opts.lowMem) {
      // the feature-scale grayscale is solve-only (SIFT + LK); on phones
      // 50 frames of it is ~140MB standing between us and the trainer
      for (const f of this.frames) f.gray = null;
    }
    this._stage({ stage: 'solved', done: 1, total: 1, detail: {
      cams: this.recon.cams.length, points: this.recon.points.length,
      rms: this.recon.rmsBA,
    } });
    return this.recon;
  }

  /** Load a reconstruction obtained elsewhere ({ cams, points, ... }). */
  useReconstruction(recon) {
    this.recon = { k1: 0, k2: 0, fScale: 1, medErr: 0, rmsBA: null, ...recon };
    return this.recon;
  }

  /** Seed Gaussians from the sparse cloud and set up the WebGPU trainer. */
  async seed(extra = {}) {
    if (!this.recon) throw new Error('solve() first');
    this._stage({ stage: 'seed', done: 0, total: 1 });
    const target = extra.initTarget || this.opts.initTarget || 60000;
    const clones = Math.min(24, Math.max(2, Math.round(target / this.recon.points.length) - 1));
    this.model = initGaussians(this.recon.points, clones);
    this._log(`initialized ${this.model.n} Gaussians (scene radius ${this.model.radius.toFixed(2)})`);

    if (!this.gpu) this.gpu = await createGpu({ device: this.opts.device });
    this.gpu.onLost = (info) => this._deviceLost(info);
    const gi = this.gpu.info || {};
    const trainerOpts = {
      maxIters: this.opts.maxIters ?? 60000,
      ...this.opts.trainer, ...extra.trainer,
      gpu: this.gpu,
    };
    this.trainer = await GSTrainer.create(trainerOpts);
    this._log(`GPU: ${gi.vendor || 'unknown'} ${gi.architecture || ''} — ` +
      `${this.trainer.tileGrad ? 'tile-shared' : 'direct'} gradient accumulation`);

    // training-resolution intrinsics (recon poses live at feature scale)
    const cams = this.recon.cams.map((c) => {
      const im = this.frames[c.imgIdx];
      const s = im.tw / im.fw;
      return { ...c, f: c.f * s, cx: im.tw / 2, cy: im.th / 2, w: im.tw, h: im.th };
    });
    // output buffers must fit the largest view the host will ever render —
    // interactive canvases are usually LARGER than the training resolution
    const maxW = Math.max(this.opts.maxViewW ?? 2560, ...cams.map((c) => c.w));
    const maxH = Math.max(this.opts.maxViewH ?? 1440, ...cams.map((c) => c.h));
    this.trainer.setup(this.model, cams, this.frames, maxW, maxH, this.model.radius);
    if (this.opts.lowMem) {
      // targets now live on the GPU (packed RGBA8); the float32 CPU copies
      // are 3x that size and nothing reads them after setup
      for (const f of this.frames) { f.rgb = null; f.sampleColor = () => [0.5, 0.5, 0.5]; }
    }

    // blur-aware training: the blurriest frames stay registered (their poses
    // hold the chain together) but are excluded from the loss so the model
    // doesn't learn their motion blur
    const sh = this.trainer.camMeta.map((m) => this.frames[m.imgIdx].sharpness);
    const med = [...sh].sort((a, b) => a - b)[sh.length >> 1];
    this.trainer.excluded = new Set();
    this.trainer.camMeta.forEach((m, i) => {
      if (sh[i] < med * 0.45) this.trainer.excluded.add(i);
    });
    if (this.trainer.excluded.size) {
      this._log(`excluding ${this.trainer.excluded.size} blurry cameras from the training loss ` +
        `(sharpness < 45% of median; poses kept)`);
    }

    // holdout: one sharp mid-sequence frame excluded from training and scored
    // as the honest novel-view metric
    const want = extra.holdout ?? this.opts.holdout ?? 'auto';
    if (want === 'auto') {
      let hi = this.trainer.camMeta.length >> 1;
      for (let k = 0; k < this.trainer.camMeta.length && this.trainer.excluded.has(hi); k++) {
        hi = (hi + 1) % this.trainer.camMeta.length;
      }
      this.holdout = hi;
    } else {
      this.holdout = (want == null || want < 0) ? -1 : want;
    }
    this.trainer.holdout = this.holdout;

    // benchmark protocol: every Nth frame joins the test set — excluded from
    // the loss (poses kept, same as the blur exclusions) and scored together
    // by evalTestPsnr(). The single chart holdout follows a mid-set test
    // frame so live progress tracks the same distribution.
    const split = extra.evalSplit ?? this.opts.evalSplit ?? 0;
    this.testCams = [];
    if (split >= 2) {
      this.trainer.camMeta.forEach((m, i) => {
        if (m.imgIdx % split === 0) { this.trainer.excluded.add(i); this.testCams.push(i); }
      });
      if (this.testCams.length && this.holdout < 0) {
        this.holdout = this.testCams[this.testCams.length >> 1];
        this.trainer.holdout = this.holdout;
      }
      this._log(`evaluation split: ${this.testCams.length} of ${this.trainer.camMeta.length} ` +
        `cameras (every ${split}th) held out of training`);
    }

    this._stage({ stage: 'seed', done: 1, total: 1, detail: { splats: this.model.n } });
    return this.model;
  }

  /**
   * Attach an ALREADY-TRAINED model instead of seeding from the sparse cloud
   * — the loading half of exportPlyBlob / a saved session.
   *
   * gaussians: { data: Float32Array (n*16, the trainer's raw parameter
   * layout), n, sh?: Float32Array (n*shK*3), shK? }.
   *
   * View-only (no frames needed): pass opts.viewOnly and the trainer is set
   * up with zero training cameras — renders, tours and exports work, training
   * does not. For a full resume call useFrames()/load() first and the normal
   * camera wiring happens exactly like seed(); pass opts.iter to continue the
   * schedules where the saved run stopped.
   */
  async seedFrom(gaussians, opts = {}) {
    if (!this.recon && !opts.viewOnly) throw new Error('useReconstruction() first');
    this._stage({ stage: 'seed', done: 0, total: 1 });
    const radius = opts.sceneRadius ?? this.recon?.sceneRadius ?? 10;
    this.model = { data: gaussians.data, n: gaussians.n, radius };

    if (!this.gpu) this.gpu = await createGpu({ device: this.opts.device });
    this.gpu.onLost = (info) => this._deviceLost(info);
    const trainerOpts = {
      maxIters: this.opts.maxIters ?? 60000,
      ...this.opts.trainer, ...opts.trainer,
      gpu: this.gpu,
    };
    if (gaussians.shK != null) trainerOpts.shDeg = { 0: 0, 3: 1, 8: 2, 15: 3 }[gaussians.shK] ?? trainerOpts.shDeg;
    this.trainer = await GSTrainer.create(trainerOpts);

    let cams = [];
    if (!opts.viewOnly) {
      cams = this.recon.cams.map((c) => {
        const im = this.frames[c.imgIdx];
        const s = im.tw / im.fw;
        return { ...c, f: c.f * s, cx: im.tw / 2, cy: im.th / 2, w: im.tw, h: im.th };
      });
    }
    const maxW = Math.max(this.opts.maxViewW ?? 2560, ...cams.map((c) => c.w));
    const maxH = Math.max(this.opts.maxViewH ?? 1440, ...cams.map((c) => c.h));
    this.trainer.setup(this.model, cams, this.frames || [], maxW, maxH, radius);
    // setup zero-fills SH (view dependence is normally learned) — a restored
    // model brings its own
    if (gaussians.sh && this.trainer.bufSH) {
      this.trainer.device.queue.writeBuffer(this.trainer.bufSH, 0,
        gaussians.sh.buffer, gaussians.sh.byteOffset, gaussians.n * this.trainer.shK * 3 * 4);
    }
    if (opts.iter) this.trainer.iter = opts.iter;
    this.trainer.excluded = new Set();
    this.holdout = -1;
    this.trainer.holdout = -1;
    this.testCams = [];
    this._log(`restored ${gaussians.n} Gaussians` +
      (opts.viewOnly ? ' (view only)' : ` at iteration ${opts.iter || 0}`));
    this._stage({ stage: 'seed', done: 1, total: 1, detail: { splats: gaussians.n } });
    return this.model;
  }

  // ---- the training loop (policy that used to live in the demo page) ----

  start() {
    if (!this.trainer) throw new Error('seed() first');
    if (this.training) return;
    this.training = true;
    this._stage({ stage: 'train', done: this.trainer.iter, total: this._maxIters() });
    this._ensureScheduler();
    this._scheduleFrame();
  }

  pause() { this.training = false; }

  /** Raise the training horizon by `moreIters` and resume. The trainer's
   *  schedules (pos-lr decay, growth stop) are horizon-relative and stretch
   *  with it — this is a real continued run, not idling at floor lr. */
  continueFor(moreIters) {
    if (!this.trainer) throw new Error('seed() first');
    this.opts.maxIters = (this.opts.maxIters ?? 60000) + moreIters;
    this.trainer.opts.maxIters = this.opts.maxIters;
    this.trainer.horizon = this.opts.maxIters;
    this.start();
    return this.opts.maxIters;
  }

  /** End the run early with the model as it stands: final metrics, then the
   *  same train-complete event an auto-stop emits. */
  async finish() {
    if (!this.trainer) return;
    this.training = false;
    this._log(`training finished early at ${this.trainer.iter} iterations`);
    await this._emitMetrics(true);
    this._em.emit('event', { kind: 'train-complete', iter: this.trainer.iter, splats: this.trainer.n });
  }

  _maxIters() { return this.opts.maxIters ?? 60000; }

  /** The GPU vanished under us (iOS reclaims WebGPU devices from backgrounded
   *  tabs; drivers reset). Training stops — the splats lived on the dead
   *  device — but frames and reconstruction are CPU-side, so recover() can
   *  rebuild and train again without redoing the solve. */
  _deviceLost(info) {
    if (this._lost) return;
    this._lost = true;
    this.training = false;
    this._fences = [];
    this._log(`GPU device lost (${(info && info.reason) || 'unknown'}) — ` +
      `${(info && info.message) || 'reclaimed by the system'}`);
    this._em.emit('event', { kind: 'device-lost', iter: this.trainer ? this.trainer.iter : 0 });
  }

  get deviceLost() { return !!this._lost; }

  /** Rebuild after device loss: a fresh device and trainer from the CPU-side
   *  reconstruction. Training restarts at iteration 0. */
  async recover() {
    if (!this.recon) throw new Error('nothing to recover — no reconstruction');
    this.training = false;
    this._lost = false;
    this._fences = [];
    this._ips = null;
    this._thrT = null;
    this._statsGap = 2000;
    this._lastStats = performance.now();
    this._metricsLock = null;    // readbacks on the dead device never resolve
    this._framePending = false;
    this.lossHistory = [];
    this._lastHold = null;
    this._lastHoldEval = 0;
    this._itersAtStats = 0;
    this.trainer = null;
    if (this.gpu && this.gpu.owned) { try { this.gpu.dispose(); } catch { /* already gone */ } }
    this.gpu = null;
    return this.seed();
  }

  _ensureScheduler() {
    if (this._sched) return;
    // Hidden tabs AND occluded windows get their rAF throttled; worker
    // messages are not. The worker tick drives the loop whenever a scheduled
    // frame is >150ms late.
    let tickWorker = null;
    try {
      const src = 'setInterval(() => postMessage(0), 33);';
      tickWorker = new Worker(URL.createObjectURL(new Blob([src], { type: 'text/javascript' })));
    } catch { /* no Worker: rAF only */ }
    this._framePending = false;
    this._frameScheduledAt = 0;
    const runFrame = () => {
      if (!this._framePending) return;
      this._framePending = false;
      // a readback rejecting mid-flight (device loss) must not become an
      // unhandled rejection — the loop stops, the device-lost event explains
      this._frameLoop().catch((e) => this._log(`frame loop: ${(e && e.message) || e}`));
    };
    if (tickWorker) {
      tickWorker.onmessage = () => {
        if (this._framePending && this.training &&
            performance.now() - this._frameScheduledAt > 150) runFrame();
      };
    }
    this._sched = {
      runFrame, tickWorker,
      raf: typeof requestAnimationFrame === 'function'
        ? (fn) => requestAnimationFrame(fn) : (fn) => setTimeout(fn, 16),
    };
    this._frameCount = 0;
    this._lastStats = performance.now();
    this._itersAtStats = 0;
    this._lastHoldEval = 0;
  }

  _scheduleFrame() {
    this._framePending = true;
    this._frameScheduledAt = performance.now();
    this._sched.raf(this._sched.runFrame);
  }

  async _frameLoop() {
    const trainer = this.trainer;
    if (!trainer || !trainer.camMeta || this._lost) return;
    this._frameCount++;

    if (this.training && trainer.iter >= this._maxIters()) {
      this.training = false;
      this._log(`training complete at ${trainer.iter} iterations`);
      await this._emitMetrics(true);
      // emitted AFTER the final readback: listeners typically call metrics()
      // right away, which must not interleave with ours on the staging buffer
      this._em.emit('event', { kind: 'train-complete', iter: trainer.iter, splats: trainer.n });
    }

    if (this.training) {
      // Batch enough iterations per frame to keep the GPU busy. The batch
      // adapts to the device: ~120ms of GPU work per frame keeps a desktop
      // GPU saturated with few submits AND keeps metrics/UI alive on a phone
      // where the same 15 iterations could take seconds.
      const batch = this.opts.itersPerFrame ??
        Math.max(4, Math.min(64, Math.round(this._batch ?? 15)));
      const t0 = performance.now();
      // true throughput, stall time included: Safari's stale fences hide the
      // GPU's real pace from per-frame dt, and a slow GPU otherwise ends up
      // with many seconds of work queued (chunky metrics, laggy pause)
      if (this._thrT == null) {
        this._thrT = t0; this._thrIter = trainer.iter;
      } else if (t0 - this._thrT > 1000) {
        const ips = (trainer.iter - this._thrIter) / (t0 - this._thrT) * 1000;
        this._ips = this._ips == null ? ips : this._ips * 0.5 + ips * 0.5;
        this._thrT = t0; this._thrIter = trainer.iter;
      }
      for (let k = 0; k < batch; k++) trainer.stepOnce();

      const tEnc = performance.now() - t0;

      // periodic refinement: relocate dead splats + grow capacity (MCMC-lite)
      if (trainer.iter > 1500 &&
          trainer.iter - (trainer.lastRefine || 0) >= (this.opts.refineEvery ?? 2500)) {
        trainer.lastRefine = trainer.iter;
        const r0 = performance.now();
        // awaited: refine reads six buffers and writes them all back — steps
        // submitted meanwhile would be silently rewound by that write-back,
        // and the six snapshots would come from six different iterations
        const r = await trainer.refine();
        if (this.perf) {
          this.perf.marks.push({ t: Math.round(r0), kind: 'refine', iter: trainer.iter,
            ms: Math.round(performance.now() - r0), moved: r.moved, grown: r.grown });
        }
        if (r.moved || r.grown) {
          this._log(`refine @${trainer.iter}: relocated ${r.moved}, grew +${r.grown} -> ${r.n} splats`);
          this._em.emit('event', { kind: 'refine', iter: trainer.iter, ...r });
        }
      }

      const now = performance.now();
      let tMet = 0;
      if (now - this._lastStats > (this._statsGap ?? 2000)) {
        const m0 = performance.now();
        await this._emitMetrics();
        tMet = performance.now() - m0;
      }

      const v0 = performance.now();
      this.view._tick(this._frameCount, this.training);
      const tView = performance.now() - v0;

      // Deep pipelining: keep a RING of fences in flight, not one. Safari
      // resolves onSubmittedWorkDone hundreds of ms late even when the GPU is
      // idle; gating each batch on a single fence made that latency the loop
      // period, and the batch adapter — reading the latency as GPU time —
      // shrank the batch into its floor (an iPhone sat at ~15 it/s on a scene
      // it can train at hundreds). With 4 fences outstanding the late fences
      // overlap; the frame period becomes ~latency/4 and the same adapter now
      // GROWS the batch until real GPU work dominates. Prompt-fence devices
      // (desktop) behave as before.
      this._fences.push(trainer.device.queue.onSubmittedWorkDone());
      const s0 = performance.now();
      if (this._fences.length > 4) await this._fences.shift();
      const tStall = performance.now() - s0;

      // adapt the batch to the measured cadence (steady-state ~= GPU time of
      // one batch); damped so it settles instead of oscillating
      const dt = Math.max(5, performance.now() - t0);
      const ideal = batch * (120 / dt);
      this._batch = Math.max(4, Math.min(64, (this._batch ?? 15) * 0.7 + ideal * 0.3));
      // throughput cap: ~0.4s of measured GPU work per batch. The dt adapter
      // reads 5ms frames when fences resolve stale and pushes the batch to
      // the clamp — fine on a desktop, ten queued seconds on a phone.
      if (this._ips != null) {
        this._batch = Math.min(this._batch, Math.max(8, Math.round(this._ips * 0.4)));
      }
      if (this.perf) {
        this.perf.frames.push([Math.round(now), trainer.iter, batch, trainer.n,
          +tEnc.toFixed(1), +tView.toFixed(1), +tStall.toFixed(1), +tMet.toFixed(1), +dt.toFixed(1)]);
      }
    } else {
      this.view._tick(this._frameCount, this.training);
      await trainer.device.queue.onSubmittedWorkDone();
      this._fences.length = 0;
    }
    if (this.training || this.view._dirty) this._scheduleFrame();
  }

  /** serialize every metric readback: they share the trainer's staging
   *  buffers, and two in flight = "outstanding map pending". On a lost
   *  device readbacks would hang forever — fail fast instead. */
  _locked(fn) {
    const call = () => {
      if (this._lost) throw new Error('GPU device lost');
      return fn();
    };
    const run = (this._metricsLock || Promise.resolve()).then(call, call);
    this._metricsLock = run.catch(() => {});
    return run;
  }

  _emitMetrics(final = false) {
    return this._locked(() => this._emitMetricsInner(final));
  }

  async _emitMetricsInner(final = false) {
    const trainer = this.trainer;
    const now = performance.now();
    const itersPerSec = (trainer.iter - this._itersAtStats) / Math.max(1, now - this._lastStats) * 1000;
    const m = { iter: trainer.iter, splats: trainer.n, itersPerSec: Math.round(itersPerSec) };
    this._lastIps = m.itersPerSec;
    const tRead = performance.now();
    const mse = await trainer.readLoss();
    if (mse != null && mse > 0) {
      m.psnrTrain = -10 * Math.log10(mse);
      this.lossHistory.push([trainer.iter, m.psnrTrain]);
    }
    const holdEvery = this.opts.evalHoldEvery ?? 4000;
    if (this.holdout >= 0 &&
        (final || trainer.iter - this._lastHoldEval >= holdEvery)) {
      this._lastHoldEval = trainer.iter;
      m.psnrHold = await trainer.evalCamPsnr(this.holdout);
      this._lastHold = m.psnrHold;
    } else if (this._lastHold != null) {
      m.psnrHold = this._lastHold;
    }
    // The readbacks above drain every queued batch first — seconds on a slow
    // GPU. Pace the next readout from COMPLETION and by what this one cost,
    // or a fixed 2s cadence turns into back-to-back queue drains.
    const end = performance.now();
    this._lastStats = end;
    this._itersAtStats = trainer.iter;
    this._statsGap = Math.min(8000, Math.max(2000, (end - tRead) * 2.5));
    this._em.emit('metrics', m);
    return m;
  }

  /** One-shot quality readout (used by tests and the done screen). */
  metrics({ refined = false } = {}) {
    return this._locked(async () => {
      const trainer = this.trainer;
      const out = { iter: trainer.iter, splats: trainer.n };
      if (this.holdout >= 0) {
        out.psnrHold = await trainer.evalCamPsnr(this.holdout);
        if (refined) out.psnrHoldRefined = await trainer.evalCamPsnrRefined(this.holdout);
      }
      return out;
    });
  }

  /** PSNR of one training camera against its own photograph (serialized with
   *  the other metric readbacks; safe to call while training). */
  evalFramePsnr(ci) {
    return this._locked(() => this.trainer.evalCamPsnr(ci));
  }

  /** Mean PSNR over the evalSplit test cameras — the number quality papers
   *  report. Resolves null when no eval split was requested. */
  evalTestPsnr() {
    if (!this.testCams || !this.testCams.length) return Promise.resolve(null);
    return this._locked(async () => {
      const frames = [];
      let sum = 0;
      for (const c of this.testCams) {
        const psnr = await this.trainer.evalCamPsnr(c);
        frames.push({ cam: c, imgIdx: this.trainer.camMeta[c].imgIdx, psnr });
        sum += psnr;
      }
      return { psnr: sum / frames.length, frames };
    });
  }

  /** Index of the camera the trainer most recently stepped on (UI pulse). */
  get activeCam() { return this.trainer ? this.trainer.lastCam : -1; }

  /** Standard 3DGS .ply with Mip opacity compensation baked (what external
   *  sorted viewers expect). */
  exportPlyBlob() {
    // shares the trainer's staging buffers with the metric readbacks
    return this._locked(async () => {
      const { data, n, sh, shK } = await this.trainer.readGaussians();
      const meta = this.trainer.camMeta[0];
      // a view-only restore has no training cameras — its model came from an
      // export that already baked the compensation, so pass it through
      const baked = meta
        ? bakeOpacityCompensation(data, n, meta.f,
            Float32Array.from(this.trainer.camMeta.flatMap(camPosition)))
        : data;
      return gaussiansToPly(baked, n, sh, shK);
    });
  }

  /** The trainer's RAW parameter state ({ data, n, sh, shK }) — the exact
   *  floats, no opacity baking or color activation: what seedFrom() takes
   *  back for a bit-identical resume. */
  exportRawState() {
    return this._locked(() => this.trainer.readGaussians());
  }

  dispose() {
    this.training = false;
    if (this._sched && this._sched.tickWorker) this._sched.tickWorker.terminate();
    if (this.gpu && this.gpu.owned) this.gpu.dispose();
    this.trainer = null;
  }
}

/** The session's render target: one canvas, one camera. */
class SessionView {
  constructor(session) {
    this.s = session;
    this.canvas = null;
    this.ctx = null;
    this.camera = null;   // { R, t, f, cx, cy, w, h } at canvas resolution
    this._dirty = false;
    this._offset = 0;     // training-target offset when looking through a frame
  }

  attach(canvas) {
    this.canvas = canvas;
    this.ctx = canvas.getContext('webgpu');
    this.ctx.configure({
      device: this.s.trainer.device,
      format: this.s.trainer.canvasFormat,
      alphaMode: 'opaque',
    });
    this._dirty = true;
    if (!this.s.training) this.s._ensureScheduler(), this.s._scheduleFrame();
  }

  /** Point the camera at training frame i (exact pose + intrinsics). Returns
   *  the camera meta so the UI can overlay the photograph. */
  lookThrough(i) {
    const meta = this.s.trainer.camMeta[i];
    this.setCamera({ ...meta });
    this._offset = meta.offset;
    return meta;
  }

  /** Free camera: { R, t, f, cx, cy, w, h }. */
  setCamera(cam) {
    this.camera = cam;
    this._offset = 0;
    this._dirty = true;
    if (!this.s.training && this.s._sched) this.s._scheduleFrame();
  }

  renderNow() {
    if (!this.ctx || !this.camera) return;
    this.s.trainer.renderView(this.camera, this.ctx, 0, this._offset);
    this._dirty = false;
    // any render satisfies the auto-refresh — a host that renders on its own
    // cadence keeps pushing this back and the tick below never double-fires
    this._lastAuto = performance.now();
  }

  _tick(frameCount, training) {
    if (!this.ctx || !this.camera) return;
    if (this._dirty) { this.renderNow(); return; }
    // during training, refresh the (unchanged-camera) view sparingly — every
    // render is a full raster pass stolen from the optimiser (~2/s at most,
    // fewer on slow devices: ~25 iterations between refreshes)
    const ips = this.s._itersAtStats && this.s._lastStats
      ? Math.max(1, this.s._lastIps || 100) : 100;
    const interval = Math.max(500, 25000 / ips);
    if (training && performance.now() - (this._lastAuto || 0) > interval) {
      this.renderNow();
    }
  }
}

/** @param {SessionOptions} [opts] */
export function createSession(opts = {}) { return new Session(opts); }