// Reading and writing .cymsim files (HDF5) with h5wasm. The layout is specified in
// cymsim/docs/format.md; this module writes exactly what the Python package writes, apart from
// the `state` group, which holds this application's own integrator state (writer "cymsimweb").
// The h5wasm module is passed in by the caller: the browser uses the in-memory build ("h5wasm")
// and writes the file image to a Blob, the tests use "h5wasm/node" on the real file system.
import type { SimulationConfig } from "../config/schema.ts";
import { AccelerationStore } from "../analysis/acceleration.ts";

/** cymsimweb's integrator state (opaque here; kept so a loaded file can be written back intact). */
export interface BackendState {
  kind: "cpu" | "gpu";
  step: number;                   // sub-steps done
  samplesDone: number;
  arrays: Record<string, Float32Array | Float64Array | Uint32Array>;
}

export const FORMAT_VERSION = 1;
export const STATE_WRITER = "cymsimweb";
const ENERGY_COLUMNS = "time kinetic membrane bending exciters max_w_over_h";
const HIT_FIELDS = ["time", "peak_force", "contact_time", "rebound_velocity", "energy_in", "energy_out"] as const;

// h5wasm's types are loose; we use the few calls we need through this minimal shape
export interface H5Module { File: new (path: string, mode: string) => any; ready: Promise<{ FS: any }> }

export interface HitRecord { time: number; kind: string; peak_force: number; contact_time: number; rebound_velocity: number; energy_in: number; energy_out: number }

/** Everything a .cymsim file holds, in memory. */
export interface CymsimRun {
  yaml: string; config: SimulationConfig;
  mesh: { xy: Float64Array; z0: Float64Array; thickness: Float64Array; normal: Float64Array; triangles: Int32Array | Uint32Array; area: Float64Array; boundary: Int32Array | Uint32Array; nNodes: number; nElements: number };
  listener: { rFront: Float64Array; rBack: Float64Array };
  fs: number; dt: number; substeps: number; duration: number; nSamples: number; nStepsDone: number; complete: boolean; wallTime: number; decimationSpan: number;
  audioBaffled: Float32Array; audioUnbaffled: Float32Array;              // the two listener signals, final samples so far
  acceleration: AccelerationStore | null; taps: Float64Array | null;  // taps: (W, m)
  displacement: { rate: number; times: number[]; frames: Float32Array[] } | null;
  energy: number[][];                                              // rows of six
  hits: HitRecord[];
  state: BackendState | null;                                      // this application's state, when written by it
  creator: string; created: string;
}

export const dtypeOf = (a: ArrayLike<number>) => a instanceof Float64Array ? "<d" : a instanceof Float32Array ? "<f" : a instanceof Int16Array ? "<h" : a instanceof Uint32Array ? "<I" : "<i";

function attr(g: any, name: string, value: number | string | Float64Array, dtype?: string) {
  if (typeof value === "string") g.create_attribute(name, value);
  else if (typeof value === "number") g.create_attribute(name, value, null, dtype ?? (Number.isInteger(value) ? "<i" : "<d"));
  else g.create_attribute(name, value, [value.length], "<d");
}

/** Create a resizable time-major dataset (rows, cols) and fill it from blocks without holding a copy of everything. */
function writeSeries(g: any, name: string, dtype: string, cols: number, rows: number, chunkRows: number, fill: (write: (r0: number, block: ArrayLike<number>) => void) => void) {
  const shape = cols > 0 ? [rows, cols] : [rows], chunks = cols > 0 ? [Math.max(1, Math.min(chunkRows, rows || chunkRows)), cols] : [Math.max(1, Math.min(chunkRows, rows || chunkRows))];
  const empty = dtype === "<h" ? new Int16Array(0) : dtype === "<d" ? new Float64Array(0) : new Float32Array(0);
  const ds = g.create_dataset({ name, data: empty, shape: cols > 0 ? [0, cols] : [0], dtype, maxshape: cols > 0 ? [null, cols] : [null], chunks });
  if (rows > 0) {
    ds.resize(shape);
    fill((r0, block) => {
      const n = cols > 0 ? block.length / cols : block.length;
      ds.write_slice(cols > 0 ? [[r0, r0 + n], [0, cols]] : [[r0, r0 + n]], block);
    });
  }
  return ds;
}

