/**
 * Scene description -> the two MATLAB files, and everything around them.
 *
 * This is the studio's hinge. A SceneSpec (src/studio/types.ts) says *what*
 * is in the scene; this module writes the model .m (init and step for every
 * object plus the air, all advanced together) and the scene .m (bodies,
 * absorption, and each object's coupling profiles), instantiates the
 * external ops each object needs with its grid geometry baked in, and lists
 * the constants, fields, state and parameters the host must offer the
 * compiler.
 *
 * The generated files are not an internal detail: the Lab view shows them,
 * and they are editable — an edited source compiles through exactly the same
 * pipeline, the generator's output being nothing more than a starting point.
 * That is why the generator works hard at readability: suffixed names
 * (`u_s1`, `cs_s1`) rather than positional arrays, one object's physics kept
 * together under its own comment, and the same idioms a hand-written model
 * (dulcimer's) uses.
 *
 * Naming scheme, per object id ID:
 *   u_ID, um_ID      the object's displacement and its history (state)
 *   xs_ID / ym_ID    node coordinates (local, 0..L), pin_ID the pin mask
 *   n_ID, L_ID       node count and length(s) — exact constants
 *   dxx_ID, lap_ID,  the object's stencil and coupling ops
 *   rad_ID, brg_ID
 *   prof_ID, bprof_ID  where in the air the object radiates / its bridge
 *                      patch drives (scene-built, air-sized)
 * The air keeps dulcimer's names: p, pm, c, sig, wall, lapw.
 */
import type { AirGrid, ObjectGrid } from '../grids.ts';
import {
  makeBridge, makeDxx, makeDxxxx, makeLap2, makeLapw, makeRadLine, makeRadRect,
  type OpInstance,
} from '../mgpu/ops.ts';
import type { StateField } from '../mgpu/model.ts';
import { BODY_TEMPLATES } from './bodies.ts';
import { templateOf } from './objects.ts';
import {
  suffixed, type ObjectSpec, type ParamSpec, type Params, type SceneSpec,
} from './types.ts';

/** Everything the compiler pipeline needs, derived from one SceneSpec. */
export interface GeneratedScene {
  modelSource: string;
  sceneSource: string;
  ops: OpInstance[];
  /** Exact scalars offered to the model .m. */
  modelConsts: Record<string, number>;
  /** Host-uploaded fields offered to the model .m, with initial contents.
   *  The scene-built ones (c, sig, wall, prof_*) are placeholders here —
   *  the session evaluates the scene and uploads the real thing. */
  modelFields: Record<string, Float32Array>;
  /** The scene-built subset of modelFields, in output order for medium(). */
  sceneOutputs: string[];
  /** Exact scalars offered to the scene .m. */
  sceneConsts: Record<string, number>;
  state: StateField[];
  /** Model parameter names (suffixed), and their slider specs by object. */
  modelParamNames: string[];
  /** Scene parameter names (placement + body), and their specs. */
  sceneParamNames: string[];
  /** Each object's grid, by id. */
  grids: Record<string, ObjectGrid>;
}

/** A parameter spec with its key suffixed for one object. */
export const suffixedSpecs = (specs: ParamSpec[], id: string): ParamSpec[] =>
  specs.map((s) => ({ ...s, key: suffixed(s.key, id) }));

/** The parameter values one object contributes, suffixed, defaults filled. */
export function suffixedParams(specs: ParamSpec[], values: Params, id: string): Params {
  return Object.fromEntries(
    specs.map((s) => [suffixed(s.key, id), values[s.key] ?? s.value]),
  );
}

const line = (s: string): string => `  ${s}`;

