micarrsim/packages/acoustics-core/realtime-localizer.js

import { computeDasMap, createDasAnalysisPlan } from "./das.js";
import {
  MAX_FREQUENCY_RESPONSE_HZ,
  MAX_FREQUENCY_RESPONSE_POINTS,
} from "./frequency-response.js";

export const REALTIME_LOCALIZATION_PLAN_KIND = "micarrsim.realtime-localization-plan/1";
export const DEFAULT_STREAM_RESULTS_PER_CALL = 1;
export const MAX_STREAM_RESULTS_PER_CALL = 32;
export const DEFAULT_MAXIMUM_STREAM_INPUT_BYTES = 256 * 1024 * 1024;

function requirePositiveFinite(value, label) {
  if (!Number.isFinite(value) || value <= 0) {
    throw new RangeError(`${label} must be a positive finite number.`);
  }
  return value;
}

function requirePositiveInteger(value, label) {
  if (!Number.isInteger(value) || value <= 0) {
    throw new RangeError(`${label} must be a positive integer.`);
  }
  return value;
}

function requirePowerOfTwo(value, label) {
  requirePositiveInteger(value, label);
  if (value < 2 || (value & (value - 1)) !== 0) {
    throw new RangeError(`${label} must be a power of two greater than one.`);
  }
  return value;
}

function freezePosition(position, label) {
  if (!Array.isArray(position) || position.length !== 3 || !position.every(Number.isFinite)) {
    throw new TypeError(`${label} must be a finite [x, y, z] array.`);
  }
  return Object.freeze([...position]);
}

function copyFrequencyResponse(value, label) {
  if (value === undefined) return undefined;
  if (
    !Array.isArray(value) ||
    value.length < 2 ||
    value.length > MAX_FREQUENCY_RESPONSE_POINTS
  ) {
    throw new RangeError(
      label + " must contain 2-" + MAX_FREQUENCY_RESPONSE_POINTS + " points.",
    );
  }
  let previousFrequency = -Infinity;
  return Object.freeze(value.map((point, index) => {
    const pointLabel = label + "[" + index + "]";
    if (!point || typeof point !== "object") {
      throw new TypeError(pointLabel + " must be an object.");
    }
    if (
      !Number.isFinite(point.frequencyHz) ||
      point.frequencyHz < 0 ||
      point.frequencyHz > MAX_FREQUENCY_RESPONSE_HZ ||
      point.frequencyHz <= previousFrequency
    ) {
      throw new RangeError(pointLabel + ".frequencyHz must be finite and strictly increasing.");
    }
    if (!Number.isFinite(point.gainDb) || point.gainDb < -120 || point.gainDb > 120) {
      throw new RangeError(pointLabel + ".gainDb must lie inside -120..120 dB.");
    }
    if (
      point.phaseDeg !== undefined &&
      (!Number.isFinite(point.phaseDeg) || point.phaseDeg < -180 || point.phaseDeg > 180)
    ) {
      throw new RangeError(pointLabel + ".phaseDeg must lie inside -180..180 degrees.");
    }
    previousFrequency = point.frequencyHz;
    return Object.freeze({ ...point });
  }));
}

function copyArrayDefinition(value) {
  if (!value || !Array.isArray(value.microphones) || value.microphones.length === 0) {
    throw new TypeError("array.microphones must be a non-empty array.");
  }
  const microphoneIds = new Set();
  const microphones = value.microphones.map((microphone, index) => {
    const id = microphone.id ?? "mic-" + index;
    if (microphoneIds.has(id)) throw new Error("duplicate microphone id: " + id);
    microphoneIds.add(id);
    if (
      microphone.sensitivityDb !== undefined &&
      !Number.isFinite(microphone.sensitivityDb)
    ) {
      throw new TypeError("array.microphones[" + index + "].sensitivityDb must be finite.");
    }
    if (
      microphone.phaseOffsetDeg !== undefined &&
      !Number.isFinite(microphone.phaseOffsetDeg)
    ) {
      throw new TypeError("array.microphones[" + index + "].phaseOffsetDeg must be finite.");
    }
    const frequencyResponse = copyFrequencyResponse(
      microphone.frequencyResponse,
      "array.microphones[" + index + "].frequencyResponse",
    );
    return Object.freeze({
      ...microphone,
      id,
      position: freezePosition(
        microphone.position,
        "array.microphones[" + index + "].position",
      ),
      ...(frequencyResponse ? { frequencyResponse: Object.freeze(frequencyResponse) } : {}),
    });
  });
  return Object.freeze({
    ...value,
    microphones: Object.freeze(microphones),
  });
}

