// gen.mjs - makes plan.bin and weights.bin for bp.c from Spotify's TF.js model.
//
//   npm i @spotify/basic-pitch@1.0.1 && npm i --legacy-peer-deps @tensorflow/tfjs@4.22.0 @tensorflow/tfjs-backend-wasm@4.22.0
//   node gen.mjs node_modules/@spotify/basic-pitch/model ../basicpitch
//
// It runs the model once in TF.js, so every shape is known. TF.js's own output goes to ref/.
import * as tf from "@tensorflow/tfjs";
import "@tensorflow/tfjs-backend-wasm";
import { readFileSync, writeFileSync, mkdirSync } from "node:fs";
import { createHash } from "node:crypto";

await tf.setBackend("wasm");
const [modelDir, out] = process.argv.slice(2);
const json = JSON.parse(readFileSync(`${modelDir}/model.json`)), bin = readFileSync(`${modelDir}/group1-shard1of1.bin`);
const model = await tf.loadGraphModel(tf.io.fromMemory({ modelTopology: json.modelTopology, weightSpecs: json.weightsManifest[0].weights, weightData: bin.buffer.slice(bin.byteOffset, bin.byteOffset + bin.length) }));
// weights.bin: every weight once
const weightAt = new Map(), kept = new Map(), parts = [];
let at = 0, size = 0;
for (const spec of json.weightsManifest[0].weights) {
  const bytes = spec.shape.reduce((a, b) => a * b, 1) * 4, data = bin.subarray(at, at + bytes), hash = createHash("sha256").update(data).digest("hex");
  at += bytes;
  if (!kept.has(hash)) { kept.set(hash, size); parts.push(data); size += bytes; }
  weightAt.set(spec.name, kept.get(hash) / 4);   // name -> offset in floats
}
if (at !== bin.length) throw new Error("the weights don't add up to the .bin");
writeFileSync(`${out}/weights.bin`, Buffer.concat(parts));
mkdirSync(`${out}/ref`, { recursive: true });

const nodes = Object.values(model.executor.graph.nodes), byName = new Map(nodes.map((n) => [n.name, n]));
const strip = (name) => name.split(":")[0];
// every node's value and shape for one window
const N = 43844;
const input = tf.tensor(Float32Array.from({ length: N }, (_, i) => Math.sin(i * 0.05) * 0.3 + Math.sin(i * 0.013) * 0.2 + (((i * 7919) % 1000) / 1000 - 0.5) * 0.05), [1, N, 1]);
const wanted = nodes.filter((n) => n.op !== "Const").map((n) => n.name);
const results = model.execute(input, wanted);
const value = new Map(wanted.map((name, i) => [name, results[i]]));
for (const n of nodes.filter((n) => n.op === "Const")) value.set(n.name, model.weights[n.name][0]);

// the outputs
const outputs = { frames: "Identity_1", onsets: "Identity_2", contours: "Identity" };
const need = new Set();
const visit = (name) => { name = strip(name); if (need.has(name)) return; need.add(name); for (const i of byName.get(name).inputNames) visit(i); };
visit(outputs.frames); visit(outputs.onsets); visit(outputs.contours);

// tensor kinds: 0 buffer, 1 weight, 2 alias (same data, new shape), 3 input
const OPS = { CONV: 1, PAD: 2, MIRRORPAD: 3, TRANSPOSE: 4, SLICE: 5, CONCAT: 6, UNARY: 7, BINARY: 8, REDUCE: 9, FREE: 10 };
const UNARY = { Neg: 1, Square: 2, Sqrt: 3, Log: 4, Sigmoid: 5, Relu: 6 };
const BINARY = { Mul: 1, Add: 2, AddV2: 2, Sub: 3, DivNoNan: 4 };
const REDUCE = { Sum: 1, Min: 2, Max: 3 };
const tensors = [], id = new Map(), ops = [];
const pad4 = (shape) => [...Array(4 - shape.length).fill(1), ...shape];
const tensorOf = (name) => {
  name = strip(name);
  if (id.has(name)) return id.get(name);
  const node = byName.get(name), v = value.get(name);
  if (node.op !== "Const") throw new Error(`${name} used before it is made`);
  if (v.dtype !== "float32") throw new Error(`int const ${name} used as a tensor`);
  if (!weightAt.has(name)) throw new Error(`no weight ${name}`);
  return add(name, { kind: 1, shape: pad4(v.shape), weight: weightAt.get(name) });
};
function add(name, t) { t.id = tensors.length; t.name = name; t.rank = value.get(name)?.shape.length ?? 4; tensors.push(t); id.set(name, t.id); return t.id; }
const ints = (name) => Array.from(value.get(strip(name)).dataSync());
const shapeOf = (name) => value.get(strip(name)).shape;