/** The generated model .m: init and step for every object plus the air. */
function generateModel(spec: SceneSpec, grids: Record<string, ObjectGrid>): string {
  const objs = spec.objects;
  const stateOut = objs.flatMap((o) => [`u_${o.id}`, `um_${o.id}`]).concat(['p', 'pm']);
  const stepOut = objs.flatMap((o) => [`un_${o.id}`, `uo_${o.id}`]).concat(['pn', 'po']);

  // ---- init ---------------------------------------------------------------
  const initArgs: string[] = [];
  const initBody: string[] = [];
  for (const o of objs) {
    const id = o.id;
    const t = templateOf(o);
    initBody.push(`% -- ${t.label.toLowerCase()} ${id}: ${t.verb.toLowerCase()} as an initial condition`);
    if (o.type === 'string') {
      initArgs.push(`xs_${id}`, `L_${id}`, `pluckpos_${id}`, `amp_${id}`);
      initBody.push(
        `a_${id} = xs_${id} / (pluckpos_${id} * L_${id});`,
        `b_${id} = (L_${id} - xs_${id}) / (L_${id} - pluckpos_${id} * L_${id});`,
        `u_${id} = amp_${id} * 0.5 * (a_${id} + b_${id} - abs(a_${id} - b_${id}));`,
        `um_${id} = u_${id};`,
      );
    } else if (o.type === 'bar') {
      initArgs.push(`xs_${id}`, `pin_${id}`, `n_${id}`, `L_${id}`, `strikepos_${id}`, `striker_${id}`, `vel_${id}`, 'dt');
      initBody.push(
        `r_${id} = (xs_${id} - strikepos_${id} * L_${id}) / striker_${id};`,
        `u_${id} = zeros(n_${id}, 1);`,
        `um_${id} = -(dt * vel_${id}) * (pin_${id} .* exp(-(r_${id} .* r_${id})));`,
      );
    } else {
      initArgs.push(
        `xm_${id}`, `ym_${id}`, `pin_${id}`, `n_${id}`, `L1_${id}`, `L2_${id}`,
        `strikex_${id}`, `strikey_${id}`, `striker_${id}`, `vel_${id}`, 'dt',
      );
      initBody.push(
        `rr_${id} = ((xm_${id} - strikex_${id} * L1_${id}).^2 + (ym_${id} - strikey_${id} * L2_${id}).^2) / (striker_${id} * striker_${id});`,
        `u_${id} = zeros(n_${id}, 1);`,
        `um_${id} = -(dt * vel_${id}) * (pin_${id} .* exp(-rr_${id}));`,
      );
    }
    initBody.push('');
  }
  initBody.push('% -- the air starts silent', 'p = zeros(npts, 1);', 'pm = zeros(npts, 1);');
  initArgs.push('npts');
  const initSig = `function [${stateOut.join(', ')}] = init(${[...new Set(initArgs)].join(', ')})`;

  // ---- step ---------------------------------------------------------------
  const stepArgs: string[] = [];
  const stepBody: string[] = [];
  for (const o of objs) {
    const id = o.id;
    const t = templateOf(o);
    stepArgs.push(`u_${id}`, `um_${id}`, `pin_${id}`);
    stepArgs.push(...(o.type === 'membrane' ? [`L1_${id}`, `L2_${id}`] : [`L_${id}`]));
    stepBody.push(`% -- ${t.label.toLowerCase()} ${id} --------------------------------------------`);
    if (o.type === 'string') {
      stepArgs.push(`f0_${id}`, `B_${id}`, `t60_${id}`, `sig1_${id}`);
      stepBody.push(
        `cs_${id} = 2 * L_${id} * f0_${id};              % tension, as a wave speed`,
        `kap2_${id} = B_${id} * (cs_${id} * L_${id} / pi)^2;  % stiffness, as inharmonicity`,
        `d_${id} = (6.91 / t60_${id}) * dt;`,
        `lu_${id} = dxx_${id}(u_${id});`,
        `lum_${id} = dxx_${id}(um_${id});`,
        `l4_${id} = dxxxx_${id}(u_${id});`,
        `un_${id} = pin_${id} .* ((2*u_${id} - (1 - d_${id})*um_${id} + (dt*dt)*((cs_${id}*cs_${id})*lu_${id} - kap2_${id}*l4_${id}) + (2*sig1_${id}*dt)*(lu_${id} - lum_${id})) ./ (1 + d_${id}));`,
      );
    } else if (o.type === 'bar') {
      stepArgs.push(`f1_${id}`, `t60_${id}`, `sig1_${id}`);
      stepBody.push(
        `kap_${id} = 2 * f1_${id} * (L_${id} * L_${id}) / pi;  % stiffness, from the first partial`,
        `d_${id} = (6.91 / t60_${id}) * dt;`,
        `lu_${id} = dxx_${id}(u_${id});`,
        `lum_${id} = dxx_${id}(um_${id});`,
        `l4_${id} = dxxxx_${id}(u_${id});`,
        `un_${id} = pin_${id} .* ((2*u_${id} - (1 - d_${id})*um_${id} - (dt*dt)*(kap_${id}*kap_${id})*l4_${id} + (2*sig1_${id}*dt)*(lu_${id} - lum_${id})) ./ (1 + d_${id}));`,
      );
    } else {
      stepArgs.push(`f0_${id}`, `t60_${id}`, `sig1_${id}`);
      stepBody.push(
        `cm_${id} = 2 * f0_${id} / sqrt(1/(L1_${id}*L1_${id}) + 1/(L2_${id}*L2_${id}));  % tension, from the (1,1) mode`,
        `d_${id} = (6.91 / t60_${id}) * dt;`,
        `lu_${id} = lap_${id}(u_${id});`,
        `lum_${id} = lap_${id}(um_${id});`,
        `un_${id} = pin_${id} .* ((2*u_${id} - (1 - d_${id})*um_${id} + (dt*dt)*(cm_${id}*cm_${id})*lu_${id} + (2*sig1_${id}*dt)*(lu_${id} - lum_${id})) ./ (1 + d_${id}));`,
      );
    }
    stepBody.push('');
  }

  stepBody.push(
    '% -- the coupling: each object\'s motion, spread where its profile says --',
  );
  let drv = '';
  let drvIdx = 0;
  for (const o of objs) {
    const id = o.id;
    stepArgs.push(`gair_${id}`);
    stepBody.push(
      `acc_${id} = (un_${id} - 2*u_${id} + um_${id}) / (dt*dt);`,
      `q_${id} = rad_${id}(acc_${id});`,
    );
    let term = `gair_${id} * (prof_${id} .* q_${id})`;
    if (o.type === 'string') {
      stepArgs.push(`gbridge_${id}`);
      stepBody.push(`fb_${id} = brg_${id}(un_${id});`);
      term += ` - (gbridge_${id} * (cs_${id}*cs_${id}) / L_${id}) * (bprof_${id} .* fb_${id})`;
    }
    const next = `drv${++drvIdx}`;
    stepBody.push(`${next} = ${drv ? `${drv} + ` : ''}${term};`);
    drv = next;
  }
  stepBody.push(
    '',
    '% -- the air ------------------------------------------------------------',
    'lp = lapw(p, wall);',
    'sd = sig * dt;',
    `pn = wall .* ((2*p - (1 - sd) .* pm + (c*dt).^2 .* lp + (dt*dt) * ${drv}) ./ (1 + sd));`,
    '',
    "% This step's fields become the next step's history. Lines of their own,",
    '% so each plans as the copy it is.',
    ...objs.map((o) => `uo_${o.id} = u_${o.id};`),
    'po = p;',
  );
  stepArgs.push('p', 'pm', 'c', 'sig', 'wall');
  for (const o of objs) {
    stepArgs.push(`prof_${o.id}`);
    if (o.type === 'string') stepArgs.push(`bprof_${o.id}`);
  }
  stepArgs.push('dt');
  const stepSig = `function [${stepOut.join(', ')}] = step(${[...new Set(stepArgs)].join(', ')})`;

  const objectSummary = objs
    .map((o) => {
      const g = grids[o.id];
      const t = templateOf(o);
      return `%   ${o.id}: ${t.label.toLowerCase()}, ${
        g.kind === 'line' ? `${g.n1} nodes over ${g.L1} m` : `${g.n1} x ${g.n2} nodes over ${g.L1} x ${g.L2} m`
      }`;
    })
    .join('\n');

  return `% Generated by soundbox from the scene you built — and editable: this file
% is the model, compiled to WebGPU kernels, and Run edits recompiles whatever
% you make of it. (Rebuilding the scene from the Build tab regenerates it.)
%
% Objects in this scene:
${objectSummary}
%
% Every object is advanced by the same explicit leapfrog family: displacement
% u and its history um, one update per timestep, pinned at the boundary by
% its pin mask. The air carries pressure p on the box grid,
%
%   p_tt + 2 sig p_t = c^2 lapw(p, wall) + s
%
% where lapw is the wall-masked Laplacian (the rigid body, at no timestep
% cost) and s collects every object's contribution: its acceleration, spread
% along the scene-built profile of where it sits (and, for a string, its
% bridge force driving a patch above the plate). The air does not push back;
% at these amplitudes that back-reaction is far below what the models already
% neglect.

${initSig}
${initBody.map(line).join('\n')}
end

${stepSig}
${stepBody.map(line).join('\n')}
end
`;
}

