worker.js

// NanoMaestro Web Worker for isolated model inference, intelligence heuristics, and parsing
try {
  // WASM-only build. WebGPU EP was removed: on most devices it was slower than WASM
  // and, worse, produced audibly broken/incoherent output because the int8-quantized
  // graph used by this model isn't reliably supported by the WebGPU EP yet. Rather than
  // ship a backend that silently corrupts generation, we standardize on the well-tested
  // multi-threaded WASM/SIMD path.
  importScripts("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js");
} catch (err) {
  console.error("Worker failed to import ONNX Runtime script:", err);
}

// ONNX runtime configuration
if (typeof ort !== "undefined") {
  ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/";
  const hwThreads = (typeof navigator !== "undefined" && navigator.hardwareConcurrency) || 4;
  ort.env.wasm.numThreads = Math.max(1, Math.min(hwThreads, 8));
  ort.env.wasm.simd = true;
  ort.env.wasm.proxy = false;
}

// State variables
let session = null;
let stoi = null;
let itos = null;
let vocabSize = 0;

let h = null;
let c = null;
let currentId = 0;
let activeModelName = "";
const currentEpUsed = "wasm";

// Model parameters
let hiddenSize = 2048;
let numLayers = 2;

// Playback settings & User BPM Override
let temperature = 0.85;
let topK = 40;
let currentBpm = 120;
let userBpmOverride = false;
let midiBpmMode = "lock";

// Seeding settings
// playingStyle: "default" | "bach" | "mozart" | "chopin" | "debussy" | "beethoven" | "satie"
//             | "midi" (continue a single uploaded track)
//             | "trained" (improvise in a style learned from one or more uploaded tracks)
let playingStyle = "default";
let midiTokens = [];
let midiStartBar = 0;
let isWarmingUp = false;
let lastSampledPitch = null;

// Rendering state
let isRendering = false;
let cancelRenderRequested = false;

// Intelligence & Heuristics Controls
let heuristicSettings = {
  dynamicTempEnabled: true,
  velocitySmoothingEnabled: true,
  chordPenaltiesEnabled: true,
  phraseMemoryEnabled: true,
  keyBiasEnabled: true,
  keyBiasStrength: 0.8,
  rhythmBalancingEnabled: true,
  cadenceGenEnabled: true,
  midiSimilarityWeight: 0.85
};

// Reused buffers for top-k sampling
const MAX_TOPK_BUFFER = 256;
const topIdxBuf = new Int32Array(MAX_TOPK_BUFFER);
const topScoreBuf = new Float32Array(MAX_TOPK_BUFFER);

// ==========================================
// MUSICAL INTELLIGENCE ENGINE IMPLEMENTATION
// ==========================================

// 1. Key Detector using Krumhansl-Schmuckler Key-Finding Algorithm
class KeyDetector {
  constructor() {
    this.pitchCounts = new Float32Array(12);
    this.recentPitches = [];
    this.maxMemory = 48;
    this.detectedRoot = 0;
    this.isMinor = false;
    this.confidence = 0;
    
    this.majorProfile = [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88];
    this.minorProfile = [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 2.69, 3.34, 3.17, 3.28];
    this.noteNames = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"];
  }

  reset() {
    this.pitchCounts.fill(0);
    this.recentPitches = [];
    this.detectedRoot = 0;
    this.isMinor = false;
    this.confidence = 0;
  }

  addPitch(midiPitch) {
    const pc = ((midiPitch % 12) + 12) % 12;
    this.recentPitches.push(pc);
    if (this.recentPitches.length > this.maxMemory) {
      this.recentPitches.shift();
    }
    
    this.pitchCounts.fill(0);
    for (let i = 0; i < this.recentPitches.length; i++) {
      this.pitchCounts[this.recentPitches[i]] += 1;
    }
    
    this.analyzeKey();
  }

  analyzeKey() {
    if (this.recentPitches.length < 6) return;

    let bestCorr = -2.0;
    let bestRoot = 0;
    let bestIsMinor = false;

    for (let root = 0; root < 12; root++) {
      const majCorr = this.correlation(this.pitchCounts, this.majorProfile, root);
      if (majCorr > bestCorr) {
        bestCorr = majCorr;
        bestRoot = root;
        bestIsMinor = false;
      }

      const minCorr = this.correlation(this.pitchCounts, this.minorProfile, root);
      if (minCorr > bestCorr) {
        bestCorr = minCorr;
        bestRoot = root;
        bestIsMinor = true;
      }
    }

    this.detectedRoot = bestRoot;
    this.isMinor = bestIsMinor;
    this.confidence = Math.max(0, Math.min(1.0, (bestCorr + 0.5) / 1.5));
  }

  correlation(counts, profile, rootShift) {
    let sumX = 0, sumY = 0, sumXY = 0, sumX2 = 0, sumY2 = 0;
    for (let i = 0; i < 12; i++) {
      const x = counts[(i + rootShift) % 12];
      const y = profile[i];
      sumX += x;
      sumY += y;
      sumXY += x * y;
      sumX2 += x * x;
      sumY2 += y * y;
    }
    const num = 12 * sumXY - sumX * sumY;
    const den = Math.sqrt((12 * sumX2 - sumX * sumX) * (12 * sumY2 - sumY * sumY));
    return den === 0 ? 0 : num / den;
  }

  getScaleDegrees() {
    const root = this.detectedRoot;
    if (this.isMinor) {
      return new Set([root, (root + 2) % 12, (root + 3) % 12, (root + 5) % 12, (root + 7) % 12, (root + 8) % 12, (root + 10) % 12, (root + 11) % 12]);
    } else {
      return new Set([root, (root + 2) % 12, (root + 4) % 12, (root + 5) % 12, (root + 7) % 12, (root + 9) % 12, (root + 11) % 12]);
    }
  }

  getKeyName() {
    return `${this.noteNames[this.detectedRoot]} ${this.isMinor ? "Minor" : "Major"}`;
  }
}

const keyDetector = new KeyDetector();