function copyScan(scan) {
  if (!scan || typeof scan !== "object") throw new TypeError("scan must be an object.");
  const numericFields = ["xMin", "xMax", "yMin", "yMax", "z"];
  for (const field of numericFields) {
    if (!Number.isFinite(scan[field])) throw new TypeError(`scan.${field} must be finite.`);
  }
  const columns = requirePositiveInteger(scan.columns, "scan.columns");
  const rows = requirePositiveInteger(scan.rows, "scan.rows");
  if (scan.xMax < scan.xMin || scan.yMax < scan.yMin) {
    throw new RangeError("scan bounds must satisfy min <= max.");
  }
  return Object.freeze({ ...scan, columns, rows });
}

function copyBands(bands, sampleRate) {
  if (!Array.isArray(bands) || bands.length === 0) {
    throw new TypeError("bands must be a non-empty array.");
  }
  const nyquist = sampleRate / 2;
  return Object.freeze(bands.map((band, index) => {
    if (
      !Number.isFinite(band?.minHz) ||
      !Number.isFinite(band?.maxHz) ||
      band.minHz < 0 ||
      band.maxHz <= band.minHz ||
      band.maxHz >= nyquist
    ) {
      throw new RangeError(`bands[${index}] must lie inside 0–${nyquist} Hz.`);
    }
    return Object.freeze({
      ...band,
      id: band.id ?? `band-${index + 1}`,
      color: band.color ?? "#65f2bf",
    });
  }));
}

function validateChannels(channels, expectedChannelCount) {
  if (!Array.isArray(channels) || channels.length !== expectedChannelCount) {
    throw new Error(
      `channel count ${channels?.length ?? 0} does not match microphone count ${expectedChannelCount}.`,
    );
  }
  if (!channels.every((channel) => channel instanceof Float32Array)) {
    throw new TypeError("every channel must be a Float32Array.");
  }
  const sampleCount = channels[0].length;
  if (channels.some((channel) => channel.length !== sampleCount)) {
    throw new Error("all channels must have the same sample count.");
  }
  return sampleCount;
}

function createAbortError(signal) {
  if (signal?.reason instanceof Error && signal.reason.name === "AbortError") return signal.reason;
  const error = new Error("localization was canceled.");
  error.name = "AbortError";
  return error;
}

function throwIfAborted(signal) {
  if (signal?.aborted) throw createAbortError(signal);
}

function locationsFromBands(bands) {
  return bands.map((band) => ({
    bandId: band.id,
    index: band.peak.index,
    position: [...band.peak.position],
    rawPower: band.peak.rawPower,
  }));
}

/**
 * Compile immutable microphone/scan/band configuration and CPU geometry caches.
 *
 * This API deliberately has no Scene or synthesis dependency: callers provide
 * sampled microphone channels directly to localizeMultichannel().
 */