for (const node of nodes) {
  if (!need.has(node.name) || node.op === "Const") continue;
  const v = value.get(node.name), a = Object.fromEntries(Object.entries(node.attrParams).map(([k, p]) => [k, p.value]));
  if (v.dtype !== "float32") continue;   // shape arithmetic, already known
  const ins = node.inputNames, shape = pad4(v.shape), off = 4 - v.shape.length;
  const op = (code, inputs, params) => { const t = add(node.name, { kind: 0, shape }); ops.push({ code, out: t, ins: inputs.map(tensorOf), params, name: node.name }); };
  switch (node.op) {
    case "Placeholder": add(node.name, { kind: 3, shape }); break;
    case "Reshape": case "ExpandDims": case "Squeeze": case "Identity":
      add(node.name, { kind: 2, shape, of: tensorOf(ins[0]) }); break;
    case "Conv2D": case "_FusedConv2D": {
      const x = shapeOf(ins[0]), w = shapeOf(ins[1]), [, sh, sw] = a.strides, [, dh, dw] = a.dilations;
      let top = 0, left = 0;
      if (a.pad === "same") {
        top = Math.floor(Math.max((v.shape[1] - 1) * sh + (w[0] - 1) * dh + 1 - x[1], 0) / 2);
        left = Math.floor(Math.max((v.shape[2] - 1) * sw + (w[1] - 1) * dw + 1 - x[2], 0) / 2);
      } else if (a.pad !== "valid") throw new Error(`pad ${a.pad}`);
      const fused = node.op === "_FusedConv2D" ? a.fusedOps : [];
      const act = fused.includes("relu") ? 1 : fused.includes("sigmoid") ? 2 : 0;
      if (fused.some((f) => !["biasadd", "relu", "sigmoid"].includes(f))) throw new Error(`fused ${fused}`);
      op(OPS.CONV, fused.includes("biasadd") ? [ins[0], ins[1], ins[2]] : [ins[0], ins[1]], [w[0], w[1], sh, sw, dh, dw, top, left, act, fused.includes("biasadd") ? 1 : 0]);
      break;
    }
    case "Pad": case "MirrorPad": {
      const p = ints(ins[1]), before = [0, 0, 0, 0];
      for (let d = 0; d < p.length / 2; d++) before[d + off] = p[2 * d];
      if (node.op === "MirrorPad" && String(a.mode ?? a.padMode).toLowerCase() !== "reflect") throw new Error(`mirror ${JSON.stringify(a)}`);
      op(node.op === "Pad" ? OPS.PAD : OPS.MIRRORPAD, [ins[0]], before);
      break;
    }
    case "Transpose": {
      const perm = ints(ins[1]), p4 = [0, 1, 2, 3];
      perm.forEach((pd, d) => { p4[d + off] = pd + off; });
      op(OPS.TRANSPOSE, [ins[0]], p4);
      break;
    }
    case "StridedSlice": {
      const x = shapeOf(ins[0]), b = ints(ins[1]), e = ints(ins[2]), s = ints(ins[3]);
      // x[:, None, :] and the like: only a new shape
      if (v.size === value.get(strip(ins[0])).size && s.every((st) => st === 1) && (a.newAxisMask || b.every((q, d) => q === 0 && (a.endMask >> d & 1 || e[d] >= x[d])))) {
        add(node.name, { kind: 2, shape, of: tensorOf(ins[0]) });
        break;
      }
      if (a.ellipsisMask || a.newAxisMask) throw new Error("slice masks");
      const xo = 4 - x.length, begin = [0, 0, 0, 0], stride = [1, 1, 1, 1], size = pad4(x);
      for (let d = 0; d < x.length; d++) {
        if (d >= b.length) continue;
        const st = s[d], dim = x[d];
        let lo = a.beginMask >> d & 1 ? (st > 0 ? 0 : dim - 1) : (b[d] < 0 ? b[d] + dim : b[d]);
        let hi = a.endMask >> d & 1 ? (st > 0 ? dim : -1) : (e[d] < 0 ? e[d] + dim : e[d]);
        lo = Math.max(Math.min(lo, st > 0 ? dim : dim - 1), st > 0 ? 0 : -1);
        hi = Math.max(Math.min(hi, st > 0 ? dim : dim - 1), st > 0 ? 0 : -1);
        if (a.shrinkAxisMask >> d & 1) hi = lo + 1;
        begin[d + xo] = lo; stride[d + xo] = st; size[d + xo] = Math.max(0, Math.ceil((hi - lo) / st));
      }
      if (size.reduce((p, q) => p * q, 1) !== v.size) throw new Error(`slice size ${node.name}: ${size} vs ${v.shape}`);
      // the node itself (shrunk axes gone) is an alias of the 4-D slice
      const t = add(node.name + "/slice", { kind: 0, shape: size });
      tensors[t].rank = 4;
      ops.push({ code: OPS.SLICE, out: t, ins: [tensorOf(ins[0])], params: [...begin, ...stride], name: node.name });
      id.delete(node.name + "/slice");
      add(node.name, { kind: 2, shape, of: t });
      break;
    }
    case "ConcatV2": case "Pack": {
      const parts = node.op === "Pack" ? ins : ins.slice(0, -1);
      let axis = node.op === "Pack" ? a.axis : ints(ins[ins.length - 1])[0];
      if (axis < 0) axis += v.shape.length;
      // pack: a new axis on each part, then a concat
      const srcs = parts.map((p) => {
        if (node.op !== "Pack") return tensorOf(p);
        const s = [...shapeOf(p)]; s.splice(axis, 0, 1);
        const t = add(strip(p) + "/packed" + node.name, { kind: 2, shape: pad4(s), of: tensorOf(p) });
        tensors[t].rank = s.length;
        return t;
      });
      const t = add(node.name, { kind: 0, shape });
      ops.push({ code: OPS.CONCAT, out: t, ins: srcs, params: [axis + off], name: node.name });
      break;
    }
    case "Neg": case "Square": case "Sqrt": case "Log": case "Sigmoid": case "Relu":
      op(OPS.UNARY, [ins[0]], [UNARY[node.op]]); break;
    case "Mul": case "Add": case "AddV2": case "Sub": case "DivNoNan":
      op(OPS.BINARY, [ins[0], ins[1]], [BINARY[node.op]]); break;
    case "Sum": case "Min": case "Max": {
      const x = shapeOf(ins[0]), xo = 4 - x.length, axes = ints(ins[1]).map((d) => (d < 0 ? d + x.length : d) + xo);
      const mask = axes.reduce((m, d) => m | 1 << d, 0), kept = pad4(x).map((n, d) => (mask >> d & 1 ? 1 : n));
      // reduced with its axes kept as 1s; without keep_dims the node's own shape is an alias of that
      const t = add(node.name + "/kept", { kind: 0, shape: kept });
      ops.push({ code: OPS.REDUCE, out: t, ins: [tensorOf(ins[0])], params: [REDUCE[node.op], mask], name: node.name });
      id.delete(node.name + "/kept");
      add(node.name, { kind: 2, shape, of: t });
      break;
    }
    default: throw new Error(`op ${node.op} (${node.name})`);
  }
}
const base = (t) => (tensors[t].kind === 2 ? base(tensors[t].of) : t);