// 2. Phrase Memory & Motif Tracker (remembers 4 bars)
class PhraseMemory {
  constructor() {
    this.history = new Array(256).fill(null);
    this.lastDurations = [];
    this.maxDurHistory = 6;
  }

  reset() {
    this.history.fill(null);
    this.lastDurations = [];
  }

  recordNote(bar, pos, pitch, durationSteps) {
    const globalStep = (Math.max(0, bar) * 64 + Math.max(0, pos)) % 256;
    this.history[globalStep] = { pitch, durationSteps, bar, pos };
  }

  recordDuration(durSteps) {
    this.lastDurations.push(durSteps);
    if (this.lastDurations.length > this.maxDurHistory) {
      this.lastDurations.shift();
    }
  }

  getMotifPitchForStep(bar, pos) {
    const currentStep = (Math.max(0, bar) * 64 + Math.max(0, pos)) % 256;
    const prev2Bar = (currentStep + 128) % 256;
    const prev4Bar = currentStep;
    
    if (this.history[prev2Bar]) return this.history[prev2Bar].pitch;
    if (this.history[prev4Bar]) return this.history[prev4Bar].pitch;
    return null;
  }

  getMonotonyPenalty(candidateDurSteps) {
    if (this.lastDurations.length < 3) return 0;
    let count = 0;
    for (let i = this.lastDurations.length - 1; i >= 0; i--) {
      if (this.lastDurations[i] === candidateDurSteps) count++;
      else break;
    }
    if (count >= 4) return -2.2;
    if (count >= 3) return -1.2;
    return 0;
  }
}

const phraseMemory = new PhraseMemory();

// 3. STYLE PROFILE — learns from one or many uploaded MIDI files and can be
// exported/imported as plain JSON so a person's trained style survives a page refresh.
//
// Rather than requiring the live generation to stay in lockstep with a single source
// bar-by-bar (which broke down quickly whenever the model's sampled rhythm drifted even
// slightly from the source file), this profile stores *statistical* habits keyed by
// grid position: how often a note lands there, which pitch-classes it tends to be, how
// long it tends to last, and how the melody tends to move (interval habits). Those
// statistics are blended into the sampler every step, so similarity degrades gracefully
// instead of collapsing the moment generation diverges from the source's bar count.
class StyleProfile {
  constructor() {
    this.name = "Untitled Style";
    this.grid = 64;
    this.posPitchClass = {};   // pos(0..grid-1) -> { pitchClass: count }
    this.posDuration = {};     // pos -> { durationSteps: count }
    this.posActive = {};       // pos -> count of note onsets landing here (rhythm habit)
    this.intervalCounts = {};  // melodic interval in semitones (-12..12) -> count
    this.velocityCounts = {};  // velocity bucket (1-8) -> count
    this.totalNotes = 0;
    this.sourceFiles = [];
  }

  static bump(obj, key, amount = 1) {
    obj[key] = (obj[key] || 0) + amount;
    return obj;
  }

  // Ingest one already-tokenized MIDI file's worth of tokens into this profile.
  // Multiple calls accumulate — this is how "train on several files" works.
  ingest(tokens, fileName) {
    let pos = 0;
    let pendingPitch = null;
    let lastPitch = null;

    for (const t of tokens) {
      if (t === "BAR") {
        pos = 0;
        continue;
      }
      if (t.startsWith("POS_")) {
        pos = Number(t.slice(4));
        if (Number.isFinite(pos)) StyleProfile.bump(this.posActive, pos);
        continue;
      }
      if (t.startsWith("NOTE_")) {
        pendingPitch = Number(t.slice(5));
        continue;
      }
      if (t.startsWith("DUR_") && pendingPitch !== null) {
        const dur = Math.max(1, Math.min(64, Number(t.slice(4)) || 1));
        const pc = ((pendingPitch % 12) + 12) % 12;
        this.posPitchClass[pos] ??= {};
        StyleProfile.bump(this.posPitchClass[pos], pc);
        this.posDuration[pos] ??= {};
        StyleProfile.bump(this.posDuration[pos], dur);
        if (lastPitch !== null) {
          const interval = Math.max(-12, Math.min(12, pendingPitch - lastPitch));
          StyleProfile.bump(this.intervalCounts, interval);
        }
        lastPitch = pendingPitch;
        this.totalNotes += 1;
        continue;
      }
      if (t.startsWith("VEL_") && pendingPitch !== null) {
        const bucket = Number(t.slice(4)) || 4;
        StyleProfile.bump(this.velocityCounts, bucket);
        pendingPitch = null;
        continue;
      }
    }

    if (fileName) this.sourceFiles.push(fileName);
  }

  isEmpty() {
    return this.totalNotes === 0;
  }

  // 0..1 — how strongly a note onset "belongs" at this grid position in the learned style
  activityBiasAt(pos) {
    const total = this._activeTotal ??= Object.values(this.posActive).reduce((a, b) => a + b, 0) || 1;
    return (this.posActive[pos] || 0) / total;
  }

  // 0..1 — likelihood this pitch-class is used when a note lands at `pos`
  pitchClassBiasAt(pos, pc) {
    const bucket = this.posPitchClass[pos];
    if (!bucket) return -1; // no data at this position
    const total = Object.values(bucket).reduce((a, b) => a + b, 0) || 1;
    return (bucket[pc] || 0) / total;
  }

  // 0..1 — likelihood a note landing at `pos` has (approximately) this duration
  durationBiasAt(pos, dur) {
    const bucket = this.posDuration[pos];
    if (!bucket) return -1;
    const total = Object.values(bucket).reduce((a, b) => a + b, 0) || 1;
    let matched = 0;
    for (const [d, count] of Object.entries(bucket)) {
      if (Math.abs(Number(d) - dur) <= 2) matched += count;
    }
    return matched / total;
  }

  // 0..1 — likelihood the melody moves by this many semitones from the previous note
  intervalBias(interval) {
    const total = Object.values(this.intervalCounts).reduce((a, b) => a + b, 0) || 1;
    return (this.intervalCounts[interval] || 0) / total;
  }

