#!/usr/bin/env node
// Play the drum kit with CHOPPED UMS instead of a soundfont.
//
//   node um-drums.mjs <hits.json> <out.wav>
//
// midi-drums.mjs says why the kit was a soundfont in the first place: "a 'd'ummm'
// can imply a beat, it cannot BE a hi-hat". That is true of a whole um. It stops
// being true once the um is cut to 50ms, pitched, filtered to one band and driven
// into clipping -- at that point what survives is an attack transient and a
// coloured burst of noise, which is all a kit part is. The vowel is the casualty
// and that is the point: nobody should be able to hear a word in the hi-hat.
//
// THE CLIPS COME FROM OUTSIDE THE MELODY PALETTE, on purpose. 879 clips have a
// measured pitch but were never accepted into accepted.json, and every one of
// them has audio in wav48/. Drums are the one voice that wants the rejects: the
// qualities that lost a clip its place in a melody (short, noisy, badly formed
// vowel, off-pitch) cost nothing here, and spending accepted clips on percussion
// would take them from the voices that do need them -- with `unique:true` the
// pool is genuinely finite.
//
//   POOL    distinct clips drawn per role (default 6)
//   GAIN    overall output gain (default 1.0)
//   MATCH   a wav to match the RMS of, so a kit swap is a fair A/B: the same
//           DRUM_GAIN downstream then lands both kits in the same place
//   SEED    clip choice and the per-hit jitter are deterministic
import { readFileSync, writeFileSync, existsSync } from "node:fs";
import { execFileSync } from "node:child_process";
import path from "node:path";
import { SONG_DATA } from "./paths.mjs";

const [HITS, OUT] = process.argv.slice(2);
if (!HITS || !OUT) { console.error("usage: node um-drums.mjs <hits.json> <out.wav>"); process.exit(1); }
const SR = 48000;
const POOL = Number(process.env.POOL ?? 6);
const GAIN = Number(process.env.GAIN ?? 1.0);
const MATCH = process.env.MATCH ?? "";
let seed = Number(process.env.SEED ?? 12345);
const rnd = () => { seed = (seed * 1103515245 + 12345) & 0x7fffffff; return seed / 0x7fffffff; };

// ---- roles -----------------------------------------------------------------
// Each role is a recipe for turning a vowel into a percussion sound: how far to
// pitch it, which single band to keep, how hard to clip it, how long to let it
// ring. `pitch` is which end of the candidate pitch range to draw clips from --
// a kick made from an already-low um needs less resampling and keeps more of its
// attack, and resampling is the step that smears attacks.
const ROLE = {
  kick:  { rate: 0.55, lp: 110,  n: 3, drive: 3.0, len: 0.22, pitch: "low"  },
  snare: { rate: 1.00, hp: 220,  lp: 3600, n: 2, drive: 6.0, len: 0.16, pitch: "any" },
  rim:   { rate: 1.40, hp: 400,  lp: 4500, n: 2, drive: 8.0, len: 0.06, pitch: "any" },
  clap:  { rate: 1.10, hp: 500,  lp: 4000, n: 2, drive: 6.0, len: 0.13, pitch: "any" },
  hatC:  { rate: 1.60, hp: 5000, n: 3, drive: 4.0, len: 0.05, pitch: "high" },
  hatO:  { rate: 1.60, hp: 5000, n: 3, drive: 4.0, len: 0.26, pitch: "high" },
  tom:   { rate: 0.75, lp: 500,  n: 2, drive: 2.5, len: 0.26, pitch: "low"  },
  crash: { rate: 1.30, hp: 3000, n: 2, drive: 3.0, len: 0.90, pitch: "high" },
  ride:  { rate: 1.50, hp: 4000, n: 2, drive: 3.0, len: 0.50, pitch: "high" },
};
const roleOf = (n) => {
  if (n === 35 || n === 36) return "kick";
  if (n === 38 || n === 40) return "snare";
  if (n === 37) return "rim";
  if (n === 39) return "clap";
  if (n === 42 || n === 44) return "hatC";
  if (n === 46) return "hatO";
  if (n === 49 || n === 52 || n === 55 || n === 57) return "crash";
  if (n === 51 || n === 53 || n === 59) return "ride";
  if (n >= 41 && n <= 50) return "tom";
  return "snare";
};

