// basicpitch.js - Spotify's Basic Pitch in the browser: audio in, notes and MIDI out.
//
//   const output = await transcribe(samples);   // mono Float32Array at SAMPLE_RATE
//   const notes = findNotes(output);            // [{ pitchMidi, startTimeSeconds, durationSeconds, amplitude }]
//   const midi = toMidi(notes);                 // the bytes of a .mid file

import { notesFromOutput } from "./notes.js";

export const SAMPLE_RATE = 22050;
const NOTES = 88, FPS = Math.floor(SAMPLE_RATE / 256);
const WINDOW = 43844, FRAMES = 172, EDGE = 15, ROWS = FRAMES - 2 * EDGE;   // as in bp.c
const OVERLAP = 2 * EDGE * 256, STEP = WINDOW - OVERLAP;

let pool = null, jobs = 0;

// Starts the workers (transcribe does too). The first call sets the number of threads.
export function preload({ threads = Math.min(8, navigator.hardwareConcurrency || 4) } = {}) {
  if (pool) return pool.ready;
  const workers = Array.from({ length: threads }, () => new Worker(import.meta.url, { type: "module", name: "basicpitch" }));
  const handlers = new Map();   // job id -> its result handler
  const ready = loadFiles().then((files) => Promise.all(workers.map((worker) => new Promise((resolve, reject) => {
    worker.onmessage = ({ data }) => {
      if (data.ready) resolve();
      else if (data.job === undefined) reject(new Error(data.error));
      else handlers.get(data.job)?.(worker, data);
    };
    worker.onerror = (e) => {
      e.preventDefault();
      const error = e.message || "Basic Pitch failed";
      reject(new Error(error));
      for (const handle of handlers.values()) handle(worker, { error });
    };
    worker.postMessage(files);
  }))));
  pool = { workers, handlers, ready };
  ready.catch(() => { for (const worker of workers) worker.terminate(); pool = null; });
  return ready;
}

async function loadFiles() {
  const simd = WebAssembly.validate(new Uint8Array([0, 97, 115, 109, 1, 0, 0, 0, 1, 5, 1, 96, 0, 1, 123, 3, 2, 1, 0, 10, 10, 1, 8, 0, 65, 0, 253, 15, 253, 98, 11]));
  const get = (name) => fetch(new URL(name, import.meta.url)).then((r) => {
    if (!r.ok) throw new Error(`can't load ${name} (${r.status})`);
    return r.arrayBuffer();
  });
  const [wasm, plan, weights] = await Promise.all([get(simd ? "bp.wasm" : "bp-plain.wasm"), get("plan.bin"), get("weights.bin")]);
  return { module: await WebAssembly.compile(wasm), plan, weights };
}

// The model's output: { frames, onsets, rows }, frames[row * 88 + key] (key 0 is MIDI note 21).
export async function transcribe(samples, { threads, onProgress, signal } = {}) {
  await preload(threads ? { threads } : undefined);
  const job = ++jobs, rows = Math.floor(samples.length * FPS / SAMPLE_RATE), count = Math.ceil((samples.length + OVERLAP / 2) / STEP);
  const frames = new Float32Array(rows * NOTES), onsets = new Float32Array(rows * NOTES);
  let next = 0, done = 0;
  try {
    await new Promise((resolve, reject) => {
      if (signal?.aborted) return reject(new DOMException("Aborted", "AbortError"));
      signal?.addEventListener("abort", () => reject(new DOMException("Aborted", "AbortError")), { once: true });
      const send = (worker) => {
        if (next >= count || signal?.aborted) return;
        const index = next++, start = index * STEP - OVERLAP / 2, window = new Float32Array(WINDOW);
        window.set(samples.subarray(Math.max(0, start), Math.min(samples.length, start + WINDOW)), Math.max(0, -start));
        worker.postMessage({ job, index, window }, [window.buffer]);
      };
      pool.handlers.set(job, (worker, data) => {
        if (data.error) return reject(new Error(data.error));
        const at = data.index * ROWS, n = Math.min(ROWS, rows - at) * NOTES;
        if (n > 0) { frames.set(data.frames.subarray(0, n), at * NOTES); onsets.set(data.onsets.subarray(0, n), at * NOTES); }
        done++;   // not inside onProgress?.(), which skips its argument without a callback
        onProgress?.(done / count);
        if (done === count) resolve();
        else send(worker);
      });
      for (const worker of pool.workers) { send(worker); send(worker); }
    });
  } finally {
    pool.handlers.delete(job);
  }
  return { frames, onsets, rows };
}