  toJSON() {
    return {
      formatVersion: 1,
      kind: "nanomaestro-style-profile",
      name: this.name,
      grid: this.grid,
      posPitchClass: this.posPitchClass,
      posDuration: this.posDuration,
      posActive: this.posActive,
      intervalCounts: this.intervalCounts,
      velocityCounts: this.velocityCounts,
      totalNotes: this.totalNotes,
      sourceFiles: this.sourceFiles,
      createdAt: new Date().toISOString()
    };
  }

  static fromJSON(obj) {
    const p = new StyleProfile();
    if (!obj || obj.kind !== "nanomaestro-style-profile") {
      throw new Error("That file doesn't look like a NanoMaestro trained style export.");
    }
    p.name = obj.name || "Imported Style";
    p.grid = obj.grid || 64;
    p.posPitchClass = obj.posPitchClass || {};
    p.posDuration = obj.posDuration || {};
    p.posActive = obj.posActive || {};
    p.intervalCounts = obj.intervalCounts || {};
    p.velocityCounts = obj.velocityCounts || {};
    p.totalNotes = obj.totalNotes || 0;
    p.sourceFiles = obj.sourceFiles || [];
    return p;
  }
}

// The currently active learned style (used by both "midi" continue-mode and "trained" mode)
let activeStyleProfile = null;
// Exact bar-indexed guide used only in "midi" (continue-a-single-track) mode for the
// high-fidelity "reproduce it almost exactly" end of the similarity slider.
let exactGuideBars = [];   // barIdx -> Map(pos -> Array<{pitch, durationSteps}>)
let exactGuideTotalBars = 0;

function buildExactGuide(tokens) {
  exactGuideBars = [];
  exactGuideTotalBars = 0;
  let currentBarMap = new Map();
  let started = false;
  let pos = 0;
  let pendingPitch = null;

  for (const t of tokens) {
    if (t === "BAR") {
      if (started) {
        exactGuideBars.push(currentBarMap);
      }
      started = true;
      currentBarMap = new Map();
      pos = 0;
      continue;
    }
    if (t.startsWith("POS_")) { pos = Number(t.slice(4)) || 0; continue; }
    if (t.startsWith("NOTE_")) { pendingPitch = Number(t.slice(5)); continue; }
    if (t.startsWith("DUR_") && pendingPitch !== null) {
      const dur = Number(t.slice(4)) || 4;
      if (!currentBarMap.has(pos)) currentBarMap.set(pos, []);
      currentBarMap.get(pos).push({ pitch: pendingPitch, durationSteps: dur });
      pendingPitch = null;
      continue;
    }
  }
  if (started && currentBarMap.size > 0) exactGuideBars.push(currentBarMap);
  exactGuideTotalBars = exactGuideBars.length;
}

function getExactGuideNotes(barIdx, pos) {
  if (exactGuideTotalBars === 0) return null;
  const mapped = ((barIdx % exactGuideTotalBars) + exactGuideTotalBars) % exactGuideTotalBars;
  const barMap = exactGuideBars[mapped];
  if (!barMap) return null;
  return barMap.get(pos) || null;
}

function tokenId(token) {
  return stoi?.[token] ?? null;
}

function makeTensorId(id) {
  return new ort.Tensor("int64", BigInt64Array.from([BigInt(id)]), [1, 1]);
}

function zeroState() {
  return new ort.Tensor("float32", new Float32Array(numLayers * 1 * hiddenSize), [numLayers, 1, hiddenSize]);
}

function updateModelDimensions(session) {
  try {
    const inputNames = session.inputNames;
    if (inputNames.includes("h") && session.handler && session.handler._model) {
      const inputs = session.handler._model.graph.inputs;
      const hInput = inputs.find(i => i.name === "h");
      if (hInput && hInput.type && hInput.type.tensorType && hInput.type.tensorType.shape) {
        const shape = hInput.type.tensorType.shape.dim;
        const layers = Number(shape[0].dimValue);
        const hidden = Number(shape[2].dimValue);
        if (layers > 0 && hidden > 0) {
          numLayers = layers;
          hiddenSize = hidden;
          console.log(`Worker set model dimensions: layers=${numLayers}, hiddenSize=${hiddenSize}`);
        }
      }
    }
  } catch (e) {
    console.error("Worker failed to detect model dimensions:", e);
  }
}

async function stepModel(record = true) {
  if (!session) return "?";
  
  const inputTensor = makeTensorId(currentId);
  const output = await session.run({ input: inputTensor, h, c });
  
  inputTensor.dispose();
  
  if (h) { try { h.dispose(); } catch (e) {} }
  if (c) { try { c.dispose(); } catch (e) {} }
  
  h = output.h_out;
  c = output.c_out;
  
  const logitsData = output.logits.data;
  currentId = sampleFromLogitsWithHeuristics(logitsData);
  output.logits.dispose();
  
  const token = itos[currentId] ?? "?";

  if (token.startsWith("NOTE_")) {
    const p = Number(token.slice(5));
    if (Number.isFinite(p)) lastSampledPitch = p;
  }
  
  if (token === "EOS") {
    if (h) { try { h.dispose(); } catch (e) {} }
    if (c) { try { c.dispose(); } catch (e) {} }
    h = zeroState();
    c = zeroState();
  }
  
  if (record && parser) {
    parser.feed(token);
  }
  return token;
}