// ---- the candidate pool ----------------------------------------------------
const acc = JSON.parse(readFileSync(path.join(SONG_DATA, "accepted.json"), "utf8"));
const accSet = new Set(Array.isArray(acc) ? acc : (acc.accepted ?? []));
const cp = JSON.parse(readFileSync(path.join(SONG_DATA, "corepitch.json"), "utf8"));
const cand = Object.entries(cp)
  .filter(([id, v]) => !accSet.has(id) && v.f0 > 0 && v.dur >= 0.12)
  .filter(([id]) => existsSync(path.join(SONG_DATA, "wav48", `${id.split("@")[0]}.wav`)))
  .map(([id, v]) => ({ id, f0: v.f0, dur: v.dur }))
  .sort((a, b) => a.f0 - b.f0);
if (cand.length < 20) throw new Error(`only ${cand.length} candidate clips outside the palette`);

const third = Math.floor(cand.length / 3);
const bucket = { low: cand.slice(0, third), high: cand.slice(2 * third), any: cand };

// ---- fetch and shape one clip ----------------------------------------------
const cache = new Map();
function clipAudio(c) {
  if (cache.has(c.id)) return cache.get(c.id);
  const [vid, start] = c.id.split("@");
  const b = execFileSync("ffmpeg", ["-v", "error", "-ss", start, "-t", String(c.dur),
    "-i", path.join(SONG_DATA, "wav48", `${vid}.wav`), "-ac", "1", "-ar", String(SR),
    "-f", "s16le", "-"], { maxBuffer: 1 << 26 });
  const a = new Int16Array(b.buffer, b.byteOffset, b.length / 2);
  const x = new Float64Array(a.length);
  let pk = 0;
  for (let i = 0; i < a.length; i += 1) { x[i] = a[i] / 32768; pk = Math.max(pk, Math.abs(x[i])); }
  // Start AT THE ONSET. A clip's window opens before the vowel does, and a drum
  // whose attack is preceded by 40ms of nothing plays late however well it is
  // placed -- the one thing a kit cannot survive.
  let s = 0;
  for (let i = 0; i < x.length; i += 1) if (Math.abs(x[i]) > 0.18 * pk) { s = i; break; }
  const out = x.slice(s);
  cache.set(c.id, out);
  return out;
}

const lowpass = (x, f, times) => {
  const a = 1 - Math.exp(-2 * Math.PI * f / SR);
  for (let t = 0; t < times; t += 1) { let y = 0; for (let i = 0; i < x.length; i += 1) { y += a * (x[i] - y); x[i] = y; } }
  return x;
};
const highpass = (x, f, times) => {
  const a = 1 - Math.exp(-2 * Math.PI * f / SR);
  for (let t = 0; t < times; t += 1) { let y = 0; for (let i = 0; i < x.length; i += 1) { y += a * (x[i] - y); x[i] -= y; } }
  return x;
};

function voice(role, c, vel, jitter) {
  const R = ROLE[role];
  const src = clipAudio(c);
  const rate = R.rate * (1 + jitter);
  const n = Math.min(Math.round(R.len * SR), Math.floor((src.length - 2) / rate));
  if (n <= 8) return null;
  const y = new Float64Array(n);
  for (let i = 0; i < n; i += 1) {              // linear-interpolated resample = pitch
    const p = i * rate, k = Math.floor(p), fr = p - k;
    y[i] = src[k] * (1 - fr) + src[k + 1] * fr;
  }
  if (R.hp) highpass(y, R.hp, R.n);
  if (R.lp) lowpass(y, R.lp, R.n);
  let pk = 0;
  for (let i = 0; i < n; i += 1) pk = Math.max(pk, Math.abs(y[i]));
  if (pk <= 0) return null;
  const atk = Math.max(2, Math.round(0.002 * SR)), tau = (R.len * SR) / 3;
  const amp = Math.pow(vel / 127, 1.2);
  for (let i = 0; i < n; i += 1) {
    let e = i < atk ? i / atk : Math.exp(-(i - atk) / tau);
    if (i > n - 240) e *= (n - i) / 240;         // 5ms tail so nothing clicks
    y[i] = Math.tanh((y[i] / pk) * R.drive) * e * amp;
  }
  return y;
}

