/**
 * The .m model, compiled and running on the GPU.
 *
 * The model file is ordinary MATLAB: it defines an `init` function that builds
 * the initial state — the plucked string and a silent air field — and a `step`
 * function that advances both one timestep. Each is specialized for the
 * current grids and compiled into a ModelPlan, and both operate on the same
 * state buffers (see HostBuffers).
 *
 * Both functions return the state, in the same order, so their signatures say
 * exactly what they produce:
 *
 *   function [u, um, p, pm] = init(xs, npts, Ls, pluckpos, amp)
 *   function [un, uold, pn, pold] = step(u, um, p, pm, ...)
 *
 * The state spans two grids — `u`, `um` are string fields (ns nodes), `p`,
 * `pm` air fields (npts cells) — and the compiled statements mix freely:
 * each line is one kernel over its own field's size, and the external ops
 * (src/mgpu/ops.ts) are where a value crosses from one grid to the other.
 *
 * The host supplies the things that are setup rather than algorithm: the grid
 * coordinates, the medium and coupling profiles the scene defines, the
 * timestep, and the parameter values. Each argument is matched to the .m's
 * declared parameter name, so the file documents its own interface.
 *
 * There is no clock here, and that is not an oversight: the pluck is an
 * initial condition, not a driven source, so nothing in the model needs to
 * know what time it is.
 */
import { HostBuffers, ModelPlan } from './plan.ts';
import type { OpPlan } from './ops.ts';
import { inFunction, inFunctionAsync, inModel } from './errors.ts';
import { CompiledModel, type Binding } from './compile.ts';
import type { AirGrid, StringGrid } from '../grid.ts';

export interface ModelParams {
  [key: string]: number;
}

/** What the scene defines, on the air grid: npts values each. */
export interface MediumFields {
  /** Sound speed, m/s. */
  c: Float32Array;
  /** Absorption rate, 1/s (the sponge and any wall absorption). */
  sig: Float32Array;
  /** 1 in air, 0 in the body's solid shell. What `lapw` masks by. */
  wall: Float32Array;
  /** Where the string radiates directly: a tube around the string line. */
  lineprof: Float32Array;
  /** Where the bridge force drives the air: a patch above the top plate. */
  boardprof: Float32Array;
}

/** One state field the .m advances, and which grid it lives on. */
export interface StateField {
  name: string;
  grid: 'air' | 'string';
}

export interface GpuModelOptions {
  device: GPUDevice;
  ops: OpPlan;
  air: AirGrid;
  string: StringGrid;
  medium: MediumFields;
  /** Model source (.m text). */
  source: string;
  /** Parameter names the .m may take as arguments. */
  paramNames: string[];
  /** State fields the .m advances, in the order its functions return them. */
  state: StateField[];
  /** Grid fields one kernel may read, overriding what the device allows.
   *  Only for tests. */
  operandBudget?: number;
}

/** Names the .m may take for the air grid coordinates. */
export const AIR_GRID_NAMES = ['x', 'y', 'z'] as const;
/** Names the .m may take for the string grid: node positions and the pin
 *  mask that terminates the ends. */
export const STRING_GRID_NAMES = ['xs', 'pin'] as const;
/** Names the .m may take for what the scene defines. */
export const MEDIUM_NAMES = ['c', 'sig', 'wall', 'lineprof', 'boardprof'] as const;

export class GpuModel {
  readonly paramNames: string[];
  readonly state: StateField[];
  readonly npts: number;
  readonly ns: number;

  #device: GPUDevice;
  #host: HostBuffers;
  #initPlan: ModelPlan;
  #stepPlan: ModelPlan;
  /** Timestep, host-owned: it follows from the grid and the medium (a CFL
   *  condition), not from anything the user types, and it is folded into every
   *  setParams so the .m's `dt` is never left with the zero a missing
   *  parameter would default to. */
  #dt = 0;
  #readback: GPUBuffer;
  /** Which function wrote the state most recently; see `read`. */
  #lastRan: 'init' | 'step' = 'init';

  private constructor(init: {
    device: GPUDevice;
    host: HostBuffers;
    initPlan: ModelPlan;
    stepPlan: ModelPlan;
    readback: GPUBuffer;
    paramNames: string[];
    state: StateField[];
    npts: number;
    ns: number;
  }) {
    this.#device = init.device;
    this.#host = init.host;
    this.#initPlan = init.initPlan;
    this.#stepPlan = init.stepPlan;
    this.#readback = init.readback;
    this.paramNames = init.paramNames;
    this.state = init.state;
    this.npts = init.npts;
    this.ns = init.ns;
  }