// Logit modifier incorporating all heuristics, exact rhythm, & trained-style guidance
function sampleFromLogitsWithHeuristics(logits) {
  let effectiveTemp = Math.max(0.05, temperature);
  const currentBar = parser ? parser.bar : 0;
  const currentPos = parser ? parser.pos : 0;
  const phraseBarIdx = Math.max(0, currentBar) % 4;
  const isCadenceStep = (phraseBarIdx === 3) && (currentPos >= 48);

  if (heuristicSettings.dynamicTempEnabled) {
    if (phraseBarIdx <= 1) {
      effectiveTemp *= 0.85;
    } else if (phraseBarIdx === 2) {
      effectiveTemp *= 1.0;
    } else if (phraseBarIdx === 3) {
      effectiveTemp *= 1.18;
    }
  }

  const k = Math.max(1, Math.min(topK, MAX_TOPK_BUFFER, logits.length));
  let filled = 0;

  const scaleDegrees = keyDetector.getScaleDegrees();
  const detectedRoot = keyDetector.detectedRoot;
  const isStrongBeat = (currentPos === 0 || currentPos === 16 || currentPos === 32 || currentPos === 48);

  const simWeight = heuristicSettings.midiSimilarityWeight;
  const useSimilarity = (playingStyle === "midi" || playingStyle === "trained") && simWeight > 0 && activeStyleProfile && !activeStyleProfile.isEmpty();
  // Exact bar-for-bar guidance only applies to single-track "continue" mode, and only
  // matters once the similarity weight is fairly high (it's the "strict reproduction" end).
  const useExactGuide = playingStyle === "midi" && simWeight > 0.5 && exactGuideTotalBars > 0;
  const exactStrength = useExactGuide ? Math.min(1, (simWeight - 0.5) / 0.5) : 0; // 0..1 above the 0.5 midpoint
  let guideNotes = null;
  if (useExactGuide) guideNotes = getExactGuideNotes(currentBar, currentPos);

  for (let i = 0; i < logits.length; i++) {
    const token = itos[i];
    if (token === undefined) continue;

    let score = logits[i] / effectiveTemp;

    // BAN EOS AND EOP WHEN PLAYING A GUIDED STYLE TO PREVENT MODEL FROM STOPPING MID-SONG
    if (token === "<EOP>" || token === "EOS") {
      if (playingStyle === "midi" || playingStyle === "trained") {
        score -= 100.0;
      } else {
        score -= 5.0;
      }
    }

    if (!isWarmingUp) {
      // 🎵 1. RHYTHM GUIDANCE (POSITION TOKENS POS_)
      if (token.startsWith("POS_")) {
        const posVal = Number(token.slice(4));
        if (Number.isFinite(posVal)) {
          if (useExactGuide) {
            const exactNotes = getExactGuideNotes(currentBar, posVal);
            if (exactNotes && exactNotes.length > 0) {
              score += 14.0 * exactStrength;
            }
          }
          if (useSimilarity) {
            const activity = activeStyleProfile.activityBiasAt(posVal);
            score += activity * 10.0 * simWeight;
          }
        }
      }

      // 🎵 2. EXACT NOTE PITCH GUIDANCE (NOTE_ TOKENS)
      if (token.startsWith("NOTE_")) {
        const pitch = Number(token.slice(5));
        if (Number.isFinite(pitch)) {
          const pc = ((pitch % 12) + 12) % 12;

          if (useExactGuide && guideNotes && guideNotes.length > 0) {
            let matchedExact = false;
            let matchedOctave = false;
            let matchedHarmony = false;

            for (const gNote of guideNotes) {
              if (pitch === gNote.pitch) { matchedExact = true; break; }
              const gPc = ((gNote.pitch % 12) + 12) % 12;
              if (pc === gPc) matchedOctave = true;
              else if (Math.abs(pc - gPc) === 3 || Math.abs(pc - gPc) === 4 || Math.abs(pc - gPc) === 7) {
                matchedHarmony = true;
              }
            }

            if (matchedExact) score += 20.0 * exactStrength;
            else if (matchedOctave) score += 6.0 * exactStrength;
            else if (matchedHarmony) score += 3.0 * exactStrength;
            else score -= 10.0 * exactStrength;
          }

          if (useSimilarity) {
            const pcBias = activeStyleProfile.pitchClassBiasAt(currentPos, pc);
            if (pcBias >= 0) {
              // Data exists at this position: pull toward what the trained style favors,
              // and softly push away from pitch-classes it never uses here.
              score += (pcBias - 0.5) * 16.0 * simWeight;
            }
            if (lastSampledPitch !== null) {
              const interval = Math.max(-12, Math.min(12, pitch - lastSampledPitch));
              const intervalBias = activeStyleProfile.intervalBias(interval);
              score += intervalBias * 6.0 * simWeight;
            }
          }

          if (heuristicSettings.keyBiasEnabled) {
            if (scaleDegrees.has(pc)) {
              score += 0.45 * heuristicSettings.keyBiasStrength;
            } else {
              score -= 0.65 * heuristicSettings.keyBiasStrength;
            }
          }

          if (heuristicSettings.chordPenaltiesEnabled && isStrongBeat) {
            const triadNotes = [detectedRoot, (detectedRoot + (keyDetector.isMinor ? 3 : 4)) % 12, (detectedRoot + 7) % 12];
            if (triadNotes.includes(pc)) {
              score += 0.5;
            } else if (!scaleDegrees.has(pc)) {
              score -= 0.9;
            }
          }

          if (heuristicSettings.phraseMemoryEnabled) {
            const motifPitch = phraseMemory.getMotifPitchForStep(currentBar, currentPos);
            if (motifPitch !== null) {
              if (pitch === motifPitch) {
                score += 0.6;
              } else if (((pitch % 12) + 12) % 12 === ((motifPitch % 12) + 12) % 12) {
                score += 0.35;
              }
            }
          }

          if (heuristicSettings.cadenceGenEnabled && isCadenceStep) {
            const tonicPitch = detectedRoot;
            const dominantPitch = (detectedRoot + 7) % 12;
            if (pc === tonicPitch) {
              score += 0.8;
            } else if (pc === dominantPitch) {
              score += 0.5;
            }
          }
        }
      }

      // 🎵 3. DURATION GUIDANCE (DUR_ TOKENS)
      if (token.startsWith("DUR_")) {
        const durSteps = Number(token.slice(4));
        if (Number.isFinite(durSteps)) {
          if (useExactGuide && guideNotes && guideNotes.length > 0) {
            for (const gNote of guideNotes) {
              if (durSteps === gNote.durationSteps) {
                score += 18.0 * exactStrength;
                break;
              }
            }
          }

          if (useSimilarity) {
            const durBias = activeStyleProfile.durationBiasAt(currentPos, durSteps);
            if (durBias >= 0) {
              score += (durBias - 0.3) * 10.0 * simWeight;
            }
          }

          if (heuristicSettings.rhythmBalancingEnabled) {
            score += phraseMemory.getMonotonyPenalty(durSteps);
          }

          if (heuristicSettings.cadenceGenEnabled && isCadenceStep) {
            if (durSteps >= 16) {
              score += 0.7;
            } else {
              score -= 0.5;
            }
          }
        }
      }
    }

    if (filled < k) {
      let pos = filled;
      while (pos > 0 && topScoreBuf[pos - 1] < score) {
        topScoreBuf[pos] = topScoreBuf[pos - 1];
        topIdxBuf[pos] = topIdxBuf[pos - 1];
        pos--;
      }
      topScoreBuf[pos] = score;
      topIdxBuf[pos] = i;
      filled++;
    } else if (score > topScoreBuf[k - 1]) {
      let pos = k - 1;
      while (pos > 0 && topScoreBuf[pos - 1] < score) {
        topScoreBuf[pos] = topScoreBuf[pos - 1];
        topIdxBuf[pos] = topIdxBuf[pos - 1];
        pos--;
      }
      topScoreBuf[pos] = score;
      topIdxBuf[pos] = i;
    }
  }

  const maxScore = topScoreBuf[0];
  let sum = 0;
  for (let j = 0; j < filled; j++) {
    const p = Math.exp(topScoreBuf[j] - maxScore);
    topScoreBuf[j] = p;
    sum += p;
  }

  let r = Math.random() * sum;
  for (let j = 0; j < filled; j++) {
    r -= topScoreBuf[j];
    if (r <= 0) return topIdxBuf[j];
  }
  return topIdxBuf[filled - 1];
}

