/**
 * The scene: bodies, absorption, and every object's coupling geometry, as a
 * .m file evaluated once on the air grid.
 *
 * A scene file is ordinary MATLAB defining one function,
 *
 *   function [c, sig, wall, prof_s1, ...] = 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
 * 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 first three outputs are fixed:
 *   c      sound speed, m/s (normally just c0 everywhere — bodies 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 a body's surfaces are given)
 *   wall   1 in air, 0 in a body's solid shell; what lapw masks by
 * After them come whatever coupling profiles this scene's objects need
 * (prof_s1, brg_s1, ...), named by the caller — the studio's code generator
 * decides the list, and the same list is offered to the compiled model as
 * uploadable fields.
 */
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 '../grids.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 {
  air: AirGrid;
  /** Scene source (.m text). */
  source: string;
  /** Output names, in order: 'c', 'sig', 'wall', then the profiles. */
  outputs: string[];
  /** Exact scalars offered by name beyond the grid's own (object lengths and
   *  the like). */
  consts: Record<string, number>;
  /** Parameter names the .m may take beyond the grid and the consts. */
  paramNames: string[];
  params: ModelParams;
}

export class Scene {
  /** Every output field, air-sized, by name. */
  readonly fields: Record<string, Float32Array>;
  readonly cmin: number;
  readonly cmax: number;

  private constructor(fields: Record<string, Float32Array>) {
    this.fields = fields;
    let lo = Infinity;
    let hi = 0;
    for (const v of fields.c) {
      if (v < lo) lo = v;
      if (v > hi) hi = v;
    }
    this.cmin = lo;
    this.cmax = hi;
  }

  get c(): Float32Array {
    return this.fields.c;
  }
  get sig(): Float32Array {
    return this.fields.sig;
  }
  get wall(): Float32Array {
    return this.fields.wall;
  }

  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 offered
 * 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<string, Float32Array> {
  const { air, source, outputs, 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 whatever consts the caller adds (object lengths —
  // the scene builds the coupling profiles around each object, so it needs
  // to know how big each one 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,
    ...opts.consts,
  };
  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: Record<string, Float32Array> = {};
  outputs.forEach((name, i) => {
    fields[name] = toGridField(
      result.variableValues[outNames[i]],
      name,
      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 });
}