// Setup and output stages of the cymbal2 render, run through numbl exactly as
// the webapp runs them. The GPU path replaces only cymbal2_sim: everything the
// time loop needs (grid operators, per-wavenumber inverses, mallet matrices,
// stroke schedule) is computed here by the unmodified MATLAB code and pulled
// out as tensors.

import { readFileSync, readdirSync } from 'node:fs';
import { fileURLToPath } from 'node:url';
import path from 'node:path';
import { executeCode, RTV } from 'numbl';

const MATLAB_DIR = path.resolve(
  path.dirname(fileURLToPath(import.meta.url)),
  '../../../simulated_instruments/cymbal2/matlab',
);

export function loadMFiles() {
  return readdirSync(MATLAB_DIR)
    .filter((f) => f.endsWith('.m'))
    .sort()
    .map((name) => ({
      name,
      source: readFileSync(path.join(MATLAB_DIR, name), 'utf8'),
    }));
}

// MATLAB lines overriding defaults from a plain params object.
function pOverrides(params) {
  const lines = [];
  for (const [k, v] of Object.entries(params ?? {})) {
    if (typeof v === 'number') lines.push(`P.${k} = ${v};`);
    else if (typeof v === 'boolean') lines.push(`P.${k} = ${v ? 1 : 0};`);
  }
  return lines.join('\n');
}

function stage(mfiles, source, inputs) {
  return executeCode(
    source,
    { initialVariableValues: inputs, displayResults: false },
    mfiles,
    'stage.m',
  );
}

function tensor(v, name) {
  if (!v || v.kind !== 'tensor') throw new Error(`expected tensor for ${name}`);
  return v; // {kind, shape, data: Float64Array} column-major
}

function scalar(v, name) {
  if (typeof v === 'number') return v;
  if (v && v.kind === 'tensor' && v.data.length === 1) return v.data[0];
  throw new Error(`expected scalar for ${name}`);
}

function f32(v) {
  const t = v.data;
  const out = new Float32Array(t.length);
  out.set(t);
  return out;
}

// The stroke-schedule block is the one in cymbal2_sim.m, verbatim, so the GPU
// render integrates exactly the schedule the CPU render would.
const SCHEDULE_SRC = `
K = max(1, round(P.strike_count));
tstr = P.pre_delay_ms * 1e-3 + (0:K - 1) / max(1e-3, P.strike_rate_hz);
vstr = zeros(1, K);
state = P.seed * 2654435761;
for i = 1:K
  state = mod(state * 1664525 + 1013904223, 4294967296);
  x = state / 4294967296;
  vstr(i) = max(0.05, P.strike_velocity * (1 + P.velocity_spread * (2 * x - 1)));
end
gap = P.strike_gap_m;
if gap <= 0
  gap = 1.05 * max(vstr)^2 / (2 * P.strike_accel) + 0.005;
end
travel = gap ./ vstr + vstr / (2 * P.strike_accel);
tlaunch = tstr - travel;
off = -min(0, min(tlaunch));
tlaunch = tlaunch + off;
duration = off + (K - 1) / max(1e-3, P.strike_rate_hz) + P.tail_s;
nsteps = floor(duration * P.sample_rate * P.oversample);
`;