function durationToSecondsFromEventSteps(steps, grid) {
  const quarterSeconds = 60 / currentBpm;
  return Math.max(0.025, (Math.max(1, steps) * 4 * quarterSeconds) / Math.max(1, grid));
}

function midiToTonePitch(midi) {
  const names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"];
  return `${names[((midi % 12) + 12) % 12]}${Math.floor(midi / 12) - 1}`;
}

function pushEvent(event) {
  if (isWarmingUp) return;
  tempQueue.push(event);
}

function getTokensUpToBar(tokens, targetBar) {
  let barCount = 0;
  const sliced = [];
  for (const t of tokens) {
    sliced.push(t);
    if (t === "BAR") {
      barCount += 1;
      if (barCount > targetBar && sliced.length >= 24) {
        break;
      }
    }
  }
  if (sliced.length < 16 && tokens.length >= 16) {
    return tokens.slice(0, Math.min(tokens.length, 128));
  }
  return sliced;
}

function makeEventParser() {
  let lastVelocity = 0.72;

  return {
    grid: 64,
    bar: -1,
    pos: 0,
    pendingPosition: null,
    pendingNotes: [],
    pendingNote: null,
    lastQ: 0,
    reset() {
      this.grid = 64;
      this.bar = -1;
      this.pos = 0;
      this.pendingPosition = null;
      this.pendingNotes = [];
      this.pendingNote = null;
      this.lastQ = 0;
      lastVelocity = 0.72;
      keyDetector.reset();
      phraseMemory.reset();
    },
    feed(token) {
      if (token.startsWith("BPM_")) {
        const bpm = Number(token.slice(4));
        if (Number.isFinite(bpm) && bpm >= 40 && bpm <= 400) {
          if (!userBpmOverride && (playingStyle !== "midi" || isWarmingUp || midiBpmMode === "model")) {
            currentBpm = bpm;
            self.postMessage({ action: "tempo", bpm: bpm });
          }
        }
        return;
      }
      if (token.startsWith("GRID_")) {
        const grid = Number(token.slice(5));
        if (Number.isFinite(grid) && grid > 0) this.grid = grid;
        return;
      }
      if (token === "BAR") {
        this.flushTo(this.absoluteQFor(this.bar + 1, 0));
        this.bar += 1;
        this.pos = 0;
        return;
      }
      if (token.startsWith("POS_")) {
        const pos = Number(token.slice(4));
        if (!Number.isFinite(pos)) return;
        this.flushTo(this.absoluteQFor(this.bar, pos));
        this.pos = pos;
        return;
      }
      if (token.startsWith("NOTE_")) {
        const midi = Number(token.slice(5));
        if (Number.isFinite(midi) && midi >= 0 && midi <= 127) {
          this.pendingNote = { midi, durationSteps: 1, velocity: 0.72 };
          keyDetector.addPitch(midi);
        }
        return;
      }
      if (token.startsWith("DUR_") && this.pendingNote) {
        const steps = Number(token.slice(4));
        if (Number.isFinite(steps)) {
          this.pendingNote.durationSteps = Math.max(1, steps);
          phraseMemory.recordDuration(steps);
        }
        return;
      }
      if (token.startsWith("VEL_") && this.pendingNote) {
        const bucket = Number(token.slice(4));
        if (Number.isFinite(bucket)) {
          let rawVel = Math.max(0.2, Math.min(0.95, bucket / 8));
          
          if (heuristicSettings.velocitySmoothingEnabled) {
            rawVel = lastVelocity * 0.55 + rawVel * 0.45;
            
            const phraseBarIdx = Math.max(0, this.bar) % 4;
            let phraseMult = 1.0;
            if (phraseBarIdx === 0) phraseMult = 0.92;
            else if (phraseBarIdx === 1) phraseMult = 1.05;
            else if (phraseBarIdx === 2) phraseMult = 1.02;
            else if (phraseBarIdx === 3) phraseMult = 0.88;
            
            rawVel = Math.max(0.2, Math.min(0.98, rawVel * phraseMult));
            lastVelocity = rawVel;
          }

          this.pendingNote.velocity = rawVel;
          phraseMemory.recordNote(this.bar, this.pos, this.pendingNote.midi, this.pendingNote.durationSteps);
        }
        const q = this.absoluteQFor(this.bar, this.pos);
        this.pendingPosition ??= q;
        this.pendingNotes.push(this.pendingNote);
        this.pendingNote = null;
      }
    },
    absoluteQFor(bar, pos) {
      return Math.max(0, bar) * 4 + (Math.max(0, pos) * 4) / Math.max(1, this.grid);
    },
    flushTo(nextQ) {
      if (this.pendingNote) {
        const q = this.absoluteQFor(this.bar, this.pos);
        this.pendingPosition ??= q;
        this.pendingNotes.push(this.pendingNote);
        this.pendingNote = null;
      }
      if (this.pendingNotes.length && this.pendingPosition !== null) {
        const gap = Math.max(0, this.pendingPosition - this.lastQ);
        if (gap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(gap) });
        pushEvent({
          type: "note",
          notes: this.pendingNotes.map((note) => midiToTonePitch(note.midi)),
          perNoteDurations: this.pendingNotes.map((note) => durationToSecondsFromEventSteps(note.durationSteps, this.grid)),
          velocities: this.pendingNotes.map((note) => note.velocity),
          duration: 0,
          advance: 0,
        });
        this.lastQ = this.pendingPosition;
      }
      const finalGap = Math.max(0, nextQ - this.lastQ);
      if (finalGap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(finalGap) });
      this.lastQ = Math.max(this.lastQ, nextQ);
      this.pendingPosition = null;
      this.pendingNotes = [];
    },
    quartersToSeconds(quarters) {
      return (quarters * 60) / currentBpm;
    },
  };
}

