/**
 * 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 — each object shaped into its excitation, silent air — and
 * a `step` function that advances everything 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. The state spans several grids — a pair of history
 * fields per vibrating object, plus the air's pressure pair — 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 another.
 *
 * Where dulcimer hardcoded its two grids by name, everything here is data the
 * caller supplies: `consts` are the exact scalars a .m may read (grid counts,
 * spacings, lengths — exact so array constructors keep static shapes),
 * `fields` are the host-uploaded arrays (coordinates, pin masks, the medium
 * and the coupling profiles), `state` names the fields the .m advances and how
 * big each is, and `paramNames` the tunable scalars. The .m names the
 * arguments it wants and the host matches each against this offer; a name
 * nobody provides is a compile error with a position, not a silently undefined
 * variable.
 *
 * There is no clock here, and that is not an oversight: every excitation 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';

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

/** One state field the .m advances, and its element count. */
export interface StateField {
  name: string;
  count: number;
}

export interface GpuModelOptions {
  device: GPUDevice;
  ops: OpPlan;
  /** Exact scalars the .m may take as arguments (grid counts, spacings). */
  consts: Record<string, number>;
  /** Host-uploaded arrays the .m may take as arguments (coordinates, masks,
   *  the medium, the coupling profiles), with their initial contents. */
  fields: Record<string, Float32Array>;
  /** Model source (.m text). */
  source: string;
  /** Tunable scalar 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;
}

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

  #device: GPUDevice;
  #host: HostBuffers;
  #initPlan: ModelPlan;
  #stepPlan: ModelPlan;
  /** Names of the uploadable fields, for uploadFields' checks. */
  #fieldNames: Set<string>;
  /** 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[];
    fieldNames: Set<string>;
  }) {
    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.#fieldNames = init.fieldNames;
  }

  static async create(opts: GpuModelOptions): Promise<GpuModel> {
    const { device, ops, consts, fields, source, paramNames, state } = opts;

    // 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> = { dt: { kind: 'param' } };
    for (const [name, value] of Object.entries(consts)) {
      bindings[name] = { kind: 'const', value };
    }
    for (const [name, data] of Object.entries(fields)) {
      bindings[name] = { kind: 'tensor', shape: [data.length, 1] };
    }
    for (const s of state) bindings[s.name] = { kind: 'tensor', shape: [s.count, 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, ops.instances));
    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, s.count);
    for (const [name, data] of Object.entries(fields)) host.ensure(name, data.length);

    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),
    );

    for (const [name, data] of Object.entries(fields)) host.upload(name, data);

    const maxCount = Math.max(...state.map((s) => s.count), ...Object.values(fields).map((f) => f.length));
    const readback = device.createBuffer({
      label: 'mgpu-readback',
      size: 4 * maxCount,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });

    return new GpuModel({
      device, host, initPlan, stepPlan, readback, paramNames, state,
      fieldNames: new Set(Object.keys(fields)),
    });
  }

  /** 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 host-uploaded fields under a running model — the medium and the
   * coupling profiles after a scene edit. They are data, not code — their
   * shapes in the bindings depend only on the grids — so changing the body's
   * geometry or moving an object is a few buffer writes and no recompile.
   */
  uploadFields(fields: Record<string, Float32Array>): void {
    for (const [name, data] of Object.entries(fields)) {
      if (!this.#fieldNames.has(name)) {
        throw new Error(`uploadFields: '${name}' is not a host field of this model`);
      }
      this.#host.upload(name, data);
    }
  }

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