/** Write ``run`` to ``path`` in h5wasm's file system (in memory in the browser, the real disk under Node). */
export function writeCymsim(h5: H5Module, path: string, run: CymsimRun, creator: string): void {
  const f = new h5.File(path, "w");
  try {
    const N = run.mesh.nNodes, E = run.mesh.nElements;
    attr(f, "format", "cymsim"); attr(f, "format_version", FORMAT_VERSION); attr(f, "creator", creator);
    attr(f, "created", new Date().toISOString().replace(/\.\d{3}Z$/, "+00:00"));
    f.create_group("config").create_dataset({ name: "yaml", data: run.yaml });

    const mesh = f.create_group("mesh");
    attr(mesh, "n_nodes", N); attr(mesh, "n_elements", E);
    mesh.create_dataset({ name: "xy", data: run.mesh.xy, shape: [N, 2], dtype: "<d" });
    mesh.create_dataset({ name: "z0", data: run.mesh.z0, shape: [N], dtype: "<d" });
    mesh.create_dataset({ name: "thickness", data: run.mesh.thickness, shape: [N], dtype: "<d" });
    mesh.create_dataset({ name: "normal", data: run.mesh.normal, shape: [N, 3], dtype: "<d" });
    mesh.create_dataset({ name: "triangles", data: Int32Array.from(run.mesh.triangles), shape: [E, 3], dtype: "<i" });
    mesh.create_dataset({ name: "area", data: run.mesh.area, shape: [E], dtype: "<d" });
    mesh.create_dataset({ name: "boundary", data: Int32Array.from(run.mesh.boundary), shape: [run.mesh.boundary.length], dtype: "<i" });

    const lis = f.create_group("listener"), L = run.config.listener;
    attr(lis, "position", Float64Array.from(L.position)); attr(lis, "c_air", L.c_air, "<d"); attr(lis, "rho_air", L.rho_air, "<d");
    attr(lis, "back_diffraction_gain", L.back_diffraction_gain, "<d");
    lis.create_dataset({ name: "r_front", data: run.listener.rFront, shape: [E], dtype: "<d" });
    lis.create_dataset({ name: "r_back", data: run.listener.rBack, shape: [E], dtype: "<d" });

    const rg = f.create_group("run");
    attr(rg, "fs", run.fs); attr(rg, "dt", run.dt, "<d"); attr(rg, "substeps", run.substeps); attr(rg, "duration", run.duration, "<d");
    attr(rg, "n_samples", run.nSamples); attr(rg, "decimation_span", run.decimationSpan); attr(rg, "n_steps_done", run.nStepsDone);
    attr(rg, "complete", run.complete ? 1 : 0); attr(rg, "wall_time", run.wallTime, "<d");

    const audio = f.create_group("audio");
    attr(audio, "fs", run.fs); attr(audio, "units", "Pa");
    for (const [name, data] of [["baffled", run.audioBaffled], ["unbaffled", run.audioUnbaffled]] as const) {
      writeSeries(audio, name, "<f", 0, data.length, 65536, (write) => write(0, data));
    }

    if (run.acceleration && run.taps) {
      const a = run.acceleration, g = f.create_group("acceleration");
      attr(g, "fs", run.fs); attr(g, "component", "normal"); attr(g, "units", "m/s^2"); attr(g, "encoding", "int16 * scale[t]");
      attr(g, "description", "acceleration of every node along the rest-surface normal (front, +z side, positive), low-pass filtered by the kernel in 'taps' and sampled at fs; sample j is centred on t = j / fs");
      const rows = Math.max(1, Math.min(1024, Math.floor(16e6 / (2 * N))));
      writeSeries(g, "normal", "<h", N, a.nSamples, rows, (write) => { for (const b of a.blocks) write(b.first, b.data); });
      writeSeries(g, "scale", "<f", 0, a.nSamples, 65536, (write) => { for (const b of a.blocks) write(b.first, b.scale); });
      const W = run.decimationSpan, m = run.substeps;
      const taps = g.create_dataset({ name: "taps", data: run.taps, shape: [W, m], dtype: "<d" });
      attr(taps, "layout", "taps[w, kk]: weight of sub-step j0 m + kk in audio sample j0 + w - W/2 + 1");
      attr(taps, "kaiser_beta", 8.0, "<d"); attr(taps, "cutoff", 0.46, "<d");
    }

    if (run.displacement) {
      const d = run.displacement, g = f.create_group("displacement");
      attr(g, "rate", d.rate, "<d"); attr(g, "units", "m"); attr(g, "description", "transverse displacement w of every node at the frame times");
      writeSeries(g, "w", "<f", N, d.frames.length, 64, (write) => d.frames.forEach((fr, i) => write(i, fr)));
      writeSeries(g, "time", "<d", 0, d.times.length, 4096, (write) => write(0, Float64Array.from(d.times)));
    }

    const eg = f.create_group("energy");
    const table = writeSeries(eg, "table", "<d", 6, run.energy.length, 1024, (write) => write(0, Float64Array.from(run.energy.flat())));
    attr(table, "columns", ENERGY_COLUMNS);

    const hg = f.create_group("hits"), K = run.hits.length;
    for (const name of HIT_FIELDS) hg.create_dataset({ name, data: Float64Array.from(run.hits, (h) => h[name]), shape: [K], dtype: "<d" });
    hg.create_dataset({ name: "kind", data: run.hits.map((h) => h.kind), shape: [K], dtype: "S16" });

    if (run.state) {
      const s = run.state, g = f.create_group("state");
      attr(g, "writer", STATE_WRITER); attr(g, "n_steps_done", s.step); attr(g, "samples_done", s.samplesDone); attr(g, "backend", s.kind);
      for (const [name, a] of Object.entries(s.arrays)) g.create_dataset({ name, data: a, shape: [a.length], dtype: dtypeOf(a) });
    }
  } finally { f.close(); }
}