function makeParser() {
  return makeEventParser();
}

async function warmPrompt() {
  h = zeroState();
  c = zeroState();
  lastSampledPitch = null;
  exactGuideBars = [];
  exactGuideTotalBars = 0;
  
  let allTokensStr;
  if (playingStyle === "midi" && midiTokens.length > 0) {
    allTokensStr = getTokensUpToBar(midiTokens, midiStartBar);
    buildExactGuide(midiTokens);
    // Also (re)build a probabilistic profile from this single track so the softer
    // similarity blending at lower slider values has statistics to draw on too.
    const singleTrackProfile = new StyleProfile();
    singleTrackProfile.name = "Current Track";
    singleTrackProfile.ingest(midiTokens, "current-track");
    activeStyleProfile = singleTrackProfile;
  } else if (playingStyle === "trained") {
    // No single sequence to replay — the model just improvises from an empty prompt,
    // steered the whole time by the learned style profile's statistics.
    allTokensStr = ["BOS", `BPM_${Math.round(currentBpm)}`, "GRID_64", "BAR", "POS_0"];
  } else if (playingStyle === "bach") {
    allTokensStr = [
      "BOS", "BPM_100", "GRID_64",
      "BAR",
      "POS_0", "NOTE_81", "DUR_4", "VEL_6",
      "POS_4", "NOTE_79", "DUR_4", "VEL_6",
      "POS_8", "NOTE_77", "DUR_4", "VEL_6",
      "POS_12", "NOTE_76", "DUR_4", "VEL_6",
      "POS_16", "NOTE_74", "DUR_4", "VEL_6",
      "POS_20", "NOTE_73", "DUR_4", "VEL_6",
      "POS_24", "NOTE_74", "DUR_16", "VEL_6",
      "BAR"
    ];
  } else if (playingStyle === "mozart") {
    allTokensStr = [
      "BOS", "BPM_132", "GRID_64",
      "BAR",
      "POS_0", "NOTE_55", "DUR_8", "VEL_6",
      "NOTE_67", "DUR_8", "VEL_6",
      "POS_8", "NOTE_74", "DUR_4", "VEL_7",
      "POS_12", "NOTE_71", "DUR_4", "VEL_6",
      "POS_16", "NOTE_79", "DUR_8", "VEL_7",
      "POS_24", "NOTE_78", "DUR_4", "VEL_6",
      "POS_28", "NOTE_76", "DUR_4", "VEL_6",
      "POS_32", "NOTE_74", "DUR_8", "VEL_7",
      "POS_40", "NOTE_72", "DUR_4", "VEL_6",
      "POS_44", "NOTE_71", "DUR_4", "VEL_6",
      "POS_48", "NOTE_69", "DUR_8", "VEL_6",
      "POS_56", "NOTE_67", "DUR_8", "VEL_7",
      "BAR",
      "POS_0", "NOTE_60", "DUR_16", "VEL_5",
      "NOTE_64", "DUR_16", "VEL_5",
      "NOTE_69", "DUR_16", "VEL_5",
      "POS_16", "NOTE_72", "DUR_4", "VEL_6",
      "POS_20", "NOTE_76", "DUR_4", "VEL_6",
      "POS_24", "NOTE_81", "DUR_8", "VEL_7",
      "POS_32", "NOTE_79", "DUR_4", "VEL_6",
      "POS_36", "NOTE_78", "DUR_4", "VEL_6",
      "POS_40", "NOTE_76", "DUR_4", "VEL_6",
      "POS_44", "NOTE_74", "DUR_4", "VEL_6",
      "POS_48", "NOTE_71", "DUR_8", "VEL_6",
      "POS_56", "NOTE_67", "DUR_8", "VEL_7",
      "BAR"
    ];
  } else if (playingStyle === "chopin") {
    allTokensStr = [
      "BOS", "BPM_68", "GRID_64",
      "BAR",
      "POS_0", "NOTE_36", "DUR_16", "VEL_4", "NOTE_72", "DUR_12", "VEL_6",
      "POS_8", "NOTE_43", "DUR_8", "VEL_4",
      "POS_12", "NOTE_75", "DUR_8", "VEL_6",
      "POS_16", "NOTE_48", "DUR_8", "VEL_4",
      "POS_24", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_12", "VEL_6",
      "POS_32", "NOTE_39", "DUR_16", "VEL_4", "NOTE_79", "DUR_16", "VEL_7",
      "POS_40", "NOTE_46", "DUR_8", "VEL_4",
      "POS_48", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_8", "VEL_6",
      "POS_56", "NOTE_55", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_5",
      "BAR"
    ];
  } else if (playingStyle === "debussy") {
    allTokensStr = [
      "BOS", "BPM_74", "GRID_64",
      "BAR",
      "POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_57", "DUR_16", "VEL_4", "NOTE_62", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_5",
      "POS_16", "NOTE_59", "DUR_16", "VEL_4", "NOTE_64", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_5",
      "POS_32", "NOTE_45", "DUR_32", "VEL_3", "NOTE_61", "DUR_16", "VEL_4", "NOTE_66", "DUR_16", "VEL_4", "NOTE_73", "DUR_16", "VEL_5",
      "POS_48", "NOTE_64", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_4", "NOTE_76", "DUR_16", "VEL_5",
      "BAR"
    ];
  } else if (playingStyle === "beethoven") {
    allTokensStr = [
      "BOS", "BPM_126", "GRID_64",
      "BAR",
      "POS_0", "NOTE_36", "DUR_8", "VEL_7", "NOTE_48", "DUR_8", "VEL_7", "NOTE_67", "DUR_4", "VEL_7",
      "POS_4", "NOTE_70", "DUR_4", "VEL_7",
      "POS_8", "NOTE_68", "DUR_8", "VEL_8",
      "POS_16", "NOTE_43", "DUR_8", "VEL_6", "NOTE_55", "DUR_8", "VEL_6", "NOTE_65", "DUR_4", "VEL_7",
      "POS_20", "NOTE_68", "DUR_4", "VEL_7",
      "POS_24", "NOTE_67", "DUR_8", "VEL_8",
      "BAR"
    ];
  } else if (playingStyle === "satie") {
    allTokensStr = [
      "BOS", "BPM_62", "GRID_64",
      "BAR",
      "POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_69", "DUR_20", "VEL_5",
      "POS_16", "NOTE_57", "DUR_24", "VEL_3", "NOTE_62", "DUR_24", "VEL_3", "NOTE_65", "DUR_24", "VEL_3",
      "BAR"
    ];
  } else {
    allTokensStr = ["BOS", "BPM_120", "GRID_64", "BAR", "POS_0"];
  }
  
  isWarmingUp = true;
  
  if (parser) {
    for (const token of allTokensStr) {
      parser.feed(token);
    }
  }
  
  const modelTokensStr = allTokensStr.slice(-256);
  const ids = modelTokensStr.map(tokenId).filter((id) => id !== null);
  
  for (const id of ids) {
    currentId = id;
    await stepModel(false);
  }
  
  isWarmingUp = false;
}