// ---- render ----------------------------------------------------------------
const raw = JSON.parse(readFileSync(HITS, "utf8"));

// COLLAPSE COINCIDENT HITS OF THE SAME ROLE. This theme puts GM 35 and GM 36 --
// two different kick sounds -- on the same tick at full velocity, 193 times out
// of 193. A soundfont takes that as reinforcement: two samples of the same drum
// sum into one bigger kick. Ums cannot, because the two clips are unrelated
// vowels; stacking them smears the one thing the hit is carrying, its attack,
// and doubles the level while it does it. One hit per role per instant.
const hits = [];
for (const h of raw.slice().sort((a, b) => a.t - b.t)) {
  const prev = hits[hits.length - 1];
  if (prev && roleOf(prev.note) === roleOf(h.note) && h.t - prev.t < 0.010) {
    if (h.vel > prev.vel) prev.vel = h.vel;
    continue;
  }
  hits.push({ ...h });
}
if (hits.length !== raw.length) {
  console.log(`collapsed ${raw.length - hits.length} coincident same-role hits (${raw.length} -> ${hits.length})`);
}
const pools = {};
for (const role of Object.keys(ROLE)) {
  const b = bucket[ROLE[role].pitch];
  pools[role] = Array.from({ length: POOL }, () => b[Math.floor(rnd() * b.length)]);
}
const end = Math.max(...hits.map((h) => h.t)) + 1.5;
const used = {};
let placed = 0;

// Render each role into its OWN buffer so the roles can be balanced against each
// other before they are summed. Matching the kit as a whole is not enough: this
// theme is 654 hats to 193 kicks, so the hats dominate the total RMS, and a
// single global match set them 3.05x the soundfont's while leaving the kicks at
// 0.86x. Level is not the same thing as balance.
const roleBuf = {};
const drawn = [];
for (let i = 0; i < hits.length; i += 1) {
  const h = hits[i];
  const role = roleOf(h.note);
  const c = pools[role][i % POOL];
  const v = voice(role, c, h.vel, (rnd() - 0.5) * 0.06);
  if (!v) continue;
  // Record WHICH clip played WHEN, in the shape run-line.mjs reads, so the kit can
  // be drawn as well as heard. It has to come from the render loop rather than be
  // reconstructed afterwards: the pool draw and the jitter both consume the same
  // PRNG, so anything that recomputes the choice out of step gets a different clip
  // than the one in the wav -- a picture that disagrees with the sound.
  drawn.push({
    slotStart: h.t, slotDur: ROLE[role].len, noteDur: ROLE[role].len,
    video: c.id.split("@")[0], srcStart: Number(c.id.split("@")[1]), srcDur: c.dur,
    role,
  });
  roleBuf[role] ??= new Float64Array(Math.ceil(end * SR));
  const b = roleBuf[role], at = Math.round(h.t * SR);
  for (let k = 0; k < v.length && at + k < b.length; k += 1) b[at + k] += v[k];
  used[role] = (used[role] ?? 0) + 1;
  placed += 1;
}