const str = (v: unknown) => (typeof v === "string" ? v : Array.isArray(v) ? String(v[0]) : String(v));
const num = (v: unknown) => Number(Array.isArray(v) ? v[0] : v);

/** Read a .cymsim file from h5wasm's file system. ``parse`` turns the YAML text into a configuration. */
export function readCymsim(h5: H5Module, path: string, parse: (yaml: string) => SimulationConfig): CymsimRun {
  const f = new h5.File(path, "r");
  try {
    if (str(f.get_attribute("format", true)) !== "cymsim") throw new Error("not a .cymsim file");
    const get = (p: string) => { const d = f.get(p); if (!d) throw new Error(`missing ${p}`); return d; };
    const has = (p: string) => f.get(p) !== null && f.get(p) !== undefined;
    const value = <T,>(p: string) => get(p).value as T;
    const a = (p: string, name: string) => get(p).get_attribute(name, true);
    const yamlText = str(value("config/yaml"));
    const config = parse(yamlText);
    const nNodes = num(a("mesh", "n_nodes")), nElements = num(a("mesh", "n_elements"));
    const f64 = (p: string) => Float64Array.from(value<ArrayLike<number>>(p));
    const f32 = (p: string) => Float32Array.from(value<ArrayLike<number>>(p));
    const fs = num(a("run", "fs")), substeps = num(a("run", "substeps"));
    let acceleration: AccelerationStore | null = null, taps: Float64Array | null = null;
    if (has("acceleration")) {
      acceleration = new AccelerationStore(nNodes, fs);
      const ds = get("acceleration/normal"), scale = f32("acceleration/scale"), T = ds.shape[0] as number;
      // read in slabs of about 64 MB so that no second copy of the whole record is needed
      const slab = Math.max(1, Math.floor(32e6 / nNodes));
      for (let r0 = 0; r0 < T; r0 += slab) {
        const r1 = Math.min(T, r0 + slab);
        const data = ds.slice([[r0, r1], [0, nNodes]]) as Int16Array;
        acceleration.pushEncoded(Int16Array.from(data), scale.slice(r0, r1));
      }
      taps = f64("acceleration/taps");
    }
    let displacement: CymsimRun["displacement"] = null;
    if (has("displacement")) {
      const w = get("displacement/w"), Td = w.shape[0] as number, all = w.value as Float32Array, times = Array.from(f64("displacement/time"));
      const frames: Float32Array[] = [];
      for (let i = 0; i < Td; i++) frames.push(Float32Array.from(all.subarray(i * nNodes, (i + 1) * nNodes)));
      displacement = { rate: num(a("displacement", "rate")), times, frames };
    }
    const et = get("energy/table"), erows = et.shape[0] as number, eflat = Array.from(et.value as ArrayLike<number>);
    const energy: number[][] = [];
    for (let i = 0; i < erows; i++) energy.push(eflat.slice(6 * i, 6 * i + 6));
    const K = get("hits/time").shape[0] as number, kinds = get("hits/kind").value as unknown as string[];
    const hits: HitRecord[] = [];
    for (let k = 0; k < K; k++) {
      const h: any = { kind: str(Array.isArray(kinds) ? kinds[k] : kinds) };
      for (const name of HIT_FIELDS) h[name] = Number((get(`hits/${name}`).value as ArrayLike<number>)[k]);
      hits.push(h);
    }
    let state: BackendState | null = null;
    if (has("state") && str(a("state", "writer")) === STATE_WRITER) {
      const g = get("state"), arrays: BackendState["arrays"] = {};
      for (const name of g.keys() as string[]) {
        const v = g.get(name).value;
        arrays[name] = v instanceof Float64Array || v instanceof Float32Array || v instanceof Uint32Array ? v : Float32Array.from(v as ArrayLike<number>);
      }
      state = { kind: str(a("state", "backend")) as "cpu" | "gpu", step: num(a("state", "n_steps_done")), samplesDone: num(a("state", "samples_done")), arrays };
    }
    const optional = (name: string, dflt: number) => { try { return num(a("run", name)); } catch { return dflt; } };
    return {
      yaml: yamlText, config,
      mesh: { xy: f64("mesh/xy"), z0: f64("mesh/z0"), thickness: f64("mesh/thickness"), normal: f64("mesh/normal"), triangles: Int32Array.from(value<ArrayLike<number>>("mesh/triangles")),
              area: f64("mesh/area"), boundary: Int32Array.from(value<ArrayLike<number>>("mesh/boundary")), nNodes, nElements },
      listener: { rFront: f64("listener/r_front"), rBack: f64("listener/r_back") },
      fs, dt: num(a("run", "dt")), substeps, duration: num(a("run", "duration")), nSamples: num(a("run", "n_samples")),
      nStepsDone: optional("n_steps_done", 0), complete: optional("complete", 0) !== 0, wallTime: optional("wall_time", 0), decimationSpan: optional("decimation_span", 32),
      audioBaffled: f32("audio/baffled"), audioUnbaffled: f32("audio/unbaffled"), acceleration, taps, displacement, energy, hits, state,
      creator: str(f.get_attribute("creator", true)), created: str(f.get_attribute("created", true)),
    };
  } finally { f.close(); }
}