// harmonic stacking (shifted copies of X, concatenated) and the conv after it become one HCONV
OPS.HCONV = 11;
const made = new Map(ops.map((o) => [o.out, o]));
for (const conv of ops.filter((o) => o.code === OPS.CONV)) {
  const [, , sh, sw, dh, dw] = conv.params;
  const cut = made.get(base(conv.ins[0]));
  if (sh !== 1 || sw !== 1 || dh !== 1 || dw !== 1 || cut?.code !== OPS.SLICE || cut.params.some((v, i) => (i < 4 ? v !== 0 : v !== 1))) continue;
  const cat = made.get(base(cut.ins[0]));
  if (cat?.code !== OPS.CONCAT || cat.params[0] !== 3) continue;
  let x = null, ok = true;
  const shifts = cat.ins.map((part) => {
    let t = base(part), shift = 0, o = made.get(t);
    if (o?.code === OPS.PAD) {   // zeros in front, along frequency
      const [p0, p1, p2, p3] = o.params;
      if (p0 || p1 || p3) ok = false;
      shift -= p2; t = base(o.ins[0]); o = made.get(t);
    }
    if (o?.code === OPS.SLICE) {   // from a bin on, along frequency
      const [b0, b1, b2, b3, ...strides] = o.params;
      if (b0 || b1 || b3 || strides.some((s) => s !== 1)) ok = false;
      shift += b2; t = base(o.ins[0]);
    }
    if (tensors[part].shape[3] !== 1) ok = false;
    if (x === null) x = t; else if (x !== t) ok = false;
    return shift;
  });
  if (!ok || tensors[x].shape[0] !== 1 || tensors[x].shape[3] !== 1) continue;
  conv.code = OPS.HCONV;
  conv.ins = [x, ...conv.ins.slice(1)];
  conv.params = [...conv.params, shifts.length, ...shifts];
}
// drop ops nothing uses any more
for (let removed = true; removed;) {
  const used = new Set([base(id.get(outputs.frames)), base(id.get(outputs.onsets)), base(id.get(outputs.contours))]);
  for (const o of ops) for (const t of o.ins) used.add(base(t));
  const before = ops.length;
  ops.splice(0, ops.length, ...ops.filter((o) => used.has(o.out)));
  removed = ops.length < before;
}