export function createLocalizationPlan(config) {
  if (!config || typeof config !== "object") throw new TypeError("config must be an object.");
  const sampleRate = requirePositiveFinite(config.sampleRate, "sampleRate");
  const arrayDefinition = copyArrayDefinition(config.array ?? config.arrayDefinition);
  const scan = copyScan(config.scan);
  const frequencyBands = copyBands(config.bands ?? config.frequencyBands, sampleRate);
  const fftSize = requirePowerOfTwo(config.stft?.fftSize ?? config.fftSize ?? 512, "fftSize");
  const hopSize = requirePositiveInteger(
    config.stft?.hopSize ?? config.stftHopSize ?? fftSize / 2,
    "stftHopSize",
  );
  if (hopSize > fftSize) throw new RangeError("stftHopSize must not exceed fftSize.");
  const speedOfSound = requirePositiveFinite(config.speedOfSound ?? 343, "speedOfSound");
  const normalization = config.normalization ?? config.beamforming?.normalization ?? "peak";
  if (normalization !== "peak" && normalization !== "none") {
    throw new RangeError('normalization must be "peak" or "none".');
  }
  const compensateFrequencyResponse =
    config.compensateFrequencyResponse ??
    config.beamforming?.compensateFrequencyResponse ??
    true;
  if (typeof compensateFrequencyResponse !== "boolean") {
    throw new TypeError("compensateFrequencyResponse must be boolean.");
  }

  const scene = Object.freeze({ speedOfSound });
  const experiment = Object.freeze({
    stft: Object.freeze({ fftSize, hopSize, window: "hann" }),
    scan,
    frequencyBands,
    beamforming: Object.freeze({ normalization, compensateFrequencyResponse }),
  });
  const dasPlan = createDasAnalysisPlan({
    sampleRate,
    arrayDefinition,
    scene,
    experiment,
    maximumSteeringCacheBytes: config.maximumSteeringCacheBytes,
    maximumPlanBytes: config.maximumPlanBytes,
  });

  return Object.freeze({
    kind: REALTIME_LOCALIZATION_PLAN_KIND,
    sampleRate,
    channelCount: arrayDefinition.microphones.length,
    arrayDefinition,
    scan,
    bands: frequencyBands,
    stft: experiment.stft,
    speedOfSound,
    normalization,
    scene,
    experiment,
    dasPlan,
    computeOptions: Object.freeze({ ...(config.computeOptions ?? {}) }),
  });
}

function requireLocalizationPlan(plan) {
  if (plan?.kind !== REALTIME_LOCALIZATION_PLAN_KIND) {
    throw new TypeError("plan must be created by createLocalizationPlan().");
  }
  return plan;
}

/**
 * Localize one already-sampled multichannel block.
 */
export async function localizeMultichannel(options) {
  if (!options || typeof options !== "object") throw new TypeError("options must be an object.");
  const plan = options.plan
    ? requireLocalizationPlan(options.plan)
    : createLocalizationPlan(options);
  const sampleCount = validateChannels(options.channels, plan.channelCount);
  if (sampleCount === 0) throw new RangeError("channels must contain at least one sample.");
  throwIfAborted(options.signal);
  const das = await computeDasMap({
    ...plan.computeOptions,
    ...(options.computeOptions ?? {}),
    channels: options.channels,
    sampleRate: plan.sampleRate,
    arrayDefinition: plan.arrayDefinition,
    scene: plan.scene,
    experiment: plan.experiment,
    analysisPlan: plan.dasPlan,
    signal: options.signal,
    onProgress: options.onProgress,
  });
  return {
    ...das,
    sampleRate: plan.sampleRate,
    sampleCount,
    duration: sampleCount / plan.sampleRate,
    locations: locationsFromBands(das.bands),
  };
}

/**
 * A fixed-capacity chunk buffer for real microphone adapters.
 *
 * append() is synchronous and remains safe while an analysis is awaiting CPU
 * batches. analyzeAvailable() copies one window before yielding, consumes its
 * hop, and returns no more than maximumResultsPerCall maps.
 */