// Balance each role against the same role in the reference kit, measured only at
// hits that role has to ITSELF (nothing else within 60ms) -- a hat measured under
// a kick is measuring the kick.
const buf = new Float64Array(Math.ceil(end * SR));
const trims = {};
if (MATCH && existsSync(MATCH)) {
  const b = execFileSync("ffmpeg", ["-v", "error", "-i", MATCH, "-ac", "1", "-ar", String(SR),
    "-f", "s16le", "-"], { maxBuffer: 1 << 28 });
  const ref = new Int16Array(b.buffer, b.byteOffset, b.length / 2);
  const rmsAt = (a, ts, d) => {
    let s = 0, n = 0;
    for (const t of ts) {
      const i0 = Math.round(t * SR), i1 = Math.round((t + d) * SR);
      for (let i = i0; i < i1 && i < a.length; i += 1) { s += a[i] ** 2; n += 1; }
    }
    return n ? Math.sqrt(s / n) : 0;
  };
  // MEASURE EACH ROLE IN ITS OWN BAND, not in a window it happens to have to
  // itself. Isolation does not work here: this groove puts a hat on every kick,
  // so the kick never gets a clean window and falls back to no correction at all
  // -- which is how the loud role ends up trimmed and the quiet one does not.
  // Filtering to the band the role occupies separates them wherever they sit.
  const BAND = {
    kick: { lp: 150 }, tom: { lp: 500 },
    snare: { hp: 200, lp: 3500 }, rim: { hp: 200, lp: 3500 }, clap: { hp: 200, lp: 3500 },
    hatC: { hp: 4000 }, hatO: { hp: 4000 }, crash: { hp: 4000 }, ride: { hp: 4000 },
  };
  const banded = (src, scale, b) => {
    const x = new Float64Array(src.length);
    for (let i = 0; i < src.length; i += 1) x[i] = src[i] * scale;
    if (b.hp) highpass(x, b.hp, 2);
    if (b.lp) lowpass(x, b.lp, 2);
    return x;
  };
  for (const role of Object.keys(roleBuf)) {
    const ts = hits.filter((h) => roleOf(h.note) === role).map((h) => h.t);
    const b = BAND[role] ?? {};
    const win = Math.min(0.10, ROLE[role].len);
    const r = rmsAt(banded(ref, 1 / 32768, b), ts, win);
    const u = rmsAt(banded(roleBuf[role], 1, b), ts, win);
    trims[role] = u > 0 ? r / u : 1;
  }
}
for (const role of Object.keys(roleBuf)) {
  const g = trims[role] ?? 1, b = roleBuf[role];
  for (let i = 0; i < buf.length; i += 1) buf[i] += b[i] * g;
}

let pk = 0;
for (let i = 0; i < buf.length; i += 1) pk = Math.max(pk, Math.abs(buf[i]));
// No global rescale on top of the per-role trims -- that would match the level
// twice and undo the balance. The trims already put every role where the
// reference kit has it, so the total follows.
let g = GAIN;
if (Object.keys(trims).length) {
  console.log(`per-role trim vs ${path.basename(MATCH)}: ` +
    Object.entries(trims).map(([k, v]) => `${k} x${v.toFixed(2)}`).join("  "));
}
if (pk * g > 0.98) { console.log(`  peak would be ${(pk * g).toFixed(2)}, holding at 0.98`); g = 0.98 / pk; }

const out = Buffer.alloc(44 + buf.length * 4);
out.write("RIFF", 0); out.writeUInt32LE(36 + buf.length * 4, 4); out.write("WAVE", 8);
out.write("fmt ", 12); out.writeUInt32LE(16, 16); out.writeUInt16LE(1, 20);
out.writeUInt16LE(2, 22); out.writeUInt32LE(SR, 24); out.writeUInt32LE(SR * 4, 28);
out.writeUInt16LE(4, 32); out.writeUInt16LE(16, 34);
out.write("data", 36); out.writeUInt32LE(buf.length * 4, 40);
for (let i = 0; i < buf.length; i += 1) {
  const s = Math.max(-1, Math.min(1, buf[i] * g));
  const v = Math.round(s * 32767);
  out.writeInt16LE(v, 44 + i * 4); out.writeInt16LE(v, 46 + i * 4);
}
writeFileSync(OUT, out);

// PLAN_OUT lets the kit be SEEN as well as heard: run-line.mjs draws these as a
// 2x2 in whichever container the tune has spare. The plan is the record of what
// actually played, taken from the render loop itself, so the picture cannot
// disagree with the sound.
if (process.env.PLAN_OUT) {
  writeFileSync(process.env.PLAN_OUT, JSON.stringify({ voices: [{ name: "drums", plan: drawn }] }, null, 1));
  const g = drawn.slice(1).map((n, i) => n.slotStart - drawn[i].slotStart).sort((a, b) => a - b);
  console.log(`  drum plan: ${drawn.length} hits, IOI p50 ${g[Math.floor(g.length / 2)].toFixed(3)}s, ` +
    `${(100 * g.filter((x) => x <= 0.15).length / g.length).toFixed(0)}% under 0.15s -> ${process.env.PLAN_OUT}`);
}
console.log(`${placed}/${hits.length} hits from ${cand.length} non-palette clips -> ${OUT}`);
console.log(`  ${Object.entries(used).map(([k, v]) => `${k} ${v}`).join("  ")}`);
console.log(`  peak ${(pk * g).toFixed(3)}  ${(buf.length / SR).toFixed(1)}s`);