/** The generated scene .m: bodies, absorption, coupling profiles. */
function generateScene(spec: SceneSpec): { source: string; outputs: string[]; paramNames: string[] } {
  const body = BODY_TEMPLATES[spec.body.type];
  const outputs = ['c', 'sig', 'wall'];
  const args = ['x', 'y', 'z', 'h', 'npts', 'Lx', 'Ly', 'Lz', 'c0'];
  const paramNames: string[] = [];

  for (const b of body.paramSpecs) {
    args.push(b.key);
    paramNames.push(b.key);
  }

  const lines: string[] = [...body.matlab, 'wall = 1 - shell;', ''];
  lines.push(
    '% Uniform air. The walls being a mask is what keeps this uniform — and',
    '% the timestep as large as air allows.',
    'c = c0;',
    '',
    "% Open boundary at the domain edge" +
      (body.paramSpecs.some((p) => p.key === 'absorb')
        ? ", plus the body's surface absorption."
        : '.'),
    body.paramSpecs.some((p) => p.key === 'absorb')
      ? 'sig = sponge3(x, y, z, Lx, Ly, Lz, 0.15*Ly, 9000) + absorb * (4 * shell .* (1 - shell));'
      : 'sig = sponge3(x, y, z, Lx, Ly, Lz, 0.15*Ly, 9000);',
    '',
  );

  for (const o of spec.objects) {
    const id = o.id;
    const t = templateOf(o);
    // The object's baked size(s) come in as exact constants; its placement
    // as scene parameters.
    args.push(...(o.type === 'membrane' ? [`L1_${id}`, `L2_${id}`] : [`L_${id}`]));
    for (const p of t.placeSpecs) {
      const k = suffixed(p.key, id);
      args.push(k);
      paramNames.push(k);
    }
    if (o.type === 'membrane') {
      outputs.push(`prof_${id}`);
      lines.push(
        `% ${t.label.toLowerCase()} ${id} radiates from its plane: a thin sheet at z = oz_${id}`,
        '% over its extent, normalized so its cross-section integrates to one.',
        `aw_${id} = 2*h;`,
        `sk_${id} = exp(-((z - oz_${id}).^2) / (aw_${id}*aw_${id})) / (sqrt(pi) * aw_${id});`,
        `exx_${id} = 0.5 * (1 + tanh((L1_${id}/2 - abs(x)) / (1.5*h)));`,
        `exy_${id} = 0.5 * (1 + tanh((L2_${id}/2 - abs(y)) / (1.5*h)));`,
        `prof_${id} = sk_${id} .* exx_${id} .* exy_${id} .* wall;`,
        '',
      );
    } else {
      outputs.push(`prof_${id}`);
      lines.push(
        `% ${t.label.toLowerCase()} ${id} radiates from its line: a tube two cells wide`,
        `% around (y = oy_${id}, z = oz_${id}), fading out at its ends, normalized so`,
        '% its cross-section integrates to one.',
        `aw_${id} = 2*h;`,
        `tube_${id} = exp(-((y - oy_${id}).^2 + (z - oz_${id}).^2) / (aw_${id}*aw_${id})) / (pi * aw_${id} * aw_${id});`,
        `span_${id} = 0.5 * (1 + tanh((L_${id}/2 - abs(x)) / (1.5*h)));`,
        `prof_${id} = tube_${id} .* span_${id} .* wall;`,
      );
      if (o.type === 'string') {
        outputs.push(`bprof_${id}`);
        lines.push(
          `% ...and its bridge force drives a patch of air just above the plate,`,
          `% around its x = +L/2 end. This stands in for a top plate moving; the`,
          '% plate\'s own modes are not modelled.',
          `pw_${id} = 1.5*h;`,
          `pr_${id} = exp(-((x - L_${id}/2).^2 + (y - oy_${id}).^2) / (0.04*0.04)) / (pi * 0.04 * 0.04);`,
          `skin_${id} = exp(-(z - pw_${id}).^2 / (pw_${id}*pw_${id})) / (sqrt(pi) * pw_${id});`,
          `bprof_${id} = pr_${id} .* skin_${id} .* wall;`,
        );
      }
      lines.push('');
    }
  }

  const source = `% Generated by soundbox — the scene: ${body.label.toLowerCase()}, and where each
% object couples into the air. Evaluated once on the CPU (numbl's
% interpreter, full MATLAB subset available), so editing it costs no
% recompile; the model .m is the compiled half.
%
% The body is a *mask*: \`wall\` is 0 inside solid and 1 in air, and the
% solver's masked Laplacian drops any flux into a masked cell — the rigid
% (Neumann) condition, at no timestep cost. What a rigid mask cannot do is
% absorb, so \`absorb\` paints a thin lossy skin over the shell's surfaces
% (4*shell*(1-shell) peaks exactly at the surface).
function [${outputs.join(', ')}] = medium(${args.join(', ')})
${lines.map(line).join('\n')}
end
`;
  return { source, outputs, paramNames };
}

