/**
 * The instrument's body and the coupling geometry: a .m scene file, evaluated
 * once on the air grid.
 *
 * A scene file is ordinary MATLAB defining one function,
 *
 *   function [c, sig, wall, lineprof, boardprof] = medium(x, y, z, <...>)
 *
 * over the solver's grid points. Unlike the model it is *not* compiled to
 * WGSL: the step runs every timestep and must lower to a fixed sequence of
 * GPU dispatches, but a scene is evaluated exactly once and survives only as
 * five arrays of numbers. So it runs through numbl's CPU interpreter instead,
 * which buys the full MATLAB subset — loops, `if`, indexing, anything in
 * tools/ — and f64 evaluation, where the step dialect is element-wise f32.
 *
 * The five outputs:
 *   c         sound speed, m/s (normally just c0 everywhere — the walls are
 *             a mask, not a material, so they cost no timestep)
 *   sig       absorption rate, 1/s (the sponge at the domain edge, plus any
 *             absorption the body's surfaces are given)
 *   wall      1 in air, 0 in the body's solid shell; what lapw masks by
 *   lineprof  where the string radiates directly (a tube around its line)
 *   boardprof where the bridge force drives the air (a patch on the top plate)
 */
import { parseMFile, type FunctionStmt } from 'numbl-src/numbl-core/parser/index.ts';
import { executeCode } from 'numbl-src/numbl-core/executeCode.ts';
import {
  RuntimeTensor,
  isRuntimeTensor,
  type RuntimeValue,
} from 'numbl-src/numbl-core/runtime/types.ts';
import type { AirGrid } from '../grid.ts';
import { C_AIR } from '../units.ts';
import { toolFiles } from '../tools.ts';
import { inFunction, inModel, ModelCompileError } from '../mgpu/errors.ts';
import type { MediumFields, ModelParams } from '../mgpu/model.ts';

/** The function a scene file must define. */
export const MEDIUM_FN = 'medium';

const OUTPUTS = ['c', 'sig', 'wall', 'lineprof', 'boardprof'] as const;

export interface SceneOptions {
  air: AirGrid;
  /** String length, metres — the scene builds the coupling profiles around
   *  the string, so it needs to know where the string is. */
  Ls: number;
  /** Scene source (.m text). */
  source: string;
  /** Parameter names the .m may take beyond the grid's own. */
  paramNames: string[];
  params: ModelParams;
}

export class Scene implements MediumFields {
  readonly c: Float32Array;
  readonly sig: Float32Array;
  readonly wall: Float32Array;
  readonly lineprof: Float32Array;
  readonly boardprof: Float32Array;
  readonly cmin: number;
  readonly cmax: number;
  /** The background speed, taken from a corner of the domain — inside the
   *  absorbing layer, where a scene has no business putting anything. */
  readonly cref: number;

  private constructor(fields: Record<(typeof OUTPUTS)[number], Float32Array>) {
    this.c = fields.c;
    this.sig = fields.sig;
    this.wall = fields.wall;
    this.lineprof = fields.lineprof;
    this.boardprof = fields.boardprof;
    let lo = Infinity;
    let hi = 0;
    for (const v of this.c) {
      if (v < lo) lo = v;
      if (v > hi) hi = v;
    }
    this.cmin = lo;
    this.cmax = hi;
    this.cref = this.c[0];
  }

  static create(opts: SceneOptions): Scene {
    const fields = evaluateScene(opts);
    const { c, sig, wall } = fields;
    for (let i = 0; i < c.length; i++) {
      if (!(c[i] > 0)) {
        throw new ModelCompileError(
          `the scene's sound speed is ${c[i]} somewhere; it must be positive ` +
            `everywhere (the timestep is set by the fastest point, and a zero ` +
            `or negative speed has no wave equation)`,
          { fn: MEDIUM_FN },
        );
      }
      if (!(sig[i] >= 0)) {
        throw new ModelCompileError(
          `the scene's absorption is ${sig[i]} somewhere; it must be zero or ` +
            `positive (a negative one would amplify rather than absorb)`,
          { fn: MEDIUM_FN },
        );
      }
      // The mask multiplies Laplacian fluxes; outside [0, 1] it would add
      // energy or invert a face. Clamp rather than refuse: a smoothed
      // difference of indicators dips a hair below zero in f64 routinely.
      if (wall[i] < 0) wall[i] = 0;
      else if (wall[i] > 1) wall[i] = 1;
    }
    return new Scene(fields);
  }
}