async function pumpTokens(targetCount = 64) {
  tempQueue = [];
  let steps = 0;
  const stepsCap = Math.max(150, targetCount * 4);
  while (tempQueue.length < targetCount && steps < stepsCap) {
    await stepModel(true);
    steps += 1;
  }

  const phraseBarIdx = (Math.max(0, parser ? parser.bar : 0) % 4) + 1;
  const isCadence = (phraseBarIdx === 4) && (parser ? parser.pos >= 48 : false);

  self.postMessage({
    action: "intelligence_state",
    detectedKey: keyDetector.getKeyName(),
    keyConfidence: keyDetector.confidence,
    phraseBar: phraseBarIdx,
    totalBar: Math.max(1, parser ? parser.bar + 1 : 1),
    isCadence: isCadence,
    epUsed: currentEpUsed
  });

  return {
    events: tempQueue,
    lastToken: itos[currentId] ?? "?"
  };
}

// Background Offline Full Song Renderer (No Audio Latency Limits)
async function renderFullSong(targetDurationSec = 60.0) {
  isRendering = true;
  cancelRenderRequested = false;
  
  const targetQuarters = Math.round((currentBpm / 60) * targetDurationSec);
  const targetBars = Math.max(4, Math.round(targetQuarters / 4));
  
  parser = makeParser();
  parser.reset();
  
  tempQueue = [];
  await warmPrompt();
  
  const renderStartTime = performance.now();
  let stepCount = 0;
  
  while (isRendering && !cancelRenderRequested) {
    await stepModel(true);
    stepCount++;
    
    const currentBar = Math.max(0, parser.bar);
    
    if (stepCount % 24 === 0) {
      const elapsedMs = performance.now() - renderStartTime;
      const progressPct = Math.min(99, Math.round((currentBar / targetBars) * 100));
      const estTotalMs = progressPct > 2 ? (elapsedMs / (progressPct / 100)) : 0;
      const estRemainingMs = Math.max(0, estTotalMs - elapsedMs);
      const noteCount = tempQueue.filter(e => e.type === "note").length;
      
      self.postMessage({
        action: "renderProgress",
        pct: progressPct,
        currentBar: currentBar + 1,
        targetBars: targetBars,
        noteCount: noteCount,
        elapsedMs: elapsedMs,
        estRemainingMs: estRemainingMs
      });
    }
    
    if (currentBar >= targetBars) {
      break;
    }
  }
  
  if (cancelRenderRequested) {
    isRendering = false;
    self.postMessage({ action: "renderCancelled" });
    return;
  }
  
  parser.flushTo(parser.absoluteQFor(parser.bar + 1, 0));
  isRendering = false;
  
  const totalNotes = tempQueue.filter(e => e.type === "note").length;
  self.postMessage({
    action: "renderComplete",
    events: tempQueue,
    totalBars: parser.bar + 1,
    bpm: currentBpm,
    noteCount: totalNotes
  });
}

