/**
 * A .m model, compiled and running on the GPU.
 *
 * A model file is ordinary MATLAB: it defines an `init` function that builds
 * the initial state and a `step` function that advances it one timestep. Each
 * is specialized for the current grid 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 [p, pm, t] = init(npts)
 *   function [pn, pold, tn] = step(p, pm, t, c, sig, x, y, dt, f, amp, ...)
 *
 * The host supplies the things that are setup rather than algorithm: the grid
 * coordinates, the medium 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.
 *
 * `t` is carried as a grid field rather than a scalar, and that is deliberate.
 * A batch of timesteps is one replay of a fixed op sequence, so nothing the
 * host writes between steps can change inside it — a clock uploaded per frame
 * would stand still for the whole batch. Making the model advance its own time
 * (`tn = t + dt`) keeps the source term correct however many steps are batched,
 * at the cost of one extra buffer and one extra kernel per step, which next to
 * the stencil is nothing.
 */
import { HostBuffers, ModelPlan } from './plan.ts';
import type { StencilPlan } from './stencil.ts';
import { inFunction, inFunctionAsync, inModel } from './errors.ts';
import { CompiledModel, type Binding } from './compile.ts';

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

/** The grid a model runs on, and the medium it runs in. */
export interface GridFields {
  nx: number;
  ny: number;
  /** Grid spacing, the same in x and y. */
  h: number;
  /** Coordinates of every grid point, npts each, x fastest. */
  x: Float32Array;
  y: Float32Array;
}

/** What the scene defines, on the grid. */
export interface MediumFields {
  /** Sound speed, npts. */
  c: Float32Array;
  /** Absorption rate (the sponge and any absorbing scatterer), npts. */
  sig: Float32Array;
}

export interface GpuModelOptions {
  device: GPUDevice;
  stencil: StencilPlan;
  grid: GridFields;
  medium: MediumFields;
  /** Model source (.m text). */
  source: string;
  /** Parameter names the .m may take as arguments. */
  paramNames: string[];
  /** State field names, in order (e.g. ['p', 'pm', 't']). */
  state: string[];
  /** Grid fields one kernel may read, overriding what the device allows.
   *  Only for tests. */
  operandBudget?: number;
}

/** Names the .m may take for the grid coordinates. */
export const GRID_NAMES = ['x', 'y'] as const;
/** Names the .m may take for the medium the scene defines. */
export const MEDIUM_NAMES = ['c', 'sig'] as const;

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

  #device: GPUDevice;
  #host: HostBuffers;
  #initPlan: ModelPlan;
  #stepPlan: ModelPlan;
  /** Timestep, host-owned: it follows from the medium and the grid (a CFL
   *  condition), not from anything the user types, and it is folded into every
   *  setParams so a .m that takes `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: string[];
    npts: 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;
  }

  static async create(opts: GpuModelOptions): Promise<GpuModel> {
    const { device, stencil, grid, medium, source, paramNames, state } = opts;
    const npts = grid.nx * grid.ny;

    // 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: grid.nx },
      ny: { kind: 'const', value: grid.ny },
      h: { kind: 'const', value: grid.h },
      dt: { kind: 'param' },
    };
    for (const g of GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
    for (const m of MEDIUM_NAMES) bindings[m] = { kind: 'tensor', shape: [npts, 1] };
    for (const s of state) bindings[s] = { kind: 'tensor', shape: [npts, 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 }));
    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];

    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 — `init` does not read
    // `p`, but it writes it, and `step` reads it back.
    for (const s of state) host.ensure(s, npts);
    for (const g of GRID_NAMES) host.ensure(g, npts);
    for (const m of MEDIUM_NAMES) host.ensure(m, npts);

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

    host.upload('x', grid.x);
    host.upload('y', grid.y);
    host.upload('c', medium.c);
    host.upload('sig', medium.sig);

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

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

  /** 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 scene is two
   * buffer writes and needs no recompile.
   */
  uploadMedium(medium: MediumFields): void {
    this.#host.upload('c', medium.c);
    this.#host.upload('sig', medium.sig);
  }

  /** 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 sound speed.
   *
   * 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();
  }
}