export function runSetup(params) {
  const mfiles = loadMFiles();
  const src = `
P = cymbal2_defaults();
${pOverrides(params)}
S = cymbal2_system(P);
MB = cymbal2_mallet(P);
G = S.G;
o_D1 = G.D1; o_D2 = G.D2;
o_F = G.F; o_Gi = G.Gi;
o_Rd = G.Rd; o_R2d = G.R2d;
o_AiB = S.AiB; o_AiC = S.AiC; o_Ai = S.Ai; o_Fi = S.Fi;
o_ASiB = MB.ASiB; o_ASiC = MB.ASiC;
o_vM = MB.vM; o_vR = MB.vR; o_vK = MB.vK; o_tipcol = MB.tipcol;
${SCHEDULE_SRC}
o_tlaunch = tlaunch; o_vstr = vstr;
c_nrt = G.nrt; c_Nr = G.Nr; c_nth = G.nth; c_hr = G.hr; c_hth = G.hth;
c_ht = S.ht; c_Eh = S.Eh; c_Rs = S.Rs; c_rhoh = S.rhoh;
c_fsc = S.forceScale; c_nlr0 = S.eomRow0; c_nDropped = S.nDropped;
c_ls = S.ls; c_js = S.js; c_lp = S.lp; c_jp = S.jp;
c_strikeR = S.strikeRadiusM; c_pickupR = S.pickupRadiusM;
c_tipi = MB.tipi; c_nfree = MB.nfree; c_ps = MB.ps;
c_KM = P.head_stiffness; c_alM = P.head_alpha;
c_aH = P.strike_accel; c_aHr = P.retract_accel;
c_holdSteps = max(1, round(P.strike_hold_ms * 1e-3 / S.ht));
c_missDepth = gap + 0.06;
c_gap = gap; c_K = K; c_nsteps = nsteps;
c_simRate = P.sample_rate * P.oversample;
c_thickness = P.thickness_m;
`;
  const r = stage(mfiles, src);
  const v = r.variableValues;
  const g = (n) => tensor(v[n], n);
  const s = (n) => scalar(v[n], n);

  const setup = {
    mfiles,
    ops: {
      D1: f32(g('o_D1')), D2: f32(g('o_D2')),
      F: f32(g('o_F')), Gi: f32(g('o_Gi')),
      Rd: f32(g('o_Rd')), R2d: f32(g('o_R2d')),
      AiB: f32(g('o_AiB')), AiC: f32(g('o_AiC')),
      Ai: f32(g('o_Ai')), Fi: f32(g('o_Fi')),
      ASiB: f32(g('o_ASiB')), ASiC: f32(g('o_ASiC')),
      vM: f32(g('o_vM')), vR: f32(g('o_vR')),
      vK: f32(g('o_vK')), tipcol: f32(g('o_tipcol')),
    },
    tlaunch: Array.from(g('o_tlaunch').data),
    vstr: Array.from(g('o_vstr').data),
    nrt: s('c_nrt'), Nr: s('c_Nr'), nth: s('c_nth'),
    hr: s('c_hr'), hth: s('c_hth'),
    ht: s('c_ht'), Eh: s('c_Eh'), Rs: s('c_Rs'), rhoh: s('c_rhoh'),
    fsc: s('c_fsc'), nlr0: s('c_nlr0'), nDropped: s('c_nDropped'),
    ls: s('c_ls'), js: s('c_js'), lp: s('c_lp'), jp: s('c_jp'),
    strikeR: s('c_strikeR'), pickupR: s('c_pickupR'),
    tipi: s('c_tipi'), nfree: s('c_nfree'), ps: s('c_ps'),
    KM: s('c_KM'), alM: s('c_alM'), aH: s('c_aH'), aHr: s('c_aHr'),
    holdSteps: s('c_holdSteps'), missDepth: s('c_missDepth'),
    gap: s('c_gap'), K: s('c_K'), nsteps: s('c_nsteps'),
    simRate: s('c_simRate'), thickness: s('c_thickness'),
  };
  if (setup.vstr.length !== setup.K) throw new Error('schedule length mismatch');
  return setup;
}

// The output stage (DC blocker, decimation, fade, normalize), run through the
// unmodified cymbal2_output.m with the GPU-computed acceleration handed in.
export function runOutput(mfiles, params, acc) {
  const accT = RTV.tensor(Float64Array.from(acc), [1, acc.length]);
  const r = stage(
    mfiles,
    `
P = cymbal2_defaults();
${pOverrides(params)}
[y, peakAccel] = cymbal2_output(P, acc);
`,
    { acc: accT },
  );
  const y = tensor(r.variableValues.y, 'y');
  return { y: f32(y), peakAccel: scalar(r.variableValues.peakAccel, 'peakAccel') };
}

// CPU reference: the unmodified cymbal2_sim through numbl (slow; for
// validation of the GPU path on reduced configurations).
export function runCpuSim(mfiles, params) {
  const t0 = performance.now();
  const r = stage(
    mfiles,
    `
P = cymbal2_defaults();
${pOverrides(params)}
S = cymbal2_system(P);
MB = cymbal2_mallet(P);
[acc, dinfo, fhist] = cymbal2_sim(P, S, MB);
d_peakF = dinfo.peakForceN;
d_cvel = dinfo.contactVel(1);
d_contactMs = dinfo.contactMs;
d_events = dinfo.contactEvents;
d_xi = dinfo.peakXiMm;
d_nsteps = dinfo.nsteps;
`,
  );
  const ms = performance.now() - t0;
  const v = r.variableValues;
  return {
    acc: f32(tensor(v.acc, 'acc')),
    fhist: f32(tensor(v.fhist, 'fhist')),
    peakForceN: scalar(v.d_peakF, 'd_peakF'),
    contactVel: scalar(v.d_cvel, 'd_cvel'),
    contactMs: scalar(v.d_contactMs, 'd_contactMs'),
    contactEvents: scalar(v.d_events, 'd_events'),
    peakXiMm: scalar(v.d_xi, 'd_xi'),
    nsteps: scalar(v.d_nsteps, 'd_nsteps'),
    ms,
  };
}