// Message Router
self.onmessage = async function (e) {
  const data = e.data;
  
  switch (data.action) {
    case "init":
      try {
        console.log(`Worker loading model: ${data.activeModelName} (WASM CPU backend)`);
        activeModelName = data.activeModelName;
        stoi = data.vocab.stoi;
        itos = Object.fromEntries(Object.entries(data.vocab.itos).map(([key, value]) => [Number(key), value]));
        vocabSize = data.vocab.vocab_size;
        
        if (data.modelConfig) {
          numLayers = data.modelConfig.numLayers || 2;
          hiddenSize = data.modelConfig.hiddenSize || 2048;
        }
        
        if (session) {
          try { session.dispose(); } catch (e) {}
          session = null;
        }
        
        session = await ort.InferenceSession.create(data.modelBuffer, {
          executionProviders: ["wasm"],
          graphOptimizationLevel: "all",
        });
        
        updateModelDimensions(session);
        
        if (h) { try { h.dispose(); } catch (e) {} h = null; }
        if (c) { try { c.dispose(); } catch (e) {} c = null; }
        
        self.postMessage({
          action: "initialized",
          activeModelName: activeModelName,
          vocabSize: vocabSize,
          epUsed: currentEpUsed
        });
      } catch (err) {
        console.error("Worker initialization failed:", err);
        self.postMessage({ action: "error", message: `Init failed: ${err.message}` });
      }
      break;
      
    case "start":
      try {
        temperature = data.temperature;
        topK = data.topK;
        if (typeof data.bpm === "number" && data.bpm >= 30) {
          currentBpm = data.bpm;
        }
        playingStyle = data.playingStyle || "default";
        midiTokens = data.midiTokens || [];
        midiStartBar = data.midiStartBar || 0;
        midiBpmMode = data.midiBpmMode || "lock";
        
        if (data.heuristicSettings) {
          heuristicSettings = { ...heuristicSettings, ...data.heuristicSettings };
        }

        if (playingStyle === "trained" && (!activeStyleProfile || activeStyleProfile.isEmpty())) {
          self.postMessage({ action: "error", message: "No trained style is loaded yet. Train one or load a saved style file first." });
          break;
        }
        
        parser = makeParser();
        parser.reset();
        
        await warmPrompt();
        
        const chunkSize = data.chunkSize || 64;
        const initialBatch = await pumpTokens(chunkSize);
        self.postMessage({
          action: "started",
          events: initialBatch.events,
          lastToken: initialBatch.lastToken
        });
      } catch (err) {
        console.error("Worker start failed:", err);
        self.postMessage({ action: "error", message: `Start failed: ${err.message}` });
      }
      break;
      
    case "pump":
      try {
        temperature = data.temperature;
        topK = data.topK;
        if (data.heuristicSettings) {
          heuristicSettings = { ...heuristicSettings, ...data.heuristicSettings };
        }
        const chunkSize = data.chunkSize || 64;
        const batch = await pumpTokens(chunkSize);
        self.postMessage({
          action: "events",
          events: batch.events,
          lastToken: batch.lastToken
        });
      } catch (err) {
        console.error("Worker pump failed:", err);
        self.postMessage({ action: "error", message: `Pump failed: ${err.message}` });
      }
      break;

    case "renderSong":
      try {
        temperature = data.temperature;
        topK = data.topK;
        if (typeof data.bpm === "number" && data.bpm >= 30) {
          currentBpm = data.bpm;
        }
        playingStyle = data.playingStyle || "default";
        midiTokens = data.midiTokens || [];
        midiStartBar = data.midiStartBar || 0;
        midiBpmMode = data.midiBpmMode || "lock";
        
        if (data.heuristicSettings) {
          heuristicSettings = { ...heuristicSettings, ...data.heuristicSettings };
        }

        if (playingStyle === "trained" && (!activeStyleProfile || activeStyleProfile.isEmpty())) {
          self.postMessage({ action: "error", message: "No trained style is loaded yet. Train one or load a saved style file first." });
          break;
        }

        await renderFullSong(data.targetDurationSec || 60.0);
      } catch (err) {
        console.error("Worker render failed:", err);
        self.postMessage({ action: "error", message: `Render failed: ${err.message}` });
      }
      break;

    case "cancelRender":
      cancelRenderRequested = true;
      break;

    case "updateBpm":
      if (typeof data.bpm === "number" && data.bpm >= 30 && data.bpm <= 400) {
        currentBpm = data.bpm;
        userBpmOverride = true;
        self.postMessage({ action: "tempo", bpm: currentBpm });
      }
      break;

    case "updateHeuristics":
      if (data.heuristicSettings) {
        heuristicSettings = { ...heuristicSettings, ...data.heuristicSettings };
      }
      break;

    // Train a new style profile from one or more uploaded (and already tokenized) MIDI files.
    case "trainStyle": {
      try {
        const profile = new StyleProfile();
        profile.name = data.name || "Untitled Style";
        const datasets = data.datasets || [];
        for (const set of datasets) {
          profile.ingest(set.tokens || [], set.fileName || "upload.mid");
        }
        if (profile.isEmpty()) {
          self.postMessage({ action: "error", message: "Couldn't find any notes in the uploaded MIDI file(s) to train on." });
          break;
        }
        activeStyleProfile = profile;
        self.postMessage({
          action: "styleTrained",
          name: profile.name,
          totalNotes: profile.totalNotes,
          fileCount: profile.sourceFiles.length,
          fileNames: profile.sourceFiles
        });
      } catch (err) {
        console.error("Worker trainStyle failed:", err);
        self.postMessage({ action: "error", message: `Training failed: ${err.message}` });
      }
      break;
    }

    // Load a previously exported style profile (JSON) back into memory, e.g. after a page refresh.
    case "loadStyleProfile": {
      try {
        const profile = StyleProfile.fromJSON(data.profile);
        activeStyleProfile = profile;
        self.postMessage({
          action: "styleLoaded",
          name: profile.name,
          totalNotes: profile.totalNotes,
          fileCount: profile.sourceFiles.length,
          fileNames: profile.sourceFiles
        });
      } catch (err) {
        console.error("Worker loadStyleProfile failed:", err);
        self.postMessage({ action: "error", message: err.message || "Couldn't load that style file." });
      }
      break;
    }

    // Serialize the currently active trained style so the main thread can offer it as a download.
    case "exportStyle": {
      if (!activeStyleProfile || activeStyleProfile.isEmpty()) {
        self.postMessage({ action: "error", message: "No trained style to export yet." });
        break;
      }
      if (typeof data.name === "string" && data.name.trim()) {
        activeStyleProfile.name = data.name.trim();
      }
      self.postMessage({
        action: "styleExported",
        profile: activeStyleProfile.toJSON()
      });
      break;
    }
      
    case "stop":
      userBpmOverride = false;
      break;
  }
};