/** Derive everything the pipeline needs from a scene description. */
export function generate(
  spec: SceneSpec,
  air: AirGrid,
  dt: number,
): GeneratedScene {
  if (spec.objects.length === 0) {
    throw new Error('a scene needs at least one object (add a string, bar or membrane)');
  }
  const seen = new Set<string>();
  for (const o of spec.objects) {
    if (seen.has(o.id)) throw new Error(`duplicate object id '${o.id}'`);
    seen.add(o.id);
  }

  const grids: Record<string, ObjectGrid> = {};
  for (const o of spec.objects) grids[o.id] = templateOf(o).makeGrid(o.size, dt);

  // Ops, consts, fields, state.
  const ops: OpInstance[] = [makeLapw('lapw', air)];
  const modelConsts: Record<string, number> = {
    npts: air.npts, nx: air.nx, ny: air.ny, nz: air.nz,
  };
  const sceneConsts: Record<string, number> = {};
  const modelFields: Record<string, Float32Array> = {
    x: air.x, y: air.y, z: air.z,
  };
  const state: StateField[] = [];
  const modelParamNames: string[] = [];

  for (const o of spec.objects) {
    const id = o.id;
    const g = grids[id];
    modelConsts[`n_${id}`] = g.npts;
    if (g.kind === 'line') {
      modelConsts[`L_${id}`] = g.L1;
      sceneConsts[`L_${id}`] = g.L1;
      modelFields[`xs_${id}`] = g.xs;
      ops.push(makeDxx(`dxx_${id}`, g), makeDxxxx(`dxxxx_${id}`, g), makeRadLine(`rad_${id}`, g, air));
      if (o.type === 'string') ops.push(makeBridge(`brg_${id}`, g, air));
    } else {
      modelConsts[`L1_${id}`] = g.L1;
      modelConsts[`L2_${id}`] = g.L2;
      sceneConsts[`L1_${id}`] = g.L1;
      sceneConsts[`L2_${id}`] = g.L2;
      modelFields[`xm_${id}`] = g.xs;
      modelFields[`ym_${id}`] = g.ys;
      ops.push(makeLap2(`lap_${id}`, g), makeRadRect(`rad_${id}`, g, air));
    }
    modelFields[`pin_${id}`] = g.pin;
    state.push({ name: `u_${id}`, count: g.npts }, { name: `um_${id}`, count: g.npts });
    for (const p of templateOf(o).paramSpecs) modelParamNames.push(suffixed(p.key, id));
  }
  state.push({ name: 'p', count: air.npts }, { name: 'pm', count: air.npts });

  const scene = generateScene(spec);
  // Scene-built fields are uploaded after evaluation; placeholders size them.
  for (const name of scene.outputs) modelFields[name] = new Float32Array(air.npts);

  return {
    modelSource: generateModel(spec, grids),
    sceneSource: scene.source,
    ops,
    modelConsts,
    modelFields,
    sceneOutputs: scene.outputs,
    sceneConsts,
    state,
    modelParamNames,
    sceneParamNames: scene.paramNames,
    grids,
  };
}