  static async create(opts: GpuModelOptions): Promise<GpuModel> {
    const { device, ops, air, string, medium, source, paramNames, state } = opts;
    const npts = air.npts;
    const ns = string.ns;
    const sizeOf = (grid: 'air' | 'string'): number => (grid === 'air' ? npts : ns);

    // What the .m may ask for by name. The grid geometry is exact, so a
    // constructor reading it (`zeros(npts, 1)`) keeps a static shape.
    const bindings: Record<string, Binding> = {
      npts: { kind: 'const', value: npts },
      nx: { kind: 'const', value: air.nx },
      ny: { kind: 'const', value: air.ny },
      nz: { kind: 'const', value: air.nz },
      h: { kind: 'const', value: air.h },
      ns: { kind: 'const', value: ns },
      hs: { kind: 'const', value: string.hs },
      Ls: { kind: 'const', value: string.Ls },
      dt: { kind: 'param' },
    };
    for (const g of AIR_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
    for (const g of STRING_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [ns, 1] };
    for (const m of MEDIUM_NAMES) bindings[m] = { kind: 'tensor', shape: [npts, 1] };
    for (const s of state) bindings[s.name] = { kind: 'tensor', shape: [sizeOf(s.grid), 1] };
    for (const p of paramNames) bindings[p] = { kind: 'param' };

    // Parsing belongs to the file, not to either function.
    const compiled = inModel(() => new CompiledModel(source, bindings, { npts, ns }));
    const nargout = state.length;
    const initFn = inFunction('init', () => compiled.specialize('init', nargout));
    const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
    compiled.finish();

    // Both functions return the state, in order, and both feed it back into
    // the shared buffers.
    const feedback = state.map((s) => s.name);

    const host = new HostBuffers(device);
    // The host owns the state and the inputs it uploads, whether or not a
    // given function happens to take them as arguments.
    for (const s of state) host.ensure(s.name, sizeOf(s.grid));
    for (const g of AIR_GRID_NAMES) host.ensure(g, npts);
    for (const g of STRING_GRID_NAMES) host.ensure(g, ns);
    for (const m of MEDIUM_NAMES) host.ensure(m, npts);

    const initPlan = await inFunctionAsync('init', () =>
      ModelPlan.create(device, ops, { fn: initFn, feedback }, host, opts.operandBudget),
    );
    const stepPlan = await inFunctionAsync('step', () =>
      ModelPlan.create(device, ops, { fn: stepFn, feedback }, host, opts.operandBudget),
    );

    host.upload('x', air.x);
    host.upload('y', air.y);
    host.upload('z', air.z);
    host.upload('xs', string.xs);
    host.upload('pin', string.pin);
    for (const m of MEDIUM_NAMES) host.upload(m, medium[m]);

    const readback = device.createBuffer({
      label: 'mgpu-readback',
      size: 4 * Math.max(npts, ns),
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });

    return new GpuModel({
      device, host, initPlan, stepPlan, readback, paramNames, state, npts, ns,
    });
  }

  /** The timestep in force. Host-owned; see `#dt`. */
  get dt(): number {
    return this.#dt;
  }

  setDt(dt: number): void {
    this.#dt = dt;
  }

  setParams(params: ModelParams): void {
    const merged = { dt: this.#dt, ...params };
    this.#initPlan.setParams(merged);
    this.#stepPlan.setParams(merged);
  }

  /**
   * Swap the medium under a running model. It is data, not code — its shape in
   * the bindings depends only on the grid — so changing the body's geometry is
   * five buffer writes and needs no recompile.
   */
  uploadMedium(medium: MediumFields): void {
    for (const m of MEDIUM_NAMES) this.#host.upload(m, medium[m]);
  }

  /** Write a host-owned value directly. Lets a test set up an exact initial
   *  condition instead of going through `init`. */
  upload(name: string, data: Float32Array): void {
    this.#host.upload(name, data);
  }

  /** Run `init`, replacing the state. */
  init(): void {
    const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
    this.#initPlan.encodeSteps(enc, 1);
    this.#device.queue.submit([enc.finish()]);
    this.#lastRan = 'init';
  }

  /**
   * Advance `steps` timesteps. Synchronous — this only records commands and
   * submits them; nothing is read back and nothing is awaited.
   *
   * `after` is recorded once per step, so anything that must see every
   * timestep (the microphone) rides along in the same submission.
   */
  step(steps = 1, after?: (encoder: GPUCommandEncoder) => void): void {
    const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
    this.#stepPlan.encodeSteps(enc, steps, after);
    this.#device.queue.submit([enc.finish()]);
    this.#lastRan = 'step';
  }

  /**
   * The buffer currently holding a named value. A field the .m computes is
   * produced by both functions, into separate buffers (only the state is
   * shared), so this resolves to whichever function ran most recently — which
   * is what makes the first frame show the initial state rather than an
   * unwritten buffer.
   */
  #locate(name: string): { buffer: GPUBuffer; count: number } | null {
    const [first, second] =
      this.#lastRan === 'init'
        ? [this.#initPlan, this.#stepPlan]
        : [this.#stepPlan, this.#initPlan];
    const buffer = first.buffer(name) ?? second.buffer(name);
    const count = first.elementCount(name) ?? second.elementCount(name);
    if (!buffer || count === undefined) return null;
    return { buffer, count };
  }

  /** The GPU buffer a named value would be read from right now. */
  valueBuffer(name: string): GPUBuffer | null {
    return this.#locate(name)?.buffer ?? null;
  }

  /**
   * The buffer a host-owned field lives in — the state between calls, or an
   * input like the wall mask.
   *
   * This is what the renderer binds, and it must be this rather than
   * `valueBuffer`: a bind group is built once and holds a particular buffer,
   * while `init` and `step` write their outputs into buffers of their own and
   * only agree here, where their feedback copies land. Binding either
   * function's private buffer would draw a stale field for half the run.
   */
  stateBuffer(name: string): GPUBuffer | null {
    return this.#host.get(name)?.buffer ?? null;
  }

  /** Read a named value back to the CPU. The only await in the whole loop. */
  async read(name: string): Promise<Float32Array> {
    const located = this.#locate(name);
    if (!located) throw new Error(`read: the model has no value named '${name}'`);
    const { buffer, count } = located;
    const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
    enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
    this.#device.queue.submit([enc.finish()]);
    await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
    const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
    this.#readback.unmap();
    return out;
  }

  /** What the .m compiled to, for display. */
  describe(): { init: string[]; step: string[] } {
    return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
  }

  destroy(): void {
    this.#initPlan.destroy();
    this.#stepPlan.destroy();
    this.#host.destroy();
    this.#readback.destroy();
  }
}