/**
 * Evaluate the scene file on the grid, through numbl's CPU interpreter.
 *
 * The .m keeps the same contract the model has: it names the arguments it
 * wants — the coordinates, the domain's numbers, and any of the registry's
 * parameters — and the host supplies them by name, so their order in the
 * signature is the .m's own business.
 */
function evaluateScene(
  opts: SceneOptions,
): Record<(typeof OUTPUTS)[number], Float32Array> {
  const { air, Ls, source, paramNames, params } = opts;
  const file = `${MEDIUM_FN}.m`;
  const ast = inModel(() => parseMFile(source, file));
  const fn = ast.body.find(
    (s): s is FunctionStmt =>
      s.type === 'Function' && (s as FunctionStmt).name === MEDIUM_FN,
  );
  if (!fn) {
    throw new ModelCompileError(`the scene defines no function named '${MEDIUM_FN}'`);
  }
  if (fn.outputs.length !== OUTPUTS.length) {
    throw new ModelCompileError(
      `'${MEDIUM_FN}' must return ${OUTPUTS.length} outputs ` +
        `[${OUTPUTS.join(', ')}], not ${fn.outputs.length}`,
      { fn: MEDIUM_FN, start: fn.span.start, end: fn.span.end },
    );
  }
  // What the grid offers a scene by name, beyond its own parameters: the
  // coordinates in metres, the numbers that describe the domain, the speed
  // of sound in air, and the string's length — the scene builds the coupling
  // profiles around the string, so it needs to know where the string is.
  const vars: Record<string, RuntimeValue> = {
    x: new RuntimeTensor(air.x64, [air.npts, 1]),
    y: new RuntimeTensor(air.y64, [air.npts, 1]),
    z: new RuntimeTensor(air.z64, [air.npts, 1]),
    Lx: air.Lx,
    Ly: air.Ly,
    Lz: air.Lz,
    h: air.h,
    c0: C_AIR,
    npts: air.npts,
    nx: air.nx,
    ny: air.ny,
    nz: air.nz,
    Ls,
  };
  const known = new Set([...Object.keys(vars), ...paramNames]);
  for (const p of fn.params) {
    if (!known.has(p)) {
      throw new ModelCompileError(
        `'${MEDIUM_FN}' takes an argument '${p}' that is neither the grid ` +
          `(${Object.keys(vars).join(', ')}) nor one of this scene's parameters` +
          (paramNames.length ? ` (${paramNames.join(', ')})` : ''),
        { fn: MEDIUM_FN, start: fn.span.start, end: fn.span.end },
      );
    }
  }

  for (const name of paramNames) {
    const v = params[name];
    // Missing parameters read as 0, as ModelPlan.setParams has it.
    vars[name] = Number.isFinite(v) ? v : 0;
  }

  const outNames = OUTPUTS.map((o) => `${o}__`);
  const driver = `[${outNames.join(', ')}] = ${MEDIUM_FN}(${fn.params.join(', ')});`;
  const result = inFunction(MEDIUM_FN, () =>
    executeCode(
      driver,
      { initialVariableValues: vars, displayResults: false, implicitCwdPath: null },
      [...toolFiles, { name: file, source }],
      'scene-driver.m',
    ),
  );

  const fields = {} as Record<(typeof OUTPUTS)[number], Float32Array>;
  OUTPUTS.forEach((name, i) => {
    fields[name] = toGridField(
      result.variableValues[outNames[i]],
      fn.outputs[i],
      air.npts,
    );
  });
  return fields;
}

/** One returned field -> npts values, rounded to the solver's f32. */
function toGridField(
  value: RuntimeValue | undefined,
  name: string,
  npts: number,
): Float32Array {
  // A uniform field stays scalar in MATLAB; spread it over the grid.
  if (typeof value === 'number') return new Float32Array(npts).fill(value);
  if (value !== undefined && isRuntimeTensor(value)) {
    if (value.imag) {
      throw new ModelCompileError(
        `the scene's '${name}' is complex; the medium must be real`,
        { fn: MEDIUM_FN },
      );
    }
    // A vector of npts values, either orientation. A reshape is refused
    // rather than reordered: the tensor's column-major layout would not match
    // the grid's x-fastest order.
    if (value.data.length === npts && value.shape.every((d) => d === 1 || d === npts)) {
      return new Float32Array(value.data);
    }
    throw new ModelCompileError(
      `the scene's '${name}' is ${value.shape.join(' x ')}, but the grid wants ` +
        `one value per point (${npts} x 1)`,
      { fn: MEDIUM_FN },
    );
  }
  throw new ModelCompileError(`the scene's '${name}' is not numeric`, { fn: MEDIUM_FN });
}