export class StreamingDasLocalizer {
  constructor(config) {
    this.plan = createLocalizationPlan(config);
    this.analysisWindowSamples = requirePositiveInteger(
      config.analysisWindowSamples ??
        Math.max(this.plan.stft.fftSize, Math.ceil(this.plan.sampleRate / 10)),
      "analysisWindowSamples",
    );
    if (this.analysisWindowSamples < this.plan.stft.fftSize) {
      throw new RangeError("analysisWindowSamples must be at least fftSize.");
    }
    this.analysisHopSamples = requirePositiveInteger(
      config.analysisHopSamples ?? this.analysisWindowSamples,
      "analysisHopSamples",
    );
    if (this.analysisHopSamples > this.analysisWindowSamples) {
      throw new RangeError("analysisHopSamples must not exceed analysisWindowSamples.");
    }
    this.maximumBufferedSamples = requirePositiveInteger(
      config.maximumBufferedSamples ?? this.analysisWindowSamples * 4,
      "maximumBufferedSamples",
    );
    if (this.maximumBufferedSamples < this.analysisWindowSamples) {
      throw new RangeError("maximumBufferedSamples must be at least analysisWindowSamples.");
    }
    this.maximumResultsPerCall = requirePositiveInteger(
      config.maximumResultsPerCall ?? DEFAULT_STREAM_RESULTS_PER_CALL,
      "maximumResultsPerCall",
    );
    if (this.maximumResultsPerCall > MAX_STREAM_RESULTS_PER_CALL) {
      throw new RangeError(
        `maximumResultsPerCall must not exceed ${MAX_STREAM_RESULTS_PER_CALL}.`,
      );
    }
    this.maximumInputBufferBytes = requirePositiveInteger(
      config.maximumInputBufferBytes ?? DEFAULT_MAXIMUM_STREAM_INPUT_BYTES,
      "maximumInputBufferBytes",
    );
    this.retainedInputBytes =
      this.plan.channelCount *
      (this.maximumBufferedSamples + this.analysisWindowSamples) *
      Float32Array.BYTES_PER_ELEMENT;
    if (
      !Number.isSafeInteger(this.retainedInputBytes) ||
      this.retainedInputBytes > this.maximumInputBufferBytes
    ) {
      throw new RangeError(
        "stream input buffers require " + this.retainedInputBytes.toLocaleString() +
        " bytes, above maximumInputBufferBytes " +
        this.maximumInputBufferBytes.toLocaleString() + ".",
      );
    }

    this.buffers = Array.from(
      { length: this.plan.channelCount },
      () => new Float32Array(this.maximumBufferedSamples),
    );
    this.analysisChannels = Array.from(
      { length: this.plan.channelCount },
      () => new Float32Array(this.analysisWindowSamples),
    );
    this.startIndex = 0;
    this.bufferedSamples = 0;
    this.bufferStartSample = 0;
    this.totalReceivedSamples = 0;
    this.droppedSamples = 0;
    this.analysisAttempts = 0;
    this.completedAnalyses = 0;
    this.failedAnalyses = 0;
    this.nextSequence = 0;
    this.busy = false;
  }

  availableAnalysisCount() {
    if (this.bufferedSamples < this.analysisWindowSamples) return 0;
    return 1 + Math.floor(
      (this.bufferedSamples - this.analysisWindowSamples) / this.analysisHopSamples,
    );
  }

  getState() {
    return {
      sampleRate: this.plan.sampleRate,
      channelCount: this.plan.channelCount,
      analysisWindowSamples: this.analysisWindowSamples,
      analysisHopSamples: this.analysisHopSamples,
      maximumBufferedSamples: this.maximumBufferedSamples,
      maximumInputBufferBytes: this.maximumInputBufferBytes,
      maximumResultsPerCall: this.maximumResultsPerCall,
      bufferedSamples: this.bufferedSamples,
      availableAnalysisCount: this.availableAnalysisCount(),
      bufferStartSample: this.bufferStartSample,
      totalReceivedSamples: this.totalReceivedSamples,
      droppedSamples: this.droppedSamples,
      analysisAttempts: this.analysisAttempts,
      completedAnalyses: this.completedAnalyses,
      failedAnalyses: this.failedAnalyses,
      nextSequence: this.nextSequence,
      busy: this.busy,
      retainedInputBytes: this.retainedInputBytes,
    };
  }

  consume(sampleCount, countAsDropped) {
    if (sampleCount <= 0) return;
    this.startIndex = (this.startIndex + sampleCount) % this.maximumBufferedSamples;
    this.bufferedSamples -= sampleCount;
    this.bufferStartSample += sampleCount;
    if (countAsDropped) this.droppedSamples += sampleCount;
  }

