/**
 * The medium: a .m scene file, evaluated once on the grid.
 *
 * A scene file is ordinary MATLAB defining one function,
 *
 *   function [c, sig] = medium(x, y, <parameters>)
 *
 * over the solver's grid points. Unlike the models it is *not* compiled to
 * WGSL: a model's step runs every frame and must lower to a fixed sequence of
 * GPU dispatches, but a scene is evaluated exactly once and survives only as
 * two arrays of numbers. So it runs through numbl's CPU interpreter instead,
 * which buys the full MATLAB subset — loops, `if`, reductions, indexing,
 * seeded randomness via `rng`/`randn`, and anything in tools/ — and f64
 * evaluation, where the step dialect is element-wise f32.
 *
 * `c` is the sound speed in metres per second and `sig` the absorption rate in
 * inverse seconds. Both are ordinary
 * fields of position, which is what lets one function describe both the
 * scatterer and the open boundary: the absorbing layer every scene puts around
 * the edge (tools/sponge.m) is just the statement that the medium swallows
 * sound out there. A scene is free to put absorption inside the domain too,
 * which makes a lossy scatterer.
 */
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 { Grid } from '../grid.ts';
import { C_AIR } from '../units.ts';
import { toolFiles } from '../tools.ts';
import { inFunction, inModel, ModelCompileError } from '../mgpu/errors.ts';
import type { ModelParams } from '../mgpu/model.ts';

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

export interface SceneOptions {
  grid: Grid;
  /** Scene source (.m text). */
  source: string;
  /** Parameter names the .m may take beyond `x` and `y`. */
  paramNames: string[];
  params: ModelParams;
}

export class Scene {
  /** Sound speed on the grid, npts. */
  readonly c: Float32Array;
  /** Absorption rate on the grid, npts. */
  readonly sig: Float32Array;
  readonly cmin: number;
  readonly cmax: number;
  /** The background speed, taken to be whatever it is in the corner of the
   *  domain — which is inside the absorbing layer, where a scene has no
   *  business putting a scatterer. What the renderer shades departures from. */
  readonly cref: number;
  /** The largest departure from `cref` anywhere, so the renderer's wash has a
   *  scale. Zero for a uniform medium, which draws nothing. */
  readonly cdev: number;

  private constructor(c: Float32Array, sig: Float32Array) {
    this.c = c;
    this.sig = sig;
    let lo = Infinity;
    let hi = 0;
    for (const v of c) {
      if (v < lo) lo = v;
      if (v > hi) hi = v;
    }
    this.cmin = lo;
    this.cmax = hi;
    this.cref = c[0];
    this.cdev = Math.max(Math.abs(hi - this.cref), Math.abs(this.cref - lo));
  }

  static create(opts: SceneOptions): Scene {
    const { grid, source, paramNames, params } = opts;
    const [c, sig] = evaluateMedium(source, paramNames, params, grid);
    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 },
        );
      }
    }
    return new Scene(c, sig);
  }
}

/**
 * Evaluate the scene file on the grid, through numbl's CPU interpreter.
 *
 * The .m keeps the same contract a model has: it names the arguments it wants
 * — `x`, `y`, 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. A
 * one-line driver script calls `medium` with exactly the arguments its
 * signature declares, with those names pre-bound in the driver's workspace.
 */
function evaluateMedium(
  source: string,
  paramNames: string[],
  params: ModelParams,
  grid: Grid,
): [Float32Array, Float32Array] {
  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 !== 2) {
    throw new ModelCompileError(
      `'${MEDIUM_FN}' must return two outputs [c, sig] — the sound speed and ` +
        `the absorption rate — 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 themselves, in metres, and the numbers that describe them. A
  // scene wants `L` to place its absorbing layer relative to the domain and
  // `h` to smooth an interface over about a cell, so both are part of the
  // contract rather than something every scene has to be told. `c0` is the
  // speed of sound in air, which is what "the background" means here.
  const vars: Record<string, RuntimeValue> = {
    x: new RuntimeTensor(grid.x64, [grid.npts, 1]),
    y: new RuntimeTensor(grid.y64, [grid.npts, 1]),
    L: grid.L,
    h: grid.h,
    c0: C_AIR,
    npts: grid.npts,
    nx: grid.nx,
    ny: grid.ny,
  };
  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 driver = `[c__, sig__] = ${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',
    ),
  );

  return [
    toGridField(result.variableValues['c__'], fn.outputs[0], grid.npts),
    toGridField(result.variableValues['sig__'], fn.outputs[1], grid.npts),
  ];
}

/** 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 2-D reshape is refused
    // rather than reordered: the tensor's column-major layout would not match
    // the grid's x-fastest rows.
    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 });
}