// The settings and defaults are basicpitch.spotify.com's (ms, Hz, 0 for no limit).
export function findNotes(output, { onsetThreshold = 0.5, frameThreshold = 0.3, minNoteLength = 128, minFrequency = 0, maxFrequency = 3000, boostOnsets = true } = {}) {
  return notesFromOutput(output, onsetThreshold, frameThreshold, Math.round(minNoteLength / 1000 * FPS), minFrequency, maxFrequency, boostOnsets);
}

// The notes as the MIDI file holds them (ticks, seconds, velocity), for playback.
export function midiNotes(notes, { tempo = 120 } = {}) {
  const ticks = 8 * tempo;   // a second, at 480 a beat
  return notes.map((n) => {
    const on = Math.round(n.startTimeSeconds * ticks), off = Math.max(on + 1, Math.round((n.startTimeSeconds + n.durationSeconds) * ticks));
    return { pitch: n.pitchMidi, on, off, start: on / ticks, end: off / ticks, velocity: Math.min(127, Math.max(1, Math.round(n.amplitude * 127))) };
  });
}

// The tempo only sets the beat grid, the notes keep their times.
export function toMidi(notes, { tempo = 120 } = {}) {
  const events = [];
  for (const n of midiNotes(notes, { tempo })) events.push([n.on, 1, n.pitch, n.velocity], [n.off, 0, n.pitch, 0]);
  events.sort((a, b) => a[0] - b[0] || a[1] - b[1]);   // offs before ons
  const us = Math.floor(60e6 / tempo);
  const track = [0x00, 0xff, 0x51, 0x03, (us >> 16) & 255, (us >> 8) & 255, us & 255];
  let last = 0;
  for (const [tick, on, pitch, velocity] of events) {
    track.push(...varint(tick - last), on ? 0x90 : 0x80, pitch, velocity);
    last = tick;
  }
  track.push(0x00, 0xff, 0x2f, 0x00);
  const n = track.length;
  return new Uint8Array([
    0x4d, 0x54, 0x68, 0x64, 0, 0, 0, 6, 0, 0, 0, 1, 0x01, 0xe0,   // MThd: format 0, 1 track, 480 ticks a beat
    0x4d, 0x54, 0x72, 0x6b, n >>> 24, (n >> 16) & 255, (n >> 8) & 255, n & 255, ...track,
  ]);
}

function varint(value) {
  const bytes = [value & 0x7f];
  while ((value >>= 7)) bytes.unshift((value & 0x7f) | 0x80);
  return bytes;
}

// In a worker

let model = null;

async function setUp({ module, plan, weights }) {
  const bp = (await WebAssembly.instantiate(module, { env: { emscripten_notify_memory_growth() {} } })).exports;
  bp._initialize?.();
  const copy = (bytes) => { const p = bp.malloc(bytes.byteLength); new Uint8Array(bp.memory.buffer, p, bytes.byteLength).set(new Uint8Array(bytes)); return p; };
  const handle = bp.bp_load(copy(plan), plan.byteLength / 4, copy(weights));
  if (!handle) throw new Error("plan.bin is not a Basic Pitch plan");
  model = { bp, handle, input: bp.malloc(WINDOW * 4), frames: bp.malloc(FRAMES * NOTES * 4), onsets: bp.malloc(FRAMES * NOTES * 4) };
  // warm up, so the browser compiles the optimised code before real work
  for (let i = 0; i < 2; i++) bp.bp_run(handle, model.input, model.frames, model.onsets, 0);
}

function run({ job, index, window }) {
  const { bp, handle, input, frames, onsets } = model;
  new Float32Array(bp.memory.buffer, input, WINDOW).set(window);
  const err = bp.bp_run(handle, input, frames, onsets, 0);
  if (err) throw new Error(`Basic Pitch error ${err}`);
  const rows = (p) => new Float32Array(bp.memory.buffer, p + EDGE * NOTES * 4, ROWS * NOTES).slice();
  const f = rows(frames), o = rows(onsets);
  self.postMessage({ job, index, frames: f, onsets: o }, [f.buffer, o.buffer]);
}

if (typeof WorkerGlobalScope !== "undefined" && self.name === "basicpitch")
  self.onmessage = ({ data }) => {
    if (data.module) setUp(data).then(() => self.postMessage({ ready: true }), (error) => self.postMessage({ error: error.message }));
    else try { run(data); } catch (error) { self.postMessage({ job: data.job, error: error.message }); }
  };