// FREE each buffer after its last use (aliases count)
const lastUse = new Map();
ops.forEach((o, i) => o.ins.forEach((t) => lastUse.set(base(t), i)));
const keep = new Set([base(id.get(outputs.frames)), base(id.get(outputs.onsets)), base(id.get(outputs.contours))]);
const program = [];
ops.forEach((o, i) => {
  program.push(o);
  for (const [t, last] of lastUse) if (last === i && tensors[t].kind === 0 && !keep.has(t)) program.push({ code: OPS.FREE, out: t, ins: [], params: [] });
});
// plan.bin: int32 words, as bp_load reads them
const words = [0x42504c4f, tensors.length, program.length, id.get(outputs.frames), id.get(outputs.onsets), id.get("input_2"), id.get(outputs.contours)];
for (const t of tensors) words.push(t.kind, t.kind === 1 ? t.weight : t.kind === 2 ? t.of : 0, ...t.shape);
for (const o of program) words.push(o.code, o.out, o.ins.length, ...o.ins, o.params.length, ...o.params);
writeFileSync(`${out}/plan.bin`, Buffer.from(Int32Array.from(words).buffer));
// TF.js's output, to check bp.c against
writeFileSync(`${out}/ref/ref_input.f32`, Buffer.from(input.dataSync().buffer));
writeFileSync(`${out}/ref/ref_frames.f32`, Buffer.from(value.get(outputs.frames).dataSync().buffer));
writeFileSync(`${out}/ref/ref_onsets.f32`, Buffer.from(value.get(outputs.onsets).dataSync().buffer));
const counts = {}; for (const o of program) counts[o.code] = (counts[o.code] || 0) + 1;
console.log(`plan: ${tensors.length} tensors, ${program.length} ops`, counts, "frames", value.get(outputs.frames).shape, "onsets", value.get(outputs.onsets).shape);
// every op's output, to find where bp.c differs
const refs = program.filter((o) => o.code !== OPS.FREE).map((o) => value.get(o.name).dataSync());
writeFileSync(`${out}/ref/ref_ops.f32`, Buffer.concat(refs.map((r) => Buffer.from(Float32Array.from(r).buffer))));
// what each op is
writeFileSync(`${out}/ref/plan.txt`, program.map((o, i) => `${i} ${o.code} ${tensors[o.out].name} ${tensors[o.out].shape} <- ${o.ins.map((t) => tensors[t].name.replace("StatefulPartitionedCall/model_1/", "")).join(", ")} | ${o.params.join(" ")}`).join("\n"));