  append(channels) {
    const chunkSamples = validateChannels(channels, this.plan.channelCount);
    if (chunkSamples === 0) return this.getState();
    const capacity = this.maximumBufferedSamples;
    const receivedBefore = this.totalReceivedSamples;

    if (chunkSamples >= capacity) {
      const retainedStart = chunkSamples - capacity;
      for (let channelIndex = 0; channelIndex < this.plan.channelCount; channelIndex += 1) {
        this.buffers[channelIndex].set(channels[channelIndex].subarray(retainedStart));
      }
      this.droppedSamples += this.bufferedSamples + retainedStart;
      this.startIndex = 0;
      this.bufferedSamples = capacity;
      this.bufferStartSample = receivedBefore + retainedStart;
      this.totalReceivedSamples += chunkSamples;
      return this.getState();
    }

    const overflow = Math.max(0, this.bufferedSamples + chunkSamples - capacity);
    this.consume(overflow, true);
    const writeIndex = (this.startIndex + this.bufferedSamples) % capacity;
    const firstLength = Math.min(chunkSamples, capacity - writeIndex);
    for (let channelIndex = 0; channelIndex < this.plan.channelCount; channelIndex += 1) {
      const source = channels[channelIndex];
      const target = this.buffers[channelIndex];
      target.set(source.subarray(0, firstLength), writeIndex);
      if (firstLength < chunkSamples) {
        target.set(source.subarray(firstLength), 0);
      }
    }
    this.bufferedSamples += chunkSamples;
    this.totalReceivedSamples += chunkSamples;
    return this.getState();
  }

  copyAnalysisWindow() {
    const firstLength = Math.min(
      this.analysisWindowSamples,
      this.maximumBufferedSamples - this.startIndex,
    );
    for (let channelIndex = 0; channelIndex < this.plan.channelCount; channelIndex += 1) {
      const source = this.buffers[channelIndex];
      const target = this.analysisChannels[channelIndex];
      target.set(source.subarray(this.startIndex, this.startIndex + firstLength), 0);
      if (firstLength < this.analysisWindowSamples) {
        target.set(
          source.subarray(0, this.analysisWindowSamples - firstLength),
          firstLength,
        );
      }
    }
  }

  async analyzeAvailable(options = {}) {
    if (this.busy) throw new Error("analyzeAvailable() is already running.");
    const maximumResults = options.maximumResults ?? this.maximumResultsPerCall;
    requirePositiveInteger(maximumResults, "maximumResults");
    if (maximumResults > this.maximumResultsPerCall) {
      throw new RangeError(
        `maximumResults cannot exceed configured bound ${this.maximumResultsPerCall}.`,
      );
    }

    this.busy = true;
    const results = [];
    try {
      while (
        results.length < maximumResults &&
        this.bufferedSamples >= this.analysisWindowSamples
      ) {
        throwIfAborted(options.signal);
        const startSample = this.bufferStartSample;
        const endSample = startSample + this.analysisWindowSamples;
        this.copyAnalysisWindow();
        this.consume(this.analysisHopSamples, false);
        this.analysisAttempts += 1;
        try {
          const localized = await localizeMultichannel({
            plan: this.plan,
            channels: this.analysisChannels,
            signal: options.signal,
            onProgress: options.onProgress,
            computeOptions: options.computeOptions,
          });
          const result = {
            sequence: this.nextSequence,
            startSample,
            endSample,
            startTime: startSample / this.plan.sampleRate,
            endTime: endSample / this.plan.sampleRate,
            centerTime: (startSample + endSample) / (2 * this.plan.sampleRate),
            ...localized,
          };
          this.nextSequence += 1;
          this.completedAnalyses += 1;
          results.push(result);
          options.onResult?.(result);
        } catch (error) {
          this.failedAnalyses += 1;
          throw error;
        }
      }
      return results;
    } finally {
      this.busy = false;
    }
  }

  async push(channels, options = {}) {
    this.append(channels);
    return this.analyzeAvailable(options);
  }

  reset() {
    if (this.busy) throw new Error("cannot reset while analysis is running.");
    this.startIndex = 0;
    this.bufferedSamples = 0;
    this.bufferStartSample = 0;
    this.totalReceivedSamples = 0;
    this.droppedSamples = 0;
    this.analysisAttempts = 0;
    this.completedAnalyses = 0;
    this.failedAnalyses = 0;
    this.nextSequence = 0;
  }
}

export function createStreamingLocalizer(config) {
  return new StreamingDasLocalizer(config);
}